fix: metastore privilege name check with privilege name all (#39476)

related: #39365 --------- Signed-off-by: shaoting-huang <[email protected]>
milvus-io · Jan 26, 2025 · 5f4bad6 · 5f4bad6
1 parent e61a841
commit 5f4bad6
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 17 deletions.
diff --git a/internal/rootcoord/rbac_task.go b/internal/rootcoord/rbac_task.go
@@ -148,16 +148,14 @@ func executeOperatePrivilegeTaskSteps(ctx context.Context, core *Core, in *milvu
 	privName := in.Entity.Grantor.Privilege.Name
 	redoTask := newBaseRedoTask(core.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 := core.getMetastorePrivilegeName(ctx, privName)
-			if err != nil {
-				return nil, err
-			}
-			in.Entity.Grantor.Privilege.Name = dbPrivName
+		// set up privilege name for metastore
+		dbPrivName, err := core.getMetastorePrivilegeName(ctx, privName)
+		if err != nil {
+			return nil, err
 		}
+		in.Entity.Grantor.Privilege.Name = dbPrivName
 
-		err := core.meta.OperatePrivilege(ctx, util.DefaultTenant, in.Entity, in.Type)
+		err = core.meta.OperatePrivilege(ctx, util.DefaultTenant, in.Entity, in.Type)
 		if err != nil && !common.IsIgnorableError(err) {
 			log.Ctx(ctx).Warn("fail to operate the privilege", zap.Any("in", in), zap.Error(err))
 			return nil, err

diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go
@@ -664,15 +664,12 @@ 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(c.ctx, util.DefaultTenant, &milvuspb.GrantEntity{
+
+			err = c.meta.OperatePrivilege(c.ctx, util.DefaultTenant, &milvuspb.GrantEntity{
 				Role:       &milvuspb.RoleEntity{Name: role},
 				Object:     &milvuspb.ObjectEntity{Name: privilege[util.RoleConfigObjectType]},
 				ObjectName: privilege[util.RoleConfigObjectName],
@@ -2775,6 +2772,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
@@ -2787,7 +2788,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
@@ -2207,6 +2207,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(mock.Anything, "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(context.Background(), 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
 }