Skip to content

Commit

Permalink
expand privilege group when list policy in rootcoord
Browse files Browse the repository at this point in the history
Signed-off-by: shaoting-huang <[email protected]>
  • Loading branch information
shaoting-huang committed Dec 25, 2024
1 parent d47372e commit aece209
Show file tree
Hide file tree
Showing 8 changed files with 94 additions and 34 deletions.
2 changes: 1 addition & 1 deletion internal/metastore/catalog.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
21 changes: 14 additions & 7 deletions internal/metastore/kv/rootcoord/kv_catalog.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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+"/", "/")
Expand All @@ -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) {
Expand Down
14 changes: 7 additions & 7 deletions internal/metastore/mocks/mock_rootcoord_catalog.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions internal/rootcoord/meta_table.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()

Expand Down
6 changes: 3 additions & 3 deletions internal/rootcoord/mock_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
}

Expand Down Expand Up @@ -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) {
Expand Down
14 changes: 7 additions & 7 deletions internal/rootcoord/mocks/meta_table.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

31 changes: 28 additions & 3 deletions internal/rootcoord/root_coord.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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"
Expand All @@ -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
Expand Down Expand Up @@ -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
}
36 changes: 32 additions & 4 deletions internal/rootcoord/root_coord_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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")
}
})
}
Expand Down

0 comments on commit aece209

Please sign in to comment.