Fix the problem of granting role not taking effect (#2495)

#### What type of PR is this?

/kind bug
/area core
/milestone 2.0

#### What this PR does / why we need it:

Fix the problem of granting role not taking effect.

#### Which issue(s) this PR fixes:

Fixes https://github.com/halo-dev/halo/issues/2490

#### Special notes for your reviewer:

Steps to test:

1. Create an user
2. Grant a role for the user
3. Check the user list

#### Does this PR introduce a user-facing change?

```release-note
None
```

src/main/java/run/halo/app/core/extension/endpoint/UserEndpoint.java

@@ -8,7 +8,6 @@ import com.fasterxml.jackson.core.type.TypeReference;
 import io.micrometer.common.util.StringUtils;
 import io.swagger.v3.oas.annotations.enums.ParameterIn;
 import io.swagger.v3.oas.annotations.media.Schema;
-import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Objects;
@@ -19,7 +18,6 @@ import org.springframework.http.MediaType;
 import org.springframework.lang.NonNull;
 import org.springframework.security.core.context.ReactiveSecurityContextHolder;
 import org.springframework.stereotype.Component;
-import org.springframework.util.CollectionUtils;
 import org.springframework.web.reactive.function.server.RouterFunction;
 import org.springframework.web.reactive.function.server.ServerRequest;
 import org.springframework.web.reactive.function.server.ServerResponse;
@@ -27,7 +25,6 @@ import org.springframework.web.server.ServerWebInputException;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import run.halo.app.core.extension.Role;
-import run.halo.app.core.extension.RoleBinding;
 import run.halo.app.core.extension.User;
 import run.halo.app.core.extension.service.UserService;
 import run.halo.app.extension.ReactiveExtensionClient;
@@ -131,45 +128,14 @@ public class UserEndpoint implements CustomEndpoint {
         return request.bodyToMono(GrantRequest.class)
             .switchIfEmpty(
                 Mono.error(() -> new ServerWebInputException("Request body is empty")))
-            .flatMap(grant -> client.get(User.class, username).thenReturn(grant))
-            .flatMap(grant -> Flux.fromIterable(grant.roles)
-                .flatMap(roleName -> client.get(Role.class, roleName))
-                .then().thenReturn(grant))
-            .zipWith(client.list(RoleBinding.class, RoleBinding.containsUser(username), null)
-                .collectList())
-            .flatMap(tuple2 -> {
-                var grant = tuple2.getT1();
-                var bindings = tuple2.getT2();
-                var bindingToUpdate = new HashSet<RoleBinding>();
-                var bindingToDelete = new HashSet<RoleBinding>();
-                var existingRoles = new HashSet<String>();
-                bindings.forEach(binding -> {
-                    var roleName = binding.getRoleRef().getName();
-                    if (grant.roles.contains(roleName)) {
-                        existingRoles.add(roleName);
-                        return;
-                    }
-                    binding.getSubjects().removeIf(RoleBinding.Subject.isUser(username));
-                    if (CollectionUtils.isEmpty(binding.getSubjects())) {
-                        // remove it if subjects is empty
-                        bindingToDelete.add(binding);
-                    } else {
-                        bindingToUpdate.add(binding);
-                    }
-                });
+            .flatMap(grantRequest -> userService.grantRoles(username, grantRequest.roles())
+                .then(ServerResponse.ok().build()));
+    }
 
-                bindingToUpdate.forEach(client::update);
-                bindingToDelete.forEach(client::delete);
-
-                // remove existing roles
-                var roles = new HashSet<>(grant.roles);
-                roles.removeAll(existingRoles);
-                roles.stream()
-                    .map(roleName -> RoleBinding.create(username, roleName))
-                    .forEach(client::create);
-
-                return ServerResponse.ok().build();
-            });
+    private Mono<GrantRequest> checkRoles(GrantRequest request) {
+        return Flux.fromIterable(request.roles)
+            .flatMap(role -> client.get(Role.class, role))
+            .then(Mono.just(request));
     }
 
     record GrantRequest(Set<String> roles) {

src/main/java/run/halo/app/core/extension/service/UserService.java

@@ -1,5 +1,6 @@
 package run.halo.app.core.extension.service;
 
+import java.util.Set;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import run.halo.app.core.extension.Role;
@@ -14,4 +15,6 @@ public interface UserService {
     Mono<User> updateWithRawPassword(String username, String rawPassword);
 
     Flux<Role> listRoles(String username);
+
+    Mono<User> grantRoles(String username, Set<String> roles);
 }

src/main/java/run/halo/app/core/extension/service/UserServiceImpl.java

@@ -2,9 +2,13 @@ package run.halo.app.core.extension.service;
 
 import static run.halo.app.core.extension.RoleBinding.containsUser;
 
+import java.util.HashSet;
 import java.util.Objects;
+import java.util.Set;
 import org.springframework.security.crypto.password.PasswordEncoder;
 import org.springframework.stereotype.Service;
+import org.springframework.transaction.annotation.Transactional;
+import org.springframework.util.CollectionUtils;
 import org.springframework.util.StringUtils;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -65,4 +69,39 @@ public class UserServiceImpl implements UserService {
             .flatMap(roleName -> client.fetch(Role.class, roleName));
     }
 
+    @Override
+    @Transactional
+    public Mono<User> grantRoles(String username, Set<String> roles) {
+        return client.get(User.class, username)
+            .flatMap(user -> {
+                var bindingsToUpdate = new HashSet<RoleBinding>();
+                var bindingsToDelete = new HashSet<RoleBinding>();
+                var existingRoles = new HashSet<String>();
+                return client.list(RoleBinding.class, RoleBinding.containsUser(username), null)
+                    .doOnNext(binding -> {
+                        var roleName = binding.getRoleRef().getName();
+                        if (roles.contains(roleName)) {
+                            existingRoles.add(roleName);
+                            return;
+                        }
+                        binding.getSubjects().removeIf(RoleBinding.Subject.isUser(username));
+                        if (CollectionUtils.isEmpty(binding.getSubjects())) {
+                            // remove it if subjects is empty
+                            bindingsToDelete.add(binding);
+                        } else {
+                            bindingsToUpdate.add(binding);
+                        }
+                    })
+                    .thenMany(Flux.fromIterable(bindingsToUpdate).flatMap(client::update))
+                    .thenMany(Flux.fromIterable(bindingsToDelete).flatMap(client::delete))
+                    .thenMany(Flux.fromStream(() -> {
+                        var mutableRoles = new HashSet<>(roles);
+                        mutableRoles.removeAll(existingRoles);
+                        return mutableRoles.stream()
+                            .map(roleName -> RoleBinding.create(username, roleName));
+                    }).flatMap(client::create))
+                    .then(Mono.just(user));
+            });
+    }
+
 }

src/test/java/run/halo/app/core/extension/endpoint/UserEndpointTest.java

@@ -1,11 +1,8 @@
 package run.halo.app.core.extension.endpoint;
 
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.ArgumentMatchers.isA;
 import static org.mockito.ArgumentMatchers.same;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -21,11 +18,9 @@ import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.autoconfigure.web.reactive.AutoConfigureWebTestClient;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.boot.test.mock.mockito.MockBean;
-import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.security.test.context.support.WithMockUser;
 import org.springframework.test.web.reactive.server.WebTestClient;
-import org.springframework.web.server.ResponseStatusException;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import run.halo.app.core.extension.Role;
@@ -169,77 +164,14 @@ class UserEndpointTest {
         }
 
         @Test
-        void shouldGetNotFoundIfUserNotFound() {
-            when(client.get(User.class, "fake-user"))
-                .thenReturn(Mono.error(new ResponseStatusException(HttpStatus.NOT_FOUND)));
-            when(client.get(Role.class, "fake-role")).thenReturn(Mono.just(mock(Role.class)));
-
-            webClient.post().uri("/apis/api.console.halo.run/v1alpha1/users/fake-user/permissions")
-                .contentType(MediaType.APPLICATION_JSON)
-                .bodyValue(new UserEndpoint.GrantRequest(Set.of("fake-role")))
-                .exchange()
-                .expectStatus().isNotFound();
-
-            verify(client, times(1)).get(same(User.class), eq("fake-user"));
-            verify(client, never()).get(same(Role.class), eq("fake-role"));
-        }
-
-        @Test
-        void shouldGetNotFoundIfRoleNotFound() {
-            when(client.get(User.class, "fake-user")).thenReturn(Mono.just(mock(User.class)));
-            when(client.get(Role.class, "fake-role"))
-                .thenReturn(Mono.error(new ResponseStatusException(HttpStatus.NOT_FOUND)));
-
-            webClient.post().uri("/apis/api.console.halo.run/v1alpha1/users/fake-user/permissions")
-                .contentType(MediaType.APPLICATION_JSON)
-                .bodyValue(new UserEndpoint.GrantRequest(Set.of("fake-role")))
-                .exchange()
-                .expectStatus().isNotFound();
-
-            verify(client).get(User.class, "fake-user");
-            verify(client).get(Role.class, "fake-role");
-        }
-
-        @Test
-        void shouldCreateRoleBindingIfNotExist() {
-            when(client.get(User.class, "fake-user")).thenReturn(Mono.just(mock(User.class)));
-            var role = mock(Role.class);
-            when(client.get(Role.class, "fake-role")).thenReturn(Mono.just(role));
+        void shouldGrantPermission() {
+            when(userService.grantRoles("fake-user", Set.of("fake-role"))).thenReturn(Mono.empty());
 
             webClient.post().uri("/apis/api.console.halo.run/v1alpha1/users/fake-user/permissions")
                 .contentType(MediaType.APPLICATION_JSON)
                 .bodyValue(new UserEndpoint.GrantRequest(Set.of("fake-role")))
                 .exchange()
                 .expectStatus().isOk();
-
-            verify(client).get(User.class, "fake-user");
-            verify(client).get(Role.class, "fake-role");
-            verify(client).create(RoleBinding.create("fake-user", "fake-role"));
-            verify(client, never()).update(isA(RoleBinding.class));
-            verify(client, never()).delete(isA(RoleBinding.class));
-        }
-
-        @Test
-        void shouldDeleteRoleBindingIfNotProvided() {
-            when(client.get(User.class, "fake-user")).thenReturn(Mono.just(mock(User.class)));
-            var role = mock(Role.class);
-            when(client.get(Role.class, "fake-role")).thenReturn(Mono.just(role));
-            var roleBinding = RoleBinding.create("fake-user", "non-provided-fake-role");
-            when(client.list(same(RoleBinding.class), any(), any()))
-                .thenReturn(Flux.fromIterable(List.of(roleBinding)));
-
-            webClient.post().uri("/apis/api.console.halo.run/v1alpha1/users/fake-user/permissions")
-                .contentType(MediaType.APPLICATION_JSON)
-                .bodyValue(new UserEndpoint.GrantRequest(Set.of("fake-role"))).exchange()
-                .expectStatus().isOk();
-
-            verify(client).get(User.class, "fake-user");
-            verify(client).get(Role.class, "fake-role");
-            verify(client).create(RoleBinding.create("fake-user", "fake-role"));
-            verify(client)
-                .delete(argThat(binding -> binding.getMetadata().getName()
-                    .equals(roleBinding.getMetadata().getName())));
-            verify(client, never()).update(isA(RoleBinding.class));
         }
 
         @Test

src/test/java/run/halo/app/core/extension/service/UserServiceImplTest.java

@@ -5,7 +5,10 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isA;
+import static org.mockito.ArgumentMatchers.same;
 import static org.mockito.Mockito.lenient;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
@@ -13,6 +16,7 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import java.util.List;
+import java.util.Set;
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
@@ -282,12 +286,93 @@ class UserServiceImplTest {
             verify(client).get(User.class, "fake-user");
         }
 
-        User createUser(String password) {
-            var user = new User();
-            user.setSpec(new User.UserSpec());
-            user.getSpec().setPassword(password);
-            return user;
-        }
     }
 
+    User createUser(String password) {
+        var user = new User();
+        user.setSpec(new User.UserSpec());
+        user.getSpec().setPassword(password);
+        return user;
+    }
+
+    @Nested
+    class GrantRolesTest {
+
+        @Test
+        void shouldGetNotFoundIfUserNotFound() {
+            when(client.get(User.class, "invalid-user"))
+                .thenReturn(Mono.error(new ExtensionNotFoundException()));
+
+            var grantRolesMono = userService.grantRoles("invalid-user", Set.of("fake-role"));
+            StepVerifier.create(grantRolesMono)
+                .expectError(ExtensionNotFoundException.class)
+                .verify();
+
+            verify(client).get(User.class, "invalid-user");
+        }
+
+        @Test
+        void shouldCreateRoleBindingIfNotExist() {
+            when(client.get(User.class, "fake-user"))
+                .thenReturn(Mono.just(createUser("fake-password")));
+            when(client.list(same(RoleBinding.class), any(), any())).thenReturn(Flux.empty());
+            when(client.create(isA(RoleBinding.class))).thenReturn(
+                Mono.just(mock(RoleBinding.class)));
+
+            var grantRolesMono = userService.grantRoles("fake-user", Set.of("fake-role"));
+            StepVerifier.create(grantRolesMono)
+                .expectNextCount(1)
+                .verifyComplete();
+
+            verify(client).get(User.class, "fake-user");
+            verify(client).list(same(RoleBinding.class), any(), any());
+            verify(client).create(isA(RoleBinding.class));
+        }
+
+        @Test
+        void shouldDeleteRoleBindingIfNotProvided() {
+            when(client.get(User.class, "fake-user")).thenReturn(Mono.just(mock(User.class)));
+            var notProvidedRoleBinding = RoleBinding.create("fake-user", "non-provided-fake-role");
+            var existingRoleBinding = RoleBinding.create("fake-user", "fake-role");
+            when(client.list(same(RoleBinding.class), any(), any())).thenReturn(
+                Flux.fromIterable(List.of(notProvidedRoleBinding, existingRoleBinding)));
+            when(client.delete(isA(RoleBinding.class)))
+                .thenReturn(Mono.just(mock(RoleBinding.class)));
+
+            StepVerifier.create(userService.grantRoles("fake-user", Set.of("fake-role")))
+                .expectNextCount(1)
+                .verifyComplete();
+
+            verify(client).get(User.class, "fake-user");
+            verify(client).list(same(RoleBinding.class), any(), any());
+            verify(client).delete(notProvidedRoleBinding);
+        }
+
+        @Test
+        void shouldUpdateRoleBindingIfExists() {
+            when(client.get(User.class, "fake-user")).thenReturn(Mono.just(mock(User.class)));
+            // add another subject
+            var anotherSubject = new RoleBinding.Subject();
+            anotherSubject.setName("another-fake-user");
+            anotherSubject.setKind(User.KIND);
+            anotherSubject.setApiGroup(User.GROUP);
+            var notProvidedRoleBinding = RoleBinding.create("fake-user", "non-provided-fake-role");
+            notProvidedRoleBinding.getSubjects().add(anotherSubject);
+
+            var existingRoleBinding = RoleBinding.create("fake-user", "fake-role");
+
+            when(client.list(same(RoleBinding.class), any(), any())).thenReturn(
+                Flux.fromIterable(List.of(notProvidedRoleBinding, existingRoleBinding)));
+            when(client.update(isA(RoleBinding.class)))
+                .thenReturn(Mono.just(mock(RoleBinding.class)));
+
+            StepVerifier.create(userService.grantRoles("fake-user", Set.of("fake-role")))
+                .expectNextCount(1)
+                .verifyComplete();
+
+            verify(client).get(User.class, "fake-user");
+            verify(client).list(same(RoleBinding.class), any(), any());
+            verify(client).update(notProvidedRoleBinding);
+        }
+    }
 }