From ecd1ea554ba4d3a970089eac0960afa8eaae5737 Mon Sep 17 00:00:00 2001
From: sthuang <167743503+shaoting-huang@users.noreply.github.com>
Date: Sun, 26 Jan 2025 12:17:19 +0800
Subject: [PATCH] fix: [2.4] metastore privilege name check with privilege name
 all (#39493)

cherry-pick from master: https://github.com/milvus-io/milvus/pull/39476
related: https://github.com/milvus-io/milvus/issues/39365

Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>
---
 internal/rootcoord/root_coord.go      | 32 ++++++++++++-------------
 internal/rootcoord/root_coord_test.go | 34 +++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 17 deletions(-)

diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go
index 92caeb219dae9..b36497615ac16 100644
--- a/internal/rootcoord/root_coord.go
+++ b/internal/rootcoord/root_coord.go
@@ -624,15 +624,11 @@ func (c *Core) initBuiltinRoles() error {
 			return errors.Wrapf(err, "failed to create a builtin role: %s", role)
 		}
 		for _, privilege := range privilegesJSON[util.RoleConfigPrivileges] {
-			privilegeName := privilege[util.RoleConfigPrivilege]
-			if !util.IsAnyWord(privilege[util.RoleConfigPrivilege]) {
-				dbPrivName, err := c.getMetastorePrivilegeName(c.ctx, privilege[util.RoleConfigPrivilege])
-				if err != nil {
-					return errors.Wrapf(err, "failed to get metastore privilege name for: %s", privilege[util.RoleConfigPrivilege])
-				}
-				privilegeName = dbPrivName
+			privilegeName, err := c.getMetastorePrivilegeName(c.ctx, privilege[util.RoleConfigPrivilege])
+			if err != nil {
+				return errors.Wrapf(err, "failed to get metastore privilege name for: %s", privilege[util.RoleConfigPrivilege])
 			}
-			err := c.meta.OperatePrivilege(util.DefaultTenant, &milvuspb.GrantEntity{
+			err = c.meta.OperatePrivilege(util.DefaultTenant, &milvuspb.GrantEntity{
 				Role:       &milvuspb.RoleEntity{Name: role},
 				Object:     &milvuspb.ObjectEntity{Name: privilege[util.RoleConfigObjectType]},
 				ObjectName: privilege[util.RoleConfigObjectName],
@@ -2664,16 +2660,14 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
 
 	redoTask := newBaseRedoTask(c.stepExecutor)
 	redoTask.AddSyncStep(NewSimpleStep("operate privilege meta data", func(ctx context.Context) ([]nestedStep, error) {
-		if !util.IsAnyWord(privName) {
-			// set up privilege name for metastore
-			dbPrivName, err := c.getMetastorePrivilegeName(ctx, privName)
-			if err != nil {
-				return nil, err
-			}
-			in.Entity.Grantor.Privilege.Name = dbPrivName
+		// set up privilege name for metastore
+		dbPrivName, err := c.getMetastorePrivilegeName(ctx, privName)
+		if err != nil {
+			return nil, err
 		}
+		in.Entity.Grantor.Privilege.Name = dbPrivName
 
-		err := c.meta.OperatePrivilege(util.DefaultTenant, in.Entity, in.Type)
+		err = c.meta.OperatePrivilege(util.DefaultTenant, in.Entity, in.Type)
 		if err != nil && !common.IsIgnorableError(err) {
 			log.Warn("fail to operate the privilege", zap.Any("in", in), zap.Error(err))
 			return nil, err
@@ -2822,6 +2816,10 @@ func (c *Core) validatePrivilegeGroupParams(ctx context.Context, entity string,
 }
 
 func (c *Core) getMetastorePrivilegeName(ctx context.Context, privName string) (string, error) {
+	// if it is '*', return directly
+	if util.IsAnyWord(privName) {
+		return privName, nil
+	}
 	// if it is built-in privilege, return the privilege name directly
 	if util.IsPrivilegeNameDefined(privName) {
 		return util.PrivilegeNameForMetastore(privName), nil
@@ -2834,7 +2832,7 @@ func (c *Core) getMetastorePrivilegeName(ctx context.Context, privName string) (
 	if customGroup {
 		return util.PrivilegeGroupNameForMetastore(privName), nil
 	}
-	return "", errors.New("not found the privilege name")
+	return "", errors.Newf("not found the privilege name [%s] from metastore", privName)
 }
 
 // SelectGrant select grant
diff --git a/internal/rootcoord/root_coord_test.go b/internal/rootcoord/root_coord_test.go
index 3aca8b79f0254..17162f00e0647 100644
--- a/internal/rootcoord/root_coord_test.go
+++ b/internal/rootcoord/root_coord_test.go
@@ -1993,6 +1993,40 @@ func TestCore_RestoreRBAC(t *testing.T) {
 	assert.False(t, merr.Ok(resp))
 }
 
+func TestCore_getMetastorePrivilegeName(t *testing.T) {
+	meta := mockrootcoord.NewIMetaTable(t)
+	c := newTestCore(withHealthyCode(), withMeta(meta))
+
+	priv, err := c.getMetastorePrivilegeName(context.Background(), util.AnyWord)
+	assert.NoError(t, err)
+	assert.Equal(t, priv, util.AnyWord)
+
+	meta.EXPECT().IsCustomPrivilegeGroup("unknown").Return(false, nil)
+	_, err = c.getMetastorePrivilegeName(context.Background(), "unknown")
+	assert.Equal(t, err.Error(), "not found the privilege name [unknown] from metastore")
+}
+
+func TestCore_expandPrivilegeGroup(t *testing.T) {
+	meta := mockrootcoord.NewIMetaTable(t)
+	c := newTestCore(withHealthyCode(), withMeta(meta))
+
+	grants := []*milvuspb.GrantEntity{
+		{
+			ObjectName: "*",
+			Object: &milvuspb.ObjectEntity{
+				Name: "Global",
+			},
+			Role:    &milvuspb.RoleEntity{Name: "role"},
+			Grantor: &milvuspb.GrantorEntity{Privilege: &milvuspb.PrivilegeEntity{Name: "*"}},
+		},
+	}
+	groups := map[string][]*milvuspb.PrivilegeEntity{}
+	expandGrants, err := c.expandPrivilegeGroups(grants, groups)
+	assert.NoError(t, err)
+	assert.Equal(t, len(expandGrants), len(grants))
+	assert.Equal(t, expandGrants[0].Grantor.Privilege.Name, grants[0].Grantor.Privilege.Name)
+}
+
 type RootCoordSuite struct {
 	suite.Suite
 }