Support token propagation during identity augmentation

extensions/oidc-token-propagation-reactive/deployment/src/main/java/io/quarkus/oidc/token/propagation/reactive/OidcTokenPropagationReactiveBuildStep.java

@@ -10,15 +10,19 @@
 import org.jboss.jandex.Type;
 
 import io.quarkus.arc.deployment.AdditionalBeanBuildItem;
+import io.quarkus.deployment.Capabilities;
+import io.quarkus.deployment.Capability;
 import io.quarkus.deployment.annotations.BuildProducer;
 import io.quarkus.deployment.annotations.BuildStep;
 import io.quarkus.deployment.annotations.BuildSteps;
 import io.quarkus.deployment.builditem.AdditionalIndexedClassesBuildItem;
 import io.quarkus.deployment.builditem.CombinedIndexBuildItem;
+import io.quarkus.deployment.builditem.SystemPropertyBuildItem;
 import io.quarkus.deployment.builditem.nativeimage.ReflectiveClassBuildItem;
 import io.quarkus.oidc.token.propagation.AccessToken;
 import io.quarkus.rest.client.reactive.deployment.DotNames;
 import io.quarkus.rest.client.reactive.deployment.RegisterProviderAnnotationInstanceBuildItem;
+import io.quarkus.runtime.configuration.ConfigurationException;
 
 @BuildSteps(onlyIf = OidcTokenPropagationReactiveBuildStep.IsEnabled.class)
 public class OidcTokenPropagationReactiveBuildStep {
@@ -51,11 +55,30 @@ void registerProvider(BuildProducer<AdditionalBeanBuildItem> additionalBeans,
 
     }
 
+    @BuildStep(onlyIf = IsEnabledDuringAuth.class)
+    SystemPropertyBuildItem activateTokenCredentialPropagationViaDuplicatedContext(Capabilities capabilities) {
+        if (capabilities.isPresent(Capability.OIDC)) {
+            return new SystemPropertyBuildItem("quarkus.oidc.propagate-token-credential-with-duplicated-context", "true");
+        }
+
+        throw new ConfigurationException(
+                "Configuration property 'quarkus.oidc-token-propagation-reactive.enable-during-auth' is set to 'true', " +
+                        "however this configuration property is only supported when Quarkus OIDC extension is present.");
+    }
+
     public static class IsEnabled implements BooleanSupplier {
         OidcTokenPropagationReactiveBuildTimeConfig config;
 
         public boolean getAsBoolean() {
             return config.enabled;
         }
     }
+
+    public static class IsEnabledDuringAuth implements BooleanSupplier {
+        OidcTokenPropagationReactiveBuildTimeConfig config;
+
+        public boolean getAsBoolean() {
+            return config.enableDuringAuth;
+        }
+    }
 }

extensions/oidc-token-propagation-reactive/deployment/src/main/java/io/quarkus/oidc/token/propagation/reactive/OidcTokenPropagationReactiveBuildTimeConfig.java

@@ -13,4 +13,17 @@ public class OidcTokenPropagationReactiveBuildTimeConfig {
      */
     @ConfigItem(defaultValue = "true")
     public boolean enabled;
+
+    /**
+     * If the OIDC Token Reactive Propagation is enabled during authentication, such as when you need
+     * to use a REST client from {@link io.quarkus.security.identity.SecurityIdentityAugmentor}.
+     * By default, access token propagation is not enabled during authentication, because
+     * current implementation requires storing of local data on Vert.X duplicated context in a safe manner.
+     * It is safe to enable this option in most cases, but please revise how your extensions work
+     * with Vert.X context.
+     *
+     * This configuration options is only supported when used together with the Quarkus OIDC extension.
+     */
+    @ConfigItem(defaultValue = "false")
+    public boolean enableDuringAuth;
 }

extensions/oidc-token-propagation-reactive/deployment/src/test/java/io/quarkus/oidc/token/propagation/reactive/FrontendResource.java

@@ -8,6 +8,8 @@
 import org.eclipse.microprofile.jwt.JsonWebToken;
 import org.eclipse.microprofile.rest.client.inject.RestClient;
 
+import io.quarkus.security.identity.CurrentIdentityAssociation;
+
 @Path("/frontend")
 public class FrontendResource {
     @Inject
@@ -17,10 +19,24 @@ public class FrontendResource {
     @Inject
     JsonWebToken jwt;
 
+    @Inject
+    CurrentIdentityAssociation identityAssociation;
+
     @GET
     @Path("token-propagation")
     @RolesAllowed("admin")
     public String userNameTokenPropagation() {
+        return getResponseWithExchangedUsername();
+    }
+
+    @GET
+    @Path("token-propagation-with-augmentor")
+    @RolesAllowed("tester") // tester role is granted by SecurityIdentityAugmentor
+    public String userNameTokenPropagationWithSecIdentityAugmentor() {
+        return getResponseWithExchangedUsername();
+    }
+
+    private String getResponseWithExchangedUsername() {
         if ("alice".equals(jwt.getName())) {
             return "Token issued to " + jwt.getName() + " has been exchanged, new user name: "
                     + accessTokenPropagationService.getUserName();

extensions/oidc-token-propagation-reactive/deployment/src/test/java/io/quarkus/oidc/token/propagation/reactive/OidcTokenPropagationWithSecurityIdentityAugmentorTest.java

@@ -0,0 +1,53 @@
+package io.quarkus.oidc.token.propagation.reactive;
+
+import static io.quarkus.oidc.token.propagation.reactive.RolesSecurityIdentityAugmentor.SUPPORTED_USER;
+import static org.hamcrest.Matchers.equalTo;
+
+import java.util.Set;
+
+import org.jboss.shrinkwrap.api.asset.StringAsset;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkus.test.QuarkusUnitTest;
+import io.quarkus.test.common.QuarkusTestResource;
+import io.quarkus.test.oidc.server.OidcWiremockTestResource;
+import io.restassured.RestAssured;
+
+@QuarkusTestResource(OidcWiremockTestResource.class)
+public class OidcTokenPropagationWithSecurityIdentityAugmentorTest {
+
+    private static Class<?>[] testClasses = {
+            FrontendResource.class,
+            ProtectedResource.class,
+            AccessTokenPropagationService.class,
+            RolesResource.class,
+            RolesService.class,
+            RolesSecurityIdentityAugmentor.class
+    };
+
+    @RegisterExtension
+    static final QuarkusUnitTest test = new QuarkusUnitTest()
+            .withApplicationRoot((jar) -> jar
+                    .addClasses(testClasses)
+                    .addAsResource("application.properties")
+                    .addAsResource(new StringAsset("quarkus.oidc-token-propagation-reactive.enable-during-auth=true\n" +
+                            "quarkus.rest-client.\"roles\".uri=http://localhost:8081/roles\n"),
+                            "META-INF/microprofile-config.properties"));
+
+    @Test
+    public void testGetUserNameWithTokenPropagation() {
+        // request only succeeds if SecurityIdentityAugmentor managed to acquire 'tester' role for user 'alice'
+        // and that is only possible if access token is propagated during augmentation
+        RestAssured.given().auth().oauth2(getBearerAccessToken())
+                .when().get("/frontend/token-propagation-with-augmentor")
+                .then()
+                .statusCode(200)
+                .body(equalTo("Token issued to alice has been exchanged, new user name: bob"));
+    }
+
+    public String getBearerAccessToken() {
+        return OidcWiremockTestResource.getAccessToken(SUPPORTED_USER, Set.of("admin"));
+    }
+
+}

extensions/oidc-token-propagation-reactive/deployment/src/test/java/io/quarkus/oidc/token/propagation/reactive/RolesResource.java

@@ -0,0 +1,26 @@
+package io.quarkus.oidc.token.propagation.reactive;
+
+import jakarta.inject.Inject;
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.Path;
+
+import org.eclipse.microprofile.jwt.JsonWebToken;
+
+import io.quarkus.security.Authenticated;
+import io.quarkus.security.ForbiddenException;
+
+@Path("/roles")
+@Authenticated
+public class RolesResource {
+
+    @Inject
+    JsonWebToken jwt;
+
+    @GET
+    public String get() {
+        if ("bob".equals(jwt.getName())) {
+            return "tester";
+        }
+        throw new ForbiddenException("Only user 'bob' is allowed to request roles");
+    }
+}

extensions/oidc-token-propagation-reactive/deployment/src/test/java/io/quarkus/oidc/token/propagation/reactive/RolesSecurityIdentityAugmentor.java

@@ -0,0 +1,34 @@
+package io.quarkus.oidc.token.propagation.reactive;
+
+import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.inject.Inject;
+
+import org.eclipse.microprofile.rest.client.inject.RestClient;
+
+import io.quarkus.security.identity.AuthenticationRequestContext;
+import io.quarkus.security.identity.SecurityIdentity;
+import io.quarkus.security.identity.SecurityIdentityAugmentor;
+import io.quarkus.security.runtime.QuarkusSecurityIdentity;
+import io.smallrye.mutiny.Uni;
+
+@ApplicationScoped
+public class RolesSecurityIdentityAugmentor implements SecurityIdentityAugmentor {
+
+    static final String SUPPORTED_USER = "alice";
+
+    @Inject
+    @RestClient
+    RolesService rolesService;
+
+    @Override
+    public Uni<SecurityIdentity> augment(SecurityIdentity securityIdentity,
+            AuthenticationRequestContext authenticationRequestContext) {
+        if (securityIdentity != null && securityIdentity.getPrincipal() != null
+                && SUPPORTED_USER.equals(securityIdentity.getPrincipal().getName())) {
+            return rolesService
+                    .getRole()
+                    .map(role -> QuarkusSecurityIdentity.builder(securityIdentity).addRole(role).build());
+        }
+        return Uni.createFrom().item(securityIdentity);
+    }
+}

extensions/oidc-token-propagation-reactive/deployment/src/test/java/io/quarkus/oidc/token/propagation/reactive/RolesService.java

@@ -0,0 +1,18 @@
+package io.quarkus.oidc.token.propagation.reactive;
+
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.Path;
+
+import org.eclipse.microprofile.rest.client.inject.RegisterRestClient;
+
+import io.quarkus.oidc.token.propagation.AccessToken;
+import io.smallrye.mutiny.Uni;
+
+@RegisterRestClient(configKey = "roles")
+@AccessToken
+@Path("/")
+public interface RolesService {
+
+    @GET
+    Uni<String> getRole();
+}

extensions/oidc-token-propagation-reactive/runtime/src/main/java/io/quarkus/oidc/token/propagation/reactive/AccessTokenRequestReactiveFilter.java

@@ -25,12 +25,15 @@
 import io.quarkus.oidc.client.runtime.DisabledOidcClientException;
 import io.quarkus.runtime.configuration.ConfigurationException;
 import io.quarkus.security.credential.TokenCredential;
+import io.quarkus.vertx.core.runtime.context.VertxContextSafetyToggle;
 import io.smallrye.mutiny.Uni;
+import io.vertx.core.Vertx;
 
 @Priority(Priorities.AUTHENTICATION)
 public class AccessTokenRequestReactiveFilter implements ResteasyReactiveClientRequestFilter {
     private static final Logger LOG = Logger.getLogger(AccessTokenRequestReactiveFilter.class);
     private static final String BEARER_SCHEME_WITH_SPACE = "Bearer ";
+    private static final String ERROR_MSG = "OIDC Token Propagation Reactive requires a safe (isolated) Vert.x sub-context because configuration property 'quarkus.oidc-token-propagation-reactive.enable-during-auth' has been set to true, but the current context hasn't been flagged as such.";
 
     @Inject
     Instance<TokenCredential> accessToken;
@@ -41,6 +44,9 @@ public class AccessTokenRequestReactiveFilter implements ResteasyReactiveClientR
     @Inject
     @ConfigProperty(name = "quarkus.oidc-token-propagation-reactive.exchange-token")
     boolean exchangeToken;
+    @Inject
+    @ConfigProperty(name = "quarkus.oidc-token-propagation-reactive.enable-during-auth", defaultValue = "false")
+    boolean enableDuringAuthentication;
 
     OidcClient exchangeTokenClient;
     String exchangeTokenProperty;
@@ -73,7 +79,7 @@ public void filter(ResteasyReactiveClientRequestContext requestContext) {
 
                 requestContext.suspend();
 
-                exchangeToken(accessToken.get().getToken()).subscribe().with(new Consumer<>() {
+                exchangeToken(getAccessToken().getToken()).subscribe().with(new Consumer<>() {
                     @Override
                     public void accept(String token) {
                         propagateToken(requestContext, token);
@@ -94,7 +100,7 @@ public void accept(Throwable t) {
                     }
                 });
             } else {
-                propagateToken(requestContext, accessToken.get().getToken());
+                propagateToken(requestContext, getAccessToken().getToken());
             }
         } else {
             abortRequest(requestContext);
@@ -111,6 +117,11 @@ public void propagateToken(ResteasyReactiveClientRequestContext requestContext,
     }
 
     protected boolean verifyTokenInstance(ResteasyReactiveClientRequestContext requestContext) {
+        if (enableDuringAuthentication && !Arc.container().requestContext().isActive()) {
+            // it is possible request context is not active when this filter is used during authentication
+            TokenCredential tokenCredential = getTokenCredentialFromContext();
+            return tokenCredential != null;
+        }
         if (!accessToken.isResolvable()) {
             LOG.debugf("Access token is not injected, aborting the request with HTTP 401 error");
             return false;
@@ -126,6 +137,18 @@ protected boolean verifyTokenInstance(ResteasyReactiveClientRequestContext reque
         return true;
     }
 
+    private TokenCredential getAccessToken() {
+        if (enableDuringAuthentication && !Arc.container().requestContext().isActive()) {
+            return getTokenCredentialFromContext();
+        }
+        return accessToken.get();
+    }
+
+    private static TokenCredential getTokenCredentialFromContext() {
+        VertxContextSafetyToggle.validateContextIfExists(ERROR_MSG, ERROR_MSG);
+        return Vertx.currentContext().getLocal(TokenCredential.class.getName());
+    }
+
     private Uni<String> exchangeToken(String token) {
         return exchangeTokenClient.getTokens(Collections.singletonMap(exchangeTokenProperty, token))
                 .onItem().transform(t -> t.getAccessToken());

extensions/oidc-token-propagation/deployment/src/main/java/io/quarkus/oidc/token/propagation/deployment/OidcTokenPropagationBuildStep.java

@@ -5,9 +5,12 @@
 import org.jboss.jandex.DotName;
 
 import io.quarkus.arc.deployment.AdditionalBeanBuildItem;
+import io.quarkus.deployment.Capabilities;
+import io.quarkus.deployment.Capability;
 import io.quarkus.deployment.annotations.BuildProducer;
 import io.quarkus.deployment.annotations.BuildStep;
 import io.quarkus.deployment.annotations.BuildSteps;
+import io.quarkus.deployment.builditem.SystemPropertyBuildItem;
 import io.quarkus.deployment.builditem.nativeimage.ReflectiveClassBuildItem;
 import io.quarkus.oidc.token.propagation.AccessToken;
 import io.quarkus.oidc.token.propagation.AccessTokenRequestFilter;
@@ -17,6 +20,7 @@
 import io.quarkus.oidc.token.propagation.runtime.OidcTokenPropagationConfig;
 import io.quarkus.restclient.deployment.RestClientAnnotationProviderBuildItem;
 import io.quarkus.resteasy.common.spi.ResteasyJaxrsProviderBuildItem;
+import io.quarkus.runtime.configuration.ConfigurationException;
 
 @BuildSteps(onlyIf = OidcTokenPropagationBuildStep.IsEnabled.class)
 public class OidcTokenPropagationBuildStep {
@@ -49,11 +53,30 @@ void registerProvider(BuildProducer<AdditionalBeanBuildItem> additionalBeans,
         }
     }
 
+    @BuildStep(onlyIf = IsEnabledDuringAuth.class)
+    SystemPropertyBuildItem activateTokenCredentialPropagationViaDuplicatedContext(Capabilities capabilities) {
+        if (capabilities.isPresent(Capability.OIDC)) {
+            return new SystemPropertyBuildItem("quarkus.oidc.propagate-token-credential-with-duplicated-context", "true");
+        }
+
+        throw new ConfigurationException(
+                "Configuration property 'quarkus.oidc-token-propagation.enable-during-auth' is set to 'true', " +
+                        "however this configuration property is only supported when Quarkus OIDC extension is present.");
+    }
+
     public static class IsEnabled implements BooleanSupplier {
         OidcTokenPropagationBuildTimeConfig config;
 
         public boolean getAsBoolean() {
             return config.enabled;
         }
     }
+
+    public static class IsEnabledDuringAuth implements BooleanSupplier {
+        OidcTokenPropagationBuildTimeConfig config;
+
+        public boolean getAsBoolean() {
+            return config.enableDuringAuth;
+        }
+    }
 }

extensions/oidc-token-propagation/deployment/src/test/java/io/quarkus/oidc/token/propagation/FrontendResource.java

@@ -21,6 +21,17 @@ public class FrontendResource {
     @Path("token-propagation")
     @RolesAllowed("admin")
     public String userNameTokenPropagation() {
+        return getResponseWithExchangedUsername();
+    }
+
+    @GET
+    @Path("token-propagation-with-augmentor")
+    @RolesAllowed("tester") // tester role is granted by SecurityIdentityAugmentor
+    public String userNameTokenPropagationWithSecIdentityAugmentor() {
+        return getResponseWithExchangedUsername();
+    }
+
+    private String getResponseWithExchangedUsername() {
         if ("alice".equals(jwt.getName())) {
             return "Token issued to " + jwt.getName() + " has been exchanged, new user name: "
                     + accessTokenPropagationService.getUserName();