From 039641eb3d797d499693f165cb7b6c602d2993df Mon Sep 17 00:00:00 2001 From: shaoting-huang Date: Fri, 1 Nov 2024 11:32:51 +0800 Subject: [PATCH] fix backup rbac Signed-off-by: shaoting-huang --- internal/metastore/kv/rootcoord/kv_catalog.go | 30 +++---- .../metastore/kv/rootcoord/kv_catalog_test.go | 7 +- tests/integration/rbac/rbac_backup_test.go | 84 ++++++++++++++++++- 3 files changed, 103 insertions(+), 18 deletions(-) diff --git a/internal/metastore/kv/rootcoord/kv_catalog.go b/internal/metastore/kv/rootcoord/kv_catalog.go index 8cbd4485a045f..c0cbd20c84ef5 100644 --- a/internal/metastore/kv/rootcoord/kv_catalog.go +++ b/internal/metastore/kv/rootcoord/kv_catalog.go @@ -1399,26 +1399,28 @@ func (kc *Catalog) RestoreRBAC(ctx context.Context, tenant string, meta *milvusp needRollbackRole = append(needRollbackRole, role) } + privGroups, err := kc.ListPrivilegeGroups(ctx) + if err != nil { + return err + } + // restore grant for _, grant := range meta.Grants { privName := grant.Grantor.Privilege.Name if util.IsPrivilegeNameDefined(privName) { grant.Grantor.Privilege.Name = util.PrivilegeNameForMetastore(privName) - } - privGroups, err := kc.ListPrivilegeGroups(ctx) - if err != nil { - return err - } - customGroup := false - for _, group := range privGroups { - if group.GroupName == privName { - grant.Grantor.Privilege.Name = util.PrivilegeGroupNameForMetastore(privName) - customGroup = true - break + } else { + customGroup := false + for _, group := range privGroups { + if group.GroupName == privName { + grant.Grantor.Privilege.Name = util.PrivilegeGroupNameForMetastore(privName) + customGroup = true + break + } + } + if !customGroup { + return errors.New("not found the privilege name") } - } - if !customGroup { - return errors.New("not found the privilege name") } err = kc.AlterGrant(ctx, tenant, grant, milvuspb.OperatePrivilegeType_Grant) if err != nil { diff --git a/internal/metastore/kv/rootcoord/kv_catalog_test.go b/internal/metastore/kv/rootcoord/kv_catalog_test.go index f6600086dae38..d62d652e81efa 100644 --- a/internal/metastore/kv/rootcoord/kv_catalog_test.go +++ b/internal/metastore/kv/rootcoord/kv_catalog_test.go @@ -2757,7 +2757,7 @@ func TestRBAC_Restore(t *testing.T) { DbName: util.DefaultDBName, Grantor: &milvuspb.GrantorEntity{ User: &milvuspb.UserEntity{Name: "user1"}, - Privilege: &milvuspb.PrivilegeEntity{Name: "PrivilegeLoad"}, + Privilege: &milvuspb.PrivilegeEntity{Name: "Load"}, }, }, }, @@ -2786,6 +2786,7 @@ func TestRBAC_Restore(t *testing.T) { assert.Equal(t, "obj_name1", grants[0].ObjectName) assert.Equal(t, "role1", grants[0].Role.Name) assert.Equal(t, "user1", grants[0].Grantor.User.Name) + assert.Equal(t, "Load", grants[0].Grantor.Privilege.Name) rbacMeta2 := &milvuspb.RBACMeta{ Users: []*milvuspb.UserInfo{ @@ -2822,7 +2823,7 @@ func TestRBAC_Restore(t *testing.T) { DbName: util.DefaultDBName, Grantor: &milvuspb.GrantorEntity{ User: &milvuspb.UserEntity{Name: "user2"}, - Privilege: &milvuspb.PrivilegeEntity{Name: "PrivilegeLoad"}, + Privilege: &milvuspb.PrivilegeEntity{Name: "Load"}, }, }, }, @@ -2847,6 +2848,7 @@ func TestRBAC_Restore(t *testing.T) { }) assert.NoError(t, err) assert.Len(t, grants, 1) + assert.Equal(t, grants[0].Grantor.Privilege.Name, "Load") } func TestRBAC_PrivilegeGroup(t *testing.T) { @@ -2896,6 +2898,7 @@ func TestRBAC_PrivilegeGroup(t *testing.T) { c = &Catalog{Txn: kvmock} ) + kvmock.EXPECT().LoadWithPrefix(funcutil.HandleTenantForEtcdKey(RolePrefix, util.DefaultTenant, "")).Return(nil, nil, nil) kvmock.EXPECT().Remove(key1).Return(nil) kvmock.EXPECT().Remove(key2).Return(errors.New("Mock remove failure")) diff --git a/tests/integration/rbac/rbac_backup_test.go b/tests/integration/rbac/rbac_backup_test.go index de4e271e9163a..5beb68ce19c15 100644 --- a/tests/integration/rbac/rbac_backup_test.go +++ b/tests/integration/rbac/rbac_backup_test.go @@ -20,6 +20,7 @@ import ( "strings" "testing" + "github.com/samber/lo" "github.com/stretchr/testify/suite" "google.golang.org/grpc/metadata" @@ -70,6 +71,7 @@ func (s *RBACBackupTestSuite) TestBackup() { s.Equal("", resp.GetRBACMeta().String()) // generate some rbac content + // create role test_role roleName := "test_role" resp1, err := s.Cluster.Proxy.CreateRole(ctx, &milvuspb.CreateRoleRequest{ Entity: &milvuspb.RoleEntity{ @@ -78,6 +80,8 @@ func (s *RBACBackupTestSuite) TestBackup() { }) s.NoError(err) s.True(merr.Ok(resp1)) + + // grant collection level search privilege to role test_role resp2, err := s.Cluster.Proxy.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{ Type: milvuspb.OperatePrivilegeType_Grant, Entity: &milvuspb.GrantEntity{ @@ -94,6 +98,42 @@ func (s *RBACBackupTestSuite) TestBackup() { s.NoError(err) s.True(merr.Ok(resp2)) s.Equal("", resp2.GetReason()) + + // create privielge group test_group + groupName := "test_group" + resp2, err = s.Cluster.Proxy.CreatePrivilegeGroup(ctx, &milvuspb.CreatePrivilegeGroupRequest{ + GroupName: groupName, + }) + s.NoError(err) + s.True(merr.Ok(resp2)) + + // add query and insert privilege to group test_group + resp2, err = s.Cluster.Proxy.OperatePrivilegeGroup(ctx, &milvuspb.OperatePrivilegeGroupRequest{ + GroupName: groupName, + Privileges: []*milvuspb.PrivilegeEntity{{Name: "Query"}, {Name: "Insert"}}, + Type: milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup, + }) + s.NoError(err) + s.True(merr.Ok(resp2)) + + // grant privilege group test_group to role test_role + resp2, err = s.Cluster.Proxy.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{ + Type: milvuspb.OperatePrivilegeType_Grant, + Entity: &milvuspb.GrantEntity{ + Role: &milvuspb.RoleEntity{Name: roleName}, + Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()}, + ObjectName: util.AnyWord, + DbName: util.AnyWord, + Grantor: &milvuspb.GrantorEntity{ + User: &milvuspb.UserEntity{Name: util.UserRoot}, + Privilege: &milvuspb.PrivilegeEntity{Name: groupName}, + }, + }, + }) + s.NoError(err) + s.True(merr.Ok(resp2)) + s.Equal("", resp2.GetReason()) + userName := "test_user" passwd := "test_passwd" resp3, err := s.Cluster.Proxy.CreateCredential(ctx, &milvuspb.CreateCredentialRequest{ @@ -109,10 +149,16 @@ func (s *RBACBackupTestSuite) TestBackup() { s.NoError(err) s.True(merr.Ok(resp4)) - // test back up rbac + // test back up rbac, grants should contain resp5, err := s.Cluster.Proxy.BackupRBAC(ctx, &milvuspb.BackupRBACMetaRequest{}) s.NoError(err) s.True(merr.Ok(resp5.GetStatus())) + s.Equal(2, len(resp5.GetRBACMeta().Grants)) + grants := lo.SliceToMap(resp5.GetRBACMeta().Grants, func(g *milvuspb.GrantEntity) (string, *milvuspb.GrantEntity) { + return g.Grantor.Privilege.Name, g + }) + s.True(grants["Search"] != nil) + s.True(grants[groupName] != nil) // test restore, expect to failed due to role/user already exist resp6, err := s.Cluster.Proxy.RestoreRBAC(ctx, &milvuspb.RestoreRBACMetaRequest{ @@ -121,7 +167,7 @@ func (s *RBACBackupTestSuite) TestBackup() { s.NoError(err) s.False(merr.Ok(resp6)) - // drop exist role/user, successful to restore + // revoke privilege search from role test_role before dropping the role resp7, err := s.Cluster.Proxy.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{ Type: milvuspb.OperatePrivilegeType_Revoke, Entity: &milvuspb.GrantEntity{ @@ -137,6 +183,24 @@ func (s *RBACBackupTestSuite) TestBackup() { }) s.NoError(err) s.True(merr.Ok(resp7)) + + // revoke privilege group test_group from role test_role before dropping the role + resp7, err = s.Cluster.Proxy.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{ + Type: milvuspb.OperatePrivilegeType_Revoke, + Entity: &milvuspb.GrantEntity{ + Role: &milvuspb.RoleEntity{Name: roleName}, + Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()}, + ObjectName: util.AnyWord, + DbName: util.AnyWord, + Grantor: &milvuspb.GrantorEntity{ + User: &milvuspb.UserEntity{Name: util.UserRoot}, + Privilege: &milvuspb.PrivilegeEntity{Name: groupName}, + }, + }, + }) + s.NoError(err) + s.True(merr.Ok(resp7)) + resp8, err := s.Cluster.Proxy.DropRole(ctx, &milvuspb.DropRoleRequest{ RoleName: roleName, }) @@ -177,6 +241,22 @@ func (s *RBACBackupTestSuite) TestBackup() { s.NoError(err) s.True(merr.Ok(resp12)) + resp12, err = s.Cluster.Proxy.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{ + Type: milvuspb.OperatePrivilegeType_Revoke, + Entity: &milvuspb.GrantEntity{ + Role: &milvuspb.RoleEntity{Name: roleName}, + Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()}, + ObjectName: util.AnyWord, + DbName: util.AnyWord, + Grantor: &milvuspb.GrantorEntity{ + User: &milvuspb.UserEntity{Name: util.UserRoot}, + Privilege: &milvuspb.PrivilegeEntity{Name: groupName}, + }, + }, + }) + s.NoError(err) + s.True(merr.Ok(resp12)) + resp13, err := s.Cluster.Proxy.DropRole(ctx, &milvuspb.DropRoleRequest{ RoleName: roleName, })