diff --git a/internal/application/repository/resource.go b/internal/application/repository/resource.go index 9ac0c431f..e06836de1 100644 --- a/internal/application/repository/resource.go +++ b/internal/application/repository/resource.go @@ -90,8 +90,11 @@ func (r *resourceRepository) GetValidGrant( return &grant, err } +// DeleteExpiredGrants drops grants that are past their expiry. A revoked grant +// is kept until then on purpose: it is the tombstone that stops a token derived +// for the same resource and window from re-creating the row and reviving access. func (r *resourceRepository) DeleteExpiredGrants(ctx context.Context, before time.Time) error { return r.db.WithContext(ctx). - Where("expires_at <= ? OR revoked_at IS NOT NULL", before). + Where("expires_at <= ?", before). Delete(&types.ResourceAccessGrant{}).Error } diff --git a/internal/application/service/resource.go b/internal/application/service/resource.go index 0a3433006..639508665 100644 --- a/internal/application/service/resource.go +++ b/internal/application/service/resource.go @@ -2,6 +2,7 @@ package service import ( "context" + "crypto/hmac" "crypto/rand" "crypto/sha256" "encoding/base64" @@ -12,6 +13,7 @@ import ( "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" + secutils "github.com/Tencent/WeKnora/internal/utils" ) const defaultResourceGrantTTL = 2 * time.Hour @@ -170,6 +172,20 @@ func (s *resourceCatalog) CreateAccessGrant(ctx context.Context, reference strin // Opportunistic cleanup keeps high-volume IM rendering from accumulating // expired capability rows; failure is non-fatal to the current grant. _ = s.repo.DeleteExpiredGrants(ctx, time.Now().UTC()) + + // Prefer reusing this resource's live grant. Rendering an answer or + // re-reading a message history would otherwise insert one row per image per + // request, which turns a read endpoint into a write-heavy one. + token, err := s.reuseOrCreateDerivedGrant(ctx, resource.ID, ttl) + if err != nil { + return "", err + } + if token != "" { + return token, nil + } + + // No reusable grant: either this deployment cannot derive tokens, or the + // derived one is not usable. Mint a fresh random token. for attempt := 0; attempt < 4; attempt++ { token, tokenErr := randomResourceToken() if tokenErr != nil { @@ -183,13 +199,89 @@ func (s *resourceCatalog) CreateAccessGrant(ctx context.Context, reference strin } if err := s.repo.CreateGrant(ctx, grant); err == nil { return token, nil - } else if !strings.Contains(strings.ToLower(err.Error()), "unique") { + } else if !isUniqueViolation(err) { return "", err } } return "", fmt.Errorf("failed to allocate unique resource access token") } +// reuseOrCreateDerivedGrant returns the token of a live grant for resourceID, +// creating the row on first use within the current window. It returns ("", nil) +// when the caller must fall back to a random token. +// +// The token is derived rather than random so it can be recomputed without ever +// storing it: the table holds only the hash, as before, and the plaintext token +// cannot be reconstructed from a database dump without SYSTEM_AES_KEY. +// Authorization still lives entirely in the row — a revoked or expired grant +// stops resolving even though the token derives to the same value. +func (s *resourceCatalog) reuseOrCreateDerivedGrant( + ctx context.Context, resourceID string, ttl time.Duration, +) (string, error) { + token, expiresAt, ok := derivedGrantToken(resourceID, ttl) + if !ok { + return "", nil + } + tokenHash := resourceLocationHash(token) + now := time.Now().UTC() + + existing, err := s.repo.GetValidGrant(ctx, tokenHash, now) + if err != nil { + return "", err + } + if existing != nil { + if existing.ResourceID != resourceID { + // A hash collision across resources is practically impossible, but + // reusing it would hand out access to the wrong file. + return "", nil + } + return token, nil + } + + err = s.repo.CreateGrant(ctx, &types.ResourceAccessGrant{ + TokenHash: tokenHash, + ResourceID: resourceID, + AccessScope: "read", + ExpiresAt: expiresAt, + }) + switch { + case err == nil: + return token, nil + case isUniqueViolation(err): + // The row exists but the lookup above rejected it, so it is revoked: + // reviving it via a fresh insert would defeat the revocation. + return "", nil + default: + return "", err + } +} + +// derivedGrantToken computes the token for resourceID in the current time +// window, together with the expiry a newly created row must carry. The window is +// half the TTL, so a reused grant always has at least ttl/2 of life left and one +// row covers every request in that window. +// +// Returns ok=false when SYSTEM_AES_KEY is not configured, in which case grants +// stay random and per-request as they were. +func derivedGrantToken(resourceID string, ttl time.Duration) (string, time.Time, bool) { + key := secutils.SystemHMACKey() + window := ttl / 2 + if key == nil || window <= 0 || resourceID == "" { + return "", time.Time{}, false + } + windowStart := time.Now().UTC().Truncate(window) + mac := hmac.New(sha256.New, key) + fmt.Fprintf(mac, "resource_grant:v1:%s:%d", resourceID, windowStart.Unix()) + token := base64.RawURLEncoding.EncodeToString(mac.Sum(nil)[:16]) + return token, windowStart.Add(ttl), true +} + +// isUniqueViolation reports whether err is a duplicate-key error. The repository +// returns driver errors verbatim, so the text is all we have. +func isUniqueViolation(err error) bool { + return err != nil && strings.Contains(strings.ToLower(err.Error()), "unique") +} + func (s *resourceCatalog) ResolveAccessGrant(ctx context.Context, token string) (*types.StoredResource, error) { grant, err := s.repo.GetValidGrant(ctx, resourceLocationHash(strings.TrimSpace(token)), time.Now().UTC()) if err != nil { diff --git a/internal/application/service/resource_test.go b/internal/application/service/resource_test.go index 17e6d9a1b..b870efa61 100644 --- a/internal/application/service/resource_test.go +++ b/internal/application/service/resource_test.go @@ -66,6 +66,90 @@ func TestResourceCatalogBindingAndAccessGrant(t *testing.T) { require.Equal(t, uint64(9), resource.TenantID) } +// Rendering one answer resolves the same image many times, and re-reading a +// message history resolves it again on every call. Each resolution used to +// insert a capability row; a live one must be reused instead. +func TestResourceCatalogReusesLiveAccessGrant(t *testing.T) { + t.Setenv("SYSTEM_AES_KEY", "weknora-test-aes-key-32bytes!!!") + catalog, db := newResourceCatalogForTest(t) + ctx := context.Background() + ref, err := catalog.Register(ctx, 9, "local://9/exports/a.png", interfaces.ResourceRegistration{}) + require.NoError(t, err) + + first, err := catalog.CreateAccessGrant(ctx, ref, time.Hour) + require.NoError(t, err) + second, err := catalog.CreateAccessGrant(ctx, ref, time.Hour) + require.NoError(t, err) + + require.Equal(t, first, second) + var grants int64 + require.NoError(t, db.Model(&types.ResourceAccessGrant{}).Count(&grants).Error) + require.EqualValues(t, 1, grants, "a live grant must be reused, not duplicated") + + resource, err := catalog.ResolveAccessGrant(ctx, second) + require.NoError(t, err) + require.Equal(t, uint64(9), resource.TenantID) +} + +// Two resources must never share a grant. +func TestResourceCatalogGrantsArePerResource(t *testing.T) { + t.Setenv("SYSTEM_AES_KEY", "weknora-test-aes-key-32bytes!!!") + catalog, _ := newResourceCatalogForTest(t) + ctx := context.Background() + first, err := catalog.Register(ctx, 9, "local://9/exports/a.png", interfaces.ResourceRegistration{}) + require.NoError(t, err) + second, err := catalog.Register(ctx, 9, "local://9/exports/b.png", interfaces.ResourceRegistration{}) + require.NoError(t, err) + + firstToken, err := catalog.CreateAccessGrant(ctx, first, time.Hour) + require.NoError(t, err) + secondToken, err := catalog.CreateAccessGrant(ctx, second, time.Hour) + require.NoError(t, err) + require.NotEqual(t, firstToken, secondToken) +} + +// Revoking a grant must stick: the derived token would otherwise recompute to +// the same value and a fresh insert would revive the access it just lost. +func TestResourceCatalogDoesNotReviveRevokedGrant(t *testing.T) { + t.Setenv("SYSTEM_AES_KEY", "weknora-test-aes-key-32bytes!!!") + catalog, db := newResourceCatalogForTest(t) + ctx := context.Background() + ref, err := catalog.Register(ctx, 9, "local://9/exports/a.png", interfaces.ResourceRegistration{}) + require.NoError(t, err) + + revoked, err := catalog.CreateAccessGrant(ctx, ref, time.Hour) + require.NoError(t, err) + now := time.Now().UTC() + require.NoError(t, db.Model(&types.ResourceAccessGrant{}). + Where("1 = 1").Update("revoked_at", &now).Error) + + fresh, err := catalog.CreateAccessGrant(ctx, ref, time.Hour) + require.NoError(t, err) + require.NotEqual(t, revoked, fresh, "a revoked token must not be handed out again") + + _, err = catalog.ResolveAccessGrant(ctx, revoked) + require.Error(t, err, "the revoked token must stay unusable") + resource, err := catalog.ResolveAccessGrant(ctx, fresh) + require.NoError(t, err) + require.Equal(t, uint64(9), resource.TenantID) +} + +// Without a signing key the deployment cannot derive tokens, so grants stay +// random and per-request — the behaviour before reuse existed. +func TestResourceCatalogWithoutSigningKeyMintsFreshGrants(t *testing.T) { + t.Setenv("SYSTEM_AES_KEY", "") + catalog, _ := newResourceCatalogForTest(t) + ctx := context.Background() + ref, err := catalog.Register(ctx, 9, "local://9/exports/a.png", interfaces.ResourceRegistration{}) + require.NoError(t, err) + + first, err := catalog.CreateAccessGrant(ctx, ref, time.Hour) + require.NoError(t, err) + second, err := catalog.CreateAccessGrant(ctx, ref, time.Hour) + require.NoError(t, err) + require.NotEqual(t, first, second) +} + func TestResourceCatalogRejectsUnsupportedPhysicalPath(t *testing.T) { catalog, _ := newResourceCatalogForTest(t) _, err := catalog.Register(context.Background(), 7, "https://example.com/a.png", interfaces.ResourceRegistration{}) diff --git a/internal/utils/presign.go b/internal/utils/presign.go index 8d31be15f..c9319fd3b 100644 --- a/internal/utils/presign.go +++ b/internal/utils/presign.go @@ -22,9 +22,10 @@ const ( presignDefaultTTL = 2 * time.Hour ) -// getPresignKey returns the HMAC key derived from SYSTEM_AES_KEY. -// Returns nil if the key is not configured or invalid. -func getPresignKey() []byte { +// SystemHMACKey returns the deployment-wide HMAC key derived from +// SYSTEM_AES_KEY, or nil when it is unset or too short to be a real secret. +// Callers must treat nil as "this deployment cannot sign", not as an empty key. +func SystemHMACKey() []byte { key := os.Getenv("SYSTEM_AES_KEY") if len(key) < 16 { return nil @@ -32,6 +33,12 @@ func getPresignKey() []byte { return []byte(key) } +// getPresignKey returns the HMAC key derived from SYSTEM_AES_KEY. +// Returns nil if the key is not configured or invalid. +func getPresignKey() []byte { + return SystemHMACKey() +} + // signPayload computes HMAC-SHA256 over the canonical payload string. func signPayload(key []byte, filePath string, tenantID uint64, expires int64) string { payload := fmt.Sprintf("file_path=%s&tenant_id=%d&expires=%d", filePath, tenantID, expires)