Add claimsFromClient API for client-originated claims

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractManagedIdentitySource.java

@@ -9,13 +9,19 @@
 import java.net.HttpURLConnection;
 import java.net.SocketException;
 import java.net.URISyntaxException;
+import java.util.HashMap;
+import java.util.Map;
 
 //base class for all sources that support managed identity
 abstract class AbstractManagedIdentitySource {
 
     private static final Logger LOG = LoggerFactory.getLogger(AbstractManagedIdentitySource.class);
     private static final String MANAGED_IDENTITY_NO_RESPONSE_RECEIVED = "[Managed Identity] Authentication unavailable. No response received from the managed identity endpoint.";
 
+    // IMDS (MSIv1) only supports this single custom claim. Any other top-level key causes IMDS to
+    // return HTTP 400 with no useful diagnostic, so it is rejected client-side before the network call.
+    private static final String XMS_AZ_NWPERIMID = "xms_az_nwperimid";
+
     protected final ManagedIdentityRequest managedIdentityRequest;
     protected final ServiceBundle serviceBundle;
     ManagedIdentitySourceType managedIdentitySourceType;
@@ -42,6 +48,7 @@ public ManagedIdentityResponse getManagedIdentityResponse(
 
         createManagedIdentityRequest(parameters.resource);
         managedIdentityRequest.addTokenRevocationParametersToQuery(parameters);
+        addClientClaimsToRequest(parameters);
         IHttpResponse response;
 
         try {
@@ -62,6 +69,59 @@ public ManagedIdentityResponse getManagedIdentityResponse(
         return handleResponse(parameters, response);
     }
 
+    /**
+     * Forwards client-originated claims (set via
+     * {@link ManagedIdentityParameters.ManagedIdentityParametersBuilder#claimsFromClient(String)}) to
+     * the managed identity endpoint. Only IMDS-based managed identity is supported; other sources fail
+     * fast rather than silently dropping the value (which would also pollute the cache with a key the
+     * endpoint never saw). For IMDS (a GET request) the claims are added as a query parameter; for any
+     * POST-based source they would be added to the body.
+     */
+    private void addClientClaimsToRequest(ManagedIdentityParameters parameters) {
+        if (StringHelper.isNullOrBlank(parameters.clientClaims)) {
+            return;
+        }
+
+        if (managedIdentitySourceType != ManagedIdentitySourceType.IMDS) {
+            throw new MsalClientException(
+                    String.format("claimsFromClient is only supported for IMDS-based managed identity sources. "
+                            + "The detected source is %s.", managedIdentitySourceType),
+                    AuthenticationErrorCode.INVALID_REQUEST);
+        }
+
+        // IMDS == MSIv1: validate that the only top-level claim key is xms_az_nwperimid.
+        validateMsiv1Claims(parameters.clientClaims);
+
+        if (managedIdentityRequest.method == HttpMethod.GET) {
+            if (managedIdentityRequest.queryParameters == null) {
+                managedIdentityRequest.queryParameters = new HashMap<>();
+            }
+            // The value is URL-encoded later by StringHelper.serializeQueryParameters.
+            managedIdentityRequest.queryParameters.put("claims", parameters.clientClaims);
+            LOG.info("[Managed Identity] Adding client claims to IMDS request as query parameter.");
+        } else {
+            if (managedIdentityRequest.bodyParameters == null) {
+                managedIdentityRequest.bodyParameters = new HashMap<>();
+            }
+            managedIdentityRequest.bodyParameters.put("claims", parameters.clientClaims);
+            LOG.info("[Managed Identity] Adding client claims to request body.");
+        }
+    }
+
+    private static void validateMsiv1Claims(String claimsJson) {
+        Map<String, Object> parsed = JsonHelper.parseJsonToMap(claimsJson);
+        for (String key : parsed.keySet()) {
+            if (!XMS_AZ_NWPERIMID.equals(key)) {
+                throw new MsalClientException(
+                        String.format("MSIv1 (IMDS v1) only supports the `%s` custom claim. "
+                                + "The claims JSON contained the unsupported key `%s`. "
+                                + "Remove all keys other than `%s` when using claimsFromClient with MSIv1.",
+                                XMS_AZ_NWPERIMID, key, XMS_AZ_NWPERIMID),
+                        AuthenticationErrorCode.INVALID_REQUEST);
+            }
+        }
+    }
+
     public ManagedIdentityResponse handleResponse(
             ManagedIdentityParameters parameters,
             IHttpResponse response) {

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenByManagedIdentitySupplier.java

@@ -67,6 +67,12 @@ AuthenticationResult execute() throws Exception {
                     context,
                     null);
 
+            // Propagate ext_cache_key_hash for cache isolation (e.g., client_claims)
+            String extCacheKeyHash = managedIdentityParameters.computeExtCacheKeyHash();
+            if (!StringHelper.isBlank(extCacheKeyHash)) {
+                silentRequest.extCacheKeyHash(extCacheKeyHash);
+            }
+
             AcquireTokenSilentSupplier supplier = new AcquireTokenSilentSupplier(
                     this.clientApplication,
                     silentRequest);

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenByOnBehalfOfSupplier.java

@@ -40,6 +40,12 @@ AuthenticationResult execute() throws Exception {
                         context,
                         onBehalfOfRequest.parameters.userAssertion());
 
+                // Propagate ext_cache_key_hash for cache isolation (e.g., client_claims)
+                String extCacheKeyHash = this.onBehalfOfRequest.parameters.computeExtCacheKeyHash();
+                if (!StringHelper.isBlank(extCacheKeyHash)) {
+                    silentRequest.extCacheKeyHash(extCacheKeyHash);
+                }
+
                 AcquireTokenSilentSupplier supplier = new AcquireTokenSilentSupplier(
                         this.clientApplication,
                         silentRequest);

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenByUserFederatedIdentityCredentialSupplier.java

@@ -48,6 +48,14 @@ AuthenticationResult execute() throws Exception {
                             context,
                             null);
 
+                    // Propagate ext_cache_key_hash for cache isolation (e.g., client_claims).
+                    // User-FIC tokens are account-scoped, so the user-token read path in TokenCache
+                    // must also filter on this hash for the isolation to take effect.
+                    String extCacheKeyHash = this.userFicRequest.parameters.computeExtCacheKeyHash();
+                    if (!StringHelper.isBlank(extCacheKeyHash)) {
+                        silentRequest.extCacheKeyHash(extCacheKeyHash);
+                    }
+
                     AcquireTokenSilentSupplier supplier = new AcquireTokenSilentSupplier(
                             this.clientApplication,
                             silentRequest);

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenSilentSupplier.java

@@ -39,7 +39,8 @@ AuthenticationResult execute() throws Exception {
                     silentRequest.parameters().account(),
                     requestAuthority,
                     silentRequest.parameters().scopes(),
-                    clientApplication.clientId());
+                    clientApplication.clientId(),
+                    silentRequest.extCacheKeyHash());
 
             if (res == null) {
                 throw new MsalClientException(AuthenticationErrorMessage.NO_TOKEN_IN_CACHE, AuthenticationErrorCode.CACHE_MISS);

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AuthenticationErrorCode.java

@@ -166,4 +166,11 @@ public class AuthenticationErrorCode {
      * This is returned by the instance discovery endpoint when the provided authority host is unknown.
      */
     public static final String INVALID_INSTANCE = "invalid_instance";
+
+    /**
+     * Indicates that the request is malformed or uses an unsupported parameter combination, for
+     * example when client-originated claims are supplied to a managed identity source that does not
+     * support them, or when an unsupported claim key is used.
+     */
+    public static final String INVALID_REQUEST = "invalid_request";
 }

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AuthorizationCodeParameters.java

@@ -6,6 +6,8 @@
 import java.net.URI;
 import java.util.Map;
 import java.util.Set;
+import java.util.SortedMap;
+import java.util.TreeMap;
 
 import static com.microsoft.aad.msal4j.ParameterValidationUtils.validateNotBlank;
 import static com.microsoft.aad.msal4j.ParameterValidationUtils.validateNotNull;
@@ -33,10 +35,20 @@ public class AuthorizationCodeParameters implements IAcquireTokenParameters {
 
     private String tenant;
 
+    private String clientClaims;
+
+    // Generic extended cache key components. The hash of these components isolates cache
+    // entries so that requests with different client-claims values do not collide.
+    private SortedMap<String, String> cacheKeyComponents;
+
+    // Memoized hash of cacheKeyComponents (computed once since parameters are immutable).
+    private String extCacheKeyHashCache;
+
     private AuthorizationCodeParameters(String authorizationCode, URI redirectUri,
                                         Set<String> scopes, ClaimsRequest claims,
                                         String codeVerifier, Map<String, String> extraHttpHeaders,
-                                        Map<String, String> extraQueryParameters, String tenant) {
+                                        Map<String, String> extraQueryParameters, String tenant,
+                                        String clientClaims) {
         this.authorizationCode = authorizationCode;
         this.redirectUri = redirectUri;
         this.scopes = scopes;
@@ -45,6 +57,10 @@ private AuthorizationCodeParameters(String authorizationCode, URI redirectUri,
         this.extraHttpHeaders = extraHttpHeaders;
         this.extraQueryParameters = extraQueryParameters;
         this.tenant = tenant;
+        this.clientClaims = clientClaims;
+
+        // Build cache key components from any parameters that require cache isolation.
+        this.cacheKeyComponents = buildCacheKeyComponents();
     }
 
     private static AuthorizationCodeParametersBuilder builder() {
@@ -104,6 +120,50 @@ public String tenant() {
         return this.tenant;
     }
 
+    /**
+     * Client-originated claims set via {@link AuthorizationCodeParametersBuilder#claimsFromClient(String)}.
+     * Forwarded to the token endpoint as the OAuth {@code claims} parameter and used as part of the
+     * extended cache key so that distinct claim values are cached separately.
+     */
+    @Override
+    public String clientClaims() {
+        return this.clientClaims;
+    }
+
+    /**
+     * Builds the sorted map of cache key components from the parameters that require cache isolation.
+     * Returns null if no components are present.
+     */
+    private SortedMap<String, String> buildCacheKeyComponents() {
+        TreeMap<String, String> components = null;
+        if (!StringHelper.isBlank(clientClaims)) {
+            components = new TreeMap<>();
+            components.put("client_claims", clientClaims);
+        }
+        return components;
+    }
+
+    /**
+     * Returns the extended cache key components for this request, if any.
+     * Used by {@link TokenCache} for both cache writes and reads.
+     */
+    SortedMap<String, String> cacheKeyComponents() {
+        return this.cacheKeyComponents;
+    }
+
+    /**
+     * Computes the extended cache key hash from all cache key components, or an empty string when
+     * there are none. The result is memoized since the parameters are immutable after construction.
+     */
+    @Override
+    public String computeExtCacheKeyHash() {
+        if (extCacheKeyHashCache != null) {
+            return extCacheKeyHashCache;
+        }
+        extCacheKeyHashCache = StringHelper.computeExtCacheKeyHash(cacheKeyComponents);
+        return extCacheKeyHashCache;
+    }
+
     public static class AuthorizationCodeParametersBuilder {
         private String authorizationCode;
         private URI redirectUri;
@@ -113,6 +173,7 @@ public static class AuthorizationCodeParametersBuilder {
         private Map<String, String> extraHttpHeaders;
         private Map<String, String> extraQueryParameters;
         private String tenant;
+        private String clientClaims;
 
         AuthorizationCodeParametersBuilder() {
         }
@@ -193,8 +254,29 @@ public AuthorizationCodeParametersBuilder tenant(String tenant) {
             return this;
         }
 
+        /**
+         * Specifies client-originated claims (a raw JSON object string) to forward to the token
+         * endpoint as the OAuth {@code claims} request parameter. Unlike {@link #claims(ClaimsRequest)}
+         * (server-issued claims challenges, which bypass the cache), tokens acquired with client claims
+         * are cached and the cache entry is keyed on the claims value, so distinct claim values produce
+         * separate cache entries. Use stable, non-dynamic values to avoid cache fragmentation.
+         * A blank value is ignored; an invalid JSON object throws {@link MsalClientException}.
+         *
+         * @param claimsJson a valid JSON object string containing the client claims
+         * @return this builder instance
+         */
+        public AuthorizationCodeParametersBuilder claimsFromClient(String claimsJson) {
+            if (StringHelper.isBlank(claimsJson)) {
+                return this;
+            }
+
+            JsonHelper.validateJsonObjectFormat(claimsJson);
+            this.clientClaims = claimsJson;
+            return this;
+        }
+
         public AuthorizationCodeParameters build() {
-            return new AuthorizationCodeParameters(this.authorizationCode, this.redirectUri, this.scopes, this.claims, this.codeVerifier, this.extraHttpHeaders, this.extraQueryParameters, this.tenant);
+            return new AuthorizationCodeParameters(this.authorizationCode, this.redirectUri, this.scopes, this.claims, this.codeVerifier, this.extraHttpHeaders, this.extraQueryParameters, this.tenant, this.clientClaims);
         }
 
         public String toString() {

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ClientCredentialParameters.java

@@ -32,6 +32,8 @@ public class ClientCredentialParameters implements IAcquireTokenParameters {
 
     private String fmiPath;
 
+    private String clientClaims;
+
     // Generic extended cache key components. Any optional or flow-specific parameters 
     // that should influence token cache isolation adds an entry here. The hash of these
     // components is used as part of the cache key in relevant scenarios entries.
@@ -40,7 +42,7 @@ public class ClientCredentialParameters implements IAcquireTokenParameters {
     // Memoized hash of cacheKeyComponents (computed once since parameters are immutable).
     private String extCacheKeyHashCache;
 
-    private ClientCredentialParameters(Set<String> scopes, Boolean skipCache, ClaimsRequest claims, Map<String, String> extraHttpHeaders, Map<String, String> extraQueryParameters, String tenant, IClientCredential clientCredential, String fmiPath) {
+    private ClientCredentialParameters(Set<String> scopes, Boolean skipCache, ClaimsRequest claims, Map<String, String> extraHttpHeaders, Map<String, String> extraQueryParameters, String tenant, IClientCredential clientCredential, String fmiPath, String clientClaims) {
         this.scopes = scopes;
         this.skipCache = skipCache;
         this.claims = claims;
@@ -49,6 +51,7 @@ private ClientCredentialParameters(Set<String> scopes, Boolean skipCache, Claims
         this.tenant = tenant;
         this.clientCredential = clientCredential;
         this.fmiPath = fmiPath;
+        this.clientClaims = clientClaims;
 
         // Build cache key components from any parameters that require cache isolation.
         this.cacheKeyComponents = buildCacheKeyComponents();
@@ -114,6 +117,16 @@ public String fmiPath() {
         return this.fmiPath;
     }
 
+    /**
+     * Client-originated claims set via {@link ClientCredentialParametersBuilder#claimsFromClient(String)}.
+     * Forwarded to the token endpoint as the OAuth {@code claims} parameter and used as part of the
+     * extended cache key so that distinct claim values are cached separately.
+     */
+    @Override
+    public String clientClaims() {
+        return this.clientClaims;
+    }
+
     /**
      * Builds the sorted map of cache key components from the parameters that require
      * cache isolation. Returns null if no components are present.
@@ -127,6 +140,12 @@ private SortedMap<String, String> buildCacheKeyComponents() {
             components = new TreeMap<>();
             components.put("fmi_path", fmiPath);
         }
+        if (!StringHelper.isBlank(clientClaims)) {
+            if (components == null) {
+                components = new TreeMap<>();
+            }
+            components.put("client_claims", clientClaims);
+        }
         return components;
     }
 
@@ -145,7 +164,8 @@ SortedMap<String, String> cacheKeyComponents() {
      * The result is memoized since ClientCredentialParameters is immutable after construction.
      * Used by both cache writes ({@link TokenCache}) and cache reads (silent lookup).
      */
-    String computeExtCacheKeyHash() {
+    @Override
+    public String computeExtCacheKeyHash() {
         if (extCacheKeyHashCache != null) {
             return extCacheKeyHashCache;
         }
@@ -162,6 +182,7 @@ public static class ClientCredentialParametersBuilder {
         private String tenant;
         private IClientCredential clientCredential;
         private String fmiPath;
+        private String clientClaims;
 
         ClientCredentialParametersBuilder() {
         }
@@ -245,8 +266,29 @@ public ClientCredentialParametersBuilder fmiPath(String fmiPath) {
             return this;
         }
 
+        /**
+         * Specifies client-originated claims (a raw JSON object string) to forward to the token
+         * endpoint as the OAuth {@code claims} request parameter. Unlike {@link #claims(ClaimsRequest)}
+         * (server-issued claims challenges, which bypass the cache), tokens acquired with client claims
+         * are cached and the cache entry is keyed on the claims value, so distinct claim values produce
+         * separate cache entries. Use stable, non-dynamic values to avoid cache fragmentation.
+         * A blank value is ignored; an invalid JSON object throws {@link MsalClientException}.
+         *
+         * @param claimsJson a valid JSON object string containing the client claims
+         * @return builder that can be used to construct ClientCredentialParameters
+         */
+        public ClientCredentialParametersBuilder claimsFromClient(String claimsJson) {
+            if (StringHelper.isBlank(claimsJson)) {
+                return this;
+            }
+
+            JsonHelper.validateJsonObjectFormat(claimsJson);
+            this.clientClaims = claimsJson;
+            return this;
+        }
+
         public ClientCredentialParameters build() {
-            return new ClientCredentialParameters(this.scopes, this.skipCache, this.claims, this.extraHttpHeaders, this.extraQueryParameters, this.tenant, this.clientCredential, this.fmiPath);
+            return new ClientCredentialParameters(this.scopes, this.skipCache, this.claims, this.extraHttpHeaders, this.extraQueryParameters, this.tenant, this.clientCredential, this.fmiPath, this.clientClaims);
         }
 
         public String toString() {