package router import ( "context" "io" "net/http" "net/http/httptest" "strings" "testing" "github.com/Tencent/WeKnora/internal/accesspolicy" "github.com/Tencent/WeKnora/internal/config" "github.com/Tencent/WeKnora/internal/handler" "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/stretchr/testify/require" ) type exportKnowledgeService struct { interfaces.KnowledgeService knowledge *types.Knowledge exported bool } func (s *exportKnowledgeService) GetKnowledgeByIDOnly(ctx context.Context, id string) (*types.Knowledge, error) { return (&downloadKnowledgeLookup{knowledge: s.knowledge}).GetKnowledgeByIDOnly(ctx, id) } func (s *exportKnowledgeService) GetKnowledgeFile(context.Context, string) (io.ReadCloser, string, error) { s.exported = true return io.NopCloser(strings.NewReader("source document")), "source.txt", nil } func (s *exportKnowledgeService) ExportFAQEntries(context.Context, string) ([]byte, error) { s.exported = true return []byte("question,answer\nq,a\n"), nil } func (s *exportKnowledgeService) ExportFAQEntriesJSON(context.Context, string) ([]byte, error) { s.exported = true return []byte(`[{"question":"q","answer":"a"}]`), nil } type exportResponseRecorder struct{ *httptest.ResponseRecorder } func (*exportResponseRecorder) CloseNotify() <-chan bool { return nil } func TestKnowledgeExportPublicationAndExistingBoundaries(t *testing.T) { gin.SetMode(gin.TestMode) retrieveKey := &types.TenantAPIKeyScope{ KeyID: 1, Capabilities: types.StringArray{"retrieve"}, KnowledgeBaseIDs: types.StringArray{"kb-export"}, } cases := []struct { name string role types.TenantRole global bool creator string systemAdmin bool native bool foreign bool grant bool scope *types.TenantAPIKeyScope want int }{ {name: "global contributor", role: types.TenantRoleContributor, global: true, want: http.StatusForbidden}, {name: "global creator", role: types.TenantRoleContributor, global: true, creator: "requester", want: http.StatusForbidden}, {name: "global grant", role: types.TenantRoleContributor, global: true, grant: true, want: http.StatusForbidden}, {name: "global viewer", role: types.TenantRoleViewer, global: true, want: http.StatusForbidden}, {name: "global admin", role: types.TenantRoleAdmin, global: true, want: http.StatusOK}, {name: "global owner role", role: types.TenantRoleOwner, global: true, want: http.StatusOK}, {name: "global system admin", role: types.TenantRoleContributor, systemAdmin: true, global: true, want: http.StatusOK}, {name: "private creator", role: types.TenantRoleContributor, creator: "requester", want: http.StatusOK}, {name: "private granted nonowner", role: types.TenantRoleContributor, creator: "owner", grant: true, want: http.StatusForbidden}, {name: "private granted unknown creator", role: types.TenantRoleContributor, grant: true, want: http.StatusForbidden}, {name: "private granted admin", role: types.TenantRoleAdmin, creator: "owner", grant: true, want: http.StatusOK}, {name: "native without snapshot", role: types.TenantRoleContributor, native: true, want: http.StatusOK}, {name: "foreign private", role: types.TenantRoleContributor, foreign: true, want: http.StatusForbidden}, {name: "foreign admin", role: types.TenantRoleAdmin, foreign: true, global: true, want: http.StatusForbidden}, {name: "foreign system admin", role: types.TenantRoleContributor, systemAdmin: true, foreign: true, global: true, want: http.StatusForbidden}, {name: "scoped retrieve key", role: types.TenantRoleViewer, global: true, scope: retrieveKey, want: http.StatusOK}, {name: "full access key", role: types.TenantRoleViewer, global: true, scope: &types.TenantAPIKeyScope{KeyID: 2, FullAccess: true}, want: http.StatusOK}, {name: "key lacks retrieve", role: types.TenantRoleOwner, global: true, scope: &types.TenantAPIKeyScope{KeyID: 3, Capabilities: types.StringArray{"ingest"}}, want: http.StatusForbidden}, {name: "key outside KB scope", role: types.TenantRoleOwner, global: true, scope: &types.TenantAPIKeyScope{KeyID: 4, Capabilities: types.StringArray{"retrieve"}, KnowledgeBaseIDs: types.StringArray{"other-kb"}}, want: http.StatusForbidden}, {name: "foreign retrieve key", role: types.TenantRoleOwner, foreign: true, global: true, scope: retrieveKey, want: http.StatusForbidden}, } for _, tc := range cases { for _, endpoint := range []struct{ name, path string }{ {"document", "/api/v1/knowledge/document-export/download"}, {"FAQ CSV", "/api/v1/knowledge-bases/kb-export/faq/entries/export"}, {"FAQ JSON", "/api/v1/knowledge-bases/kb-export/faq/entries/export?format=json"}, } { t.Run(tc.name+"/"+endpoint.name, func(t *testing.T) { enabled := true cfg := &config.Config{Tenant: &config.TenantConfig{EnableRBAC: &enabled}} kb := &types.KnowledgeBase{ID: "kb-export", TenantID: 1, CreatorID: tc.creator} if tc.foreign { kb.TenantID = 2 } svc := &exportKnowledgeService{knowledge: &types.Knowledge{ID: "document-export", KnowledgeBaseID: kb.ID, TenantID: kb.TenantID}} guards := &rbacGuards{ cfg: cfg, knowledgeService: svc, kbService: &stubWikiKBLookup{kbs: map[string]*types.KnowledgeBase{kb.ID: kb}}, apiKeyAuthorizer: middleware.NewAPIKeyRouteAuthorizer(), } 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, "requester") ctx = context.WithValue(ctx, types.TenantRoleContextKey, tc.role) ctx = context.WithValue(ctx, types.SystemAdminContextKey, tc.systemAdmin) if !tc.native { ctx = accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{ Publications: map[string]bool{"knowledge:" + kb.ID: tc.global}, Grants: map[string]bool{"knowledge:" + kb.ID: tc.grant}, }) } if tc.scope != nil { ctx = types.WithTenantAPIKeyScope(ctx, *tc.scope) } c.Request = c.Request.WithContext(ctx) c.Set(types.TenantIDContextKey.String(), uint64(1)) c.Set(types.UserIDContextKey.String(), "requester") c.Next() }) r.Use(guards.apiKeyAuthorizer.Middleware()) v1 := r.Group("/api/v1") RegisterKnowledgeRoutes(v1, handler.NewKnowledgeHandler(cfg, svc, nil, nil, nil, nil, nil), guards) RegisterFAQRoutes(v1, handler.NewFAQHandler(svc, nil), guards) rec := &exportResponseRecorder{httptest.NewRecorder()} r.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, endpoint.path, nil)) require.Equal(t, tc.want, rec.Code, "body=%s", rec.Body.String()) require.Equal(t, tc.want == http.StatusOK, svc.exported, "export must not reach the service after rejection") }) } } }