expand privilege group when list policy in rootcoord

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

internal/metastore/catalog.go

@@ -77,7 +77,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

@@ -1354,13 +1354,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(ctx, granteeKey)
 	if err != nil {
 		log.Ctx(ctx).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 {
@@ -1373,7 +1373,7 @@ func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]string, err
 		idKeys, _, err := kc.Txn.LoadWithPrefix(ctx, granteeIDKey)
 		if err != nil {
 			log.Ctx(ctx).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+"/", "/")
@@ -1382,11 +1382,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/rootcoord/meta_table.go

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

internal/rootcoord/mock_test.go

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

internal/rootcoord/root_coord.go

@@ -2712,8 +2712,7 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
 		}
 		grants := []*milvuspb.GrantEntity{in.Entity}
 
-		allGroups, err := c.meta.ListPrivilegeGroups(ctx)
-		allGroups = append(allGroups, Params.RbacConfig.GetDefaultPrivilegeGroups()...)
+		allGroups, err := c.getPrivilegeGroups(ctx)
 		if err != nil {
 			return nil, err
 		}
@@ -2901,6 +2900,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(ctx)
+	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(ctx, 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(ctx, util.DefaultTenant)
 	if err != nil {
 		errMsg := "fail to list user-role"
@@ -2923,7 +2938,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
@@ -3397,3 +3412,13 @@ func (c *Core) expandPrivilegeGroups(ctx context.Context, grants []*milvuspb.Gra
 		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(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
+	allGroups, err := c.meta.ListPrivilegeGroups(ctx)
+	allGroups = append(allGroups, Params.RbacConfig.GetDefaultPrivilegeGroups()...)
+	if err != nil {
+		return nil, err
+	}
+	return allGroups, nil
+}

internal/rootcoord/root_coord_test.go

@@ -1054,6 +1054,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(ctx, 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(ctx).Return([]*milvuspb.PrivilegeGroupInfo{}, nil)
+
+		meta.EXPECT().ListUserRole(ctx, 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()
@@ -1917,14 +1945,14 @@ func TestRootCoord_RBACError(t *testing.T) {
 		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
 
 		mockMeta := c.meta.(*mockMetaTable)
-		mockMeta.ListPolicyFunc = func(ctx context.Context, tenant string) ([]string, error) {
-			return []string{}, nil
+		mockMeta.ListPolicyFunc = func(ctx context.Context, tenant string) ([]*milvuspb.GrantEntity, error) {
+			return []*milvuspb.GrantEntity{}, nil
 		}
 		resp, err = c.ListPolicy(ctx, &internalpb.ListPolicyRequest{})
 		assert.NoError(t, err)
 		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
-		mockMeta.ListPolicyFunc = func(ctx context.Context, tenant string) ([]string, error) {
-			return []string{}, errors.New("mock error")
+		mockMeta.ListPolicyFunc = func(ctx context.Context, tenant string) ([]*milvuspb.GrantEntity, error) {
+			return []*milvuspb.GrantEntity{}, errors.New("mock error")
 		}
 	})
 }

tests/integration/rbac/privilege_group_test.go

@@ -180,6 +180,12 @@ func (s *PrivilegeGroupTestSuite) TestGrantV2BuiltinPrivilegeGroup() {
 	s.True(merr.Ok(resp))
 	resp, _ = s.operatePrivilegeV2(ctx, roleName, "CollectionAdmin", util.AnyWord, "col1", milvuspb.OperatePrivilegeType_Grant)
 	s.False(merr.Ok(resp))
+
+	s.Cluster.Proxy.Stop()
+	s.Cluster.Proxy.Run()
+
+	_, err = s.Cluster.Proxy.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{})
+	s.NoError(err)
 }
 
 func (s *PrivilegeGroupTestSuite) TestGrantV2CustomPrivilegeGroup() {