expand privilege group when list policy in rootcoord

Signed-off-by: shaoting-huang <[email protected]>

internal/metastore/catalog.go

@@ -76,7 +76,7 @@ type RootCoordCatalog interface {
 	// ListGrant lists all grant infos accoording to entity for the tenant
 	// Please make sure entity valid before calling this API
 	ListGrant(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error)
-	ListPolicy(ctx context.Context, tenant string) ([]string, error)
+	ListPolicy(ctx context.Context, tenant string) ([]*milvuspb.GrantEntity, error)
 	// List all user role pair in string for the tenant
 	// For example []string{"user1/role1"}
 	ListUserRole(ctx context.Context, tenant string) ([]string, error)

internal/metastore/kv/rootcoord/kv_catalog.go

@@ -1283,13 +1283,13 @@ func (kc *Catalog) DeleteGrant(ctx context.Context, tenant string, role *milvusp
 	return err
 }
 
-func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]string, error) {
-	var grantInfoStrs []string
+func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]*milvuspb.GrantEntity, error) {
+	var grants []*milvuspb.GrantEntity
 	granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "")
 	keys, values, err := kc.Txn.LoadWithPrefix(granteeKey)
 	if err != nil {
 		log.Error("fail to load all grant privilege entities", zap.String("key", granteeKey), zap.Error(err))
-		return []string{}, err
+		return []*milvuspb.GrantEntity{}, err
 	}
 
 	for i, key := range keys {
@@ -1302,7 +1302,7 @@ func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]string, err
 		idKeys, _, err := kc.Txn.LoadWithPrefix(granteeIDKey)
 		if err != nil {
 			log.Error("fail to load the grantee ids", zap.String("key", granteeIDKey), zap.Error(err))
-			return []string{}, err
+			return []*milvuspb.GrantEntity{}, err
 		}
 		for _, idKey := range idKeys {
 			granteeIDInfos := typeutil.AfterN(idKey, granteeIDKey+"/", "/")
@@ -1311,11 +1311,18 @@ func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]string, err
 				continue
 			}
 			dbName, objectName := funcutil.SplitObjectName(grantInfos[2])
-			grantInfoStrs = append(grantInfoStrs,
-				funcutil.PolicyForPrivilege(grantInfos[0], grantInfos[1], objectName, granteeIDInfos[0], dbName))
+			grants = append(grants, &milvuspb.GrantEntity{
+				Role:       &milvuspb.RoleEntity{Name: grantInfos[0]},
+				Object:     &milvuspb.ObjectEntity{Name: grantInfos[1]},
+				ObjectName: objectName,
+				DbName:     dbName,
+				Grantor: &milvuspb.GrantorEntity{
+					Privilege: &milvuspb.PrivilegeEntity{Name: util.PrivilegeNameForAPI(granteeIDInfos[0])},
+				},
+			})
 		}
 	}
-	return grantInfoStrs, nil
+	return grants, nil
 }
 
 func (kc *Catalog) ListUserRole(ctx context.Context, tenant string) ([]string, error) {

internal/metastore/kv/rootcoord/kv_catalog_test.go

@@ -2476,6 +2476,18 @@ func TestRBAC_Grant(t *testing.T) {
 			secondLoadWithPrefixReturn atomic.Bool
 		)
 
+		grant := func(role, obj, objName, privilege, dbName string) *milvuspb.GrantEntity {
+			return &milvuspb.GrantEntity{
+				Role:       &milvuspb.RoleEntity{Name: role},
+				Object:     &milvuspb.ObjectEntity{Name: obj},
+				ObjectName: objName,
+				DbName:     dbName,
+				Grantor: &milvuspb.GrantorEntity{
+					Privilege: &milvuspb.PrivilegeEntity{Name: util.PrivilegeNameForAPI(privilege)},
+				},
+			}
+		}
+
 		kvmock.EXPECT().LoadWithPrefix(mock.Anything).Call.Return(
 			func(key string) []string {
 				contains := strings.Contains(key, GranteeIDPrefix)
@@ -2548,11 +2560,11 @@ func TestRBAC_Grant(t *testing.T) {
 				if test.isValid {
 					assert.NoError(t, err)
 					assert.Equal(t, 4, len(policy))
-					ps := []string{
-						funcutil.PolicyForPrivilege("role1", "obj1", "obj_name1", "PrivilegeLoad", "default"),
-						funcutil.PolicyForPrivilege("role1", "obj1", "obj_name1", "PrivilegeRelease", "default"),
-						funcutil.PolicyForPrivilege("role2", "obj2", "obj_name2", "PrivilegeLoad", "default"),
-						funcutil.PolicyForPrivilege("role2", "obj2", "obj_name2", "PrivilegeRelease", "default"),
+					ps := []*milvuspb.GrantEntity{
+						grant("role1", "obj1", "obj_name1", "PrivilegeLoad", "default"),
+						grant("role1", "obj1", "obj_name1", "PrivilegeRelease", "default"),
+						grant("role2", "obj2", "obj_name2", "PrivilegeLoad", "default"),
+						grant("role2", "obj2", "obj_name2", "PrivilegeRelease", "default"),
 					}
 					assert.ElementsMatch(t, ps, policy)
 				} else {

internal/rootcoord/meta_table.go

@@ -94,7 +94,7 @@ type IMetaTable interface {
 	OperatePrivilege(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error
 	SelectGrant(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error)
 	DropGrant(tenant string, role *milvuspb.RoleEntity) error
-	ListPolicy(tenant string) ([]string, error)
+	ListPolicy(tenant string) ([]*milvuspb.GrantEntity, error)
 	ListUserRole(tenant string) ([]string, error)
 	BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error)
 	RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error
@@ -1430,7 +1430,7 @@ func (mt *MetaTable) DropGrant(tenant string, role *milvuspb.RoleEntity) error {
 	return mt.catalog.DeleteGrant(mt.ctx, tenant, role)
 }
 
-func (mt *MetaTable) ListPolicy(tenant string) ([]string, error) {
+func (mt *MetaTable) ListPolicy(tenant string) ([]*milvuspb.GrantEntity, error) {
 	mt.permissionLock.RLock()
 	defer mt.permissionLock.RUnlock()
 

internal/rootcoord/mock_test.go

@@ -93,7 +93,7 @@ type mockMetaTable struct {
 	OperatePrivilegeFunc             func(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error
 	SelectGrantFunc                  func(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error)
 	DropGrantFunc                    func(tenant string, role *milvuspb.RoleEntity) error
-	ListPolicyFunc                   func(tenant string) ([]string, error)
+	ListPolicyFunc                   func(tenant string) ([]*milvuspb.GrantEntity, error)
 	ListUserRoleFunc                 func(tenant string) ([]string, error)
 	DescribeDatabaseFunc             func(ctx context.Context, dbName string) (*model.Database, error)
 	CreatePrivilegeGroupFunc         func(groupName string) error
@@ -248,7 +248,7 @@ func (m mockMetaTable) DropGrant(tenant string, role *milvuspb.RoleEntity) error
 	return m.DropGrantFunc(tenant, role)
 }
 
-func (m mockMetaTable) ListPolicy(tenant string) ([]string, error) {
+func (m mockMetaTable) ListPolicy(tenant string) ([]*milvuspb.GrantEntity, error) {
 	return m.ListPolicyFunc(tenant)
 }
 
@@ -540,7 +540,7 @@ func withInvalidMeta() Opt {
 	meta.DropGrantFunc = func(tenant string, role *milvuspb.RoleEntity) error {
 		return errors.New("error mock DropGrant")
 	}
-	meta.ListPolicyFunc = func(tenant string) ([]string, error) {
+	meta.ListPolicyFunc = func(tenant string) ([]*milvuspb.GrantEntity, error) {
 		return nil, errors.New("error mock ListPolicy")
 	}
 	meta.ListUserRoleFunc = func(tenant string) ([]string, error) {

internal/rootcoord/root_coord.go

@@ -2686,8 +2686,7 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
 		}
 		grants := []*milvuspb.GrantEntity{in.Entity}
 
-		allGroups, err := c.meta.ListPrivilegeGroups()
-		allGroups = append(allGroups, Params.RbacConfig.GetDefaultPrivilegeGroups()...)
+		allGroups, err := c.getPrivilegeGroups()
 		if err != nil {
 			return nil, err
 		}
@@ -2875,6 +2874,22 @@ func (c *Core) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest)
 			Status: merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_ListPolicyFailure),
 		}, nil
 	}
+	// expand privilege groups and turn to policies
+	allGroups, err := c.getPrivilegeGroups()
+	if err != nil {
+		return nil, err
+	}
+	groups := lo.SliceToMap(allGroups, func(group *milvuspb.PrivilegeGroupInfo) (string, []*milvuspb.PrivilegeEntity) {
+		return group.GroupName, group.Privileges
+	})
+	expandGrants, err := c.expandPrivilegeGroups(policies, groups)
+	if err != nil {
+		return nil, err
+	}
+	expandPolicies := lo.Map(expandGrants, func(r *milvuspb.GrantEntity, _ int) string {
+		return funcutil.PolicyForPrivilege(r.Role.Name, r.Object.Name, r.ObjectName, r.Grantor.Privilege.Name, r.DbName)
+	})
+
 	userRoles, err := c.meta.ListUserRole(util.DefaultTenant)
 	if err != nil {
 		errMsg := "fail to list user-role"
@@ -2897,7 +2912,7 @@ func (c *Core) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest)
 	metrics.RootCoordDDLReqLatency.WithLabelValues(method).Observe(float64(tr.ElapseSpan().Milliseconds()))
 	return &internalpb.ListPolicyResponse{
 		Status:          merr.Success(),
-		PolicyInfos:     policies,
+		PolicyInfos:     expandPolicies,
 		UserRoles:       userRoles,
 		PrivilegeGroups: privGroups,
 	}, nil
@@ -3358,3 +3373,13 @@ func (c *Core) expandPrivilegeGroups(grants []*milvuspb.GrantEntity, groups map[
 		return fmt.Sprintf("%s-%s-%s-%s-%s-%s", g.Role, g.Object, g.ObjectName, g.Grantor.User, g.Grantor.Privilege.Name, g.DbName)
 	}), nil
 }
+
+// getPrivilegeGroups returns default privilege groups and user-defined privilege groups.
+func (c *Core) getPrivilegeGroups() ([]*milvuspb.PrivilegeGroupInfo, error) {
+	allGroups, err := c.meta.ListPrivilegeGroups()
+	allGroups = append(allGroups, Params.RbacConfig.GetDefaultPrivilegeGroups()...)
+	if err != nil {
+		return nil, err
+	}
+	return allGroups, nil
+}

internal/rootcoord/root_coord_test.go

@@ -1027,6 +1027,34 @@ func TestRootCoord_RenameCollection(t *testing.T) {
 	})
 }
 
+func TestRootCoord_ListPolicy(t *testing.T) {
+	t.Run("expand privilege groups", func(t *testing.T) {
+		meta := mockrootcoord.NewIMetaTable(t)
+		c := newTestCore(withHealthyCode(), withMeta(meta))
+		ctx := context.Background()
+
+		meta.EXPECT().ListPolicy(util.DefaultTenant).Return([]*milvuspb.GrantEntity{
+			{
+				ObjectName: "*",
+				Object: &milvuspb.ObjectEntity{
+					Name: "Global",
+				},
+				Role:    &milvuspb.RoleEntity{Name: "role"},
+				Grantor: &milvuspb.GrantorEntity{Privilege: &milvuspb.PrivilegeEntity{Name: "CollectionAdmin"}},
+			},
+		}, nil)
+
+		meta.EXPECT().ListPrivilegeGroups().Return([]*milvuspb.PrivilegeGroupInfo{}, nil)
+
+		meta.EXPECT().ListUserRole(util.DefaultTenant).Return([]string{}, nil)
+
+		resp, err := c.ListPolicy(ctx, &internalpb.ListPolicyRequest{})
+		assert.Equal(t, len(Params.RbacConfig.GetDefaultPrivilegeGroup("CollectionAdmin").Privileges), len(resp.PolicyInfos))
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
+	})
+}
+
 func TestRootCoord_ShowConfigurations(t *testing.T) {
 	t.Run("not healthy", func(t *testing.T) {
 		ctx := context.Background()
@@ -1778,14 +1806,25 @@ func TestRootCoord_RBACError(t *testing.T) {
 		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
 
 		mockMeta := c.meta.(*mockMetaTable)
-		mockMeta.ListPolicyFunc = func(tenant string) ([]string, error) {
-			return []string{}, nil
+		mockMeta.ListPolicyFunc = func(tenant string) ([]*milvuspb.GrantEntity, error) {
+			return []*milvuspb.GrantEntity{}, nil
+		}
+		mockMeta.ListPrivilegeGroupsFunc = func() ([]*milvuspb.PrivilegeGroupInfo, error) {
+			return []*milvuspb.PrivilegeGroupInfo{
+				{
+					GroupName:  "CollectionAdmin",
+					Privileges: []*milvuspb.PrivilegeEntity{{Name: "CreateCollection"}},
+				},
+			}, nil
 		}
 		resp, err = c.ListPolicy(ctx, &internalpb.ListPolicyRequest{})
 		assert.NoError(t, err)
 		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
-		mockMeta.ListPolicyFunc = func(tenant string) ([]string, error) {
-			return []string{}, errors.New("mock error")
+		mockMeta.ListPolicyFunc = func(tenant string) ([]*milvuspb.GrantEntity, error) {
+			return []*milvuspb.GrantEntity{}, errors.New("mock error")
+		}
+		mockMeta.ListPrivilegeGroupsFunc = func() ([]*milvuspb.PrivilegeGroupInfo, error) {
+			return []*milvuspb.PrivilegeGroupInfo{}, errors.New("mock error")
 		}
 	})
 }