From b81b248972cd1110a24f834e99c7b07c31303af2 Mon Sep 17 00:00:00 2001 From: shaoting-huang Date: Wed, 25 Dec 2024 21:05:10 +0800 Subject: [PATCH] expand privilege group when list policy in rootcoord Signed-off-by: shaoting-huang --- internal/metastore/catalog.go | 2 +- internal/metastore/kv/rootcoord/kv_catalog.go | 21 ++++-- .../metastore/kv/rootcoord/kv_catalog_test.go | 22 +++++-- .../metastore/mocks/mock_rootcoord_catalog.go | 14 ++-- internal/rootcoord/meta_table.go | 4 +- internal/rootcoord/mock_test.go | 6 +- internal/rootcoord/mocks/meta_table.go | 14 ++-- internal/rootcoord/root_coord.go | 43 +++++++++--- internal/rootcoord/root_coord_test.go | 66 +++++++++++++++++-- 9 files changed, 147 insertions(+), 45 deletions(-) diff --git a/internal/metastore/catalog.go b/internal/metastore/catalog.go index 2d23104912ee5..0248aca02f6c2 100644 --- a/internal/metastore/catalog.go +++ b/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) diff --git a/internal/metastore/kv/rootcoord/kv_catalog.go b/internal/metastore/kv/rootcoord/kv_catalog.go index 6b50f5d38788c..d03da2ab431b5 100644 --- a/internal/metastore/kv/rootcoord/kv_catalog.go +++ b/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) { diff --git a/internal/metastore/kv/rootcoord/kv_catalog_test.go b/internal/metastore/kv/rootcoord/kv_catalog_test.go index abe0f1f46be96..e116201d3bc11 100644 --- a/internal/metastore/kv/rootcoord/kv_catalog_test.go +++ b/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 { diff --git a/internal/metastore/mocks/mock_rootcoord_catalog.go b/internal/metastore/mocks/mock_rootcoord_catalog.go index 8c35d288c1143..8bc119e07b671 100644 --- a/internal/metastore/mocks/mock_rootcoord_catalog.go +++ b/internal/metastore/mocks/mock_rootcoord_catalog.go @@ -1548,19 +1548,19 @@ func (_c *RootCoordCatalog_ListGrant_Call) RunAndReturn(run func(context.Context } // ListPolicy provides a mock function with given fields: ctx, tenant -func (_m *RootCoordCatalog) ListPolicy(ctx context.Context, tenant string) ([]string, error) { +func (_m *RootCoordCatalog) ListPolicy(ctx context.Context, tenant string) ([]*milvuspb.GrantEntity, error) { ret := _m.Called(ctx, tenant) - var r0 []string + var r0 []*milvuspb.GrantEntity var r1 error - if rf, ok := ret.Get(0).(func(context.Context, string) ([]string, error)); ok { + if rf, ok := ret.Get(0).(func(context.Context, string) ([]*milvuspb.GrantEntity, error)); ok { return rf(ctx, tenant) } - if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok { + if rf, ok := ret.Get(0).(func(context.Context, string) []*milvuspb.GrantEntity); ok { r0 = rf(ctx, tenant) } else { if ret.Get(0) != nil { - r0 = ret.Get(0).([]string) + r0 = ret.Get(0).([]*milvuspb.GrantEntity) } } @@ -1592,12 +1592,12 @@ func (_c *RootCoordCatalog_ListPolicy_Call) Run(run func(ctx context.Context, te return _c } -func (_c *RootCoordCatalog_ListPolicy_Call) Return(_a0 []string, _a1 error) *RootCoordCatalog_ListPolicy_Call { +func (_c *RootCoordCatalog_ListPolicy_Call) Return(_a0 []*milvuspb.GrantEntity, _a1 error) *RootCoordCatalog_ListPolicy_Call { _c.Call.Return(_a0, _a1) return _c } -func (_c *RootCoordCatalog_ListPolicy_Call) RunAndReturn(run func(context.Context, string) ([]string, error)) *RootCoordCatalog_ListPolicy_Call { +func (_c *RootCoordCatalog_ListPolicy_Call) RunAndReturn(run func(context.Context, string) ([]*milvuspb.GrantEntity, error)) *RootCoordCatalog_ListPolicy_Call { _c.Call.Return(run) return _c } diff --git a/internal/rootcoord/meta_table.go b/internal/rootcoord/meta_table.go index 473e49ac9ac7f..acafebad70fb0 100644 --- a/internal/rootcoord/meta_table.go +++ b/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() diff --git a/internal/rootcoord/mock_test.go b/internal/rootcoord/mock_test.go index c03a887bf4d6c..08a6c5691c092 100644 --- a/internal/rootcoord/mock_test.go +++ b/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) { diff --git a/internal/rootcoord/mocks/meta_table.go b/internal/rootcoord/mocks/meta_table.go index 0272510d1bdf5..a378960d828bb 100644 --- a/internal/rootcoord/mocks/meta_table.go +++ b/internal/rootcoord/mocks/meta_table.go @@ -1894,19 +1894,19 @@ func (_c *IMetaTable_ListDatabases_Call) RunAndReturn(run func(context.Context, } // ListPolicy provides a mock function with given fields: tenant -func (_m *IMetaTable) ListPolicy(tenant string) ([]string, error) { +func (_m *IMetaTable) ListPolicy(tenant string) ([]*milvuspb.GrantEntity, error) { ret := _m.Called(tenant) - var r0 []string + var r0 []*milvuspb.GrantEntity var r1 error - if rf, ok := ret.Get(0).(func(string) ([]string, error)); ok { + if rf, ok := ret.Get(0).(func(string) ([]*milvuspb.GrantEntity, error)); ok { return rf(tenant) } - if rf, ok := ret.Get(0).(func(string) []string); ok { + if rf, ok := ret.Get(0).(func(string) []*milvuspb.GrantEntity); ok { r0 = rf(tenant) } else { if ret.Get(0) != nil { - r0 = ret.Get(0).([]string) + r0 = ret.Get(0).([]*milvuspb.GrantEntity) } } @@ -1937,12 +1937,12 @@ func (_c *IMetaTable_ListPolicy_Call) Run(run func(tenant string)) *IMetaTable_L return _c } -func (_c *IMetaTable_ListPolicy_Call) Return(_a0 []string, _a1 error) *IMetaTable_ListPolicy_Call { +func (_c *IMetaTable_ListPolicy_Call) Return(_a0 []*milvuspb.GrantEntity, _a1 error) *IMetaTable_ListPolicy_Call { _c.Call.Return(_a0, _a1) return _c } -func (_c *IMetaTable_ListPolicy_Call) RunAndReturn(run func(string) ([]string, error)) *IMetaTable_ListPolicy_Call { +func (_c *IMetaTable_ListPolicy_Call) RunAndReturn(run func(string) ([]*milvuspb.GrantEntity, error)) *IMetaTable_ListPolicy_Call { _c.Call.Return(run) return _c } diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go index 6456d481f7109..e8d09a37a66ea 100644 --- a/internal/rootcoord/root_coord.go +++ b/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,31 +2874,47 @@ func (c *Core) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest) Status: merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_ListPolicyFailure), }, nil } - userRoles, err := c.meta.ListUserRole(util.DefaultTenant) + // expand privilege groups and turn to policies + allGroups, err := c.getPrivilegeGroups() if err != nil { - errMsg := "fail to list user-role" - ctxLog.Warn(errMsg, zap.Any("in", in), zap.Error(err)) + errMsg := "fail to get privilege groups" + ctxLog.Warn(errMsg, zap.Error(err)) return &internalpb.ListPolicyResponse{ Status: merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_ListPolicyFailure), }, nil } - privGroups, err := c.meta.ListPrivilegeGroups() + 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 { - errMsg := "fail to list privilege groups" + errMsg := "fail to expand privilege groups" ctxLog.Warn(errMsg, zap.Error(err)) return &internalpb.ListPolicyResponse{ Status: merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_ListPolicyFailure), }, nil } + 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" + ctxLog.Warn(errMsg, zap.Any("in", in), zap.Error(err)) + return &internalpb.ListPolicyResponse{ + Status: merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_ListPolicyFailure), + }, nil + } ctxLog.Debug(method + " success") metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.SuccessLabel).Inc() metrics.RootCoordDDLReqLatency.WithLabelValues(method).Observe(float64(tr.ElapseSpan().Milliseconds())) return &internalpb.ListPolicyResponse{ Status: merr.Success(), - PolicyInfos: policies, + PolicyInfos: expandPolicies, UserRoles: userRoles, - PrivilegeGroups: privGroups, + PrivilegeGroups: allGroups, }, 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 +} diff --git a/internal/rootcoord/root_coord_test.go b/internal/rootcoord/root_coord_test.go index 32428925e9f6a..110c3875a9f7a 100644 --- a/internal/rootcoord/root_coord_test.go +++ b/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,44 @@ 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{{ + ObjectName: "*", + Object: &milvuspb.ObjectEntity{ + Name: "Global", + }, + Role: &milvuspb.RoleEntity{Name: "role"}, + Grantor: &milvuspb.GrantorEntity{Privilege: &milvuspb.PrivilegeEntity{Name: "CustomGroup"}}, + }}, 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.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.IsCustomPrivilegeGroupFunc = func(groupName string) (bool, error) { + return true, 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) ([]*milvuspb.GrantEntity, error) { + return []*milvuspb.GrantEntity{}, errors.New("mock error") + } + mockMeta.ListPrivilegeGroupsFunc = func() ([]*milvuspb.PrivilegeGroupInfo, error) { + return []*milvuspb.PrivilegeGroupInfo{}, errors.New("mock error") + } + mockMeta.IsCustomPrivilegeGroupFunc = func(groupName string) (bool, error) { + return false, errors.New("mock error") } }) }