Merge pull request #42599 from Malandril/fix-augment-bug

Fix augmentors called multiple times for each identity provider
quarkusio · Aug 28, 2024 · a0906bb · a0906bb
2 parents 4b1ae89 + 1609bee
commit a0906bb
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 10 deletions.
diff --git a/extensions/security/runtime/pom.xml b/extensions/security/runtime/pom.xml
@@ -49,6 +49,11 @@
             <artifactId>quarkus-junit5-internal</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

diff --git a/...runtime/src/main/java/io/quarkus/security/runtime/QuarkusIdentityProviderManagerImpl.java b/...runtime/src/main/java/io/quarkus/security/runtime/QuarkusIdentityProviderManagerImpl.java
@@ -64,7 +64,7 @@ public Uni<SecurityIdentity> authenticate(AuthenticationRequest request) {
             if (providers.size() == 1) {
                 return handleSingleProvider(providers.get(0), request);
             }
-            return handleProvider(0, (List) providers, request);
+            return handleProviders((List) providers, request);
         } catch (Throwable t) {
             return Uni.createFrom().failure(t);
         }
@@ -106,18 +106,30 @@ public SecurityIdentity authenticateBlocking(AuthenticationRequest request) {
             throw new IllegalArgumentException(
                     "No IdentityProviders were registered to handle AuthenticationRequest " + request);
         }
-        return (SecurityIdentity) handleProvider(0, (List) providers, request).await().indefinitely();
+        return (SecurityIdentity) handleProviders((List) providers, request).await().indefinitely();
     }
 
-    private <T extends AuthenticationRequest> Uni<SecurityIdentity> handleProvider(int pos,
+    private <T extends AuthenticationRequest> Uni<SecurityIdentity> handleProviders(
             List<IdentityProvider<T>> providers, T request) {
+        return handleProvider(0, providers, request)
+                .onItem()
+                .transformToUni(new Function<SecurityIdentity, Uni<? extends SecurityIdentity>>() {
+                    @Override
+                    public Uni<? extends SecurityIdentity> apply(SecurityIdentity securityIdentity) {
+                        return handleIdentityFromProvider(0, securityIdentity, request.getAttributes());
+                    }
+                });
+    }
+
+    private <T extends AuthenticationRequest> Uni<SecurityIdentity> handleProvider(int pos, List<IdentityProvider<T>> providers,
+            T request) {
         if (pos == providers.size()) {
             //we failed to authentication
             log.debug("Authentication failed as providers would authenticate the request");
             return Uni.createFrom().failure(new AuthenticationFailedException());
         }
         IdentityProvider<T> current = providers.get(pos);
-        Uni<SecurityIdentity> cs = current.authenticate(request, blockingRequestContext)
+        return current.authenticate(request, blockingRequestContext)
                 .onItem().transformToUni(new Function<SecurityIdentity, Uni<? extends SecurityIdentity>>() {
                     @Override
                     public Uni<SecurityIdentity> apply(SecurityIdentity securityIdentity) {
@@ -127,12 +139,6 @@ public Uni<SecurityIdentity> apply(SecurityIdentity securityIdentity) {
                         return handleProvider(pos + 1, providers, request);
                     }
                 });
-        return cs.onItem().transformToUni(new Function<SecurityIdentity, Uni<? extends SecurityIdentity>>() {
-            @Override
-            public Uni<? extends SecurityIdentity> apply(SecurityIdentity securityIdentity) {
-                return handleIdentityFromProvider(0, securityIdentity, request.getAttributes());
-            }
-        });
     }
 
     private Uni<SecurityIdentity> handleIdentityFromProvider(int pos, SecurityIdentity identity,

diff --git a/...ime/src/test/java/io/quarkus/security/runtime/QuarkusIdentityProviderManagerImplTest.java b/...ime/src/test/java/io/quarkus/security/runtime/QuarkusIdentityProviderManagerImplTest.java
@@ -1,6 +1,11 @@
 package io.quarkus.security.runtime;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import java.util.concurrent.Executors;
 
@@ -11,6 +16,7 @@
 import io.quarkus.security.identity.IdentityProvider;
 import io.quarkus.security.identity.IdentityProviderManager;
 import io.quarkus.security.identity.SecurityIdentity;
+import io.quarkus.security.identity.SecurityIdentityAugmentor;
 import io.quarkus.security.identity.request.BaseAuthenticationRequest;
 import io.smallrye.mutiny.Uni;
 
@@ -32,6 +38,21 @@ void testIdentityProviderPriority() {
         assertEquals(new QuarkusPrincipal("Bob"), identity.getPrincipal());
     }
 
+    @Test
+    void testIdentityProviderAugmentOnlyOnce() {
+        TestSecurityAugmentor augmentor = spy(new TestSecurityAugmentor());
+        IdentityProviderManager identityProviderManager = QuarkusIdentityProviderManagerImpl.builder()
+                .addProvider(new TestIdentityProviderSystemFirstPriorityNoop())
+                .addProvider(new TestIdentityProviderSystemLastPriorityUser())
+                .addProvider(new AnonymousIdentityProvider())
+                .addSecurityIdentityAugmentor(augmentor)
+                .setBlockingExecutor(Executors.newSingleThreadExecutor()).build();
+        SecurityIdentity identity = identityProviderManager.authenticateBlocking(new TestAuthenticationRequest());
+        assertEquals(new QuarkusPrincipal("Bob"), identity.getPrincipal());
+        assertTrue(identity.getRoles().contains("role"));
+        verify(augmentor, times(1)).augment(any(), any());
+    }
+
     static class TestAuthenticationRequest extends BaseAuthenticationRequest {
     }
 
@@ -78,4 +99,26 @@ public int priority() {
             return SYSTEM_LAST;
         }
     }
+
+    static class TestIdentityProviderSystemFirstPriorityNoop extends TestIdentityProviderSystemLastPriority {
+        @Override
+        public Uni<SecurityIdentity> authenticate(TestAuthenticationRequest request, AuthenticationRequestContext context) {
+            return Uni.createFrom().nullItem();
+        }
+    }
+
+    static class TestIdentityProviderSystemLastPriorityUser extends TestIdentityProviderUserFirstPriority {
+        @Override
+        public int priority() {
+            return SYSTEM_LAST;
+        }
+    }
+
+    private static class TestSecurityAugmentor implements SecurityIdentityAugmentor {
+        @Override
+        public Uni<SecurityIdentity> augment(SecurityIdentity securityIdentity,
+                AuthenticationRequestContext authenticationRequestContext) {
+            return Uni.createFrom().item(QuarkusSecurityIdentity.builder(securityIdentity).addRole("role").build());
+        }
+    }
 }