enhance: encrypt all dbs when defaultKey is given (#47049)

when default key is provided, all the new database will be encrypted by
default and use the
default key as the key.
See also: #40013

---------

Signed-off-by: yangxuan <xuan.yang@zilliz.com>

internal/rootcoord/ddl_callbacks_alter_collection_properties_test.go

@@ -55,7 +55,7 @@ func TestDDLCallbacksAlterCollectionProperties(t *testing.T) {
 	})
 	require.ErrorIs(t, merr.CheckRPCCall(resp, err), merr.ErrParameterInvalid)
 
-	// hook related properties are not allowed to be altered.
+	// cipher related properties are not allowed to be altered.
 	resp, err = core.AlterCollection(ctx, &milvuspb.AlterCollectionRequest{
 		DbName:         dbName,
 		CollectionName: collectionName,
@@ -63,6 +63,22 @@ func TestDDLCallbacksAlterCollectionProperties(t *testing.T) {
 	})
 	require.ErrorIs(t, merr.CheckRPCCall(resp, err), merr.ErrParameterInvalid)
 
+	// cipher related properties are not allowed to be altered.
+	resp, err = core.AlterCollection(ctx, &milvuspb.AlterCollectionRequest{
+		DbName:         dbName,
+		CollectionName: collectionName,
+		Properties:     []*commonpb.KeyValuePair{{Key: common.EncryptionEzIDKey, Value: "1"}},
+	})
+	require.ErrorIs(t, merr.CheckRPCCall(resp, err), merr.ErrParameterInvalid)
+
+	// cipher related properties are not allowed to be altered.
+	resp, err = core.AlterCollection(ctx, &milvuspb.AlterCollectionRequest{
+		DbName:         dbName,
+		CollectionName: collectionName,
+		Properties:     []*commonpb.KeyValuePair{{Key: common.EncryptionRootKeyKey, Value: "1"}},
+	})
+	require.ErrorIs(t, merr.CheckRPCCall(resp, err), merr.ErrParameterInvalid)
+
 	// Alter a database that does not exist should return error.
 	resp, err = core.AlterCollection(ctx, &milvuspb.AlterCollectionRequest{
 		DbName:         dbName,

internal/rootcoord/key_manager_test.go

@@ -59,10 +59,6 @@ func TestKeyManager_GetDatabaseByEzID(t *testing.T) {
 					Key:   common.EncryptionEzIDKey,
 					Value: "123", // the same as the dbID
 				},
-				{
-					Key:   common.EncryptionEnabledKey,
-					Value: "true",
-				},
 			},
 		}
 
@@ -90,10 +86,6 @@ func TestKeyManager_GetDatabaseByEzID(t *testing.T) {
 					Key:   common.EncryptionEzIDKey,
 					Value: strconv.FormatInt(ezID, 10),
 				},
-				{
-					Key:   common.EncryptionEnabledKey,
-					Value: "true",
-				},
 			},
 		}
 

internal/rootcoord/meta_table.go

@@ -45,7 +45,6 @@ import (
 	"github.com/milvus-io/milvus/pkg/v2/util/crypto"
 	"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
 	"github.com/milvus-io/milvus/pkg/v2/util/merr"
-	"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/v2/util/timerecord"
 	"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
 )
@@ -353,7 +352,9 @@ func (mt *MetaTable) reloadWithNonDatabase() error {
 }
 
 func (mt *MetaTable) createDefaultDb() error {
-	ts, err := mt.tsoAllocator.GenerateTSO(1)
+	// Generate ezID and db ts for default database
+	// Use unique ID as ezID because the default dbID(1) for each cluster is the same
+	ts, err := mt.tsoAllocator.GenerateTSO(2)
 	if err != nil {
 		return err
 	}
@@ -364,21 +365,16 @@ func (mt *MetaTable) createDefaultDb() error {
 		return err
 	}
 
-	defaultRootKey := paramtable.GetCipherParams().DefaultRootKey.GetValue()
-	if hookutil.IsClusterEncryptionEnabled() && len(defaultRootKey) > 0 {
-		// Set unique ID as ezID because the default dbID for each cluster
-		// is the same
-		ezID, err := mt.tsoAllocator.GenerateTSO(1)
-		if err != nil {
-			return err
-		}
+	// Apply same encryption logic as regular database creation
+	// This respects the defaultKey setting
+	defaultProperties, err = hookutil.TidyDBCipherProperties(int64(ts-1), defaultProperties)
+	if err != nil {
+		return err
+	}
 
-		cipherProps := hookutil.GetDBCipherProperties(ezID, defaultRootKey)
-		defaultProperties = append(defaultProperties, cipherProps...)
-
-		if err := hookutil.CreateEZByDBProperties(defaultProperties); err != nil {
-			return err
-		}
+	// Create EZ if encryption is enabled
+	if err := hookutil.CreateEZByDBProperties(defaultProperties); err != nil {
+		return err
 	}
 
 	return mt.createDatabasePrivate(mt.ctx, model.NewDefaultDatabase(defaultProperties), ts)

internal/rootcoord/meta_table_test.go

@@ -35,6 +35,7 @@ import (
 	"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/channel"
 	mocktso "github.com/milvus-io/milvus/internal/tso/mocks"
 	kvfactory "github.com/milvus-io/milvus/internal/util/dependency/kv"
+	"github.com/milvus-io/milvus/internal/util/hookutil"
 	"github.com/milvus-io/milvus/pkg/v2/common"
 	pb "github.com/milvus-io/milvus/pkg/v2/proto/etcdpb"
 	"github.com/milvus-io/milvus/pkg/v2/proto/internalpb"
@@ -1949,6 +1950,139 @@ func TestMetaTable_CreateDatabase(t *testing.T) {
 	})
 }
 
+func TestCreateDefaultDb(t *testing.T) {
+	hookutil.InitTestCipher()
+
+	// Save original config and restore after test
+	originalDefaultKey := paramtable.GetCipherParams().DefaultRootKey.GetValue()
+	defer func() {
+		paramtable.GetCipherParams().Save("cipherPlugin.kms.defaultKey", originalDefaultKey)
+	}()
+
+	t.Run("default db without encryption when defaultKey is empty", func(t *testing.T) {
+		paramtable.GetCipherParams().Save("cipherPlugin.kms.defaultKey", "")
+
+		catalog := mocks.NewRootCoordCatalog(t)
+		catalog.On("CreateDatabase",
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+		).Return(nil)
+
+		tsoAllocator := mocktso.NewAllocator(t)
+		tsoAllocator.On("GenerateTSO", mock.Anything).Return(uint64(100), nil)
+
+		meta := &MetaTable{
+			ctx:          context.Background(),
+			dbName2Meta:  make(map[string]*model.Database),
+			names:        newNameDb(),
+			aliases:      newNameDb(),
+			catalog:      catalog,
+			tsoAllocator: tsoAllocator,
+		}
+
+		err := meta.createDefaultDb()
+		assert.NoError(t, err)
+
+		// Verify default database was created
+		db, ok := meta.dbName2Meta[util.DefaultDBName]
+		assert.True(t, ok)
+		assert.Equal(t, util.DefaultDBName, db.Name)
+
+		// Verify no encryption properties
+		hasEncryption := hookutil.IsDBEncrypted(db.Properties)
+		assert.False(t, hasEncryption, "default DB should not be encrypted when defaultKey is empty")
+	})
+
+	t.Run("default db with encryption when defaultKey is set", func(t *testing.T) {
+		paramtable.GetCipherParams().Save("cipherPlugin.kms.defaultKey", "default-test-key")
+
+		catalog := mocks.NewRootCoordCatalog(t)
+		catalog.On("CreateDatabase",
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+		).Return(nil)
+
+		tsoAllocator := mocktso.NewAllocator(t)
+		tsoAllocator.On("GenerateTSO", mock.Anything).Return(uint64(200), nil)
+
+		meta := &MetaTable{
+			ctx:          context.Background(),
+			dbName2Meta:  make(map[string]*model.Database),
+			names:        newNameDb(),
+			aliases:      newNameDb(),
+			catalog:      catalog,
+			tsoAllocator: tsoAllocator,
+		}
+
+		err := meta.createDefaultDb()
+		assert.NoError(t, err)
+
+		// Verify default database was created
+		db, ok := meta.dbName2Meta[util.DefaultDBName]
+		assert.True(t, ok)
+		assert.Equal(t, util.DefaultDBName, db.Name)
+
+		// Verify encryption properties are present
+		hasEzID := false
+		hasRootKey := false
+		for _, prop := range db.Properties {
+			if prop.Key == common.EncryptionEzIDKey {
+				hasEzID = true
+				assert.Equal(t, prop.GetValue(), "199")
+			}
+			if prop.Key == common.EncryptionRootKeyKey && prop.Value == "default-test-key" {
+				hasRootKey = true
+			}
+		}
+		assert.True(t, hasRootKey, "default DB should have root key when encrypted")
+		assert.True(t, hasEzID, "default DB should have ezID when encrypted")
+	})
+
+	t.Run("TSO allocation failure", func(t *testing.T) {
+		tsoAllocator := mocktso.NewAllocator(t)
+		tsoAllocator.On("GenerateTSO", mock.Anything).Return(uint64(0), errors.New("TSO allocation failed"))
+
+		meta := &MetaTable{
+			ctx:          context.Background(),
+			dbName2Meta:  make(map[string]*model.Database),
+			tsoAllocator: tsoAllocator,
+		}
+
+		err := meta.createDefaultDb()
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "TSO allocation failed")
+	})
+
+	t.Run("catalog CreateDatabase failure", func(t *testing.T) {
+		paramtable.GetCipherParams().Save("cipherPlugin.kms.defaultKey", "")
+
+		catalog := mocks.NewRootCoordCatalog(t)
+		catalog.On("CreateDatabase",
+			mock.Anything,
+			mock.Anything,
+			mock.Anything,
+		).Return(errors.New("catalog error"))
+
+		tsoAllocator := mocktso.NewAllocator(t)
+		tsoAllocator.On("GenerateTSO", mock.Anything).Return(uint64(300), nil)
+
+		meta := &MetaTable{
+			ctx:          context.Background(),
+			dbName2Meta:  make(map[string]*model.Database),
+			names:        newNameDb(),
+			aliases:      newNameDb(),
+			catalog:      catalog,
+			tsoAllocator: tsoAllocator,
+		}
+
+		err := meta.createDefaultDb()
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "catalog error")
+	})
+}
+
 func TestAlterDatabase(t *testing.T) {
 	t.Run("normal case", func(t *testing.T) {
 		catalog := mocks.NewRootCoordCatalog(t)

internal/util/hookutil/cipher.go

@@ -166,23 +166,6 @@ func GetEzByCollProperties(collProperties []*commonpb.KeyValuePair, collectionID
 	return nil
 }
 
-func GetDBCipherProperties(ezID uint64, kmsKey string) []*commonpb.KeyValuePair {
-	return []*commonpb.KeyValuePair{
-		{
-			Key:   common.EncryptionEnabledKey,
-			Value: "true",
-		},
-		{
-			Key:   common.EncryptionEzIDKey,
-			Value: strconv.FormatUint(ezID, 10),
-		},
-		{
-			Key:   common.EncryptionRootKeyKey,
-			Value: kmsKey,
-		},
-	}
-}
-
 func CreateEZByDBProperties(dbProperties []*commonpb.KeyValuePair) error {
 	ezID, hasEzID := ParseEzIDFromProperties(dbProperties)
 	if !hasEzID {
@@ -198,46 +181,76 @@ func CreateEZByDBProperties(dbProperties []*commonpb.KeyValuePair) error {
 	return nil
 }
 
-// When creating a new database with encryption eabled, set the ezID to the dbProperties,
-// and try to use the the config rootKey if rootKey not provided
-// An encrypted DB's properties will contain three properties:
-// cipher.enabled, cipher.ezID, cipher.key
+// When creating a new database
+// System will encrypt the DB If defaultKey is not empty.
+// However, if user controls with "cipher.enabled" and "cipher.key", respect the user's choice.
+// An encrypted DB's properties will contain two properties:
+// cipher.ezID, cipher.key
+//
+// Property consistency: DB without cipher.ezID and cipher.Key is NOT encrypted
+// - Encrypted DB: cipher.ezID and cipher.key are not empty
+// - Non-encrypted DB: no cipher.* key in properties
 func TidyDBCipherProperties(ezID int64, dbProperties []*commonpb.KeyValuePair) ([]*commonpb.KeyValuePair, error) {
-	dbEncryptionEnabled := IsDBEncrypted(dbProperties)
-	if GetCipher() == nil {
-		if dbEncryptionEnabled {
-			return nil, ErrCipherPluginMissing
+	defaultRootKey := paramtable.GetCipherParams().DefaultRootKey.GetValue()
+	defaultEncrypt := len(defaultRootKey) > 0
+
+	for _, property := range dbProperties {
+		switch property.Key {
+		case common.EncryptionEnabledKey:
+			value := strings.ToLower(property.Value)
+			if value != "true" && value != "false" {
+				return nil, fmt.Errorf("invalid value for %s: %q, must be \"true\" or \"false\"",
+					common.EncryptionEnabledKey, property.Value)
+			}
+			defaultEncrypt = value == "true"
+			log.Info("User explicitly controls encryption", zap.Int64("ezID", ezID), zap.Bool("value", defaultEncrypt))
+
+		case common.EncryptionRootKeyKey:
+			defaultRootKey = property.Value
+			log.Info("User explicitly set rootKey", zap.Int64("ezID", ezID), zap.String("value", property.GetValue()))
 		}
-		return dbProperties, nil
 	}
 
-	if dbEncryptionEnabled {
-		ezIDKv := &commonpb.KeyValuePair{
+	cipher := GetCipher()
+	// If cipher plugin is missing but encryption is requested, return error
+	if cipher == nil && defaultEncrypt {
+		return nil, ErrCipherPluginMissing
+	}
+
+	// No plugin or not encrypted
+	if cipher == nil || !defaultEncrypt {
+		return removeCipherProperties(dbProperties), nil
+	}
+
+	if defaultRootKey == "" {
+		return nil, fmt.Errorf("encryption enabled but no key provided and no default key configured")
+	}
+
+	result := removeCipherProperties(dbProperties)
+	result = append(result,
+		&commonpb.KeyValuePair{
 			Key:   common.EncryptionEzIDKey,
 			Value: strconv.FormatInt(ezID, 10),
-		}
-		// kmsKey already in the properties
-		for _, property := range dbProperties {
-			if property.Key == common.EncryptionRootKeyKey {
-				dbProperties = append(dbProperties, ezIDKv)
-				return dbProperties, nil
-			}
-		}
+		},
+		&commonpb.KeyValuePair{
+			Key:   common.EncryptionRootKeyKey,
+			Value: defaultRootKey,
+		},
+	)
 
-		if defaultRootKey := paramtable.GetCipherParams().DefaultRootKey.GetValue(); defaultRootKey != "" {
-			// set default root key from config if EncryuptionRootKeyKey left empty
-			dbProperties = append(dbProperties,
-				ezIDKv,
-				&commonpb.KeyValuePair{
-					Key:   common.EncryptionRootKeyKey,
-					Value: defaultRootKey,
-				},
-			)
-			return dbProperties, nil
+	return result, nil
+}
+
+func removeCipherProperties(properties []*commonpb.KeyValuePair) []*commonpb.KeyValuePair {
+	result := make([]*commonpb.KeyValuePair, 0, len(properties))
+	for _, property := range properties {
+		if property.Key != common.EncryptionEnabledKey &&
+			property.Key != common.EncryptionEzIDKey &&
+			property.Key != common.EncryptionRootKeyKey {
+			result = append(result, property)
 		}
-		return nil, fmt.Errorf("Empty default root key for encrypted database without kms key")
 	}
-	return dbProperties, nil
+	return result
 }
 
 func TidyCollPropsByDBProps(collProps, dbProps []*commonpb.KeyValuePair) []*commonpb.KeyValuePair {
@@ -270,12 +283,15 @@ func getCollEzPropsByDBProps(dbProperties []*commonpb.KeyValuePair) *commonpb.Ke
 }
 
 func IsDBEncrypted(dbProperties []*commonpb.KeyValuePair) bool {
+	var hasEzID, hasKey bool = false, false
 	for _, property := range dbProperties {
-		if property.Key == common.EncryptionEnabledKey && strings.ToLower(property.Value) == "true" {
-			return true
+		if property.Key == common.EncryptionEzIDKey && property.Value != "" {
+			hasEzID = true
+		} else if property.Key == common.EncryptionRootKeyKey && property.Value != "" {
+			hasKey = true
 		}
 	}
-	return false
+	return hasEzID && hasKey
 }
 
 func RemoveEZByDBProperties(dbProperties []*commonpb.KeyValuePair) error {

internal/util/hookutil/cipher_test.go

@@ -101,11 +101,11 @@ func (s *CipherSuite) TestTidyDBCipherProperties() {
 	}
 	result, err := TidyDBCipherProperties(1, dbPropertiesWithRootKey)
 	s.NoError(err)
-	s.Equal(3, len(result))
+	s.Equal(2, len(result))
 	for _, kv := range result {
 		switch kv.Key {
 		case common.EncryptionEnabledKey:
-			s.Equal(kv.Value, "true")
+			s.Fail("unexpected key")
 		case common.EncryptionEzIDKey:
 			s.Equal(kv.Value, "1")
 		case common.EncryptionRootKeyKey:
@@ -137,10 +137,16 @@ func (s *CipherSuite) TestIsDBEncryptionEnabled() {
 	dbProperties := []*commonpb.KeyValuePair{
 		{Key: common.EncryptionEnabledKey, Value: "true"},
 	}
-	s.True(IsDBEncrypted(dbProperties))
+	s.False(IsDBEncrypted(dbProperties))
 
 	dbProperties = []*commonpb.KeyValuePair{}
 	s.False(IsDBEncrypted(dbProperties))
+
+	dbProperties = []*commonpb.KeyValuePair{
+		{Key: common.EncryptionEzIDKey, Value: "1"},
+		{Key: common.EncryptionRootKeyKey, Value: "abc"},
+	}
+	s.True(IsDBEncrypted(dbProperties))
 }
 
 func (s *CipherSuite) TestTidyDBCipherPropertiesError() {
@@ -154,6 +160,33 @@ func (s *CipherSuite) TestTidyDBCipherPropertiesError() {
 	s.Equal(ErrCipherPluginMissing, err)
 }
 
+func (s *CipherSuite) TestTidyDBCipherPropertiesInvalidValue() {
+	InitTestCipher()
+
+	// Test invalid cipher.enabled values
+	invalidValues := []string{"1", "0", "yes", "no", "enabled", "disabled", "on", "off", ""}
+	for _, invalidValue := range invalidValues {
+		props := []*commonpb.KeyValuePair{
+			{Key: common.EncryptionEnabledKey, Value: invalidValue},
+		}
+		_, err := TidyDBCipherProperties(1, props)
+		s.Error(err, "expected error for invalid value: %q", invalidValue)
+		s.Contains(err.Error(), "invalid value")
+		s.Contains(err.Error(), invalidValue)
+	}
+
+	// Test valid values should work (case-insensitive)
+	validValues := []string{"true", "false", "True", "False", "TRUE", "FALSE", "TrUe", "FaLsE"}
+	for _, validValue := range validValues {
+		props := []*commonpb.KeyValuePair{
+			{Key: common.EncryptionEnabledKey, Value: validValue},
+			{Key: common.EncryptionRootKeyKey, Value: "some-key"},
+		}
+		_, err := TidyDBCipherProperties(1, props)
+		s.NoError(err, "expected no error for valid value: %q", validValue)
+	}
+}
+
 func (s *CipherSuite) TestTestCipherInit() {
 	cipher := testCipher{}
 	err := cipher.Init(map[string]string{"key": "value"})
@@ -254,17 +287,6 @@ func (s *CipherSuite) TestGetStoragePluginContext() {
 	s.Nil(result)
 }
 
-func (s *CipherSuite) TestGetDBCipherProperties() {
-	props := GetDBCipherProperties(123, "test-kms-key")
-	s.Equal(3, len(props))
-	s.Equal(common.EncryptionEnabledKey, props[0].Key)
-	s.Equal("true", props[0].Value)
-	s.Equal(common.EncryptionEzIDKey, props[1].Key)
-	s.Equal("123", props[1].Value)
-	s.Equal(common.EncryptionRootKeyKey, props[2].Key)
-	s.Equal("test-kms-key", props[2].Value)
-}
-
 func (s *CipherSuite) TestRemoveEZByDBProperties() {
 	storeCipher(nil)
 	err := RemoveEZByDBProperties([]*commonpb.KeyValuePair{{Key: common.EncryptionEzIDKey, Value: "1"}})
@@ -309,15 +331,146 @@ func (s *CipherSuite) TestGetEzByCollPropertiesInvalidValue() {
 	s.Equal(int64(456), result.CollectionID)
 }
 
-func (s *CipherSuite) TestIsDBEncryptionEnabledCaseInsensitive() {
-	props := []*commonpb.KeyValuePair{{Key: common.EncryptionEnabledKey, Value: "True"}}
-	s.True(IsDBEncrypted(props))
+func (s *CipherSuite) TestTidyDBCipherPropertiesWithDefaultKey() {
+	InitTestCipher()
 
-	props = []*commonpb.KeyValuePair{{Key: common.EncryptionEnabledKey, Value: "TRUE"}}
-	s.True(IsDBEncrypted(props))
+	// Save original config and restore after test
+	originalDefaultKey := paramtable.GetCipherParams().DefaultRootKey.GetValue()
+	defer func() {
+		paramtable.GetCipherParams().Save("cipherPlugin.kms.defaultKey", originalDefaultKey)
+	}()
 
-	props = []*commonpb.KeyValuePair{{Key: common.EncryptionEnabledKey, Value: "false"}}
-	s.False(IsDBEncrypted(props))
+	// Test 1: defaultKey is empty, no cipher.enabled in request -> not encrypted
+	paramtable.GetCipherParams().Save("cipherPlugin.kms.defaultKey", "")
+	props := []*commonpb.KeyValuePair{
+		{Key: "other.key", Value: "other.value"},
+	}
+	result, err := TidyDBCipherProperties(100, props)
+	s.NoError(err)
+	s.Equal(1, len(result))
+	s.Equal("other.key", result[0].Key)
+	// Verify no cipher properties
+	for _, prop := range result {
+		s.NotEqual(common.EncryptionEnabledKey, prop.Key)
+		s.NotEqual(common.EncryptionEzIDKey, prop.Key)
+		s.NotEqual(common.EncryptionRootKeyKey, prop.Key)
+	}
+
+	// Test 2: defaultKey is empty, cipher.enabled=false -> not encrypted
+	props = []*commonpb.KeyValuePair{
+		{Key: common.EncryptionEnabledKey, Value: "false"},
+		{Key: "other.key", Value: "other.value"},
+	}
+	result, err = TidyDBCipherProperties(101, props)
+	s.NoError(err)
+	// Verify no cipher properties
+	s.False(IsDBEncrypted(result))
+	for _, prop := range result {
+		s.NotEqual(common.EncryptionEnabledKey, prop.Key)
+		s.NotEqual(common.EncryptionEzIDKey, prop.Key)
+		s.NotEqual(common.EncryptionRootKeyKey, prop.Key)
+	}
+
+	// Test 3: defaultKey is empty, cipher.enabled=true, user provides key -> encrypted with user key
+	props = []*commonpb.KeyValuePair{
+		{Key: common.EncryptionEnabledKey, Value: "true"},
+		{Key: common.EncryptionRootKeyKey, Value: "user-key"},
+		{Key: "other.key", Value: "other.value"},
+	}
+	result, err = TidyDBCipherProperties(102, props)
+	s.NoError(err)
+	// only has ezID and key
+	s.False(hasKeyValue(result, common.EncryptionEnabledKey, "true"))
+	s.True(hasKeyValue(result, common.EncryptionEzIDKey, "102"))
+	s.True(hasKeyValue(result, common.EncryptionRootKeyKey, "user-key"))
+	s.True(hasKeyValue(result, "other.key", "other.value"))
+
+	// Test 4: defaultKey is empty, cipher.enabled=true, no user key -> error
+	props = []*commonpb.KeyValuePair{
+		{Key: common.EncryptionEnabledKey, Value: "true"},
+	}
+	_, err = TidyDBCipherProperties(103, props)
+	s.Error(err)
+	s.Contains(err.Error(), "no key provided")
+
+	// Test 5: defaultKey is set, no cipher.enabled in request -> encrypted with defaultKey
+	paramtable.GetCipherParams().Save("cipherPlugin.kms.defaultKey", "default-kms-key")
+	props = []*commonpb.KeyValuePair{
+		{Key: "other.key", Value: "other.value"},
+	}
+	result, err = TidyDBCipherProperties(104, props)
+	s.NoError(err)
+	s.True(IsDBEncrypted(result))
+	s.False(hasKeyValue(result, common.EncryptionEnabledKey, "true"))
+	s.True(hasKeyValue(result, common.EncryptionEzIDKey, "104"))
+	s.True(hasKeyValue(result, common.EncryptionRootKeyKey, "default-kms-key"))
+	s.True(hasKeyValue(result, "other.key", "other.value"))
+
+	// Test 6: defaultKey is set, cipher.enabled=false -> not encrypted
+	props = []*commonpb.KeyValuePair{
+		{Key: common.EncryptionEnabledKey, Value: "false"},
+		{Key: "other.key", Value: "other.value"},
+	}
+	result, err = TidyDBCipherProperties(105, props)
+	s.NoError(err)
+	// Verify no cipher properties
+	s.False(IsDBEncrypted(result))
+	for _, prop := range result {
+		s.NotEqual(common.EncryptionEzIDKey, prop.Key)
+		s.NotEqual(common.EncryptionRootKeyKey, prop.Key)
+	}
+	s.True(hasKeyValue(result, "other.key", "other.value"))
+
+	// Test 7: defaultKey is set, cipher.enabled=true with user key -> encrypted with user key
+	props = []*commonpb.KeyValuePair{
+		{Key: common.EncryptionEnabledKey, Value: "true"},
+		{Key: common.EncryptionRootKeyKey, Value: "user-override-key"},
+		{Key: "other.key", Value: "other.value"},
+	}
+	result, err = TidyDBCipherProperties(106, props)
+	s.NoError(err)
+	s.True(IsDBEncrypted(result))
+	s.False(hasKeyValue(result, common.EncryptionEnabledKey, "true"))
+	s.True(hasKeyValue(result, common.EncryptionEzIDKey, "106"))
+	s.True(hasKeyValue(result, common.EncryptionRootKeyKey, "user-override-key"))
+	s.True(hasKeyValue(result, "other.key", "other.value"))
+
+	// Test 8: defaultKey is set, cipher.enabled=true without user key -> encrypted with defaultKey
+	props = []*commonpb.KeyValuePair{
+		{Key: common.EncryptionEnabledKey, Value: "true"},
+		{Key: "other.key", Value: "other.value"},
+	}
+	result, err = TidyDBCipherProperties(107, props)
+	s.NoError(err)
+	s.True(IsDBEncrypted(result))
+	s.False(hasKeyValue(result, common.EncryptionEnabledKey, "true"))
+	s.True(hasKeyValue(result, common.EncryptionEzIDKey, "107"))
+	s.True(hasKeyValue(result, common.EncryptionRootKeyKey, "default-kms-key"))
+	s.True(hasKeyValue(result, "other.key", "other.value"))
+
+	// Test 9: defaultKey is set, user provides key, no cipher.enabled -> encrypted with user key
+	// This tests that user key overrides defaultKey even when cipher.enabled is not specified
+	props = []*commonpb.KeyValuePair{
+		{Key: common.EncryptionRootKeyKey, Value: "user-custom-key"},
+		{Key: "other.key", Value: "other.value"},
+	}
+	result, err = TidyDBCipherProperties(108, props)
+	s.NoError(err)
+	s.True(IsDBEncrypted(result))
+	s.False(hasKeyValue(result, common.EncryptionEnabledKey, "true"))
+	s.True(hasKeyValue(result, common.EncryptionEzIDKey, "108"))
+	s.True(hasKeyValue(result, common.EncryptionRootKeyKey, "user-custom-key"))
+	s.True(hasKeyValue(result, "other.key", "other.value"))
+}
+
+// Helper function to check if a key-value pair exists in the result
+func hasKeyValue(props []*commonpb.KeyValuePair, key, value string) bool {
+	for _, prop := range props {
+		if prop.Key == key && prop.Value == value {
+			return true
+		}
+	}
+	return false
 }
 
 func (s *CipherSuite) TestRegisterCallback() {
@@ -332,40 +485,6 @@ func (s *CipherSuite) TestRegisterCallback() {
 	s.NoError(err)
 }
 
-func (s *CipherSuite) TestTidyDBCipherPropertiesWithDefaultKey() {
-	initCipherOnce = sync.Once{}
-	storeCipher(nil)
-	InitTestCipher()
-
-	oldValue := paramtable.GetCipherParams().DefaultRootKey.GetValue()
-	defer paramtable.GetCipherParams().DefaultRootKey.SwapTempValue(oldValue)
-
-	paramtable.GetCipherParams().DefaultRootKey.SwapTempValue("default-test-key")
-
-	dbProperties := []*commonpb.KeyValuePair{
-		{Key: common.EncryptionEnabledKey, Value: "true"},
-	}
-	result, err := TidyDBCipherProperties(1, dbProperties)
-	s.NoError(err)
-	s.NotNil(result)
-	s.Equal(3, len(result))
-
-	foundEzID := false
-	foundRootKey := false
-	for _, kv := range result {
-		if kv.Key == common.EncryptionEzIDKey {
-			s.Equal("1", kv.Value)
-			foundEzID = true
-		}
-		if kv.Key == common.EncryptionRootKeyKey {
-			s.Equal("default-test-key", kv.Value)
-			foundRootKey = true
-		}
-	}
-	s.True(foundEzID)
-	s.True(foundRootKey)
-}
-
 func (s *CipherSuite) TestGetEzByCollPropertiesNotFound() {
 	props := []*commonpb.KeyValuePair{
 		{Key: "other_key", Value: "value"},
@@ -389,22 +508,23 @@ func (s *CipherSuite) TestBackupEZ() {
 func (s *CipherSuite) TestBackupEZKFromDBProperties() {
 	storeCipher(nil)
 	props := []*commonpb.KeyValuePair{
-		{Key: common.EncryptionEnabledKey, Value: "true"},
 		{Key: common.EncryptionEzIDKey, Value: "123"},
+		{Key: common.EncryptionRootKeyKey, Value: "abc-key"},
 	}
 	_, err := BackupEZKFromDBProperties(props)
 	s.Error(err)
 	s.Equal(ErrCipherPluginMissing, err)
 
 	props = []*commonpb.KeyValuePair{
-		{Key: common.EncryptionEnabledKey, Value: "false"},
+		{Key: common.EncryptionRootKeyKey, Value: "abc-key"},
 	}
 	_, err = BackupEZKFromDBProperties(props)
 	s.Error(err)
 	s.Contains(err.Error(), "not an encryption zone")
 
 	props = []*commonpb.KeyValuePair{
-		{Key: common.EncryptionEnabledKey, Value: "true"},
+		{Key: common.EncryptionRootKeyKey, Value: "abc-key"},
+		{Key: common.EncryptionEzIDKey, Value: "fack"},
 	}
 	_, err = BackupEZKFromDBProperties(props)
 	s.Error(err)

pkg/streaming/util/message/cipher.go

@@ -2,8 +2,13 @@ package message
 
 import (
 	"sync"
+	"time"
+
+	"github.com/cockroachdb/errors"
+	"go.uber.org/zap"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/hook"
+	"github.com/milvus-io/milvus/pkg/v2/log"
 )
 
 // cipher is a global variable that is used to encrypt and decrypt messages.
@@ -29,6 +34,65 @@ func mustGetCipher() hook.Cipher {
 	return cipher
 }
 
+// ErrKmsKeyInvalid is the error returned when a KMS key is invalid or revoked.
+// This error is also defined in the milvus-cloud-plugin. It is checked using `errors.Is`
+// to allow for proper error wrapping and reliable error handling.
+var ErrKmsKeyInvalid = errors.New("kms key invalid")
+
+func isKmsKeyInvalidError(err error) bool {
+	return errors.Is(err, ErrKmsKeyInvalid)
+}
+
+// getDecryptorWithRetry wraps cipher.GetDecryptor with retry logic for streaming node consumption.
+// It retries with exponential backoff if the error is KmsKeyInvalid (retriable).
+// For other errors, it returns immediately without retry.
+func getDecryptorWithRetry(ezID, collectionID int64, safeKey []byte) (hook.Decryptor, error) {
+	cipher := mustGetCipher()
+
+	const (
+		initialBackoff = 100 * time.Millisecond
+		maxBackoff     = 3 * time.Second
+		backoffFactor  = 2.0
+	)
+
+	backoff := initialBackoff
+	attempt := 0
+
+	for {
+		attempt++
+		decryptor, err := cipher.GetDecryptor(ezID, collectionID, safeKey)
+		if err == nil {
+			return decryptor, nil
+		}
+
+		// If it's NOT a KMS key invalid error, fail immediately (non-retriable)
+		if !isKmsKeyInvalidError(err) {
+			log.Error("failed to get decryptor with non-retriable error",
+				zap.Int64("ezID", ezID),
+				zap.Int64("collectionID", collectionID),
+				zap.Int("attempt", attempt),
+				zap.Error(err))
+			return nil, err
+		}
+
+		// KMS key invalid error - log and retry
+		log.Warn("KMS key invalid, will retry",
+			zap.Int64("ezID", ezID),
+			zap.Int64("collectionID", collectionID),
+			zap.Int("attempt", attempt),
+			zap.Duration("backoff", backoff),
+			zap.Error(err))
+
+		time.Sleep(backoff)
+
+		// Exponential backoff with max cap
+		backoff = time.Duration(float64(backoff) * backoffFactor)
+		if backoff > maxBackoff {
+			backoff = maxBackoff
+		}
+	}
+}
+
 // CipherConfig is the configuration for cipher that is used to encrypt and decrypt messages.
 type CipherConfig struct {
 	// EzID is the encryption zone ID.

pkg/streaming/util/message/cipher_test.go

@@ -0,0 +1,277 @@
+package message
+
+import (
+	"fmt"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/milvus-io/milvus-proto/go-api/v2/hook"
+)
+
+// mockCipher is a simple mock implementation for testing
+type mockCipher struct {
+	getDecryptorFunc func(ezID, collectionID int64, safeKey []byte) (hook.Decryptor, error)
+	getEncryptorFunc func(ezID, collectionID int64) (hook.Encryptor, []byte, error)
+	getUnsafeKeyFunc func(ezID, collectionID int64) []byte
+	initFunc         func(params map[string]string) error
+}
+
+func (m *mockCipher) Init(params map[string]string) error {
+	if m.initFunc != nil {
+		return m.initFunc(params)
+	}
+	return nil
+}
+
+func (m *mockCipher) GetEncryptor(ezID, collectionID int64) (hook.Encryptor, []byte, error) {
+	if m.getEncryptorFunc != nil {
+		return m.getEncryptorFunc(ezID, collectionID)
+	}
+	return nil, nil, nil
+}
+
+func (m *mockCipher) GetDecryptor(ezID, collectionID int64, safeKey []byte) (hook.Decryptor, error) {
+	if m.getDecryptorFunc != nil {
+		return m.getDecryptorFunc(ezID, collectionID, safeKey)
+	}
+	return nil, nil
+}
+
+func (m *mockCipher) GetUnsafeKey(ezID, collectionID int64) []byte {
+	if m.getUnsafeKeyFunc != nil {
+		return m.getUnsafeKeyFunc(ezID, collectionID)
+	}
+	return nil
+}
+
+// mockDecryptor is a simple mock decryptor
+type mockDecryptor struct {
+	decryptFunc func([]byte) ([]byte, error)
+}
+
+func (m *mockDecryptor) Decrypt(data []byte) ([]byte, error) {
+	if m.decryptFunc != nil {
+		return m.decryptFunc(data)
+	}
+	return data, nil
+}
+
+func TestGetDecryptorWithRetry_Success(t *testing.T) {
+	// Setup: register a mock cipher that succeeds immediately
+	origCipher := cipher
+	defer func() { cipher = origCipher }()
+
+	mockDec := &mockDecryptor{}
+	cipher = &mockCipher{
+		getDecryptorFunc: func(ezID, collectionID int64, safeKey []byte) (hook.Decryptor, error) {
+			return mockDec, nil
+		},
+	}
+
+	// Test
+	decryptor, err := getDecryptorWithRetry(123, 456, []byte("safe-key"))
+
+	// Assert
+	assert.NoError(t, err)
+	assert.Equal(t, mockDec, decryptor)
+}
+
+func TestGetDecryptorWithRetry_NonRetriableError(t *testing.T) {
+	// Setup: register a mock cipher that returns a non-retriable error
+	origCipher := cipher
+	defer func() { cipher = origCipher }()
+
+	nonRetriableErr := errors.New("non-retriable error")
+	cipher = &mockCipher{
+		getDecryptorFunc: func(ezID, collectionID int64, safeKey []byte) (hook.Decryptor, error) {
+			return nil, nonRetriableErr
+		},
+	}
+
+	// Test
+	start := time.Now()
+	decryptor, err := getDecryptorWithRetry(123, 456, []byte("safe-key"))
+	duration := time.Since(start)
+
+	// Assert - should fail immediately without retry
+	assert.Error(t, err)
+	assert.Nil(t, decryptor)
+	assert.Contains(t, err.Error(), "non-retriable error")
+	assert.Less(t, duration.Milliseconds(), int64(50), "should return quickly without retry")
+}
+
+func TestGetDecryptorWithRetry_KmsKeyInvalidWithRetry(t *testing.T) {
+	// Setup: register a mock cipher that fails 3 times with KmsKeyInvalid, then succeeds
+	origCipher := cipher
+	defer func() { cipher = origCipher }()
+
+	var attemptCount int32
+	mockDec := &mockDecryptor{}
+	cipher = &mockCipher{
+		getDecryptorFunc: func(ezID, collectionID int64, safeKey []byte) (hook.Decryptor, error) {
+			attempt := atomic.AddInt32(&attemptCount, 1)
+			if attempt < 4 {
+				return nil, fmt.Errorf("%w: ezID=%d, state=revoked", ErrKmsKeyInvalid, ezID)
+			}
+			return mockDec, nil
+		},
+	}
+
+	// Test
+	start := time.Now()
+	decryptor, err := getDecryptorWithRetry(123, 456, []byte("safe-key"))
+	duration := time.Since(start)
+
+	// Assert
+	assert.NoError(t, err)
+	assert.Equal(t, mockDec, decryptor)
+	assert.Equal(t, int32(4), atomic.LoadInt32(&attemptCount))
+	// Should have retried at least 3 times with exponential backoff
+	// Min expected duration: 100ms + 200ms + 400ms = 700ms
+	assert.Greater(t, duration.Milliseconds(), int64(300), "should have retried with delay")
+}
+
+func TestGetDecryptorWithRetry_ExponentialBackoff(t *testing.T) {
+	// Setup: register a mock cipher that tracks timing between attempts
+	origCipher := cipher
+	defer func() { cipher = origCipher }()
+
+	var attemptCount int32
+	var lastAttemptTime time.Time
+	var backoffs []time.Duration
+	mockDec := &mockDecryptor{}
+
+	cipher = &mockCipher{
+		getDecryptorFunc: func(ezID, collectionID int64, safeKey []byte) (hook.Decryptor, error) {
+			attempt := atomic.AddInt32(&attemptCount, 1)
+			now := time.Now()
+
+			if !lastAttemptTime.IsZero() && attempt > 1 {
+				backoff := now.Sub(lastAttemptTime)
+				backoffs = append(backoffs, backoff)
+			}
+			lastAttemptTime = now
+
+			if attempt < 5 {
+				return nil, ErrKmsKeyInvalid
+			}
+			return mockDec, nil
+		},
+	}
+
+	// Test
+	_, err := getDecryptorWithRetry(123, 456, []byte("safe-key"))
+
+	// Assert
+	assert.NoError(t, err)
+	assert.Equal(t, int32(5), atomic.LoadInt32(&attemptCount))
+	assert.Len(t, backoffs, 4)
+
+	// Verify exponential backoff (with some tolerance for timing variance)
+	// Expected: ~100ms, ~200ms, ~400ms, ~800ms
+	expectedBackoffs := []time.Duration{
+		100 * time.Millisecond,
+		200 * time.Millisecond,
+		400 * time.Millisecond,
+		800 * time.Millisecond,
+	}
+
+	for i, backoff := range backoffs {
+		expected := expectedBackoffs[i]
+		// Allow 50% tolerance due to timing variance
+		minBackoff := time.Duration(float64(expected) * 0.5)
+		maxBackoff := time.Duration(float64(expected) * 1.5)
+		assert.GreaterOrEqual(t, backoff, minBackoff,
+			"backoff %d should be at least %v, got %v", i, minBackoff, backoff)
+		assert.LessOrEqual(t, backoff, maxBackoff,
+			"backoff %d should be at most %v, got %v", i, maxBackoff, backoff)
+	}
+}
+
+func TestGetDecryptorWithRetry_MaxBackoffCap(t *testing.T) {
+	// Setup: register a mock cipher that fails many times to test max backoff
+	origCipher := cipher
+	defer func() { cipher = origCipher }()
+
+	var attemptCount int32
+	var lastAttemptTime time.Time
+	var backoffs []time.Duration
+	mockDec := &mockDecryptor{}
+
+	cipher = &mockCipher{
+		getDecryptorFunc: func(ezID, collectionID int64, safeKey []byte) (hook.Decryptor, error) {
+			attempt := atomic.AddInt32(&attemptCount, 1)
+			now := time.Now()
+
+			if !lastAttemptTime.IsZero() {
+				backoff := now.Sub(lastAttemptTime)
+				backoffs = append(backoffs, backoff)
+			}
+			lastAttemptTime = now
+
+			// Succeed after enough retries to test max backoff
+			if attempt < 10 {
+				return nil, ErrKmsKeyInvalid
+			}
+			return mockDec, nil
+		},
+	}
+
+	// Test
+	_, err := getDecryptorWithRetry(123, 456, []byte("safe-key"))
+
+	// Assert
+	assert.NoError(t, err)
+
+	// Verify that backoff is capped at 3 seconds
+	// After several retries, all backoffs should be at max (3s)
+	if len(backoffs) > 5 {
+		for i := 5; i < len(backoffs); i++ {
+			// Allow some tolerance
+			assert.LessOrEqual(t, backoffs[i], 3500*time.Millisecond,
+				"backoff %d should be capped at ~3s, got %v", i, backoffs[i])
+			assert.GreaterOrEqual(t, backoffs[i], 2500*time.Millisecond,
+				"backoff %d should be around 3s, got %v", i, backoffs[i])
+		}
+	}
+}
+
+func TestIsKmsKeyInvalidError(t *testing.T) {
+	tests := []struct {
+		name     string
+		err      error
+		expected bool
+	}{
+		{
+			name:     "exact error",
+			err:      ErrKmsKeyInvalid,
+			expected: true,
+		},
+		{
+			name:     "wrapped error",
+			err:      fmt.Errorf("%w: additional context", ErrKmsKeyInvalid),
+			expected: true,
+		},
+		{
+			name:     "different error",
+			err:      errors.New("some other error"),
+			expected: false,
+		},
+		{
+			name:     "nil error",
+			err:      nil,
+			expected: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := isKmsKeyInvalidError(tt.err)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}

pkg/streaming/util/message/message_impl.go

@@ -39,10 +39,12 @@ func (m *messageImpl) MessageTypeWithVersion() MessageTypeWithVersion {
 }
 
 // Payload returns payload of current message.
+// If the message is encrypted, it will be decrypted with automatic retry for transient KMS errors.
 func (m *messageImpl) Payload() []byte {
 	if ch := m.cipherHeader(); ch != nil {
-		cipher := mustGetCipher()
-		decryptor, err := cipher.GetDecryptor(ch.EzId, ch.CollectionId, ch.SafeKey)
+		// Use getDecryptorWithRetry for resilient decryption with automatic retry
+		// for transient KMS key errors (e.g., temporarily invalid/revoked keys)
+		decryptor, err := getDecryptorWithRetry(ch.EzId, ch.CollectionId, ch.SafeKey)
 		if err != nil {
 			panic(fmt.Sprintf("can not get decryptor for message: %s", err))
 		}