Skip to content

Commit

Permalink
RBAC custom privilege group
Browse files Browse the repository at this point in the history 
Signed-off-by: shaoting-huang <[email protected]>
  • Loading branch information
@shaoting-huang
shaoting-huang committed Oct 23, 2024
1 parent 3d4afcb commit 6853e8c
Show file tree
Hide file tree
Showing 3 changed files with 91 additions and 69 deletions.
5 changes: 2 additions & 3 deletions internal/metastore/catalog.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,11 +86,10 @@ type RootCoordCatalog interface {
BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error)
RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error

CreatePrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error
GetPrivilegeGroup(ctx context.Context, groupName string) ([]*milvuspb.PrivilegeEntity, error)
DropPrivilegeGroup(ctx context.Context, groupName string) error
AlterPrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error
ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupEntity, error)
AddPrivilegesToGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error
DropPrivilegesFromGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error

Close()
}
Expand Down
113 changes: 50 additions & 63 deletions internal/metastore/kv/rootcoord/kv_catalog.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1439,98 +1439,85 @@ func (kc *Catalog) RestoreRBAC(ctx context.Context, tenant string, meta *milvusp
return err
}

func (kc *Catalog) CreatePrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeGroupEntity) error {
func (kc *Catalog) GetPrivilegeGroup(ctx context.Context, groupName string) ([]*milvuspb.PrivilegeEntity, error) {
k := BuildPrivilegeGroupkey(groupName)
err := kc.save(k)
if err != nil && !common.IsIgnorableError(err) {
log.Warn("fail to save the privilege group", zap.String("key", k), zap.Error(err))
}
return err
}

func (kc *Catalog) DropPrivilegeGroup(ctx context.Context, groupName string) error {
k := BuildPrivilegeGroupkey(groupName)
err := kc.Txn.Remove(k)
data, err := kc.Txn.Load(k)
if err != nil {
log.Warn("fail to drop privilege group", zap.String("key", k), zap.Error(err))
return err
if errors.Is(err, merr.ErrIoKeyNotFound) {
return nil, fmt.Errorf("privilege group [%s] does not exist", groupName)
}
log.Error("failed to load privilege group", zap.String("group", groupName), zap.Error(err))
return nil, err
}
return nil
}

func (kc *Catalog) ListPrivilegeGroup(ctx context.Context) ([]*milvuspb.PrivilegeGroupEntity, error) {
var results []*milvuspb.PrivilegeGroupEntity
keys, _, err := kc.Txn.LoadWithPrefix(PrivilegeGroupPrefix)
var privilegeNames []string
err = json.Unmarshal([]byte(data), &privilegeNames)
if err != nil {
log.Error("list all privilege groups fail", zap.String("prefix", PrivilegeGroupPrefix), zap.Error(err))
log.Error("failed to unmarshal privilege group data", zap.String("group", groupName), zap.Error(err))
return nil, err
}
for _, key := range keys {
groupName := typeutil.AfterN(key, PrivilegeGroupPrefix+"/", "/")
if len(groupName) != 1 {
log.Warn("invalid privilege group key", zap.String("string", key), zap.String("sub_string", PrivilegeGroupPrefix))
continue
}
results = append(results, &milvuspb.PrivilegeGroupEntity{
Name: groupName[0],
})

var privileges []*milvuspb.PrivilegeEntity
for _, name := range privilegeNames {
privileges = append(privileges, &milvuspb.PrivilegeEntity{Name: name})
}
return results, nil
return privileges, nil
}

func (kc *Catalog) AddPrivilegesToGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error {
func (kc *Catalog) DropPrivilegeGroup(ctx context.Context, groupName string) error {
k := BuildPrivilegeGroupkey(groupName)
data, err := kc.Txn.Load(k)
if err != nil {
if errors.Is(err, merr.ErrIoKeyNotFound) {
return fmt.Errorf("privilege group [%s] does not exist", groupName)
}
log.Error("failed to load privilege group", zap.String("group", groupName), zap.Error(err))
return err
}

var existingGroup milvuspb.PrivilegeGroupEntity
err = proto.Unmarshal([]byte(data), &existingGroup)
err := kc.Txn.Remove(k)
if err != nil {
log.Error("failed to unmarshal privilege group data", zap.String("group", groupName), zap.Error(err))
log.Warn("fail to drop privilege group", zap.String("key", k), zap.Error(err))
return err
}
return nil
}

func (kc *Catalog) AlterPrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error {
var privilegeNames []string
for _, privilege := range privileges {
if !containsPrivilege(existingGroup.Privileges, privilege) {
existingGroup.Privileges = append(existingGroup.Privileges, privilege)
}
privilegeNames = append(privilegeNames, privilege.Name)
}

updatedData, err := proto.Marshal(&existingGroup)
privilegeData, err := json.Marshal(privilegeNames)
if err != nil {
log.Error("failed to marshal updated privilege group", zap.String("group", groupName), zap.Error(err))
log.Error("failed to marshal privileges", zap.String("group", groupName), zap.Error(err))
return err
}

err = kc.Txn.Save(k, string(updatedData))
k := BuildPrivilegeGroupkey(groupName)
hasKey, err := kc.Txn.Has(k)
if hasKey {
err = kc.Txn.Remove(k)
}
err = kc.Txn.Save(k, string(privilegeData))
if err != nil {
log.Error("failed to save updated privilege group", zap.String("group", groupName), zap.Error(err))
log.Warn("fail to put privilege group", zap.String("key", k), zap.Error(err))
return err
}

log.Info("successfully added privileges to group", zap.String("group", groupName))
return nil
}

func (kc *Catalog) DropPrivilegesFromGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error {
k := BuildPrivilegeGroupkey(groupName)
privilegesBytes, err := proto.Marshal(privileges)
func (kc *Catalog) ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupEntity, error) {
var privilegeGroups []*milvuspb.PrivilegeGroupEntity

keys, _, err := kc.Txn.LoadWithPrefix(PrivilegeGroupPrefix)
if err != nil {
log.Error("marshal privileges fail", zap.Error(err))
return err
log.Error("failed to list privilege groups", zap.String("prefix", PrivilegeGroupPrefix), zap.Error(err))
return nil, err
}
err = kc.Txn.Remove(k, string(privilegesBytes))
if err != nil {
log.Error("delete privileges from group fail", zap.String("key", k), zap.Error(err))
return err

for _, key := range keys {
groupName := typeutil.AfterN(key, PrivilegeGroupPrefix+"/", "/")
if len(groupName) != 1 {
log.Warn("invalid privilege group key", zap.String("string", key), zap.String("sub_string", PrivilegeGroupPrefix))
continue
}

privilegeGroups = append(privilegeGroups, &milvuspb.PrivilegeGroupEntity{
Name: groupName[0],
})
}
return nil
return privilegeGroups, nil
}

func (kc *Catalog) Close() {
Expand Down
42 changes: 39 additions & 3 deletions internal/rootcoord/meta_table.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1481,7 +1481,7 @@ func (mt *MetaTable) CreatePrivilegeGroup(ctx context.Context, groupName string,
}
}

return mt.catalog.CreatePrivilegeGroup(mt.ctx, groupName, privileges)
return mt.catalog.AlterPrivilegeGroup(mt.ctx, groupName, privileges)
}

func (mt *MetaTable) DropPrivilegeGroup(ctx context.Context, groupName string) error {
Expand All @@ -1502,12 +1502,48 @@ func (mt *MetaTable) AddPrivilegesToGroup(ctx context.Context, groupName string,
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()

return mt.catalog.AddPrivilegesToGroup(mt.ctx, groupName, privileges)
toMergePrivileges, err := mt.catalog.GetPrivilegeGroup(mt.ctx, groupName)
if err != nil {
log.Warn("fail to get privilege group", zap.String("privilege_group", groupName), zap.Error(err))
return err
}
privilegeSet := make(map[string]struct{})
for _, p := range toMergePrivileges {
privilegeSet[p.Name] = struct{}{}
}
for _, p := range privileges {
privilegeSet[p.Name] = struct{}{}
}
var mergedPrivileges []*milvuspb.PrivilegeEntity
for privilegeName := range privilegeSet {
mergedPrivileges = append(mergedPrivileges, &milvuspb.PrivilegeEntity{Name: privilegeName})
}

return mt.catalog.AlterPrivilegeGroup(mt.ctx, groupName, mergedPrivileges)
}

func (mt *MetaTable) DropPrivilegesFromGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error {
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()

return mt.catalog.DropPrivilegesFromGroup(mt.ctx, groupName, privileges)
toUpdatePrivileges, err := mt.catalog.GetPrivilegeGroup(mt.ctx, groupName)
if err != nil {
log.Warn("fail to get privilege group", zap.String("privilege_group", groupName), zap.Error(err))
return err
}

privilegesToRemove := make(map[string]struct{})
for _, p := range privileges {
privilegesToRemove[p.Name] = struct{}{}
}

var updatedPrivileges []*milvuspb.PrivilegeEntity
for _, p := range toUpdatePrivileges {
// If the privilege is not in the removal set, keep it
if _, exists := privilegesToRemove[p.Name]; !exists {
updatedPrivileges = append(updatedPrivileges, p)
}
}

return mt.catalog.AlterPrivilegeGroup(mt.ctx, groupName, updatedPrivileges)
}

0 comments on commit 6853e8c

Please sign in to comment.