package handler import ( "context" "net/http" "net/http/httptest" "strings" "testing" "github.com/Tencent/WeKnora/internal/config" "github.com/Tencent/WeKnora/internal/middleware" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" "github.com/gin-gonic/gin" "github.com/hibiken/asynq" "github.com/stretchr/testify/require" ) type batchOwnershipKBService struct { interfaces.KnowledgeBaseService } func (*batchOwnershipKBService) GetKnowledgeBaseByID(context.Context, string) (*types.KnowledgeBase, error) { return &types.KnowledgeBase{ID: "kb-owner", TenantID: 1, CreatorID: "owner"}, nil } type batchOwnershipKnowledgeService struct { interfaces.KnowledgeService updated bool fetched bool } func (*batchOwnershipKnowledgeService) GetKnowledgeByIDOnly(context.Context, string) (*types.Knowledge, error) { return &types.Knowledge{ID: "doc-owner", TenantID: 1, KnowledgeBaseID: "kb-owner"}, nil } func (s *batchOwnershipKnowledgeService) UpdateKnowledgeTagBatch(_ context.Context, kbID string, _ map[string][]string) error { s.updated = kbID == "kb-owner" return nil } func (s *batchOwnershipKnowledgeService) GetKnowledgeBatch(context.Context, uint64, []string) ([]*types.Knowledge, error) { s.fetched = true return []*types.Knowledge{{ID: "doc-owner", TenantID: 1, KnowledgeBaseID: "kb-owner"}}, nil } type batchOwnershipQueue struct{ enqueued bool } func (q *batchOwnershipQueue) Enqueue(*asynq.Task, ...asynq.Option) (*asynq.TaskInfo, error) { q.enqueued = true return &asynq.TaskInfo{ID: "reparse-test"}, nil } func TestKnowledgeBatchMutationsRequireKBOwnerOrAdmin(t *testing.T) { gin.SetMode(gin.TestMode) for _, principal := range []struct { name string user string role types.TenantRole want int }{ {"other contributor", "other", types.TenantRoleContributor, http.StatusForbidden}, {"other viewer", "other", types.TenantRoleViewer, http.StatusForbidden}, {"KB creator", "owner", types.TenantRoleContributor, http.StatusOK}, {"tenant admin", "other", types.TenantRoleAdmin, http.StatusOK}, } { for _, operation := range []struct { name string body string }{ {"tags with KB", `{"kb_id":"kb-owner","updates":{"doc-owner":["tag-1"]}}`}, {"tags inferred KB", `{"updates":{"doc-owner":["tag-1"]}}`}, {"batch reparse", `{"kb_id":"kb-owner","ids":["doc-owner"]}`}, } { t.Run(principal.name+"/"+operation.name, func(t *testing.T) { enabled := true cfg := &config.Config{Tenant: &config.TenantConfig{EnableRBAC: &enabled}} kg := &batchOwnershipKnowledgeService{} queue := &batchOwnershipQueue{} h := &KnowledgeHandler{cfg: cfg, kbService: &batchOwnershipKBService{}, kgService: kg, asynqClient: queue} r := gin.New() r.Use(middleware.ErrorHandler()) r.Use(func(c *gin.Context) { ctx := context.WithValue(c.Request.Context(), types.TenantIDContextKey, uint64(1)) ctx = context.WithValue(ctx, types.UserIDContextKey, principal.user) ctx = context.WithValue(ctx, types.TenantRoleContextKey, principal.role) c.Request = c.Request.WithContext(ctx) c.Set(types.TenantIDContextKey.String(), uint64(1)) c.Set(types.UserIDContextKey.String(), principal.user) c.Next() }) if operation.name == "batch reparse" { r.POST("/batch", h.BatchReparseKnowledge) } else { r.POST("/batch", h.UpdateKnowledgeTagBatch) } req := httptest.NewRequest(http.MethodPost, "/batch", strings.NewReader(operation.body)) req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() r.ServeHTTP(w, req) require.Equal(t, principal.want, w.Code, "body=%s", w.Body.String()) allowed := principal.want == http.StatusOK if operation.name == "batch reparse" { require.Equal(t, allowed, kg.fetched, "reject before reading batch") require.Equal(t, allowed, queue.enqueued, "reject before scheduling reparse") } else { require.Equal(t, allowed, kg.updated, "reject before changing tags") } }) } } }