feat: implement KMS key revocation deny read/write (milvus-io#45936)

Adds KMS key state monitoring and deny reading/writing for abnormal kms

Key Changes:
- Add KeyManager in RootCoord for periodic KMS state polling
- Integrate KeyManager with QuotaCenter for access denial
- Implement revocation checks in Proxy SimpleLimiter

Access Denial:
- Revoked keys: Release collections + deny DML/DQL (DDL still allowed)
- Manual LoadCollection required after key recovery

See also: milvus-io#45117, milvus-io#44981, milvus-io#45242

---------

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

Author: XuanYang-cn

internal/mocks/flushcommon/mock_util/mock_MsgHandler.go

@@ -162,23 +162,27 @@ func isNotCollectionLevelLimitRequest(rt internalpb.RateType) bool {
 func (m *SimpleLimiter) GetQuotaStates() ([]milvuspb.QuotaState, []string) {
 	m.quotaStatesMu.RLock()
 	defer m.quotaStatesMu.RUnlock()
-	serviceStates := make(map[milvuspb.QuotaState]typeutil.Set[commonpb.ErrorCode])
+	type stateReasonKey struct {
+		ErrorCode commonpb.ErrorCode
+		Reason    string
+	}
+	serviceStates := make(map[milvuspb.QuotaState]typeutil.Set[stateReasonKey])
 
 	rlinternal.TraverseRateLimiterTree(m.rateLimiter.GetRootLimiters(), nil,
-		func(node *rlinternal.RateLimiterNode, state milvuspb.QuotaState, errCode commonpb.ErrorCode) bool {
+		func(node *rlinternal.RateLimiterNode, state milvuspb.QuotaState, errCode commonpb.ErrorCode, reason string) bool {
 			if serviceStates[state] == nil {
-				serviceStates[state] = typeutil.NewSet[commonpb.ErrorCode]()
+				serviceStates[state] = typeutil.NewSet[stateReasonKey]()
 			}
-			serviceStates[state].Insert(errCode)
+			serviceStates[state].Insert(stateReasonKey{ErrorCode: errCode, Reason: reason})
 			return true
 		})
 
 	states := make([]milvuspb.QuotaState, 0)
 	reasons := make([]string, 0)
-	for state, errCodes := range serviceStates {
-		for errCode := range errCodes {
+	for state, stateReasonKeys := range serviceStates {
+		for key := range stateReasonKeys {
 			states = append(states, state)
-			reasons = append(reasons, ratelimitutil.GetQuotaErrorString(errCode))
+			reasons = append(reasons, ratelimitutil.GetQuotaErrorStringWithReason(key.ErrorCode, key.Reason))
 		}
 	}
 
@@ -285,11 +289,19 @@ func (m *SimpleLimiter) updateLimiterNode(req *proxypb.Limiter, node *rlinternal
 		limit.SetLimit(ratelimitutil.Limit(rate.GetR()))
 		setRateGaugeByRateType(rate.GetRt(), paramtable.GetNodeID(), sourceID, rate.GetR())
 	}
-	quotaStates := typeutil.NewConcurrentMap[milvuspb.QuotaState, commonpb.ErrorCode]()
+	quotaStates := typeutil.NewConcurrentMap[milvuspb.QuotaState, *rlinternal.QuotaStateInfo]()
 	states := req.GetStates()
 	codes := req.GetCodes()
+	reasons := req.GetReasons()
 	for i, state := range states {
-		quotaStates.Insert(state, codes[i])
+		reason := ""
+		if i < len(reasons) {
+			reason = reasons[i]
+		}
+		quotaStates.Insert(state, &rlinternal.QuotaStateInfo{
+			ErrorCode: codes[i],
+			Reason:    reason,
+		})
 	}
 	node.SetQuotaStates(quotaStates)
 	return nil

internal/proxy/simple_rate_limiter.go

@@ -25,7 +25,6 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
-	"github.com/milvus-io/milvus/internal/util/hookutil"
 	"github.com/milvus-io/milvus/pkg/v2/common"
 	"github.com/milvus-io/milvus/pkg/v2/util"
 	"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
@@ -60,7 +59,7 @@ func TestDDLCallbacksAlterCollectionProperties(t *testing.T) {
 	resp, err = core.AlterCollection(ctx, &milvuspb.AlterCollectionRequest{
 		DbName:         dbName,
 		CollectionName: collectionName,
-		Properties:     []*commonpb.KeyValuePair{{Key: hookutil.EncryptionEnabledKey, Value: "1"}},
+		Properties:     []*commonpb.KeyValuePair{{Key: common.EncryptionEnabledKey, Value: "1"}},
 	})
 	require.ErrorIs(t, merr.CheckRPCCall(resp, err), merr.ErrParameterInvalid)
 

internal/rootcoord/ddl_callbacks_alter_collection_properties_test.go

@@ -137,7 +137,7 @@ func (c *DDLCallback) alterDatabaseV1AckCallback(ctx context.Context, result mes
 	header := result.Message.Header()
 	body := result.Message.MustBody()
 
-	db := model.NewDatabase(header.DbId, header.DbName, etcdpb.DatabaseState_DatabaseCreated, result.Message.MustBody().Properties)
+	db := model.NewDatabase(header.DbId, header.DbName, etcdpb.DatabaseState_DatabaseCreated, body.Properties)
 	if err := c.meta.AlterDatabase(ctx, db, result.GetControlChannelResult().TimeTick); err != nil {
 		return errors.Wrap(err, "failed to alter database")
 	}

internal/rootcoord/ddl_callbacks_alter_database.go

@@ -23,7 +23,6 @@ import (
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
-	"github.com/milvus-io/milvus/internal/util/hookutil"
 	"github.com/milvus-io/milvus/pkg/v2/common"
 	"github.com/milvus-io/milvus/pkg/v2/proto/rootcoordpb"
 	"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
@@ -54,7 +53,7 @@ func TestDDLCallbacksAlterDatabase(t *testing.T) {
 	// hook related properties are not allowed to be altered.
 	resp, err = core.AlterDatabase(ctx, &rootcoordpb.AlterDatabaseRequest{
 		DbName:     dbName,
-		Properties: []*commonpb.KeyValuePair{{Key: hookutil.EncryptionEnabledKey, Value: "1"}},
+		Properties: []*commonpb.KeyValuePair{{Key: common.EncryptionEnabledKey, Value: "1"}},
 	})
 	require.ErrorIs(t, merr.CheckRPCCall(resp, err), merr.ErrParameterInvalid)
 

internal/rootcoord/ddl_callbacks_alter_database_test.go

@@ -0,0 +1,95 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package rootcoord
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+
+	"github.com/samber/lo"
+	"go.uber.org/zap"
+
+	"github.com/milvus-io/milvus/internal/metastore/model"
+	"github.com/milvus-io/milvus/internal/util/hookutil"
+	"github.com/milvus-io/milvus/pkg/v2/common"
+	"github.com/milvus-io/milvus/pkg/v2/log"
+	"github.com/milvus-io/milvus/pkg/v2/util"
+)
+
+type KeyManager struct {
+	ctx     context.Context
+	meta    IMetaTable
+	enabled bool
+}
+
+func NewKeyManager(
+	ctx context.Context,
+	meta IMetaTable,
+) *KeyManager {
+	if hookutil.GetCipherWithState() == nil {
+		log.Info("KeyManager disabled (cipher plugin not loaded)")
+		return nil
+	}
+	log.Info("KeyManager enabled")
+	return &KeyManager{
+		ctx:  ctx,
+		meta: meta,
+	}
+}
+
+func (km *KeyManager) GetRevokedDatabases() ([]int64, error) {
+	currentStates, err := hookutil.GetEzStates()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get cipher states: %w", err)
+	}
+
+	abnormalDB := make(map[int64]struct{})
+	for ezID, currentState := range currentStates {
+		if currentState != hookutil.KeyStateEnabled {
+			db, err := km.getDatabaseByEzID(ezID)
+			if err != nil {
+				log.Warn("KeyManager: failed to get database for ezID", zap.Int64("ezID", ezID), zap.Error(err))
+				continue
+			}
+
+			abnormalDB[db.ID] = struct{}{}
+		}
+	}
+
+	revokedDBIDs := lo.Keys(abnormalDB)
+	return revokedDBIDs, nil
+}
+
+func (km *KeyManager) getDatabaseByEzID(ezID int64) (*model.Database, error) {
+	// use ezID as dbID to get database
+	db, err := km.meta.GetDatabaseByID(km.ctx, ezID, 0)
+	if err != nil {
+		// fallback to default database(dbID=1)
+		db, err = km.meta.GetDatabaseByID(km.ctx, util.DefaultDBID, 0)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	// verify the ezID matches the retrieved DB
+	if db.GetProperty(common.EncryptionEzIDKey) != strconv.FormatInt(ezID, 10) {
+		return nil, fmt.Errorf("db for ezID %d not found", ezID)
+	}
+
+	return db, nil
+}

internal/rootcoord/key_manager.go

@@ -0,0 +1,112 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package rootcoord
+
+import (
+	"context"
+	"strconv"
+	"testing"
+
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
+	"github.com/milvus-io/milvus/internal/metastore/model"
+	mockrootcoord "github.com/milvus-io/milvus/internal/rootcoord/mocks"
+	"github.com/milvus-io/milvus/internal/util/hookutil"
+	"github.com/milvus-io/milvus/pkg/v2/common"
+	"github.com/milvus-io/milvus/pkg/v2/util"
+)
+
+func TestNewKeyManager(t *testing.T) {
+	ctx := context.Background()
+	meta := mockrootcoord.NewIMetaTable(t)
+	hookutil.InitTestCipher()
+
+	km := NewKeyManager(ctx, meta)
+
+	assert.NotNil(t, km)
+	assert.Equal(t, ctx, km.ctx)
+	assert.Equal(t, meta, km.meta)
+}
+
+func TestKeyManager_GetDatabaseByEzID(t *testing.T) {
+	ctx := context.Background()
+	hookutil.InitTestCipher()
+
+	t.Run("success get database", func(t *testing.T) {
+		meta := mockrootcoord.NewIMetaTable(t)
+
+		expectedDB := &model.Database{
+			ID:   123,
+			Name: "test_db",
+			Properties: []*commonpb.KeyValuePair{
+				{
+					Key:   common.EncryptionEzIDKey,
+					Value: "123", // the same as the dbID
+				},
+				{
+					Key:   common.EncryptionEnabledKey,
+					Value: "true",
+				},
+			},
+		}
+
+		meta.EXPECT().GetDatabaseByID(ctx, int64(123), uint64(0)).Return(expectedDB, nil).Once()
+
+		km := &KeyManager{
+			ctx:  ctx,
+			meta: meta,
+		}
+
+		db, err := km.getDatabaseByEzID(123)
+		assert.NoError(t, err)
+		assert.Equal(t, expectedDB, db)
+	})
+
+	t.Run("fallback to default database", func(t *testing.T) {
+		meta := mockrootcoord.NewIMetaTable(t)
+
+		ezID := int64(19530)
+		defaultDB := &model.Database{
+			ID:   util.DefaultDBID,
+			Name: util.DefaultDBName,
+			Properties: []*commonpb.KeyValuePair{
+				{
+					Key:   common.EncryptionEzIDKey,
+					Value: strconv.FormatInt(ezID, 10),
+				},
+				{
+					Key:   common.EncryptionEnabledKey,
+					Value: "true",
+				},
+			},
+		}
+
+		meta.EXPECT().GetDatabaseByID(ctx, ezID, uint64(0)).Return(nil, errors.New("db not found")).Once()
+		meta.EXPECT().GetDatabaseByID(ctx, util.DefaultDBID, uint64(0)).Return(defaultDB, nil).Once()
+
+		km := &KeyManager{
+			ctx:  ctx,
+			meta: meta,
+		}
+
+		db, err := km.getDatabaseByEzID(ezID)
+		assert.NoError(t, err)
+		assert.Equal(t, defaultDB, db)
+	})
+}