xinmi/internal/router/router_knowledge_export_test.go

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")
			})
		}
	}
}