diff --git a/core/backup_impl_restore_backup.go b/core/backup_impl_restore_backup.go index b5b5012..34bd475 100644 --- a/core/backup_impl_restore_backup.go +++ b/core/backup_impl_restore_backup.go @@ -6,15 +6,12 @@ import ( "strings" "time" - "github.com/zilliztech/milvus-backup/core/meta" - "github.com/zilliztech/milvus-backup/core/restore" - jsoniter "github.com/json-iterator/go" - "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" - "github.com/samber/lo" "go.uber.org/zap" + "github.com/zilliztech/milvus-backup/core/meta" "github.com/zilliztech/milvus-backup/core/proto/backuppb" + "github.com/zilliztech/milvus-backup/core/restore" "github.com/zilliztech/milvus-backup/core/utils" "github.com/zilliztech/milvus-backup/internal/common" "github.com/zilliztech/milvus-backup/internal/log" @@ -122,13 +119,15 @@ func (b *BackupContext) RestoreBackup(ctx context.Context, request *backuppb.Res // restore rbac if request.GetRbac() { - err := b.restoreRBAC(ctx, backup) - if err != nil { + rt := restore.NewRBACTask(b.getMilvusClient(), backup.GetRbacMeta()) + if err := rt.Execute(ctx); err != nil { log.Error("fail to restore RBAC", zap.Error(err)) resp.Code = backuppb.ResponseCode_Fail resp.Msg = fmt.Sprintf("fail to restore RBAC, err: %s", err) return resp } + } else { + log.Info("skip restore RBAC") } // 2, initial restoreCollectionTasks @@ -425,84 +424,3 @@ func (b *BackupContext) executeRestoreCollectionTask(ctx context.Context, backup return task, err } - -func (b *BackupContext) restoreRBAC(ctx context.Context, backupInfo *backuppb.BackupInfo) error { - log.Info("restore RBAC") - - resp, err := b.getMilvusClient().BackupRBAC(ctx) - if err != nil { - return fmt.Errorf("fail to get current RBAC, err: %s", err) - } - curRBAC := resp.RBACMeta - - rbacBackup := backupInfo.GetRbacMeta() - curUsers := lo.SliceToMap(curRBAC.Users, func(user *milvuspb.UserInfo) (string, struct{}) { - return user.User, struct{}{} - }) - users := make([]*milvuspb.UserInfo, 0, len(rbacBackup.GetUsers())) - for _, user := range rbacBackup.GetUsers() { - // skip if user already exist - if _, ok := curUsers[user.GetUser()]; ok { - continue - } - - ur := lo.Map(user.GetRoles(), func(role *backuppb.RoleEntity, index int) *milvuspb.RoleEntity { - return &milvuspb.RoleEntity{Name: role.Name} - }) - userEntity := &milvuspb.UserInfo{ - User: user.User, - Password: user.Password, - Roles: ur, - } - users = append(users, userEntity) - } - - curRoles := lo.SliceToMap(curRBAC.Roles, func(role *milvuspb.RoleEntity) (string, struct{}) { - return role.Name, struct{}{} - }) - roles := make([]*milvuspb.RoleEntity, 0, len(rbacBackup.GetRoles())) - for _, role := range rbacBackup.GetRoles() { - // skip if role already exist - if _, ok := curRoles[role.GetName()]; ok { - continue - } - - roleEntity := &milvuspb.RoleEntity{ - Name: role.GetName(), - } - roles = append(roles, roleEntity) - } - - grants := make([]*milvuspb.GrantEntity, 0, len(rbacBackup.GetGrants())) - curGrants := lo.SliceToMap(curRBAC.Grants, func(grant *milvuspb.GrantEntity) (string, struct{}) { - return fmt.Sprintf("%s/%s/%s/%s/%s/%s", grant.Object, grant.ObjectName, grant.Role.Name, grant.Grantor.User, grant.Grantor.Privilege.Name, grant.DbName), struct{}{} - }) - for _, roleGrant := range rbacBackup.GetGrants() { - key := fmt.Sprintf("%s/%s/%s/%s/%s/%s", roleGrant.Object.GetName(), roleGrant.GetObjectName(), roleGrant.GetRole().GetName(), roleGrant.GetGrantor().GetUser().GetName(), roleGrant.GetGrantor().GetPrivilege().GetName(), roleGrant.GetDbName()) - // skip if grant already exist - if _, ok := curGrants[key]; ok { - continue - } - roleGrantEntity := &milvuspb.GrantEntity{ - Object: &milvuspb.ObjectEntity{Name: roleGrant.Object.Name}, - ObjectName: roleGrant.GetObjectName(), - Role: &milvuspb.RoleEntity{Name: roleGrant.Role.Name}, - Grantor: &milvuspb.GrantorEntity{ - User: &milvuspb.UserEntity{Name: roleGrant.Grantor.User.Name}, - Privilege: &milvuspb.PrivilegeEntity{Name: roleGrant.Grantor.Privilege.Name}, - }, - DbName: roleGrant.GetDbName(), - } - grants = append(grants, roleGrantEntity) - } - - rbacMeta := &milvuspb.RBACMeta{ - Users: users, - Roles: roles, - Grants: grants, - } - - log.Info("restore RBAC", zap.Int("users", len(users)), zap.Int("roles", len(roles)), zap.Int("grants", len(grants))) - err = b.getMilvusClient().RestoreRBAC(ctx, rbacMeta) - return err -} diff --git a/core/restore/rbac.go b/core/restore/rbac.go new file mode 100644 index 0000000..6ddaadb --- /dev/null +++ b/core/restore/rbac.go @@ -0,0 +1,146 @@ +package restore + +import ( + "context" + "fmt" + + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/samber/lo" + "go.uber.org/zap" + + "github.com/zilliztech/milvus-backup/core/client" + "github.com/zilliztech/milvus-backup/core/proto/backuppb" + "github.com/zilliztech/milvus-backup/internal/log" +) + +type RBACTask struct { + grpcCli client.Grpc + + bakRBAC *backuppb.RBACMeta +} + +func (rt *RBACTask) Execute(ctx context.Context) error { + curRBAC, err := rt.grpcCli.BackupRBAC(ctx) + if err != nil { + return fmt.Errorf("restore_rbac: get current rbac: %w", err) + } + + users := rt.users(curRBAC.GetRBACMeta().GetUsers()) + roles := rt.roles(curRBAC.GetRBACMeta().GetRoles()) + grants := rt.grants(curRBAC.GetRBACMeta().GetGrants()) + + rbacMeta := &milvuspb.RBACMeta{ + Users: users, + Roles: roles, + Grants: grants, + } + + log.Info("insert rbac to milvus", zap.Int("users", len(users)), + zap.Int("roles", len(roles)), + zap.Int("grants", len(grants))) + if err := rt.grpcCli.RestoreRBAC(ctx, rbacMeta); err != nil { + return fmt.Errorf("restore_rbac: restore rbac: %w", err) + } + + return nil +} + +func NewRBACTask(grpcCli client.Grpc, bakRBAC *backuppb.RBACMeta) *RBACTask { + return &RBACTask{ + grpcCli: grpcCli, + + bakRBAC: bakRBAC, + } +} + +func (rt *RBACTask) users(curUsers []*milvuspb.UserInfo) []*milvuspb.UserInfo { + um := lo.SliceToMap(curUsers, func(user *milvuspb.UserInfo) (string, struct{}) { + return user.User, struct{}{} + }) + users := make([]*milvuspb.UserInfo, 0, len(rt.bakRBAC.GetUsers())) + for _, user := range rt.bakRBAC.GetUsers() { + // skip if user already exist + if _, ok := um[user.GetUser()]; ok { + continue + } + + ur := lo.Map(user.GetRoles(), func(role *backuppb.RoleEntity, index int) *milvuspb.RoleEntity { + return &milvuspb.RoleEntity{Name: role.Name} + }) + userEntity := &milvuspb.UserInfo{ + User: user.User, + Password: user.Password, + Roles: ur, + } + users = append(users, userEntity) + } + + return users +} + +func (rt *RBACTask) roles(curRoles []*milvuspb.RoleEntity) []*milvuspb.RoleEntity { + rm := lo.SliceToMap(curRoles, func(role *milvuspb.RoleEntity) (string, struct{}) { + return role.Name, struct{}{} + }) + roles := make([]*milvuspb.RoleEntity, 0, len(rt.bakRBAC.GetRoles())) + for _, role := range rt.bakRBAC.GetRoles() { + // skip if role already exist + if _, ok := rm[role.GetName()]; ok { + continue + } + + roleEntity := &milvuspb.RoleEntity{ + Name: role.GetName(), + } + roles = append(roles, roleEntity) + } + + return roles +} + +func backupGrantKey(grant *backuppb.GrantEntity) string { + return fmt.Sprintf("%s/%s/%s/%s/%s/%s", + grant.GetObject().GetName(), + grant.GetObjectName(), + grant.GetRole().GetName(), + grant.GetGrantor().GetUser().GetName(), + grant.GetGrantor().GetPrivilege().GetName(), + grant.GetDbName()) +} + +func milvusGrantKey(grant *milvuspb.GrantEntity) string { + return fmt.Sprintf("%s/%s/%s/%s/%s/%s", + grant.GetObject().GetName(), + grant.GetObjectName(), + grant.GetRole().GetName(), + grant.GetGrantor().GetUser().GetName(), + grant.GetGrantor().GetPrivilege().GetName(), + grant.GetDbName()) +} + +func (rt *RBACTask) grants(curGrants []*milvuspb.GrantEntity) []*milvuspb.GrantEntity { + gm := lo.SliceToMap(curGrants, func(grant *milvuspb.GrantEntity) (string, struct{}) { + return milvusGrantKey(grant), struct{}{} + }) + grants := make([]*milvuspb.GrantEntity, 0, len(rt.bakRBAC.GetGrants())) + for _, grant := range rt.bakRBAC.GetGrants() { + // skip if grant already exist + if _, ok := gm[backupGrantKey(grant)]; ok { + continue + } + + grantEntity := &milvuspb.GrantEntity{ + Role: &milvuspb.RoleEntity{Name: grant.GetRole().GetName()}, + Object: &milvuspb.ObjectEntity{Name: grant.GetObject().GetName()}, + ObjectName: grant.GetObjectName(), + Grantor: &milvuspb.GrantorEntity{ + User: &milvuspb.UserEntity{Name: grant.GetGrantor().GetUser().GetName()}, + Privilege: &milvuspb.PrivilegeEntity{Name: grant.GetGrantor().GetPrivilege().GetName()}, + }, + DbName: grant.GetDbName(), + } + grants = append(grants, grantEntity) + } + + return grants +} diff --git a/core/restore/rbac_test.go b/core/restore/rbac_test.go new file mode 100644 index 0000000..f3d8512 --- /dev/null +++ b/core/restore/rbac_test.go @@ -0,0 +1,118 @@ +package restore + +import ( + "testing" + + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/stretchr/testify/assert" + + "github.com/zilliztech/milvus-backup/core/proto/backuppb" +) + +func TestBackupGrantKey(t *testing.T) { + grant := &backuppb.GrantEntity{ + Role: &backuppb.RoleEntity{Name: "role.name"}, + Object: &backuppb.ObjectEntity{Name: "object.name"}, + ObjectName: "objectName", + Grantor: &backuppb.GrantorEntity{ + User: &backuppb.UserEntity{Name: "grantor.user.name"}, + Privilege: &backuppb.PrivilegeEntity{Name: "grantor.privilege.name"}, + }, + DbName: "dbName", + } + + grantKey := backupGrantKey(grant) + expected := "object.name/objectName/role.name/grantor.user.name/grantor.privilege.name/dbName" + assert.Equal(t, expected, grantKey) +} + +func TestMilvusGrantKey(t *testing.T) { + grant := &milvuspb.GrantEntity{ + Role: &milvuspb.RoleEntity{Name: "role.name"}, + Object: &milvuspb.ObjectEntity{Name: "object.name"}, + ObjectName: "objectName", + Grantor: &milvuspb.GrantorEntity{ + User: &milvuspb.UserEntity{Name: "grantor.user.name"}, + Privilege: &milvuspb.PrivilegeEntity{Name: "grantor.privilege.name"}, + }, + DbName: "dbName", + } + + grantKey := milvusGrantKey(grant) + expected := "object.name/objectName/role.name/grantor.user.name/grantor.privilege.name/dbName" + assert.Equal(t, expected, grantKey) +} + +func TestRBACTask_Grants(t *testing.T) { + bakGrants := []*backuppb.GrantEntity{ + { + Role: &backuppb.RoleEntity{Name: "role.name"}, + Object: &backuppb.ObjectEntity{Name: "object.name"}, + ObjectName: "objectName", + Grantor: &backuppb.GrantorEntity{ + User: &backuppb.UserEntity{Name: "grantor.user.name"}, + Privilege: &backuppb.PrivilegeEntity{Name: "grantor.privilege.name"}, + }, + DbName: "dbName", + }, + { + Role: &backuppb.RoleEntity{Name: "role1.name"}, // different role + Object: &backuppb.ObjectEntity{Name: "object.name"}, + ObjectName: "objectName", + Grantor: &backuppb.GrantorEntity{ + User: &backuppb.UserEntity{Name: "grantor.user.name"}, + Privilege: &backuppb.PrivilegeEntity{Name: "grantor.privilege.name"}, + }, + DbName: "dbName", + }, + } + + curGrants := []*milvuspb.GrantEntity{ + { + Role: &milvuspb.RoleEntity{Name: "role.name"}, + Object: &milvuspb.ObjectEntity{Name: "object.name"}, + ObjectName: "objectName", + Grantor: &milvuspb.GrantorEntity{ + User: &milvuspb.UserEntity{Name: "grantor.user.name"}, + Privilege: &milvuspb.PrivilegeEntity{Name: "grantor.privilege.name"}, + }, + DbName: "dbName", + }, + } + + rt := &RBACTask{bakRBAC: &backuppb.RBACMeta{Grants: bakGrants}} + restoreGrants := rt.grants(curGrants) + assert.Len(t, restoreGrants, 1) +} + +func TestRBACTask_Roles(t *testing.T) { + bakRoles := []*backuppb.RoleEntity{ + {Name: "role.name"}, + {Name: "role1.name"}, + } + + curRoles := []*milvuspb.RoleEntity{ + {Name: "role.name"}, + } + + rt := &RBACTask{bakRBAC: &backuppb.RBACMeta{Roles: bakRoles}} + restoreRoles := rt.roles(curRoles) + assert.Len(t, restoreRoles, 1) + assert.Equal(t, "role1.name", restoreRoles[0].Name) +} + +func TestRBACTask_Users(t *testing.T) { + bakUsers := []*backuppb.UserInfo{ + {User: "user.user"}, + {User: "user1.user"}, + } + + curUsers := []*milvuspb.UserInfo{ + {User: "user.user"}, + } + + rt := &RBACTask{bakRBAC: &backuppb.RBACMeta{Users: bakUsers}} + restoreUsers := rt.users(curUsers) + assert.Len(t, restoreUsers, 1) + assert.Equal(t, "user1.user", restoreUsers[0].User) +}