RBAC custom privilege group

Signed-off-by: shaoting-huang <[email protected]>
milvus-io · Oct 23, 2024 · 370e705 · 370e705
1 parent 7da2c9f
commit 370e705
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 6 deletions.
diff --git a/internal/rootcoord/meta_table.go b/internal/rootcoord/meta_table.go
@@ -92,7 +92,7 @@ type IMetaTable interface {
 	OperateUserRole(tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error
 	SelectRole(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error)
 	SelectUser(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error)
-	OperatePrivilege(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error
+	OperatePrivilege(ctx context.Context, 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)
@@ -1370,7 +1370,7 @@ func (mt *MetaTable) SelectUser(tenant string, entity *milvuspb.UserEntity, incl
 }
 
 // OperatePrivilege grant or revoke privilege by setting the operateType param
-func (mt *MetaTable) OperatePrivilege(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
+func (mt *MetaTable) OperatePrivilege(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
 	if funcutil.IsEmptyString(entity.ObjectName) {
 		return fmt.Errorf("the object name in the grant entity is empty")
 	}
@@ -1386,6 +1386,10 @@ func (mt *MetaTable) OperatePrivilege(tenant string, entity *milvuspb.GrantEntit
 	if entity.Grantor.Privilege == nil || funcutil.IsEmptyString(entity.Grantor.Privilege.Name) {
 		return fmt.Errorf("the privilege name in the grant entity is empty")
 	}
+	conflict, err := mt.CheckWithPrivilegeGroupNames(ctx, entity.Grantor.Privilege.Name)
+	if err != nil || conflict {
+		return fmt.Errorf("the privilege name in the grant entity is conflict with any of the privilege group names")
+	}
 	if entity.Grantor.User == nil || funcutil.IsEmptyString(entity.Grantor.User.Name) {
 		return fmt.Errorf("the grantor name in the grant entity is empty")
 	}
@@ -1402,6 +1406,23 @@ func (mt *MetaTable) OperatePrivilege(tenant string, entity *milvuspb.GrantEntit
 	return mt.catalog.AlterGrant(mt.ctx, tenant, entity, operateType)
 }
 
+// IsPrivilegeInAnyGroup checks if the given privilege name is conflict with any of the privilege group names.
+func (mt *MetaTable) CheckWithPrivilegeGroupNames(ctx context.Context, privilegeName string) (bool, error) {
+	privilegeGroups, err := mt.catalog.ListPrivilegeGroups(ctx)
+	if err != nil {
+		log.Warn("fail to list privilege groups", zap.Error(err))
+		return false, err
+	}
+
+	privilegeGroupNames := make(map[string]struct{})
+	for _, group := range privilegeGroups {
+		privilegeGroupNames[group.Name] = struct{}{}
+	}
+
+	_, exists := privilegeGroupNames[privilegeName]
+	return exists, nil
+}
+
 // SelectGrant select grant
 // The principal entity MUST be not empty in the grant entity
 // The resource entity and the resource name are optional, and the two params should be not empty together when you select some grants about the resource kind.

diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go
@@ -570,7 +570,7 @@ func (c *Core) initPublicRolePrivilege() error {
 
 	var err error
 	for _, globalPrivilege := range globalPrivileges {
-		err = c.meta.OperatePrivilege(util.DefaultTenant, &milvuspb.GrantEntity{
+		err = c.meta.OperatePrivilege(c.ctx, util.DefaultTenant, &milvuspb.GrantEntity{
 			Role:       &milvuspb.RoleEntity{Name: util.RolePublic},
 			Object:     &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Global.String()},
 			ObjectName: util.AnyWord,
@@ -585,7 +585,7 @@ func (c *Core) initPublicRolePrivilege() error {
 		}
 	}
 	for _, collectionPrivilege := range collectionPrivileges {
-		err = c.meta.OperatePrivilege(util.DefaultTenant, &milvuspb.GrantEntity{
+		err = c.meta.OperatePrivilege(c.ctx, util.DefaultTenant, &milvuspb.GrantEntity{
 			Role:       &milvuspb.RoleEntity{Name: util.RolePublic},
 			Object:     &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
 			ObjectName: util.AnyWord,
@@ -615,7 +615,7 @@ func (c *Core) initBuiltinRoles() error {
 			if !util.IsAnyWord(privilege[util.RoleConfigPrivilege]) {
 				privilegeName = util.PrivilegeNameForMetastore(privilege[util.RoleConfigPrivilege])
 			}
-			err := c.meta.OperatePrivilege(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],
@@ -2579,7 +2579,7 @@ 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) {
-		err := c.meta.OperatePrivilege(util.DefaultTenant, in.Entity, in.Type)
+		err := c.meta.OperatePrivilege(c.ctx, 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