package middleware import ( "context" "net/http" "net/http/httptest" "testing" "github.com/Tencent/WeKnora/internal/accesspolicy" "github.com/Tencent/WeKnora/internal/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) func TestRequireKnowledgeExportFailsClosedWithoutResolvedKnowledgeBase(t *testing.T) { gin.SetMode(gin.TestMode) for _, tc := range []struct { name string access any }{ {name: "absent"}, {name: "nil access", access: (*KBAccess)(nil)}, {name: "missing KB", access: &KBAccess{}}, {name: "invalid type", access: "kb-id"}, } { t.Run(tc.name, func(t *testing.T) { r := gin.New() r.Use(ErrorHandler()) r.Use(func(c *gin.Context) { ctx := accesspolicy.WithSnapshot(c.Request.Context(), &accesspolicy.Snapshot{ Publications: map[string]bool{"knowledge:kb-id": true}, }) c.Request = c.Request.WithContext(ctx) if tc.access != nil { c.Set(KBAccessContextKey, tc.access) } c.Next() }) r.GET("/export", RequireKnowledgeExport(), func(c *gin.Context) { t.Fatal("export handler ran without an authorized KB") }) rec := httptest.NewRecorder() r.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/export", nil)) require.Equal(t, http.StatusServiceUnavailable, rec.Code) }) } } func TestRequireKnowledgeExportPreservesSystemAdministrator(t *testing.T) { gin.SetMode(gin.TestMode) r := gin.New() r.Use(func(c *gin.Context) { ctx := context.WithValue(c.Request.Context(), types.SystemAdminContextKey, true) ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleViewer) ctx = accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-id": true}}) c.Request = c.Request.WithContext(ctx) c.Set(KBAccessContextKey, &KBAccess{KnowledgeBase: &types.KnowledgeBase{ID: "kb-id"}}) c.Next() }) r.GET("/export", RequireKnowledgeExport(), func(c *gin.Context) { c.Status(http.StatusNoContent) }) rec := httptest.NewRecorder() r.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/export", nil)) require.Equal(t, http.StatusNoContent, rec.Code) }