mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 02:05:41 +00:00
fix: avoid repeated grantee scans in ListPolicy (#51404)
Related: #50236 - Reuse the initially loaded grantee keys/values when ListPolicy checks legacy grantee ID ownership. - Avoid reloading the full grantee-privileges tree once per legacy grant. - Add a regression test that verifies ListPolicy performs only one full grantee-prefix load for legacy grants. Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>
This commit is contained in:
@@ -1491,11 +1491,24 @@ func shouldRemoveGranteeIDSubtree(ctx context.Context, granteePrefix string, idS
|
||||
}
|
||||
|
||||
func (kc *Catalog) loadGranteeIDPrefix(ctx context.Context, tenant string, granteeKey string, idStr string) ([]string, []string, string, error) {
|
||||
return kc.loadGranteeIDPrefixWithLoadedGrantees(ctx, tenant, granteeKey, idStr, nil, nil)
|
||||
}
|
||||
|
||||
func (kc *Catalog) loadGranteeIDPrefixWithLoadedGrantees(ctx context.Context, tenant string, granteeKey string, idStr string, granteeKeys []string, granteeValues []string) ([]string, []string, string, error) {
|
||||
var firstPrefix string
|
||||
newID := crypto.GranteeID(granteeKey)
|
||||
if isLegacyGranteeID(idStr) && idStr != newID {
|
||||
granteeIDKey := funcutil.HandleTenantForEtcdPrefix(GranteeIDPrefix, tenant, idStr)
|
||||
otherGranteeKey, err := kc.findOtherGranteeWithID(ctx, tenant, granteeKey, idStr)
|
||||
var (
|
||||
otherGranteeKey string
|
||||
err error
|
||||
)
|
||||
if granteeKeys != nil {
|
||||
granteePrefix := funcutil.HandleTenantForEtcdPrefix(GranteePrefix, tenant)
|
||||
otherGranteeKey, err = findOtherGranteeWithIDFromKeys(ctx, granteePrefix, granteeKey, idStr, granteeKeys, granteeValues)
|
||||
} else {
|
||||
otherGranteeKey, err = kc.findOtherGranteeWithID(ctx, tenant, granteeKey, idStr)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, nil, granteeIDKey, err
|
||||
}
|
||||
@@ -1981,7 +1994,7 @@ func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]*milvuspb.G
|
||||
continue
|
||||
}
|
||||
logicalGranteeKey := fmt.Sprintf("%s/%s/%s/%s", GranteePrefix, grantInfos[0], grantInfos[1], grantInfos[2])
|
||||
idKeys, _, granteeIDKey, err := kc.loadGranteeIDPrefix(ctx, tenant, logicalGranteeKey, values[i])
|
||||
idKeys, _, granteeIDKey, err := kc.loadGranteeIDPrefixWithLoadedGrantees(ctx, tenant, logicalGranteeKey, values[i], keys, values)
|
||||
if err != nil {
|
||||
mlog.Error(ctx, "fail to load the grantee ids", mlog.String("key", granteeIDKey), mlog.Err(err))
|
||||
return []*milvuspb.GrantEntity{}, err
|
||||
|
||||
@@ -3254,6 +3254,56 @@ func TestRBACGrantSharedLegacyGranteeIDMigrationFailsClosed(t *testing.T) {
|
||||
assert.Equal(t, "donor-user", donorUser)
|
||||
}
|
||||
|
||||
func TestRBACListPolicyLegacyGranteeIDReusesLoadedGrantKeys(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tenant := util.DefaultTenant
|
||||
kvmock := mocks.NewTxnKV(t)
|
||||
c := NewCatalog(kvmock)
|
||||
|
||||
granteePrefix := funcutil.HandleTenantForEtcdPrefix(GranteePrefix, tenant)
|
||||
granteeIDPrefix := funcutil.HandleTenantForEtcdPrefix(GranteeIDPrefix, tenant)
|
||||
objectType := commonpb.ObjectType_Collection.String()
|
||||
firstLogicalKey := fmt.Sprintf("%s/%s/%s/%s", GranteePrefix, "role1", objectType, funcutil.CombineObjectName(util.DefaultDBName, "coll1"))
|
||||
secondLogicalKey := fmt.Sprintf("%s/%s/%s/%s", GranteePrefix, "role2", objectType, funcutil.CombineObjectName(util.DefaultDBName, "coll2"))
|
||||
firstEtcdKey := fmt.Sprintf("%s%s/%s/%s", granteePrefix, "role1", objectType, funcutil.CombineObjectName(util.DefaultDBName, "coll1"))
|
||||
secondEtcdKey := fmt.Sprintf("%s%s/%s/%s", granteePrefix, "role2", objectType, funcutil.CombineObjectName(util.DefaultDBName, "coll2"))
|
||||
firstLegacyID := crypto.MD5(firstLogicalKey)
|
||||
secondLegacyID := crypto.MD5(secondLogicalKey)
|
||||
var granteePrefixLoads atomic.Int32
|
||||
|
||||
kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
|
||||
func(ctx context.Context, key string) []string {
|
||||
switch {
|
||||
case key == granteePrefix:
|
||||
granteePrefixLoads.Inc()
|
||||
return []string{firstEtcdKey, secondEtcdKey}
|
||||
case strings.HasPrefix(key, granteeIDPrefix):
|
||||
return []string{key + "PrivilegeLoad"}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
},
|
||||
func(ctx context.Context, key string) []string {
|
||||
switch {
|
||||
case key == granteePrefix:
|
||||
return []string{firstLegacyID, secondLegacyID}
|
||||
case strings.HasPrefix(key, granteeIDPrefix):
|
||||
return []string{"root"}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
},
|
||||
func(ctx context.Context, key string) error {
|
||||
return nil
|
||||
},
|
||||
)
|
||||
|
||||
policies, err := c.ListPolicy(ctx, tenant)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, policies, 2)
|
||||
assert.Equal(t, int32(1), granteePrefixLoads.Load(), "ListPolicy should reuse the initially loaded grantee keys for legacy ID collision checks")
|
||||
}
|
||||
|
||||
func TestRBACGrantMigrationIgnoresUnreferencedComputedLegacyID(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
etcdCli, _ := etcd.GetEtcdClient(
|
||||
|
||||
Reference in New Issue
Block a user