Weknora/Enterprise-Knowledge-Base
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(api-keys): implement throttled last_used_at updates for API key authentication

Browse Source 
Enhance the tenant API key service by introducing asynchronous throttling for the last_used_at updates during API key authentication. This change reduces database write frequency under high request rates, improving performance. Additionally, update tests to validate the throttling behavior, ensuring that last_used_at is only updated once per key within the defined interval. This implementation aligns with the goal of optimizing API key management and maintaining system efficiency.
This commit is contained in:
wizardchen
2026-07-06 19:38:16 +08:00
committed by lyingbug
parent 97c36f634a
commit b5cc7d2a28
4 changed files with 175 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/application/service/tenant_api_key.go
+30 -2
View File
@@ -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)
}
internal/application/service/tenant_api_key_test.go
+33 -2
View File
@@ -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)
}
}
internal/router/router.go
+14 -13
View File
@@ -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)
}
}
internal/router/router_wiki_test.go
+98 -26
View File
@@ -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())
})
}
}