Support list of TokenCustomizers

quarkusio · Jul 7, 2024 · 3fd082f · 3fd082f
1 parent 7736fe5
commit 3fd082f
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 33 deletions.
diff --git a/extensions/oidc/runtime/src/main/java/io/quarkus/oidc/runtime/OidcProvider.java b/extensions/oidc/runtime/src/main/java/io/quarkus/oidc/runtime/OidcProvider.java
@@ -72,7 +72,7 @@ public class OidcProvider implements Closeable {
     final RefreshableVerificationKeyResolver asymmetricKeyResolver;
     final DynamicVerificationKeyResolver keyResolverProvider;
     final OidcTenantConfig oidcConfig;
-    final TokenCustomizer tokenCustomizer;
+    final List<TokenCustomizer> tokenCustomizers;
     final String issuer;
     final String[] audience;
     final Map<String, String> requiredClaims;
@@ -85,10 +85,10 @@ public OidcProvider(OidcProviderClient client, OidcTenantConfig oidcConfig, Json
     }
 
     public OidcProvider(OidcProviderClient client, OidcTenantConfig oidcConfig, JsonWebKeySet jwks,
-            TokenCustomizer tokenCustomizer, Key tokenDecryptionKey, List<Validator> customValidators) {
+            List<TokenCustomizer> tokenCustomizers, Key tokenDecryptionKey, List<Validator> customValidators) {
         this.client = client;
         this.oidcConfig = oidcConfig;
-        this.tokenCustomizer = tokenCustomizer;
+        this.tokenCustomizers = tokenCustomizers;
         if (jwks != null) {
             this.asymmetricKeyResolver = new JsonWebKeyResolver(jwks, oidcConfig.token.forcedJwkRefreshInterval);
         } else if (oidcConfig != null && oidcConfig.certificateChain.trustStoreFile.isPresent()) {
@@ -113,7 +113,7 @@ public OidcProvider(OidcProviderClient client, OidcTenantConfig oidcConfig, Json
     public OidcProvider(String publicKeyEnc, OidcTenantConfig oidcConfig, Key tokenDecryptionKey) {
         this.client = null;
         this.oidcConfig = oidcConfig;
-        this.tokenCustomizer = TenantFeatureFinder.find(oidcConfig);
+        this.tokenCustomizers = TenantFeatureFinder.find(oidcConfig);
         if (publicKeyEnc != null) {
             this.asymmetricKeyResolver = new LocalPublicKeyResolver(publicKeyEnc);
         } else if (oidcConfig.certificateChain.trustStoreFile.isPresent()) {
@@ -274,17 +274,18 @@ private TokenVerificationResult verifyJwtTokenInternal(String token,
     }
 
     private String customizeJwtToken(String token) {
-        if (tokenCustomizer != null) {
-            JsonObject headers = AbstractJsonObjectResponse.toJsonObject(
-                    OidcUtils.decodeJwtHeadersAsString(token));
-            headers = tokenCustomizer.customizeHeaders(headers);
-            if (headers != null) {
-                String newHeaders = new String(
-                        Base64.getUrlEncoder().withoutPadding().encode(headers.toString().getBytes()),
-                        StandardCharsets.UTF_8);
-                int dotIndex = token.indexOf('.');
-                String newToken = newHeaders + token.substring(dotIndex);
-                return newToken;
+        if (tokenCustomizers != null) {
+            for (TokenCustomizer tokenCustomizer : tokenCustomizers) {
+                JsonObject headers = AbstractJsonObjectResponse.toJsonObject(OidcUtils.decodeJwtHeadersAsString(token));
+                headers = tokenCustomizer.customizeHeaders(headers);
+                if (headers != null) {
+                    String newHeaders = new String(
+                            Base64.getUrlEncoder().withoutPadding().encode(headers.toString().getBytes()),
+                            StandardCharsets.UTF_8);
+                    int dotIndex = token.indexOf('.');
+                    String newToken = newHeaders + token.substring(dotIndex);
+                    return newToken;
+                }
             }
         }
         return token;

diff --git a/extensions/oidc/runtime/src/main/java/io/quarkus/oidc/runtime/TenantFeatureFinder.java b/extensions/oidc/runtime/src/main/java/io/quarkus/oidc/runtime/TenantFeatureFinder.java
@@ -21,35 +21,25 @@ private TenantFeatureFinder() {
 
     }
 
-    public static TokenCustomizer find(OidcTenantConfig oidcConfig) {
+    public static List<TokenCustomizer> find(OidcTenantConfig oidcConfig) {
         if (oidcConfig == null) {
-            return null;
+            return List.of();
         }
         ArcContainer container = Arc.container();
         if (container != null) {
             String customizerName = oidcConfig.token.customizerName.orElse(null);
             if (customizerName != null && !customizerName.isEmpty()) {
                 InstanceHandle<TokenCustomizer> tokenCustomizer = container.instance(customizerName);
                 if (tokenCustomizer.isAvailable()) {
-                    return tokenCustomizer.get();
+                    return List.of(tokenCustomizer.get());
                 } else {
                     throw new OIDCException("Unable to find TokenCustomizer " + customizerName);
                 }
-            } else if (oidcConfig.tenantId.isPresent()) {
-                String tenantId = oidcConfig.tenantId.get();
-                List<TokenCustomizer> list = findTenantFeaturesByTenantId(TokenCustomizer.class, tenantId, container);
-                if (!list.isEmpty()) {
-                    if (list.size() >= 2) {
-                        throw new OIDCException(
-                                "Found multiple TokenCustomizers that are annotated with @TenantFeature that has tenantId ("
-                                        + tenantId + ")");
-                    }
-                    return list.get(0);
-                }
-
+            } else {
+                return find(oidcConfig, TokenCustomizer.class);
             }
         }
-        return null;
+        return List.of();
     }
 
     public static <T> List<T> find(OidcTenantConfig oidcTenantConfig, Class<T> tenantFeatureClass) {

diff --git a/extensions/oidc/runtime/src/test/java/io/quarkus/oidc/runtime/OidcProviderTest.java b/extensions/oidc/runtime/src/test/java/io/quarkus/oidc/runtime/OidcProviderTest.java
@@ -54,14 +54,14 @@ public void testAlgorithmCustomizer() throws Exception {
             }
         }
 
-        try (OidcProvider provider = new OidcProvider(null, oidcConfig, jwkSet, new TokenCustomizer() {
+        try (OidcProvider provider = new OidcProvider(null, oidcConfig, jwkSet, List.of(new TokenCustomizer() {
 
             @Override
             public JsonObject customizeHeaders(JsonObject headers) {
                 return Json.createObjectBuilder(headers).add("alg", "RS256").build();
             }
 
-        }, null, null)) {
+        }), null, null)) {
             TokenVerificationResult result = provider.verifyJwtToken(newToken, false, false, null);
             assertEquals("http://keycloak/realm", result.localVerificationResult.getString("iss"));
         }