use the cache in most places

pkg/registry/apis/secret/contracts/encryption.go

@@ -14,6 +14,9 @@ type EncryptionManager interface {
 	// implementation present at manager.EncryptionService.
 	Encrypt(ctx context.Context, namespace xkube.Namespace, payload []byte) (EncryptedPayload, error)
 	Decrypt(ctx context.Context, namespace xkube.Namespace, payload EncryptedPayload) ([]byte, error)
+
+	// Since consolidation occurs at a level above the EncryptionManager, we need to allow that process to manually flush the cache
+	FlushCache(namespace xkube.Namespace)
 }
 
 type EncryptedPayload struct {

pkg/registry/apis/secret/encryption/manager/manager.go

@@ -65,6 +65,7 @@ func ProvideEncryptionManager(
 		cipher:         enc,
 		log:            log.New("encryption"),
 		providerConfig: providerConfig,
+		dataKeyCache:   dataKeyCache,
 		cfg:            cfg,
 	}
 
@@ -182,6 +183,11 @@ func (s *EncryptionManager) currentDataKey(ctx context.Context, namespace xkube.
 // dataKeyByLabel looks up for data key in cache by label.
 // Otherwise, it fetches it from database, decrypts it and caches it decrypted.
 func (s *EncryptionManager) dataKeyByLabel(ctx context.Context, namespace, label string) (string, []byte, error) {
+	// 0. Get data key from in-memory cache.
+	if entry, exists := s.dataKeyCache.GetByLabel(namespace, label); exists && entry.Active {
+		return entry.Id, entry.DataKey, nil
+	}
+
 	// 1. Get data key from database.
 	dataKey, err := s.store.GetCurrentDataKey(ctx, namespace, label)
 	if err != nil {
@@ -203,6 +209,9 @@ func (s *EncryptionManager) dataKeyByLabel(ctx context.Context, namespace, label
 		return "", nil, err
 	}
 
+	// 3. Store the decrypted data key into the in-memory cache.
+	s.cacheDataKey(namespace, dataKey, decrypted)
+
 	return dataKey.UID, decrypted, nil
 }
 
@@ -249,6 +258,9 @@ func (s *EncryptionManager) newDataKey(ctx context.Context, namespace string, la
 		return "", nil, err
 	}
 
+	// 4. Store the decrypted data key into the in-memory cache.
+	s.cacheDataKey(namespace, &dbDataKey, dataKey)
+
 	return id, dataKey, nil
 }
 
@@ -312,6 +324,11 @@ func (s *EncryptionManager) dataKeyById(ctx context.Context, namespace, id strin
 	))
 	defer span.End()
 
+	// 0. Get data key from in-memory cache.
+	if entry, exists := s.dataKeyCache.GetById(namespace, id); exists && entry.Active {
+		return entry.DataKey, nil
+	}
+
 	// 1. Get encrypted data key from database.
 	dataKey, err := s.store.GetDataKey(ctx, namespace, id)
 	if err != nil {
@@ -330,6 +347,9 @@ func (s *EncryptionManager) dataKeyById(ctx context.Context, namespace, id strin
 		return nil, err
 	}
 
+	// 3. Store the decrypted data key into the in-memory cache.
+	s.cacheDataKey(namespace, dataKey, decrypted)
+
 	return decrypted, nil
 }
 
@@ -337,10 +357,12 @@ func (s *EncryptionManager) GetProviders() encryption.ProviderConfig {
 	return s.providerConfig
 }
 
+func (s *EncryptionManager) FlushCache(namespace xkube.Namespace) {
+	s.dataKeyCache.Flush(namespace.String())
+}
+
 func (s *EncryptionManager) Run(ctx context.Context) error {
-	gc := time.NewTicker(
-		s.cfg.SecretsManagement.DataKeysCacheCleanupInterval,
-	)
+	gc := time.NewTicker(s.cfg.SecretsManagement.DataKeysCacheCleanupInterval)
 
 	grp, gCtx := errgroup.WithContext(ctx)
 
@@ -362,3 +384,45 @@ func (s *EncryptionManager) Run(ctx context.Context) error {
 		}
 	}
 }
+
+// NB: Much of this was copied or derived from the original implementation in the legacy SecretsService.
+//
+// Caching a data key is tricky, because at SecretsService level we cannot guarantee
+// that a newly created data key has actually been persisted, depending on the different
+// use cases that rely on SecretsService encryption and different database engines that
+// we have support for, because the data key creation may have happened within a DB TX,
+// that may fail afterwards.
+//
+// Therefore, if we cache a data key that hasn't been persisted with success (and won't),
+// and later that one is used for a encryption operation (aside from the DB TX that created
+// it), we may end up with data encrypted by a non-persisted data key, which could end up
+// in (unrecoverable) data corruption.
+//
+// So, we cache the data key by id and/or by label, depending on the data key's lifetime,
+// assuming that a data key older than a "caution period" should have been persisted.
+//
+// Look at the comments inline for further details.
+// You can also take a look at the issue below for more context:
+// https://github.com/grafana/grafana-enterprise/issues/4252
+func (s *EncryptionManager) cacheDataKey(namespace string, dataKey *contracts.SecretDataKey, decrypted []byte) {
+	// First, we cache the data key by id, because cache "by id" is
+	// only used by decrypt operations, so no risk of corrupting data.
+	entry := &encryption.DataKeyCacheEntry{
+		Namespace: namespace,
+		Id:        dataKey.UID,
+		Label:     dataKey.Label,
+		DataKey:   decrypted,
+		Active:    dataKey.Active,
+	}
+
+	s.dataKeyCache.AddById(namespace, entry)
+
+	// Then, we cache the data key by label, ONLY if data key's lifetime
+	// is longer than a certain "caution period", because cache "by label"
+	// is used (only) by encrypt operations, and we want to ensure that
+	// no data key is cached for encryption ops before being persisted.
+	nowMinusCautionPeriod := time.Now().Add(-s.cfg.SecretsManagement.DataKeysCacheCautionPeriod)
+	if dataKey.Created.Before(nowMinusCautionPeriod) {
+		s.dataKeyCache.AddByLabel(namespace, entry)
+	}
+}

pkg/registry/apis/secret/encryption/manager/manager_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"testing"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
@@ -646,3 +647,88 @@ func TestEncryptionService_ThirdPartyProviders(t *testing.T) {
 	require.Len(t, encMgr.providerConfig.AvailableProviders, 1)
 	require.Contains(t, encMgr.providerConfig.AvailableProviders, encryption.ProviderID("fakeProvider.v1"))
 }
+
+func TestEncryptionService_FlushCache(t *testing.T) {
+	ctx := context.Background()
+	namespace := xkube.Namespace("test-namespace")
+	plaintext := []byte("secret data to encrypt")
+
+	// Set up the encryption manager with a real OSS DEK cache
+	testDB := sqlstore.NewTestStore(t, sqlstore.WithMigrator(migrator.New()))
+	tracer := noop.NewTracerProvider().Tracer("test")
+	database := database.ProvideDatabase(testDB, tracer)
+
+	cfg := &setting.Cfg{
+		SecretsManagement: setting.SecretsManagerSettings{
+			CurrentEncryptionProvider:  "secret_key.v1",
+			ConfiguredKMSProviders:     map[string]map[string]string{"secret_key.v1": {"secret_key": "SW2YcwTIb9zpOOhoPsMm"}},
+			DataKeysCacheTTL:           time.Hour,       // Long TTL to ensure keys don't expire during test
+			DataKeysCacheCautionPeriod: 0 * time.Second, // Override the caution period for testing
+		},
+	}
+
+	store, err := encryptionstorage.ProvideDataKeyStorage(database, tracer, nil)
+	require.NoError(t, err)
+
+	usageStats := &usagestats.UsageStatsMock{T: t}
+	enc, err := service.ProvideAESGCMCipherService(tracer, usageStats)
+	require.NoError(t, err)
+
+	ossProviders, err := osskmsproviders.ProvideOSSKMSProviders(cfg, enc)
+	require.NoError(t, err)
+
+	// Create a real OSS DEK cache
+	dekCache := ProvideOSSDataKeyCache(cfg)
+
+	encMgr, err := ProvideEncryptionManager(
+		tracer,
+		store,
+		usageStats,
+		enc,
+		ossProviders,
+		dekCache,
+		cfg,
+	)
+	require.NoError(t, err)
+
+	svc := encMgr.(*EncryptionManager)
+
+	// Encrypt some data - this will create a DEK and cache it
+	encrypted, err := svc.Encrypt(ctx, namespace, plaintext)
+	require.NoError(t, err)
+
+	// Verify we can decrypt - this should use the cached key
+	decrypted, err := svc.Decrypt(ctx, namespace, encrypted)
+	require.NoError(t, err)
+	assert.Equal(t, plaintext, decrypted)
+
+	// Get the data key ID from the encrypted payload
+	dataKeyID := encrypted.DataKeyID
+
+	// Verify the key is in the cache by checking both by ID and by label
+	label := encryption.KeyLabel(svc.providerConfig.CurrentProvider)
+	_, existsById := dekCache.GetById(namespace.String(), dataKeyID)
+	assert.True(t, existsById, "DEK should be cached by ID before flush")
+
+	_, existsByLabel := dekCache.GetByLabel(namespace.String(), label)
+	assert.True(t, existsByLabel, "DEK should be cached by label before flush")
+
+	// Flush the cache for this namespace
+	svc.FlushCache(namespace)
+
+	// Verify the cache is empty for this namespace
+	_, existsById = dekCache.GetById(namespace.String(), dataKeyID)
+	assert.False(t, existsById, "DEK should not be in cache by ID after flush")
+
+	_, existsByLabel = dekCache.GetByLabel(namespace.String(), label)
+	assert.False(t, existsByLabel, "DEK should not be in cache by label after flush")
+
+	// Verify we can still decrypt - this should fetch from DB and re-cache
+	decrypted, err = svc.Decrypt(ctx, namespace, encrypted)
+	require.NoError(t, err)
+	assert.Equal(t, plaintext, decrypted)
+
+	// Verify the key is back in the cache after the decrypt operation
+	_, existsById = dekCache.GetById(namespace.String(), dataKeyID)
+	assert.True(t, existsById, "DEK should be re-cached by ID after decrypt")
+}

pkg/setting/setting_secrets_manager.go

@@ -18,6 +18,8 @@ type SecretsManagerSettings struct {
 	DataKeysCacheTTL time.Duration
 	// The interval to remove expired data keys from the cache
 	DataKeysCacheCleanupInterval time.Duration
+	// The caution period is the time after which a data key is assumed to be persisted in the worst case scenario.
+	DataKeysCacheCautionPeriod time.Duration
 	// Whether to use a Redis cache for data keys instead of the in-memory cache
 	DataKeysCacheUseRedis bool
 
@@ -77,6 +79,8 @@ func (cfg *Cfg) readSecretsManagerSettings() {
 	cfg.SecretsManagement.DataKeysCacheUseRedis = secretsMgmt.Key("data_keys_cache_use_redis").MustBool(false)
 	cfg.SecretsManagement.DataKeysCacheTTL = secretsMgmt.Key("data_keys_cache_ttl").MustDuration(15 * time.Minute)
 	cfg.SecretsManagement.DataKeysCacheCleanupInterval = secretsMgmt.Key("data_keys_cache_cleanup_interval").MustDuration(1 * time.Minute)
+	// We consider a "caution period" of 10m to be long enough for any database transaction that implied a data key creation to have finished successfully.
+	cfg.SecretsManagement.DataKeysCacheCautionPeriod = secretsMgmt.Key("data_keys_cache_caution_period").MustDuration(10 * time.Minute)
 
 	// Extract available KMS providers from configuration sections
 	providers := make(map[string]map[string]string)