From 430d4b5394a470821db464a27a6369aed957018a Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Tue, 30 Jun 2026 21:23:53 -0700 Subject: [PATCH] kms: cache decrypted data keys so cache_enabled/cache_ttl take effect (#10173) * kms: cache decrypted data keys so cache_enabled/cache_ttl take effect cache_enabled/cache_ttl were parsed into KMSConfig but never consulted, so every SSE-KMS read repeated a KMS Decrypt round-trip. Add a CachedKMSProvider decorator keyed on (ciphertext, encryption context) and wrap providers with it wherever one is created. Decrypt is deterministic for that pair; GenerateDataKey/DescribeKey/GetKeyID pass through untouched. Cache hits and stores return private copies so the read path's ClearSensitiveData can't wipe the cached key. TTL and max_cache_size bound the cache. * kms: zero superseded data keys and pool cache-key hash states Wipe the old plaintext when a cache entry is overwritten so a key material buffer left behind by two readers racing on the same miss does not linger in memory. Reuse sha256 states via a sync.Pool so the read hot path stops allocating a fresh hash on every Decrypt. * kms: drop cache writes after Close An in-flight Decrypt miss can call set after Close scrubbed the map, repopulating it with a data key that would then never be cleared. Guard set with a closed flag so post-Close stores wipe their copy and no-op. --- weed/kms/cache.go | 245 ++++++++++++++++++++++++++++++++++++ weed/kms/cache_test.go | 273 +++++++++++++++++++++++++++++++++++++++++ weed/kms/config.go | 5 + 3 files changed, 523 insertions(+) create mode 100644 weed/kms/cache.go create mode 100644 weed/kms/cache_test.go diff --git a/weed/kms/cache.go b/weed/kms/cache.go new file mode 100644 index 0000000000..8a4a694aa5 --- /dev/null +++ b/weed/kms/cache.go @@ -0,0 +1,245 @@ +package kms + +import ( + "context" + "crypto/sha256" + "encoding/binary" + "hash" + "sort" + "sync" + "time" + + "github.com/seaweedfs/seaweedfs/weed/glog" +) + +const ( + // DefaultDataKeyCacheTTL is used when caching is enabled without an explicit TTL. + DefaultDataKeyCacheTTL = time.Hour + // DefaultDataKeyCacheSize bounds the number of cached decrypted data keys. + DefaultDataKeyCacheSize = 1000 +) + +// CachedKMSProvider wraps a KMSProvider with an in-memory cache of decrypted +// data keys. A KMS Decrypt call is deterministic for a given (ciphertext, +// encryption context) pair, so every read of the same SSE-KMS object would +// otherwise repeat an identical, network-bound KMS round-trip. Caching the +// plaintext data key removes that round-trip on the read hot path — this is +// what the cache_enabled / cache_ttl provider settings control. +// +// Only Decrypt is cached. GenerateDataKey must return a unique key per call, +// and DescribeKey / GetKeyID are cheap metadata lookups, so all three pass +// straight through to the wrapped provider via the embedded interface. +type CachedKMSProvider struct { + KMSProvider // wrapped provider; supplies the un-cached methods + + ttl time.Duration + maxEntries int + + mu sync.Mutex + entries map[string]*dataKeyCacheEntry + closed bool +} + +type dataKeyCacheEntry struct { + keyID string + plaintext []byte + expiresAt time.Time +} + +// NewCachedKMSProvider wraps provider with a decrypted-data-key cache. A +// non-positive ttl or maxEntries falls back to the package defaults. +func NewCachedKMSProvider(provider KMSProvider, ttl time.Duration, maxEntries int) *CachedKMSProvider { + if ttl <= 0 { + ttl = DefaultDataKeyCacheTTL + } + if maxEntries <= 0 { + maxEntries = DefaultDataKeyCacheSize + } + return &CachedKMSProvider{ + KMSProvider: provider, + ttl: ttl, + maxEntries: maxEntries, + entries: make(map[string]*dataKeyCacheEntry), + } +} + +// Decrypt returns the plaintext data key from cache when present, otherwise +// delegates to the wrapped provider and caches the result. +func (c *CachedKMSProvider) Decrypt(ctx context.Context, req *DecryptRequest) (*DecryptResponse, error) { + if req == nil || len(req.CiphertextBlob) == 0 { + // Let the wrapped provider produce its normal validation error. + return c.KMSProvider.Decrypt(ctx, req) + } + + cacheKey := dataKeyCacheKey(req) + if resp := c.get(cacheKey); resp != nil { + glog.V(4).Infof("KMS data key cache hit") + return resp, nil + } + + resp, err := c.KMSProvider.Decrypt(ctx, req) + if err != nil { + return nil, err + } + + // Store our own copy, then hand the original back to the caller. The + // caller clears the plaintext it receives (ClearSensitiveData); keeping a + // private copy means later cache hits still return valid key material. + c.set(cacheKey, resp) + return resp, nil +} + +// get returns a fresh copy of the cached response, or nil on miss/expiry. The +// copy is essential: callers clear the returned plaintext after use, which +// would otherwise wipe the cached key material for every subsequent reader. +func (c *CachedKMSProvider) get(cacheKey string) *DecryptResponse { + c.mu.Lock() + defer c.mu.Unlock() + + entry, ok := c.entries[cacheKey] + if !ok { + return nil + } + if time.Now().After(entry.expiresAt) { + ClearSensitiveData(entry.plaintext) + delete(c.entries, cacheKey) + return nil + } + + plaintext := make([]byte, len(entry.plaintext)) + copy(plaintext, entry.plaintext) + return &DecryptResponse{KeyID: entry.keyID, Plaintext: plaintext} +} + +// set stores a private copy of the plaintext data key. +func (c *CachedKMSProvider) set(cacheKey string, resp *DecryptResponse) { + if resp == nil { + return + } + + plaintext := make([]byte, len(resp.Plaintext)) + copy(plaintext, resp.Plaintext) + + c.mu.Lock() + defer c.mu.Unlock() + + if c.closed { + // A miss can finish its underlying Decrypt after Close scrubbed the + // map; don't repopulate it with key material that won't be cleared. + ClearSensitiveData(plaintext) + return + } + + if old, exists := c.entries[cacheKey]; exists { + // Two readers can race on the same miss and both store; wipe the + // superseded key material instead of leaving it to linger in memory. + ClearSensitiveData(old.plaintext) + } else { + c.evictIfFullLocked() + } + c.entries[cacheKey] = &dataKeyCacheEntry{ + keyID: resp.KeyID, + plaintext: plaintext, + expiresAt: time.Now().Add(c.ttl), + } +} + +// evictIfFullLocked drops expired entries and, if still at capacity, the entry +// closest to expiry. Callers must hold c.mu. +func (c *CachedKMSProvider) evictIfFullLocked() { + if len(c.entries) < c.maxEntries { + return + } + + now := time.Now() + for k, entry := range c.entries { + if now.After(entry.expiresAt) { + ClearSensitiveData(entry.plaintext) + delete(c.entries, k) + } + } + + if len(c.entries) < c.maxEntries { + return + } + + // Still full: evict the soonest-to-expire entry. + var oldestKey string + var oldestAt time.Time + for k, entry := range c.entries { + if oldestKey == "" || entry.expiresAt.Before(oldestAt) { + oldestKey = k + oldestAt = entry.expiresAt + } + } + if oldestKey != "" { + ClearSensitiveData(c.entries[oldestKey].plaintext) + delete(c.entries, oldestKey) + } +} + +// Close clears cached key material and closes the wrapped provider. +func (c *CachedKMSProvider) Close() error { + c.mu.Lock() + c.closed = true + for k, entry := range c.entries { + ClearSensitiveData(entry.plaintext) + delete(c.entries, k) + } + c.mu.Unlock() + return c.KMSProvider.Close() +} + +// sha256Pool reuses hash states across cache-key computations so the read hot +// path doesn't allocate a fresh one on every Decrypt. +var sha256Pool = sync.Pool{ + New: func() interface{} { return sha256.New() }, +} + +// dataKeyCacheKey derives a stable cache key from the ciphertext blob and the +// encryption context. The context is part of the key because the same +// ciphertext decrypted under a different context is a different request, and +// caching by ciphertext alone would bypass the provider's context check. +func dataKeyCacheKey(req *DecryptRequest) string { + h := sha256Pool.Get().(hash.Hash) + + // Length-prefix each field so distinct inputs cannot collide by concatenation. + writeLengthPrefixed(h, req.CiphertextBlob) + + keys := make([]string, 0, len(req.EncryptionContext)) + for k := range req.EncryptionContext { + keys = append(keys, k) + } + sort.Strings(keys) + for _, k := range keys { + writeLengthPrefixed(h, []byte(k)) + writeLengthPrefixed(h, []byte(req.EncryptionContext[k])) + } + + key := string(h.Sum(nil)) + h.Reset() + sha256Pool.Put(h) + return key +} + +func writeLengthPrefixed(h hash.Hash, b []byte) { + var lenBuf [8]byte + binary.BigEndian.PutUint64(lenBuf[:], uint64(len(b))) + h.Write(lenBuf[:]) + h.Write(b) +} + +// wrapWithCacheIfEnabled returns provider wrapped in a data-key cache when the +// configuration enables caching, otherwise provider unchanged. +func wrapWithCacheIfEnabled(provider KMSProvider, config *KMSConfig) KMSProvider { + if provider == nil || config == nil || !config.CacheEnabled { + return provider + } + // Avoid double-wrapping if a cached provider is somehow reused. + if _, ok := provider.(*CachedKMSProvider); ok { + return provider + } + cached := NewCachedKMSProvider(provider, config.CacheTTL, config.MaxCacheSize) + glog.V(1).Infof("KMS data key caching enabled (ttl=%s, max=%d)", cached.ttl, cached.maxEntries) + return cached +} diff --git a/weed/kms/cache_test.go b/weed/kms/cache_test.go new file mode 100644 index 0000000000..41dab7d7e8 --- /dev/null +++ b/weed/kms/cache_test.go @@ -0,0 +1,273 @@ +package kms + +import ( + "bytes" + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/seaweedfs/seaweedfs/weed/util" +) + +// countingKMSProvider is a minimal KMSProvider that records how many times each +// method is invoked so tests can assert the cache actually avoids round-trips. +type countingKMSProvider struct { + decryptCalls int64 + generateCalls int64 + + mu sync.Mutex + store map[string][]byte // ciphertext(string) -> plaintext data key +} + +func newCountingKMSProvider() *countingKMSProvider { + return &countingKMSProvider{store: make(map[string][]byte)} +} + +func (p *countingKMSProvider) GenerateDataKey(ctx context.Context, req *GenerateDataKeyRequest) (*GenerateDataKeyResponse, error) { + n := atomic.AddInt64(&p.generateCalls, 1) + plaintext := []byte(fmt.Sprintf("datakey-%d-abcdefghijklmnopqrstuv", n))[:32] + ciphertext := []byte(fmt.Sprintf("cipher-%d", n)) + + p.mu.Lock() + p.store[string(ciphertext)] = append([]byte(nil), plaintext...) + p.mu.Unlock() + + return &GenerateDataKeyResponse{ + KeyID: req.KeyID, + Plaintext: plaintext, + CiphertextBlob: ciphertext, + }, nil +} + +func (p *countingKMSProvider) Decrypt(ctx context.Context, req *DecryptRequest) (*DecryptResponse, error) { + atomic.AddInt64(&p.decryptCalls, 1) + + p.mu.Lock() + plaintext, ok := p.store[string(req.CiphertextBlob)] + p.mu.Unlock() + if !ok { + return nil, fmt.Errorf("ciphertext not found") + } + + // Return a fresh copy, mirroring real providers that allocate per call. + return &DecryptResponse{ + KeyID: "test-key", + Plaintext: append([]byte(nil), plaintext...), + }, nil +} + +func (p *countingKMSProvider) DescribeKey(ctx context.Context, req *DescribeKeyRequest) (*DescribeKeyResponse, error) { + return &DescribeKeyResponse{KeyID: req.KeyID, KeyState: KeyStateEnabled, KeyUsage: KeyUsageEncryptDecrypt}, nil +} + +func (p *countingKMSProvider) GetKeyID(ctx context.Context, keyIdentifier string) (string, error) { + return keyIdentifier, nil +} + +func (p *countingKMSProvider) Close() error { return nil } + +// TestCachedProviderDecryptHitsProviderOnce is the regression test for the +// reported bug: with caching enabled, repeated Decrypt of the same data key +// must consult the underlying KMS only once. +func TestCachedProviderDecryptHitsProviderOnce(t *testing.T) { + inner := newCountingKMSProvider() + cached := NewCachedKMSProvider(inner, time.Hour, 100) + + ctx := context.Background() + dk, err := cached.GenerateDataKey(ctx, &GenerateDataKeyRequest{KeyID: "k1", KeySpec: KeySpecAES256}) + if err != nil { + t.Fatalf("GenerateDataKey: %v", err) + } + + req := &DecryptRequest{CiphertextBlob: dk.CiphertextBlob} + first, err := cached.Decrypt(ctx, req) + if err != nil { + t.Fatalf("first Decrypt: %v", err) + } + + // Emulate the S3 read path clearing the returned plaintext after use. A + // correct cache must hand back an independent copy so this does not wipe + // the cached key material. + ClearSensitiveData(first.Plaintext) + + second, err := cached.Decrypt(ctx, req) + if err != nil { + t.Fatalf("second Decrypt: %v", err) + } + + if got := atomic.LoadInt64(&inner.decryptCalls); got != 1 { + t.Fatalf("expected underlying Decrypt to be called once, got %d", got) + } + + // The cached result must still be the real data key. If the cache handed + // back the same buffer it stored, the ClearSensitiveData above would have + // zeroed it. + if bytes.Equal(second.Plaintext, make([]byte, len(second.Plaintext))) { + t.Fatalf("cached plaintext was zeroed by a previous caller clearing its copy") + } +} + +// TestCachedProviderContextSensitivity ensures the encryption context is part +// of the cache key: a different context must not return a stale plaintext. +func TestCachedProviderContextSensitivity(t *testing.T) { + inner := newCountingKMSProvider() + cached := NewCachedKMSProvider(inner, time.Hour, 100) + ctx := context.Background() + + dk, err := cached.GenerateDataKey(ctx, &GenerateDataKeyRequest{KeyID: "k1", KeySpec: KeySpecAES256}) + if err != nil { + t.Fatalf("GenerateDataKey: %v", err) + } + + // Same ciphertext, two different contexts. The counting provider ignores + // context, but the cache must still treat these as distinct entries so a + // real provider's context check is never short-circuited. + if _, err := cached.Decrypt(ctx, &DecryptRequest{CiphertextBlob: dk.CiphertextBlob, EncryptionContext: map[string]string{"a": "1"}}); err != nil { + t.Fatalf("Decrypt ctx a: %v", err) + } + if _, err := cached.Decrypt(ctx, &DecryptRequest{CiphertextBlob: dk.CiphertextBlob, EncryptionContext: map[string]string{"a": "2"}}); err != nil { + t.Fatalf("Decrypt ctx b: %v", err) + } + if got := atomic.LoadInt64(&inner.decryptCalls); got != 2 { + t.Fatalf("expected 2 underlying Decrypt calls for distinct contexts, got %d", got) + } +} + +// TestCachedProviderTTLExpiry verifies expired entries fall through to the +// underlying provider. +func TestCachedProviderTTLExpiry(t *testing.T) { + inner := newCountingKMSProvider() + cached := NewCachedKMSProvider(inner, 20*time.Millisecond, 100) + ctx := context.Background() + + dk, err := cached.GenerateDataKey(ctx, &GenerateDataKeyRequest{KeyID: "k1", KeySpec: KeySpecAES256}) + if err != nil { + t.Fatalf("GenerateDataKey: %v", err) + } + req := &DecryptRequest{CiphertextBlob: dk.CiphertextBlob} + + if _, err := cached.Decrypt(ctx, req); err != nil { + t.Fatalf("Decrypt: %v", err) + } + time.Sleep(40 * time.Millisecond) + if _, err := cached.Decrypt(ctx, req); err != nil { + t.Fatalf("Decrypt after expiry: %v", err) + } + if got := atomic.LoadInt64(&inner.decryptCalls); got != 2 { + t.Fatalf("expected 2 underlying Decrypt calls across TTL expiry, got %d", got) + } +} + +// TestCachedProviderMaxEntries verifies the cache is size-bounded. +func TestCachedProviderMaxEntries(t *testing.T) { + inner := newCountingKMSProvider() + cached := NewCachedKMSProvider(inner, time.Hour, 2) + ctx := context.Background() + + var blobs [][]byte + for i := 0; i < 3; i++ { + dk, err := cached.GenerateDataKey(ctx, &GenerateDataKeyRequest{KeyID: "k1", KeySpec: KeySpecAES256}) + if err != nil { + t.Fatalf("GenerateDataKey: %v", err) + } + blobs = append(blobs, dk.CiphertextBlob) + if _, err := cached.Decrypt(ctx, &DecryptRequest{CiphertextBlob: dk.CiphertextBlob}); err != nil { + t.Fatalf("Decrypt: %v", err) + } + } + + cached.mu.Lock() + size := len(cached.entries) + cached.mu.Unlock() + if size > 2 { + t.Fatalf("cache exceeded max entries: got %d, want <= 2", size) + } +} + +// TestCachedProviderClearsOverwrittenPlaintext verifies that when an entry is +// replaced (two readers racing on the same miss), the superseded key material +// is zeroed rather than left lingering in memory. +func TestCachedProviderClearsOverwrittenPlaintext(t *testing.T) { + inner := newCountingKMSProvider() + cached := NewCachedKMSProvider(inner, time.Hour, 10) + + cached.set("dupkey", &DecryptResponse{KeyID: "k", Plaintext: []byte("0123456789abcdef0123456789abcdef")}) + + cached.mu.Lock() + oldBuf := cached.entries["dupkey"].plaintext + cached.mu.Unlock() + + cached.set("dupkey", &DecryptResponse{KeyID: "k", Plaintext: []byte("ffffffffffffffffffffffffffffffff")}) + + if !bytes.Equal(oldBuf, make([]byte, len(oldBuf))) { + t.Fatalf("superseded cache plaintext was not zeroed: %v", oldBuf) + } +} + +// TestCachedProviderNoWriteAfterClose verifies a store that lands after Close +// does not repopulate the scrubbed cache with key material. +func TestCachedProviderNoWriteAfterClose(t *testing.T) { + inner := newCountingKMSProvider() + cached := NewCachedKMSProvider(inner, time.Hour, 10) + + if err := cached.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + cached.set("k", &DecryptResponse{KeyID: "k", Plaintext: []byte("0123456789abcdef0123456789abcdef")}) + + cached.mu.Lock() + n := len(cached.entries) + cached.mu.Unlock() + if n != 0 { + t.Fatalf("cache repopulated after Close: %d entries", n) + } +} + +// TestAddKMSProviderWiresCache is the wiring regression test: a provider +// configured with cache_enabled must actually be handed back wrapped in a +// caching provider (the bug was that the setting was parsed but never applied). +func TestAddKMSProviderWiresCache(t *testing.T) { + RegisterProvider("cache-wiring-test", func(config util.Configuration) (KMSProvider, error) { + return newCountingKMSProvider(), nil + }) + + mgr := &KMSManager{ + providers: make(map[string]KMSProvider), + configs: make(map[string]*KMSConfig), + bucketKMS: make(map[string]string), + } + + if err := mgr.AddKMSProvider("p-cached", &KMSConfig{ + Provider: "cache-wiring-test", + CacheEnabled: true, + CacheTTL: time.Hour, + MaxCacheSize: 10, + }); err != nil { + t.Fatalf("AddKMSProvider (cached): %v", err) + } + p, err := mgr.GetKMSProviderByName("p-cached") + if err != nil { + t.Fatalf("GetKMSProviderByName: %v", err) + } + if _, ok := p.(*CachedKMSProvider); !ok { + t.Fatalf("expected *CachedKMSProvider when cache enabled, got %T", p) + } + + if err := mgr.AddKMSProvider("p-plain", &KMSConfig{ + Provider: "cache-wiring-test", + CacheEnabled: false, + }); err != nil { + t.Fatalf("AddKMSProvider (plain): %v", err) + } + p2, err := mgr.GetKMSProviderByName("p-plain") + if err != nil { + t.Fatalf("GetKMSProviderByName: %v", err) + } + if _, ok := p2.(*CachedKMSProvider); ok { + t.Fatalf("expected unwrapped provider when cache disabled, got %T", p2) + } +} diff --git a/weed/kms/config.go b/weed/kms/config.go index 8f3146c285..f1c66c58fb 100644 --- a/weed/kms/config.go +++ b/weed/kms/config.go @@ -128,6 +128,7 @@ func InitializeGlobalKMS(config *KMSConfig) error { if err != nil { return err } + provider = wrapWithCacheIfEnabled(provider, config) globalKMSMutex.Lock() defer globalKMSMutex.Unlock() @@ -224,6 +225,10 @@ func (km *KMSManager) AddKMSProvider(name string, config *KMSConfig) error { return fmt.Errorf("failed to create KMS provider %s: %w", name, err) } + // Wrap with a decrypted-data-key cache when enabled so repeated reads of + // the same SSE-KMS object avoid a KMS round-trip per read. + provider = wrapWithCacheIfEnabled(provider, config) + // Store provider and configuration km.providers[name] = provider km.configs[name] = config