fix restore rbac empty meta crash

Signed-off-by: shaoting-huang <[email protected]>
milvus-io · Jan 10, 2025 · cb486db · cb486db
1 parent 95809ca
commit cb486db
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/internal/metastore/kv/rootcoord/kv_catalog.go b/internal/metastore/kv/rootcoord/kv_catalog.go
@@ -1527,7 +1527,7 @@ func (kc *Catalog) RestoreRBAC(ctx context.Context, tenant string, meta *milvusp
 		return err
 	}
 	existRoleMap := lo.SliceToMap(existRoles, func(entity *milvuspb.RoleResult) (string, struct{}) { return entity.GetRole().GetName(), struct{}{} })
-	for _, role := range meta.Roles {
+	for _, role := range meta.GetRoles() {
 		if _, ok := existRoleMap[role.GetName()]; ok {
 			log.Ctx(ctx).Warn("failed to restore, role already exists", zap.String("role", role.GetName()))
 			err = errors.Newf("role [%s] already exists", role.GetName())
@@ -1609,7 +1609,7 @@ func (kc *Catalog) RestoreRBAC(ctx context.Context, tenant string, meta *milvusp
 		entity := &milvuspb.UserEntity{
 			Name: user.User,
 		}
-		for _, role := range user.Roles {
+		for _, role := range user.GetRoles() {
 			err = kc.AlterUserRole(ctx, tenant, entity, role, milvuspb.OperateUserRoleType_AddUserToRole)
 			if err != nil {
 				return err

diff --git a/internal/proxy/impl.go b/internal/proxy/impl.go
@@ -5656,6 +5656,9 @@ func (node *Proxy) RestoreRBAC(ctx context.Context, req *milvuspb.RestoreRBACMet
 	if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
 		return merr.Status(err), nil
 	}
+	if req.RBACMeta == nil {
+		return merr.Status(fmt.Errorf("should provide rbac meta to restore rbac")), nil
+	}
 
 	result, err := node.rootCoord.RestoreRBAC(ctx, req)
 	if err != nil {

diff --git a/tests/integration/rbac/rbac_backup_test.go b/tests/integration/rbac/rbac_backup_test.go
@@ -154,8 +154,13 @@ func (s *RBACBackupTestSuite) TestBackup() {
 	s.Equal(groupName, backupRBACResp.GetRBACMeta().PrivilegeGroups[0].GroupName)
 	s.Equal(2, len(backupRBACResp.GetRBACMeta().PrivilegeGroups[0].Privileges))
 
+	// should pass rbac meta to restore
+	restoreRBACResp, err := s.Cluster.Proxy.RestoreRBAC(ctx, &milvuspb.RestoreRBACMetaRequest{})
+	s.NoError(err)
+	s.False(merr.Ok(restoreRBACResp))
+
 	// test restore, expect to failed due to role/user already exist
-	restoreRBACResp, err := s.Cluster.Proxy.RestoreRBAC(ctx, &milvuspb.RestoreRBACMetaRequest{
+	restoreRBACResp, err = s.Cluster.Proxy.RestoreRBAC(ctx, &milvuspb.RestoreRBACMetaRequest{
 		RBACMeta: backupRBACResp.GetRBACMeta(),
 	})
 	s.NoError(err)