Skip to content

Commit

Permalink
fix backup rbac
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 Nov 1, 2024
1 parent fdb1715 commit 039641e
Show file tree
Hide file tree
Showing 3 changed files with 103 additions and 18 deletions.
30 changes: 16 additions & 14 deletions internal/metastore/kv/rootcoord/kv_catalog.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
7 changes: 5 additions & 2 deletions internal/metastore/kv/rootcoord/kv_catalog_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"},
},
},
},
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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"},
},
},
},
Expand All @@ -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) {
Expand Down Expand Up @@ -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"))

Expand Down
84 changes: 82 additions & 2 deletions tests/integration/rbac/rbac_backup_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import (
"strings"
"testing"

"github.com/samber/lo"
"github.com/stretchr/testify/suite"
"google.golang.org/grpc/metadata"

Expand Down Expand Up @@ -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{
Expand All @@ -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{
Expand All @@ -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{
Expand All @@ -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{
Expand All @@ -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{
Expand All @@ -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,
})
Expand Down Expand Up @@ -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,
})
Expand Down

0 comments on commit 039641e

Please sign in to comment.