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

@@ -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
+}

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)
+	}
+}

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