diff --git a/internal/application/service/tenant_api_key.go b/internal/application/service/tenant_api_key.go index fe35b5c4..ffd17aa8 100644 --- a/internal/application/service/tenant_api_key.go +++ b/internal/application/service/tenant_api_key.go @@ -8,15 +8,23 @@ import ( "encoding/hex" "errors" "strings" + "sync" "time" apprepo "github.com/Tencent/WeKnora/internal/application/repository" + "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" ) +// apiKeyLastUsedMinInterval bounds how often we persist last_used_at per key. +// The UI only needs minute-level freshness; throttling avoids a DB write on +// every authenticated request under high QPS. +const apiKeyLastUsedMinInterval = time.Minute + type tenantAPIKeyService struct { - repo interfaces.TenantAPIKeyRepository + repo interfaces.TenantAPIKeyRepository + lastUsedTouch sync.Map // key ID (uint64) -> time.Time of last persisted touch } func NewTenantAPIKeyService(repo interfaces.TenantAPIKeyRepository) interfaces.TenantAPIKeyService { @@ -70,10 +78,30 @@ func (s *tenantAPIKeyService) AuthenticateAPIKey(ctx context.Context, token stri if key.ExpiresAt != nil && time.Now().After(*key.ExpiresAt) { return nil, apprepo.ErrTenantAPIKeyNotFound } - _ = s.repo.UpdateAPIKeyLastUsed(ctx, key.ID, time.Now()) + s.touchAPIKeyLastUsedAsync(key.ID) return key, nil } +// touchAPIKeyLastUsedAsync persists last_used_at at most once per key per +// apiKeyLastUsedMinInterval. The write runs in a detached goroutine so auth +// latency is not tied to an UPDATE on the hot path. +func (s *tenantAPIKeyService) touchAPIKeyLastUsedAsync(keyID uint64) { + now := time.Now() + if v, ok := s.lastUsedTouch.Load(keyID); ok { + if now.Sub(v.(time.Time)) < apiKeyLastUsedMinInterval { + return + } + } + s.lastUsedTouch.Store(keyID, now) + go func(id uint64, at time.Time) { + if err := s.repo.UpdateAPIKeyLastUsed(context.Background(), id, at); err != nil { + logger.Warnf(context.Background(), + "failed to update tenant api key last_used_at (id=%d): %v", id, err) + s.lastUsedTouch.Delete(id) + } + }(keyID, now) +} + func (s *tenantAPIKeyService) ListAPIKeys(ctx context.Context, tenantID uint64) ([]*types.TenantAPIKey, error) { return s.repo.ListAPIKeys(ctx, tenantID) } diff --git a/internal/application/service/tenant_api_key_test.go b/internal/application/service/tenant_api_key_test.go index 132eeadb..ae8d35bb 100644 --- a/internal/application/service/tenant_api_key_test.go +++ b/internal/application/service/tenant_api_key_test.go @@ -13,8 +13,9 @@ import ( ) type fakeTenantAPIKeyRepo struct { - byHash map[string]*types.TenantAPIKey - nextID uint64 + byHash map[string]*types.TenantAPIKey + nextID uint64 + lastUsedUpdateCount int } func TestTenantAPIKeyServiceCreateAPIKeyUsesSKPrefix(t *testing.T) { @@ -118,6 +119,7 @@ func (r *fakeTenantAPIKeyRepo) ListKeysWithPlaceholderHash(_ context.Context) ([ } func (r *fakeTenantAPIKeyRepo) UpdateAPIKeyLastUsed(_ context.Context, id uint64, at time.Time) error { + r.lastUsedUpdateCount++ for _, key := range r.byHash { if key.ID == id && key.RevokedAt == nil { key.LastUsedAt = &at @@ -178,3 +180,32 @@ func TestTenantAPIKeyServiceRevokeAPIKey(t *testing.T) { t.Fatal("revoked key should not authenticate") } } + +func TestTenantAPIKeyServiceAuthenticateThrottlesLastUsedUpdates(t *testing.T) { + ctx := context.Background() + repo := newFakeTenantAPIKeyRepo() + svc := NewTenantAPIKeyService(repo) + + created, err := svc.CreateAPIKey(ctx, interfaces.TenantAPIKeyCreateRequest{ + TenantID: 42, + Name: "integration", + Role: types.TenantRoleViewer, + }) + if err != nil { + t.Fatalf("CreateAPIKey returned error: %v", err) + } + + for i := 0; i < 5; i++ { + if _, err := svc.AuthenticateAPIKey(ctx, created.Token); err != nil { + t.Fatalf("AuthenticateAPIKey #%d returned error: %v", i+1, err) + } + } + + deadline := time.Now().Add(500 * time.Millisecond) + for repo.lastUsedUpdateCount == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if repo.lastUsedUpdateCount != 1 { + t.Fatalf("last_used update count = %d, want 1 (throttled async write)", repo.lastUsedUpdateCount) + } +} diff --git a/internal/router/router.go b/internal/router/router.go index 365c1704..74aad8bb 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -751,11 +751,12 @@ func RegisterInitializationRoutes(r *gin.RouterGroup, handler *handler.Initializ g.apiKeyRoute(r, http.MethodGet, "/initialization/config/:kbId", apiKeyViewer(), g.Viewer(), g.KBAccessRead("kbId"), handler.GetCurrentConfigByKB) // InitializeByKB / UpdateKBConfig 都是改 KB 的核心模型/storage 配置 — - // 跟 PUT /knowledge-bases/:id 同等敏感，挂同款 OwnedKB 矩阵。 + // 跟 PUT /knowledge-bases/:id 同等敏感，挂同款 OwnedKB 矩阵 + KBAccessWrite + //（API-key 主体短路 Owned* 守卫，KB allow-list 只能靠 KBAccess 兜底）。 g.apiKeyRoute(r, http.MethodPost, "/initialization/initialize/:kbId", - apiKeyContributor(), g.OwnedKBOrAdminFromKbIDParam(), handler.InitializeByKB) + apiKeyContributor(), g.OwnedKBOrAdminFromKbIDParam(), g.KBAccessWrite("kbId"), handler.InitializeByKB) g.apiKeyRoute(r, http.MethodPut, "/initialization/config/:kbId", - apiKeyContributor(), g.OwnedKBOrAdminFromKbIDParam(), handler.UpdateKBConfig) + apiKeyContributor(), g.OwnedKBOrAdminFromKbIDParam(), g.KBAccessWrite("kbId"), handler.UpdateKBConfig) // Ollama / 远程 API / 抽取等系统级检测/下载操作。这些不绑某个 KB， // 会改租户级模型配置或拉远端模型；JWT 侧只读探测 Viewer+、变更 Admin+。 @@ -1797,17 +1798,17 @@ func RegisterWikiPageRoutes(r *gin.RouterGroup, wikiHandler *handler.WikiPageHan { // Page CRUD wiki.GET("/pages", g.Viewer(), g.KBAccessRead("kb_id"), wikiHandler.ListPages) - wiki.POST("/pages", g.OwnedWikiKBOrAdmin(), wikiHandler.CreatePage) - wiki.PUT("/move-page", g.OwnedWikiKBOrAdmin(), wikiHandler.MovePage) + wiki.POST("/pages", g.OwnedWikiKBOrAdmin(), g.KBAccessWrite("kb_id"), wikiHandler.CreatePage) + wiki.PUT("/move-page", g.OwnedWikiKBOrAdmin(), g.KBAccessWrite("kb_id"), wikiHandler.MovePage) wiki.GET("/pages/*slug", g.Viewer(), g.KBAccessRead("kb_id"), wikiHandler.GetPage) - wiki.PUT("/pages/*slug", g.OwnedWikiKBOrAdmin(), wikiHandler.UpdatePage) - wiki.DELETE("/pages/*slug", g.OwnedWikiKBOrAdmin(), wikiHandler.DeletePage) + wiki.PUT("/pages/*slug", g.OwnedWikiKBOrAdmin(), g.KBAccessWrite("kb_id"), wikiHandler.UpdatePage) + wiki.DELETE("/pages/*slug", g.OwnedWikiKBOrAdmin(), g.KBAccessWrite("kb_id"), wikiHandler.DeletePage) // Folder tree (directory nodes) wiki.GET("/folders", g.Viewer(), g.KBAccessRead("kb_id"), wikiHandler.ListFolders) - wiki.POST("/folders", g.OwnedWikiKBOrAdmin(), wikiHandler.CreateFolder) - wiki.PUT("/folders/:folder_id", g.OwnedWikiKBOrAdmin(), wikiHandler.UpdateFolder) - wiki.DELETE("/folders/:folder_id", g.OwnedWikiKBOrAdmin(), wikiHandler.DeleteFolder) + wiki.POST("/folders", g.OwnedWikiKBOrAdmin(), g.KBAccessWrite("kb_id"), wikiHandler.CreateFolder) + wiki.PUT("/folders/:folder_id", g.OwnedWikiKBOrAdmin(), g.KBAccessWrite("kb_id"), wikiHandler.UpdateFolder) + wiki.DELETE("/folders/:folder_id", g.OwnedWikiKBOrAdmin(), g.KBAccessWrite("kb_id"), wikiHandler.DeleteFolder) // Special pages wiki.GET("/index", g.Viewer(), g.KBAccessRead("kb_id"), wikiHandler.GetIndex) @@ -1819,12 +1820,12 @@ func RegisterWikiPageRoutes(r *gin.RouterGroup, wikiHandler *handler.WikiPageHan // Search and maintenance wiki.GET("/search", g.Viewer(), g.KBAccessRead("kb_id"), wikiHandler.SearchPages) - wiki.POST("/rebuild-links", g.OwnedWikiKBOrAdmin(), wikiHandler.RebuildLinks) + wiki.POST("/rebuild-links", g.OwnedWikiKBOrAdmin(), g.KBAccessWrite("kb_id"), wikiHandler.RebuildLinks) wiki.GET("/lint", g.Viewer(), g.KBAccessRead("kb_id"), wikiHandler.Lint) - wiki.POST("/auto-fix", g.OwnedWikiKBOrAdmin(), wikiHandler.AutoFix) + wiki.POST("/auto-fix", g.OwnedWikiKBOrAdmin(), g.KBAccessWrite("kb_id"), wikiHandler.AutoFix) // Issues wiki.GET("/issues", g.Viewer(), g.KBAccessRead("kb_id"), wikiHandler.ListIssues) - wiki.PUT("/issues/:issue_id/status", g.OwnedWikiKBOrAdmin(), wikiHandler.UpdateIssueStatus) + wiki.PUT("/issues/:issue_id/status", g.OwnedWikiKBOrAdmin(), g.KBAccessWrite("kb_id"), wikiHandler.UpdateIssueStatus) } } diff --git a/internal/router/router_wiki_test.go b/internal/router/router_wiki_test.go index f14939ec..e3769c5b 100644 --- a/internal/router/router_wiki_test.go +++ b/internal/router/router_wiki_test.go @@ -4,6 +4,7 @@ import ( "context" "net/http" "net/http/httptest" + "strings" "testing" apprepo "github.com/Tencent/WeKnora/internal/application/repository" @@ -27,34 +28,24 @@ func (s *stubWikiKBLookup) GetKnowledgeBaseByID(_ context.Context, id string) (* } func newWikiRouteTestEngine(t *testing.T, callerTenantID uint64, kbLookup *stubWikiKBLookup) *gin.Engine { - t.Helper() - gin.SetMode(gin.TestMode) - - enabled := true - cfg := &config.Config{ - Tenant: &config.TenantConfig{EnableRBAC: &enabled}, - } - guards := &rbacGuards{ - cfg: cfg, - kbService: kbLookup, - } - - r := gin.New() - r.Use(middleware.ErrorHandler()) - r.Use(func(c *gin.Context) { - ctx := c.Request.Context() - ctx = context.WithValue(ctx, types.TenantIDContextKey, callerTenantID) - ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleViewer) - c.Request = c.Request.WithContext(ctx) - c.Set(types.TenantIDContextKey.String(), callerTenantID) - c.Next() + return newKBRouteTestEngine(t, callerTenantID, kbLookup, nil, func(r *gin.RouterGroup, guards *rbacGuards) { + RegisterWikiPageRoutes(r, &handler.WikiPageHandler{}, guards) }) - - RegisterWikiPageRoutes(r.Group("/api/v1"), &handler.WikiPageHandler{}, guards) - return r } func newInitializationRouteTestEngine(t *testing.T, callerTenantID uint64, kbLookup *stubWikiKBLookup) *gin.Engine { + return newKBRouteTestEngine(t, callerTenantID, kbLookup, nil, func(r *gin.RouterGroup, guards *rbacGuards) { + RegisterInitializationRoutes(r, &handler.InitializationHandler{}, guards) + }) +} + +func newKBRouteTestEngine( + t *testing.T, + callerTenantID uint64, + kbLookup *stubWikiKBLookup, + apiKeyScope *types.TenantAPIKeyScope, + register func(r *gin.RouterGroup, guards *rbacGuards), +) *gin.Engine { t.Helper() gin.SetMode(gin.TestMode) @@ -72,16 +63,33 @@ func newInitializationRouteTestEngine(t *testing.T, callerTenantID uint64, kbLoo r.Use(func(c *gin.Context) { ctx := c.Request.Context() ctx = context.WithValue(ctx, types.TenantIDContextKey, callerTenantID) - ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleViewer) + role := types.TenantRoleViewer + if apiKeyScope != nil { + ctx = types.WithTenantAPIKeyScope(ctx, *apiKeyScope) + role = apiKeyScope.Role + if role == "" { + role = types.TenantRoleViewer + } + } + ctx = context.WithValue(ctx, types.TenantRoleContextKey, role) c.Request = c.Request.WithContext(ctx) c.Set(types.TenantIDContextKey.String(), callerTenantID) c.Next() }) - RegisterInitializationRoutes(r.Group("/api/v1"), &handler.InitializationHandler{}, guards) + register(r.Group("/api/v1"), guards) return r } +func tenantKBLookupFixture() *stubWikiKBLookup { + return &stubWikiKBLookup{ + kbs: map[string]*types.KnowledgeBase{ + "kb-allowed": {ID: "kb-allowed", TenantID: 1, Type: types.KnowledgeBaseTypeWiki}, + "kb-other": {ID: "kb-other", TenantID: 1, Type: types.KnowledgeBaseTypeWiki}, + }, + } +} + func TestInitializationConfigRouteDenyCrossTenantKB(t *testing.T) { kbLookup := &stubWikiKBLookup{ kbs: map[string]*types.KnowledgeBase{ @@ -130,3 +138,67 @@ func TestWikiReadRoutesDenyCrossTenantKB(t *testing.T) { }) } } + +func TestInitializationWriteRoutesDenyOutOfScopeAPIKeyKB(t *testing.T) { + kbLookup := tenantKBLookupFixture() + scope := &types.TenantAPIKeyScope{ + Role: types.TenantRoleContributor, + KnowledgeBaseIDs: types.StringArray{"kb-allowed"}, + } + engine := newKBRouteTestEngine(t, 1, kbLookup, scope, func(r *gin.RouterGroup, guards *rbacGuards) { + RegisterInitializationRoutes(r, &handler.InitializationHandler{}, guards) + }) + + cases := []struct { + method string + path string + body string + }{ + {http.MethodPut, "/api/v1/initialization/config/kb-other", `{}`}, + {http.MethodPost, "/api/v1/initialization/initialize/kb-other", `{}`}, + } + + for _, tc := range cases { + t.Run(tc.method+" "+tc.path, func(t *testing.T) { + rec := httptest.NewRecorder() + req := httptest.NewRequest(tc.method, tc.path, strings.NewReader(tc.body)) + req.Header.Set("Content-Type", "application/json") + engine.ServeHTTP(rec, req) + require.Equal(t, http.StatusForbidden, rec.Code, "body=%s", rec.Body.String()) + }) + } +} + +func TestWikiWriteRoutesDenyOutOfScopeAPIKeyKB(t *testing.T) { + kbLookup := tenantKBLookupFixture() + scope := &types.TenantAPIKeyScope{ + Role: types.TenantRoleContributor, + KnowledgeBaseIDs: types.StringArray{"kb-allowed"}, + } + engine := newKBRouteTestEngine(t, 1, kbLookup, scope, func(r *gin.RouterGroup, guards *rbacGuards) { + RegisterWikiPageRoutes(r, &handler.WikiPageHandler{}, guards) + }) + + cases := []struct { + method string + path string + }{ + {http.MethodPost, "/api/v1/knowledgebase/kb-other/wiki/pages"}, + {http.MethodPut, "/api/v1/knowledgebase/kb-other/wiki/pages/some-page"}, + {http.MethodDelete, "/api/v1/knowledgebase/kb-other/wiki/pages/some-page"}, + {http.MethodPost, "/api/v1/knowledgebase/kb-other/wiki/folders"}, + {http.MethodPost, "/api/v1/knowledgebase/kb-other/wiki/rebuild-links"}, + {http.MethodPost, "/api/v1/knowledgebase/kb-other/wiki/auto-fix"}, + {http.MethodPut, "/api/v1/knowledgebase/kb-other/wiki/issues/1/status"}, + } + + for _, tc := range cases { + t.Run(tc.method+" "+tc.path, func(t *testing.T) { + rec := httptest.NewRecorder() + req := httptest.NewRequest(tc.method, tc.path, strings.NewReader(`{}`)) + req.Header.Set("Content-Type", "application/json") + engine.ServeHTTP(rec, req) + require.Equal(t, http.StatusForbidden, rec.Code, "body=%s", rec.Body.String()) + }) + } +}