Merge pull request #988 from AzureAD/avdunn/improve-credentials

Avery-Dunn · web-flow · commit cabae29b0217 · 2025-09-11T10:44:37.000-07:00
Improve behavior related to assertions
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/Authority.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/Authority.java
@@ -131,6 +131,40 @@ static void validateAuthority(URL authorityUrl) {
         }
     }
 
+    /**
+     * Creates a new Authority instance with a different tenant.
+     * This is useful when overriding the tenant at request level.
+     *
+     * @param originalAuthority The original authority to base the new one on
+     * @param newTenant The new tenant to use in the authority URL
+     * @return A new Authority instance with the specified tenant
+     * @throws MalformedURLException If the new authority URL is invalid
+     * @throws NullPointerException If originalAuthority or newTenant is null
+     */
+    static Authority replaceTenant(Authority originalAuthority, String newTenant) throws MalformedURLException {
+        if (originalAuthority == null) {
+            throw new NullPointerException("originalAuthority");
+        }
+        if (StringHelper.isBlank(newTenant)) {
+            throw new NullPointerException("newTenant");
+        }
+
+        URL originalUrl = originalAuthority.canonicalAuthorityUrl();
+        String host = originalUrl.getHost();
+        String protocol = originalUrl.getProtocol();
+        int port = originalUrl.getPort();
+
+        // Build path with new tenant
+        String newAuthority = String.format("%s://%s%s/%s/",
+                protocol,
+                host,
+                (port == -1 ? "" : ":" + port),
+                newTenant);
+
+        // Create proper authority instance with the tenant-specific URL
+        return createAuthority(new URL(newAuthority));
+    }
+
     static String getTenant(URL authorityUrl, AuthorityType authorityType) {
         String[] segments = authorityUrl.getPath().substring(1).split("/");
         if (authorityType == AuthorityType.B2C) {
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ClientAssertion.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ClientAssertion.java
@@ -4,21 +4,66 @@
 package com.microsoft.aad.msal4j;
 
 import java.util.Objects;
+import java.util.concurrent.Callable;
 
 final class ClientAssertion implements IClientAssertion {
 
     static final String ASSERTION_TYPE_JWT_BEARER = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer";
     private final String assertion;
+    private final Callable<String> assertionProvider;
 
+    /**
+     * Constructor that accepts a static assertion string
+     *
+     * @param assertion The JWT assertion string to use
+     * @throws NullPointerException if assertion is null or empty
+     */
     ClientAssertion(final String assertion) {
         if (StringHelper.isBlank(assertion)) {
             throw new NullPointerException("assertion");
         }
 
         this.assertion = assertion;
+        this.assertionProvider = null;
     }
 
+    /**
+     * Constructor that accepts a callable that provides the assertion string
+     *
+     * @param assertionProvider A callable that returns a JWT assertion string
+     * @throws NullPointerException if assertionProvider is null
+     */
+    ClientAssertion(final Callable<String> assertionProvider) {
+        if (assertionProvider == null) {
+            throw new NullPointerException("assertionProvider");
+        }
+
+        this.assertion = null;
+        this.assertionProvider = assertionProvider;
+    }
+
+    /**
+     * Gets the JWT assertion for client authentication.
+     * If this ClientAssertion was created with a Callable, the callable will be
+     * invoked each time this method is called to generate a fresh assertion.
+     *
+     * @return A JWT assertion string
+     * @throws MsalClientException if the assertion provider returns null/empty or throws an exception
+     */
     public String assertion() {
+        if (assertionProvider != null) {
+            try {
+                String generatedAssertion = assertionProvider.call();
+                if (StringHelper.isBlank(generatedAssertion)) {
+                    throw new MsalClientException("Assertion provider returned null or empty assertion",
+                            AuthenticationErrorCode.INVALID_JWT);
+                }
+                return generatedAssertion;
+            } catch (Exception ex) {
+                throw new MsalClientException(ex);
+            }
+        }
+
         return this.assertion;
     }
 
@@ -30,11 +75,24 @@ public boolean equals(Object o) {
         if (!(o instanceof ClientAssertion)) return false;
 
         ClientAssertion other = (ClientAssertion) o;
+
+        // For assertion providers, we consider them equal if they're the same object
+        if (this.assertionProvider != null && other.assertionProvider != null) {
+            return this.assertionProvider == other.assertionProvider;
+        }
+
+        // For static assertions, compare the assertion strings
         return Objects.equals(assertion(), other.assertion());
     }
 
     @Override
     public int hashCode() {
+        // For assertion providers, use the provider's identity hash code
+        if (assertionProvider != null) {
+            return System.identityHashCode(assertionProvider);
+        }
+
+        // For static assertions, hash the assertion string
         int result = 1;
         result = result * 59 + (this.assertion == null ? 43 : this.assertion.hashCode());
         return result;
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ClientCredentialFactory.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ClientCredentialFactory.java
@@ -91,15 +91,18 @@ public static IClientAssertion createFromClientAssertion(String clientAssertion)
 
     /**
      * Static method to create a {@link ClientAssertion} instance from a provided Callable.
+     * The callable will be invoked each time the assertion is needed, allowing for dynamic
+     * generation of assertions.
      *
      * @param callable Callable that produces a JWT token encoded as a base64 URL encoded string
-     * @return {@link ClientAssertion}
+     * @return {@link ClientAssertion} that will invoke the callable each time assertion() is called
+     * @throws NullPointerException if callable is null
      */
-    public static IClientAssertion createFromCallback(Callable<String> callable) throws ExecutionException, InterruptedException {
-        ExecutorService executor = Executors.newSingleThreadExecutor();
-
-        Future<String> future = executor.submit(callable);
+    public static IClientAssertion createFromCallback(Callable<String> callable) {
+        if (callable == null) {
+            throw new NullPointerException("callable");
+        }
 
-        return new ClientAssertion(future.get());
+        return new ClientAssertion(callable);
     }
 }
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/TokenRequestExecutor.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/TokenRequestExecutor.java
@@ -90,64 +90,63 @@ private void addQueryParameters(OAuthHttpRequest oauthHttpRequest) {
         if (msalRequest.application() instanceof ConfidentialClientApplication) {
             ConfidentialClientApplication application = (ConfidentialClientApplication) msalRequest.application();
 
-            // Determine which credential to use - either from the request or from the application
-            IClientCredential credential = getCredentialToUse(application);
-
-            // Add appropriate authentication parameters based on the credential type
-            addCredentialToRequest(queryParameters, credential, application);
+            // Consolidated credential and tenant override handling
+            addCredentialToRequest(queryParameters, application);
         }
 
         oauthHttpRequest.setQuery(StringHelper.serializeQueryParameters(queryParameters));
     }
 
-    /**
-     * Determines which credential to use for authentication:
-     * - If the request is a ClientCredentialRequest with a specified credential, use that
-     * - Otherwise use the application's credential
-     *
-     * @param application The confidential client application
-     * @return The credential to use, may be null if no credential is available
-     */
-    private IClientCredential getCredentialToUse(ConfidentialClientApplication application) {
-        if (msalRequest instanceof ClientCredentialRequest &&
-            ((ClientCredentialRequest) msalRequest).parameters.clientCredential() != null) {
-            return ((ClientCredentialRequest) msalRequest).parameters.clientCredential();
-        }
-        return application.clientCredential;
-    }
-
     /**
      * Adds the appropriate authentication parameters to the request based on credential type.
      * Handles different credential types (secret, assertion, certificate) by adding the appropriate
      * parameters to the request.
      *
      * @param queryParameters The map of query parameters to add to
-     * @param credential The credential to use for authentication, may be null
      * @param application The confidential client application
      */
     private void addCredentialToRequest(Map<String, String> queryParameters,
-                                       IClientCredential credential,
                                        ConfidentialClientApplication application) {
-        if (credential == null) {
+        IClientCredential credentialToUse = application.clientCredential;
+        Authority authorityToUse = application.authenticationAuthority;
+
+        // A ClientCredentialRequest may have parameters which override the credentials used to build the application.
+        if (msalRequest instanceof ClientCredentialRequest) {
+            ClientCredentialParameters parameters = ((ClientCredentialRequest) msalRequest).parameters;
+
+            if (parameters.clientCredential() != null) {
+                credentialToUse = parameters.clientCredential();
+            }
+
+            if (parameters.tenant() != null) {
+                try {
+                    authorityToUse = Authority.replaceTenant(authorityToUse, parameters.tenant());
+                } catch (MalformedURLException e) {
+                    log.warn("Could not create authority with tenant override: {}", e.getMessage());
+                }
+            }
+        }
+
+        // Quick return if no credential is provided
+        if (credentialToUse == null) {
             return;
         }
 
-        if (credential instanceof ClientSecret) {
+        if (credentialToUse instanceof ClientSecret) {
             // For client secret, add client_secret parameter
-            queryParameters.put("client_secret", ((ClientSecret) credential).clientSecret());
-        } else if (credential instanceof ClientAssertion) {
+            queryParameters.put("client_secret", ((ClientSecret) credentialToUse).clientSecret());
+        } else if (credentialToUse instanceof ClientAssertion) {
             // For client assertion, add client_assertion and client_assertion_type parameters
-            addJWTBearerAssertionParams(queryParameters, ((ClientAssertion) credential).assertion());
-        } else if (credential instanceof ClientCertificate) {
+            addJWTBearerAssertionParams(queryParameters, ((ClientAssertion) credentialToUse).assertion());
+        } else if (credentialToUse instanceof ClientCertificate) {
             // For client certificate, generate a new assertion and add it to the request
-            ClientCertificate certificate = (ClientCertificate) credential;
+            ClientCertificate certificate = (ClientCertificate) credentialToUse;
             String assertion = certificate.getAssertion(
-                application.authenticationAuthority,
+                authorityToUse,
                 application.clientId(),
                 application.sendX5c());
             addJWTBearerAssertionParams(queryParameters, assertion);
         }
-        // If credential is of an unknown type, no additional parameters are added
     }
 
     /**
diff --git a/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ClientCertificateTest.java b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ClientCertificateTest.java
@@ -12,6 +12,7 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.*;
 
@@ -156,6 +157,108 @@ void testClientCertificate_GeneratesNewAssertionEachTime() throws Exception {
                 "The access tokens from each request should be different");
     }
 
+    @Test
+    void testClientCertificate_TenantOverride() throws Exception {
+        DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class);
+        Map<String, String> capturedTenants = new HashMap<>();
+
+        ConfidentialClientApplication cca =
+                ConfidentialClientApplication.builder("clientId", ClientCredentialFactory.createFromCertificate(TestHelper.getPrivateKey(), TestHelper.getX509Cert()))
+                        .authority("https://login.microsoftonline.com/default-tenant")
+                        .instanceDiscovery(false)
+                        .validateAuthority(false)
+                        .httpClient(httpClientMock)
+                        .build();
+
+        // Mock the HTTP client to capture and analyze assertions from each request
+        when(httpClientMock.send(any(HttpRequest.class))).thenAnswer(parameters -> {
+            HttpRequest request = parameters.getArgument(0);
+            String requestBody = request.body();
+            String url = request.url().toString();
+
+            // Capture which tenant was used in the authority
+            String tenant = extractTenantFromUrl(url);
+
+            // Extract the assertion to verify its audience claim
+            String clientAssertion = extractClientAssertion(requestBody);
+            if (clientAssertion != null) {
+                SignedJWT signedJWT = SignedJWT.parse(clientAssertion);
+
+                // Get the audience claim to verify it matches the tenant
+                String audience = signedJWT.getJWTClaimsSet().getAudience().get(0);
+
+                // Store the tenant and audience for verification
+                capturedTenants.put(tenant, audience);
+
+                // Verify it's a valid JWT with proper headers
+                if (signedJWT.getHeader().toJSONObject().containsKey("x5t#S256")) {
+                    HashMap<String, String> tokenResponseValues = new HashMap<>();
+                    tokenResponseValues.put("access_token", "access_token_for_" + tenant);
+                    return TestHelper.expectedResponse(200, TestHelper.getSuccessfulTokenResponse(tokenResponseValues));
+                }
+            }
+            return null;
+        });
+
+        // First request with default tenant
+        ClientCredentialParameters defaultParameters = ClientCredentialParameters.builder(Collections.singleton("scopes"))
+                .skipCache(true)
+                .build();
+        IAuthenticationResult resultDefault = cca.acquireToken(defaultParameters).get();
+
+        // Second request with override tenant
+        String overrideTenant = "override-tenant";
+        ClientCredentialParameters overrideParameters = ClientCredentialParameters.builder(Collections.singleton("scopes"))
+                .skipCache(true)
+                .tenant(overrideTenant)
+                .build();
+        IAuthenticationResult resultOverride = cca.acquireToken(overrideParameters).get();
+
+        // Verify both requests were processed
+        assertEquals(2, capturedTenants.size(), "Two requests with different tenants should have been processed");
+
+        // Verify both tenants were used
+        assertTrue(capturedTenants.containsKey("default-tenant"), "Default tenant should have been used");
+        assertTrue(capturedTenants.containsKey(overrideTenant), "Override tenant should have been used");
+
+        // Verify the audience in the JWT assertions reflects the different tenants
+        assertNotEquals(
+            capturedTenants.get("default-tenant"),
+            capturedTenants.get(overrideTenant),
+            "JWT audience should differ between default and override tenant"
+        );
+
+        // Verify the audience claims match the expected format with the correct tenant
+        assertTrue(
+            capturedTenants.get("default-tenant").contains("default-tenant"),
+            "Audience for default tenant should contain the default tenant name"
+        );
+        assertTrue(
+            capturedTenants.get(overrideTenant).contains(overrideTenant),
+            "Audience for override tenant should contain the override tenant name"
+        );
+
+        // Verify different access tokens were returned
+        assertNotEquals(resultDefault.accessToken(), resultOverride.accessToken(),
+            "Access tokens should differ when using different tenants");
+    }
+
+    /**
+     * Extracts the tenant name from an authority URL
+     * @param url The full URL containing the tenant
+     * @return The tenant name
+     */
+    private String extractTenantFromUrl(String url) {
+        // Authority URL format is typically https://login.microsoftonline.com/tenant/...
+        String[] parts = url.split("/");
+        for (int i = 0; i < parts.length; i++) {
+            if (parts[i].equalsIgnoreCase("login.microsoftonline.com") && i + 1 < parts.length) {
+                return parts[i + 1];
+            }
+        }
+        return null;
+    }
+
     /**
      * Extracts the client_assertion value from a URL-encoded request body
      * @param requestBody The request body string
diff --git a/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ClientCredentialTest.java b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ClientCredentialTest.java
