Improve assertion callback options

Avery-Dunn · Avery-Dunn · commit a5b5b014f2c0 · 2026-05-06T14:17:18.000-07:00
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AssertionResponse.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AssertionResponse.java
@@ -0,0 +1,63 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.aad.msal4j;
+
+import java.security.cert.X509Certificate;
+
+/**
+ * Container returned from context-aware assertion provider callbacks.
+ * Allows the callback to supply both the client assertion JWT and an optional
+ * token-binding certificate for mutual-TLS Proof-of-Possession (mTLS PoP) scenarios.
+ *
+ * <p>When a {@link #tokenBindingCertificate()} is provided, MSAL sets the
+ * {@code client_assertion_type} to {@code urn:ietf:params:oauth:client-assertion-type:jwt-pop}
+ * instead of the default {@code jwt-bearer}.</p>
+ *
+ * @see ClientCredentialFactory#createFromCallback(java.util.function.Function)
+ * @see AssertionRequestOptions
+ */
+public class AssertionResponse {
+
+    private final String assertion;
+    private final X509Certificate tokenBindingCertificate;
+
+    /**
+     * Creates an AssertionResponse with just an assertion string (no token binding certificate).
+     *
+     * @param assertion the JWT assertion string
+     */
+    public AssertionResponse(String assertion) {
+        this(assertion, null);
+    }
+
+    /**
+     * Creates an AssertionResponse with an assertion string and an optional token-binding certificate.
+     *
+     * @param assertion               the JWT assertion string
+     * @param tokenBindingCertificate optional certificate for mTLS PoP binding, or null for standard jwt-bearer
+     */
+    public AssertionResponse(String assertion, X509Certificate tokenBindingCertificate) {
+        this.assertion = assertion;
+        this.tokenBindingCertificate = tokenBindingCertificate;
+    }
+
+    /**
+     * Gets the JWT assertion string to use as the {@code client_assertion} parameter.
+     *
+     * @return the JWT assertion string
+     */
+    public String assertion() {
+        return assertion;
+    }
+
+    /**
+     * Gets the optional token-binding certificate for mutual-TLS Proof-of-Possession (mTLS PoP).
+     * When present, MSAL uses {@code client_assertion_type=jwt-pop} instead of {@code jwt-bearer}.
+     *
+     * @return the binding certificate, or null if not using mTLS PoP
+     */
+    public X509Certificate tokenBindingCertificate() {
+        return tokenBindingCertificate;
+    }
+}
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
@@ -10,9 +10,11 @@
 final class ClientAssertion implements IClientAssertion {
 
     static final String ASSERTION_TYPE_JWT_BEARER = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer";
+    static final String ASSERTION_TYPE_JWT_POP = "urn:ietf:params:oauth:client-assertion-type:jwt-pop";
     private final String assertion;
     private final Callable<String> assertionProvider;
     private final Function<AssertionRequestOptions, String> contextAwareAssertionProvider;
+    private final Function<AssertionRequestOptions, AssertionResponse> contextAwareResponseProvider;
 
     /**
      * Constructor that accepts a static assertion string
@@ -28,6 +30,7 @@ final class ClientAssertion implements IClientAssertion {
         this.assertion = assertion;
         this.assertionProvider = null;
         this.contextAwareAssertionProvider = null;
+        this.contextAwareResponseProvider = null;
     }
 
     /**
@@ -44,6 +47,7 @@ final class ClientAssertion implements IClientAssertion {
         this.assertion = null;
         this.assertionProvider = assertionProvider;
         this.contextAwareAssertionProvider = null;
+        this.contextAwareResponseProvider = null;
     }
 
     /**
@@ -62,6 +66,27 @@ final class ClientAssertion implements IClientAssertion {
         this.assertion = null;
         this.assertionProvider = null;
         this.contextAwareAssertionProvider = contextAwareAssertionProvider;
+        this.contextAwareResponseProvider = null;
+    }
+
+    /**
+     * Constructor that accepts a context-aware function returning an {@link AssertionResponse}.
+     * This allows the callback to supply both the assertion JWT and an optional token-binding
+     * certificate for mTLS PoP scenarios.
+     *
+     * @param contextAwareResponseProvider A function that receives context and returns an AssertionResponse
+     * @throws NullPointerException if contextAwareResponseProvider is null
+     */
+    ClientAssertion(final Function<AssertionRequestOptions, AssertionResponse> contextAwareResponseProvider,
+                    boolean responseProvider) {
+        if (contextAwareResponseProvider == null) {
+            throw new NullPointerException("contextAwareResponseProvider");
+        }
+
+        this.assertion = null;
+        this.assertionProvider = null;
+        this.contextAwareAssertionProvider = null;
+        this.contextAwareResponseProvider = contextAwareResponseProvider;
     }
 
     /**
@@ -74,6 +99,11 @@ final class ClientAssertion implements IClientAssertion {
      * @throws MsalClientException if the assertion provider returns null/empty or throws an exception
      */
     public String assertion() {
+        if (contextAwareResponseProvider != null) {
+            AssertionResponse response = assertionResponse(new AssertionRequestOptions(null, null, null));
+            return response.assertion();
+        }
+
         if (contextAwareAssertionProvider != null) {
             return assertion(new AssertionRequestOptions(null, null, null));
         }
@@ -95,6 +125,11 @@ public String assertion() {
      * @throws MsalClientException if the assertion provider returns null/empty or throws an exception
      */
     String assertion(AssertionRequestOptions options) {
+        if (contextAwareResponseProvider != null) {
+            AssertionResponse response = assertionResponse(options);
+            return response.assertion();
+        }
+
         if (contextAwareAssertionProvider != null) {
             try {
                 String generatedAssertion = contextAwareAssertionProvider.apply(options);
@@ -118,10 +153,40 @@ String assertion(AssertionRequestOptions options) {
     }
 
     /**
-     * Returns true if this assertion uses a context-aware provider.
+     * Gets the full AssertionResponse from the context-aware response provider.
+     * Returns null if this ClientAssertion does not use a response provider.
+     *
+     * @param options context information for the assertion request
+     * @return An AssertionResponse, or null if not using a response provider
+     * @throws MsalClientException if the provider returns null or throws an exception
+     */
+    AssertionResponse assertionResponse(AssertionRequestOptions options) {
+        if (contextAwareResponseProvider == null) {
+            return null;
+        }
+
+        try {
+            AssertionResponse response = contextAwareResponseProvider.apply(options);
+
+            if (response == null || StringHelper.isBlank(response.assertion())) {
+                throw new MsalClientException(
+                    "Assertion provider returned null or empty assertion",
+                    AuthenticationErrorCode.INVALID_JWT);
+            }
+
+            return response;
+        } catch (MsalClientException ex) {
+            throw ex;
+        } catch (Exception ex) {
+            throw new MsalClientException(ex);
+        }
+    }
+
+    /**
+     * Returns true if this assertion uses a context-aware provider (either string or response).
      */
     boolean isContextAware() {
-        return contextAwareAssertionProvider != null;
+        return contextAwareAssertionProvider != null || contextAwareResponseProvider != null;
     }
 
     private String invokeCallable() {
@@ -151,6 +216,11 @@ public boolean equals(Object o) {
 
         ClientAssertion other = (ClientAssertion) o;
 
+        // For context-aware response providers, we consider them equal if they're the same object
+        if (this.contextAwareResponseProvider != null && other.contextAwareResponseProvider != null) {
+            return this.contextAwareResponseProvider == other.contextAwareResponseProvider;
+        }
+
         // For context-aware providers, we consider them equal if they're the same object
         if (this.contextAwareAssertionProvider != null && other.contextAwareAssertionProvider != null) {
             return this.contextAwareAssertionProvider == other.contextAwareAssertionProvider;
@@ -167,6 +237,11 @@ public boolean equals(Object o) {
 
     @Override
     public int hashCode() {
+        // For context-aware response providers, use the provider's identity hash code
+        if (contextAwareResponseProvider != null) {
+            return System.identityHashCode(contextAwareResponseProvider);
+        }
+
         // For context-aware providers, use the provider's identity hash code
         if (contextAwareAssertionProvider != null) {
             return System.identityHashCode(contextAwareAssertionProvider);
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
@@ -125,4 +125,29 @@ public static IClientAssertion createFromCallback(Function<AssertionRequestOptio
 
         return new ClientAssertion(assertionProvider);
     }
+
+    /**
+     * Static method to create a {@link ClientAssertion} instance from a provided Function that
+     * receives {@link AssertionRequestOptions} context and returns an {@link AssertionResponse}.
+     * This overload allows the callback to supply both the assertion JWT and an optional
+     * token-binding certificate for mTLS PoP scenarios.
+     *
+     * <p>When the returned {@link AssertionResponse} includes a
+     * {@link AssertionResponse#tokenBindingCertificate()}, MSAL uses
+     * {@code client_assertion_type=urn:ietf:params:oauth:client-assertion-type:jwt-pop}
+     * instead of the default {@code jwt-bearer}.</p>
+     *
+     * @param assertionProvider Function that receives {@link AssertionRequestOptions} and produces
+     *                          an {@link AssertionResponse} containing the assertion and optional certificate
+     * @return {@link ClientAssertion} that will invoke the function each time assertion() is called
+     * @throws NullPointerException if assertionProvider is null
+     */
+    public static IClientAssertion createFromAssertionResponseCallback(
+            Function<AssertionRequestOptions, AssertionResponse> assertionProvider) {
+        if (assertionProvider == null) {
+            throw new NullPointerException("assertionProvider");
+        }
+
+        return new ClientAssertion(assertionProvider, true);
+    }
 }
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
@@ -138,7 +138,6 @@ private void addCredentialToRequest(Map<String, String> queryParameters,
         } else if (credentialToUse instanceof ClientAssertion) {
             // For client assertion, add client_assertion and client_assertion_type parameters
             ClientAssertion clientAssertion = (ClientAssertion) credentialToUse;
-            String assertionValue;
             if (clientAssertion.isContextAware()) {
                 // Build assertion context with fmi_path if available
                 String fmiPath = null;
@@ -156,11 +155,17 @@ private void addCredentialToRequest(Map<String, String> queryParameters,
                         application.clientId(),
                         tokenEndpoint,
                         fmiPath);
-                assertionValue = clientAssertion.assertion(options);
+
+                // Try to get the full AssertionResponse first
+                AssertionResponse response = clientAssertion.assertionResponse(options);
+                if (response != null) {
+                    addAssertionResponseParams(queryParameters, response);
+                } else {
+                    addJWTBearerAssertionParams(queryParameters, clientAssertion.assertion(options));
+                }
             } else {
-                assertionValue = clientAssertion.assertion();
+                addJWTBearerAssertionParams(queryParameters, clientAssertion.assertion());
             }
-            addJWTBearerAssertionParams(queryParameters, assertionValue);
         } else if (credentialToUse instanceof ClientCertificate) {
             // For client certificate, generate a new assertion and add it to the request
             ClientCertificate certificate = (ClientCertificate) credentialToUse;
@@ -183,6 +188,23 @@ private void addJWTBearerAssertionParams(Map<String, String> queryParameters, St
         queryParameters.put("client_assertion_type", ClientAssertion.ASSERTION_TYPE_JWT_BEARER);
     }
 
+    /**
+     * Adds assertion parameters from an AssertionResponse, using jwt-pop assertion type
+     * when a token-binding certificate is present, or jwt-bearer otherwise.
+     *
+     * @param queryParameters The map of query parameters to add to
+     * @param response The AssertionResponse containing the assertion and optional certificate
+     */
+    private void addAssertionResponseParams(Map<String, String> queryParameters, AssertionResponse response) {
+        queryParameters.put("client_assertion", response.assertion());
+
+        if (response.tokenBindingCertificate() != null) {
+            queryParameters.put("client_assertion_type", ClientAssertion.ASSERTION_TYPE_JWT_POP);
+        } else {
+            queryParameters.put("client_assertion_type", ClientAssertion.ASSERTION_TYPE_JWT_BEARER);
+        }
+    }
+
     private AuthenticationResult createAuthenticationResultFromOauthHttpResponse(HttpResponse oauthHttpResponse) {
         AuthenticationResult result;
 
diff --git a/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/FmiTest.java b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/FmiTest.java
