From 9752d9943235ed97880ab8095b67eb315390fac0 Mon Sep 17 00:00:00 2001 From: avdunn Date: Wed, 15 Apr 2026 14:34:42 -0700 Subject: [PATCH 1/2] First draft agent ID support --- .../aad/msal4j/AbstractApplicationBase.java | 4 + .../AcquireTokenForAgentParameters.java | 176 +++++++++ .../msal4j/AcquireTokenForAgentRequest.java | 21 + .../msal4j/AcquireTokenForAgentSupplier.java | 282 +++++++++++++ .../microsoft/aad/msal4j/AgentIdentity.java | 102 +++++ .../msal4j/ConfidentialClientApplication.java | 25 ++ .../IConfidentialClientApplication.java | 14 + .../com/microsoft/aad/msal4j/PublicApi.java | 1 + .../aad/msal4j/AcquireTokenForAgentTest.java | 374 ++++++++++++++++++ .../aad/msal4j/AgentIdentityTest.java | 140 +++++++ 10 files changed, 1139 insertions(+) create mode 100644 msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentParameters.java create mode 100644 msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentRequest.java create mode 100644 msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentSupplier.java create mode 100644 msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AgentIdentity.java create mode 100644 msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AcquireTokenForAgentTest.java create mode 100644 msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AgentIdentityTest.java diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractApplicationBase.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractApplicationBase.java index 5cc20f1a..19064264 100644 --- a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractApplicationBase.java +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractApplicationBase.java @@ -115,6 +115,10 @@ private AuthenticationResultSupplier getAuthenticationResultSupplier(MsalRequest supplier = new AcquireTokenByUserFederatedIdentityCredentialSupplier( (ConfidentialClientApplication) this, (UserFederatedIdentityCredentialRequest) msalRequest); + } else if (msalRequest instanceof AcquireTokenForAgentRequest) { + supplier = new AcquireTokenForAgentSupplier( + (ConfidentialClientApplication) this, + (AcquireTokenForAgentRequest) msalRequest); } else if (msalRequest instanceof ManagedIdentityRequest) { supplier = new AcquireTokenByManagedIdentitySupplier( (ManagedIdentityApplication) this, diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentParameters.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentParameters.java new file mode 100644 index 00000000..4dad1a69 --- /dev/null +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentParameters.java @@ -0,0 +1,176 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.microsoft.aad.msal4j; + +import java.util.Map; +import java.util.Set; + +import static com.microsoft.aad.msal4j.ParameterValidationUtils.validateNotNull; + +/** + * Object containing parameters for the composite agent token acquisition flow. + * This orchestrates the full three-leg FMI/FIC token exchange: the developer passes + * scopes and an {@link AgentIdentity}, and MSAL handles Legs 1-3 internally. + *

+ * Can be used as parameter to + * {@link ConfidentialClientApplication#acquireTokenForAgent(AcquireTokenForAgentParameters)} + */ +public class AcquireTokenForAgentParameters implements IAcquireTokenParameters { + + private Set scopes; + private AgentIdentity agentIdentity; + private boolean forceRefresh; + private ClaimsRequest claims; + private Map extraHttpHeaders; + private Map extraQueryParameters; + private String tenant; + + private AcquireTokenForAgentParameters( + Set scopes, + AgentIdentity agentIdentity, + boolean forceRefresh, + ClaimsRequest claims, + Map extraHttpHeaders, + Map extraQueryParameters, + String tenant) { + this.scopes = scopes; + this.agentIdentity = agentIdentity; + this.forceRefresh = forceRefresh; + this.claims = claims; + this.extraHttpHeaders = extraHttpHeaders; + this.extraQueryParameters = extraQueryParameters; + this.tenant = tenant; + } + + /** + * Builder for {@link AcquireTokenForAgentParameters}. + * + * @param scopes scopes application is requesting access to + * @param agentIdentity the identity of the agent and (optionally) the target user + * @return builder that can be used to construct AcquireTokenForAgentParameters + */ + public static AcquireTokenForAgentParametersBuilder builder( + Set scopes, AgentIdentity agentIdentity) { + validateNotNull("scopes", scopes); + validateNotNull("agentIdentity", agentIdentity); + + return new AcquireTokenForAgentParametersBuilder() + .scopes(scopes) + .agentIdentity(agentIdentity); + } + + public Set scopes() { + return this.scopes; + } + + public AgentIdentity agentIdentity() { + return this.agentIdentity; + } + + public boolean forceRefresh() { + return this.forceRefresh; + } + + public ClaimsRequest claims() { + return this.claims; + } + + public Map extraHttpHeaders() { + return this.extraHttpHeaders; + } + + public Map extraQueryParameters() { + return this.extraQueryParameters; + } + + public String tenant() { + return this.tenant; + } + + public static class AcquireTokenForAgentParametersBuilder { + private Set scopes; + private AgentIdentity agentIdentity; + private boolean forceRefresh; + private ClaimsRequest claims; + private Map extraHttpHeaders; + private Map extraQueryParameters; + private String tenant; + + AcquireTokenForAgentParametersBuilder() { + } + + AcquireTokenForAgentParametersBuilder scopes(Set scopes) { + this.scopes = scopes; + return this; + } + + AcquireTokenForAgentParametersBuilder agentIdentity(AgentIdentity agentIdentity) { + this.agentIdentity = agentIdentity; + return this; + } + + /** + * If true, the request will ignore cached access tokens on read, but will still write + * them to the cache once obtained from the identity provider. The default is false. + * + * @param forceRefresh whether to bypass the user token cache + * @return this builder + */ + public AcquireTokenForAgentParametersBuilder forceRefresh(boolean forceRefresh) { + this.forceRefresh = forceRefresh; + return this; + } + + /** + * Claims to be requested through the OIDC claims request parameter, allowing requests + * for standard and custom claims. + * + * @param claims {@link ClaimsRequest} + * @return this builder + */ + public AcquireTokenForAgentParametersBuilder claims(ClaimsRequest claims) { + this.claims = claims; + return this; + } + + /** + * Adds additional headers to the token request. + * + * @param extraHttpHeaders headers to include + * @return this builder + */ + public AcquireTokenForAgentParametersBuilder extraHttpHeaders(Map extraHttpHeaders) { + this.extraHttpHeaders = extraHttpHeaders; + return this; + } + + /** + * Adds additional query parameters to the token request. + * + * @param extraQueryParameters query parameters to include + * @return this builder + */ + public AcquireTokenForAgentParametersBuilder extraQueryParameters(Map extraQueryParameters) { + this.extraQueryParameters = extraQueryParameters; + return this; + } + + /** + * Sets the tenant for the request, overriding the application's configured authority. + * + * @param tenant tenant ID or domain + * @return this builder + */ + public AcquireTokenForAgentParametersBuilder tenant(String tenant) { + this.tenant = tenant; + return this; + } + + public AcquireTokenForAgentParameters build() { + return new AcquireTokenForAgentParameters( + scopes, agentIdentity, forceRefresh, claims, + extraHttpHeaders, extraQueryParameters, tenant); + } + } +} diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentRequest.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentRequest.java new file mode 100644 index 00000000..e73480af --- /dev/null +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentRequest.java @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.microsoft.aad.msal4j; + +/** + * Internal request class for the composite agent token acquisition flow. + * This request does not create its own grant; actual grants are produced + * by the inner CCA calls orchestrated by {@link AcquireTokenForAgentSupplier}. + */ +class AcquireTokenForAgentRequest extends MsalRequest { + + AcquireTokenForAgentParameters parameters; + + AcquireTokenForAgentRequest(AcquireTokenForAgentParameters parameters, + ConfidentialClientApplication application, + RequestContext requestContext) { + super(application, null, requestContext); + this.parameters = parameters; + } +} diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentSupplier.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentSupplier.java new file mode 100644 index 00000000..47b67edf --- /dev/null +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentSupplier.java @@ -0,0 +1,282 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.microsoft.aad.msal4j; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.net.MalformedURLException; +import java.util.Collections; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; + +/** + * Orchestrates a multi-leg token acquisition for agent scenarios. + *

+ * Two CCA instances are involved: + *

    + *
  1. Blueprint CCA — the developer-created CCA that holds the real credential + * (certificate, secret, etc.). It only participates in Leg 1: acquiring an FMI + * credential via AcquireTokenForClient + WithFmiPath. Its app token cache stores + * the FMI credential.
    2. + *
  2. Agent CCA — an internal CCA keyed by the agent's app ID, created and cached + * by this class. Its client assertion callback delegates to the Blueprint for FMI + * credentials (Leg 1). It handles both Leg 2 (AcquireTokenForClient for the assertion + * token) and Leg 3 (AcquireTokenByUserFederatedIdentityCredential for the user token).
    3. + *
+ *

+ * Caching behavior: + *

+ */ +class AcquireTokenForAgentSupplier extends AuthenticationResultSupplier { + + private static final Logger LOG = LoggerFactory.getLogger(AcquireTokenForAgentSupplier.class); + private static final Set TOKEN_EXCHANGE_SCOPE = + Collections.singleton("api://AzureADTokenExchange/.default"); + private static final String AGENT_CCA_KEY_PREFIX = "agent_"; + + private final AcquireTokenForAgentRequest agentRequest; + private final ConfidentialClientApplication blueprintApplication; + + AcquireTokenForAgentSupplier(ConfidentialClientApplication clientApplication, + AcquireTokenForAgentRequest agentRequest) { + super(clientApplication, agentRequest); + this.agentRequest = agentRequest; + this.blueprintApplication = clientApplication; + } + + @Override + AuthenticationResult execute() throws Exception { + AgentIdentity agentIdentity = agentRequest.parameters.agentIdentity(); + String agentAppId = agentIdentity.agentApplicationId(); + Set callerScopes = agentRequest.parameters.scopes(); + + // Retrieve (or create) the internal Agent CCA for this agent app ID. + ConfidentialClientApplication agentCca = getOrCreateAgentCca(agentAppId); + + if (!agentIdentity.hasUserIdentifier()) { + // App-only flow: AcquireTokenForClient has built-in cache-first logic, + // so no explicit silent pre-check is needed. + LOG.debug("App-only agent flow for agent app ID: {}", agentAppId); + return (AuthenticationResult) joinAndUnwrap( + agentCca.acquireToken( + ClientCredentialParameters.builder(callerScopes).build())); + } + + // --- User identity flow --- + + // Check the Agent CCA's user token cache for a previously-acquired token for this user. + // ForceRefresh skips this check so a fresh user token is always obtained from the network. + if (!agentRequest.parameters.forceRefresh()) { + AuthenticationResult cachedResult = tryAcquireTokenSilent(agentCca, agentIdentity, callerScopes); + if (cachedResult != null) { + LOG.debug("Returning cached user token for agent app ID: {}", agentAppId); + return cachedResult; + } + } + + // Cache miss (or ForceRefresh) — execute Leg 2 + Leg 3. + + // Leg 2: Acquire an assertion token from the Agent CCA's app token cache (or network). + // The Agent CCA's assertion callback will invoke Leg 1 (FMI credential from Blueprint), + // but AcquireTokenForClient's built-in cache handles repeat calls. + LOG.debug("Executing Leg 2 (assertion token) for agent app ID: {}", agentAppId); + IAuthenticationResult assertionResult = joinAndUnwrap( + agentCca.acquireToken( + ClientCredentialParameters.builder(TOKEN_EXCHANGE_SCOPE).build())); + + String assertion = assertionResult.accessToken(); + + // Leg 3: Exchange the assertion for a user-scoped token via UserFIC. + // The result is written to the Agent CCA's user token cache for future silent retrieval. + LOG.debug("Executing Leg 3 (user FIC token) for agent app ID: {}", agentAppId); + UserFederatedIdentityCredentialParameters ficParams; + if (agentIdentity.userObjectId() != null) { + ficParams = UserFederatedIdentityCredentialParameters + .builder(callerScopes, agentIdentity.userObjectId(), assertion) + .forceRefresh(true) // always fetch from network (we already checked the cache above) + .build(); + } else { + ficParams = UserFederatedIdentityCredentialParameters + .builder(callerScopes, agentIdentity.username(), assertion) + .forceRefresh(true) + .build(); + } + + return (AuthenticationResult) joinAndUnwrap(agentCca.acquireToken(ficParams)); + } + + /** + * Searches the Agent CCA's user token cache for a previously-acquired token + * matching the specified user identity (by OID or UPN). + * Returns null if no matching account exists or the cached token is expired. + */ + private AuthenticationResult tryAcquireTokenSilent( + ConfidentialClientApplication agentCca, + AgentIdentity agentIdentity, + Set scopes) { + try { + Set accounts = joinAndUnwrap(agentCca.getAccounts()); + IAccount matchedAccount = findMatchingAccount(accounts, agentIdentity); + if (matchedAccount == null) { + return null; + } + + SilentParameters silentParams = SilentParameters + .builder(scopes, matchedAccount) + .build(); + + return (AuthenticationResult) joinAndUnwrap( + agentCca.acquireTokenSilently(silentParams)); + } catch (Exception ex) { + // Token expired or requires interaction — fall through to full Leg 2 + Leg 3 flow + LOG.debug("Silent token acquisition failed for agent: {}", ex.getMessage()); + return null; + } + } + + /** + * Finds an account in the Agent CCA's cache that matches the user identity. + * Matches by OID (HomeAccountId objectId) if the caller specified a UUID, + * otherwise by UPN (Account.username). Both comparisons are case-insensitive. + */ + private static IAccount findMatchingAccount(Set accounts, AgentIdentity agentIdentity) { + if (agentIdentity.userObjectId() != null) { + String targetOid = agentIdentity.userObjectId().toString(); + return accounts.stream() + .filter(a -> a.homeAccountId() != null && + targetOid.equalsIgnoreCase(extractOid(a.homeAccountId()))) + .findFirst() + .orElse(null); + } + + return accounts.stream() + .filter(a -> agentIdentity.username().equalsIgnoreCase(a.username())) + .findFirst() + .orElse(null); + } + + /** + * Extracts the OID portion from a homeAccountId (format: "oid.tid"). + */ + private static String extractOid(String homeAccountId) { + int dotIndex = homeAccountId.indexOf('.'); + return dotIndex >= 0 ? homeAccountId.substring(0, dotIndex) : homeAccountId; + } + + // ======================================================================== + // Agent CCA Construction and Configuration + // ======================================================================== + + /** + * Retrieves the cached internal Agent CCA for the given agent app ID, or creates one + * if this is the first call. The Agent CCA is stored in the Blueprint's agentCcaCache + * so its app and user token caches persist across calls. + */ + private ConfidentialClientApplication getOrCreateAgentCca(String agentAppId) { + String key = AGENT_CCA_KEY_PREFIX + agentAppId; + return blueprintApplication.agentCcaCache.computeIfAbsent(key, k -> { + try { + return buildAgentCca(agentAppId); + } catch (MalformedURLException e) { + throw new MsalClientException(e); + } + }); + } + + /** + * Builds a new internal Agent CCA configured with: + *
    + *
  • Client ID = the agent's app ID
    • + *
  • Authority = the Blueprint's resolved authority
    • + *
  • Client assertion callback = Leg 1 (FMI credential from Blueprint)
    • + *
  • App-level config = propagated from the Blueprint
    • + *
+ */ + private ConfidentialClientApplication buildAgentCca(String agentAppId) + throws MalformedURLException { + // Capture only the blueprint reference (long-lived) in the assertion callback. + // Do NOT capture 'this' (per-request state) to avoid pinning stale request data. + final ConfidentialClientApplication blueprint = blueprintApplication; + + IClientCredential assertionCredential = ClientCredentialFactory.createFromCallback( + (AssertionRequestOptions opts) -> { + try { + // Leg 1: Acquire an FMI credential from the Blueprint CCA. + // AcquireTokenForClient has built-in cache-first logic — only the + // first call hits the network; subsequent calls return cached credential. + IAuthenticationResult result = joinAndUnwrap( + blueprint.acquireToken( + ClientCredentialParameters.builder(TOKEN_EXCHANGE_SCOPE) + .fmiPath(agentAppId) + .build())); + return result.accessToken(); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new MsalClientException(e); + } + }); + + ConfidentialClientApplication.Builder builder = + ConfidentialClientApplication.builder(agentAppId, assertionCredential) + .authority(blueprint.authority()); + + propagateBlueprintConfig(builder, blueprint); + return builder.build(); + } + + /** + * Propagates app-level configuration from the Blueprint CCA to the Agent CCA builder. + * This ensures the Agent CCA shares the Blueprint's HTTP behavior, logging, instance + * discovery settings, and telemetry identity. + */ + private static void propagateBlueprintConfig( + ConfidentialClientApplication.Builder builder, + ConfidentialClientApplication blueprint) { + // HTTP: share the same HTTP client + if (blueprint.httpClient() != null) { + builder.httpClient(blueprint.httpClient()); + } + + // Logging + builder.logPii(blueprint.logPii()); + + // Instance discovery: honor the Blueprint's settings + builder.instanceDiscovery(blueprint.instanceDiscovery()); + builder.validateAuthority(blueprint.validateAuthority()); + + // Telemetry: attribute network calls to the same caller + if (blueprint.applicationName() != null) { + builder.applicationName(blueprint.applicationName()); + } + if (blueprint.applicationVersion() != null) { + builder.applicationVersion(blueprint.applicationVersion()); + } + } + + /** + * Calls {@link CompletableFuture#join()} and unwraps any {@link CompletionException} + * so the original exception propagates with its correct type. + */ + private static T joinAndUnwrap(CompletableFuture future) throws Exception { + try { + return future.join(); + } catch (CompletionException e) { + Throwable cause = e.getCause(); + if (cause instanceof Exception) { + throw (Exception) cause; + } + throw e; + } + } +} diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AgentIdentity.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AgentIdentity.java new file mode 100644 index 00000000..81f71dfe --- /dev/null +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AgentIdentity.java @@ -0,0 +1,102 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.microsoft.aad.msal4j; + +import java.util.UUID; + +import static com.microsoft.aad.msal4j.ParameterValidationUtils.validateNotBlank; +import static com.microsoft.aad.msal4j.ParameterValidationUtils.validateNotNull; + +/** + * Represents the identity of an agent application and the user it acts on behalf of. + * Used with {@link IConfidentialClientApplication#acquireTokenForAgent(AcquireTokenForAgentParameters)} + * to acquire tokens for agent scenarios using Federated Managed Identity (FMI) and + * User Federated Identity Credentials (UserFIC). + */ +public final class AgentIdentity { + + private final String agentApplicationId; + private UUID userObjectId; + private String username; + + private AgentIdentity(String agentApplicationId) { + validateNotBlank("agentApplicationId", agentApplicationId); + this.agentApplicationId = agentApplicationId; + } + + /** + * Creates an {@link AgentIdentity} that identifies the user by their object ID (OID). + * This is the recommended approach for identifying users in agent scenarios. + * + * @param agentApplicationId the client ID of the agent application + * @param userObjectId the object ID (OID) of the user the agent acts on behalf of + */ + public AgentIdentity(String agentApplicationId, UUID userObjectId) { + this(agentApplicationId); + validateNotNull("userObjectId", userObjectId); + this.userObjectId = userObjectId; + } + + /** + * Creates an {@link AgentIdentity} that identifies the user by their UPN (User Principal Name). + * + * @param agentApplicationId the client ID of the agent application + * @param username the UPN of the user the agent acts on behalf of + * @return an {@link AgentIdentity} configured with the user's UPN + */ + public static AgentIdentity withUsername(String agentApplicationId, String username) { + validateNotBlank("username", username); + AgentIdentity identity = new AgentIdentity(agentApplicationId); + identity.username = username; + return identity; + } + + /** + * Creates an {@link AgentIdentity} for app-only (no user) scenarios, where only Legs 1-2 + * of the agent token acquisition are performed. + * + * @param agentApplicationId the client ID of the agent application + * @return an {@link AgentIdentity} configured for app-only access + */ + public static AgentIdentity appOnly(String agentApplicationId) { + return new AgentIdentity(agentApplicationId); + } + + /** + * Gets the client ID of the agent application. + * + * @return the agent application's client ID + */ + public String agentApplicationId() { + return agentApplicationId; + } + + /** + * Gets the object ID (OID) of the user, if specified. + * + * @return the user's OID, or null if not specified + */ + public UUID userObjectId() { + return userObjectId; + } + + /** + * Gets the UPN of the user, if specified. + * + * @return the user's UPN, or null if not specified + */ + public String username() { + return username; + } + + /** + * Returns whether this identity includes a user identifier (OID or UPN). + * When false, the agent flow is app-only (Legs 1-2 only). + * + * @return true if a user OID or UPN is present + */ + public boolean hasUserIdentifier() { + return userObjectId != null || !StringHelper.isBlank(username); + } +} diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ConfidentialClientApplication.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ConfidentialClientApplication.java index 50df4466..d4824db3 100644 --- a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ConfidentialClientApplication.java +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ConfidentialClientApplication.java @@ -6,6 +6,7 @@ import org.slf4j.LoggerFactory; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Function; import static com.microsoft.aad.msal4j.ParameterValidationUtils.validateNotNull; @@ -22,6 +23,12 @@ public class ConfidentialClientApplication extends AbstractClientApplicationBase IClientCredential clientCredential; private boolean sendX5c; + /** + * Cache of internal Agent CCA instances keyed by "agent_" + agentAppId. + * Each agent CCA has its own app and user token caches that persist across calls. + */ + final ConcurrentHashMap agentCcaCache = new ConcurrentHashMap<>(); + /** AppTokenProvider creates a Credential from a function that provides access tokens. The function must be concurrency safe. This is intended only to allow the Azure SDK to cache MSI tokens. It isn't useful to applications in general because the token provider must implement all authentication logic. */ @@ -81,6 +88,24 @@ public CompletableFuture acquireToken(UserFederatedIdenti return this.executeRequest(userFicRequest); } + @Override + public CompletableFuture acquireTokenForAgent(AcquireTokenForAgentParameters parameters) { + validateNotNull("parameters", parameters); + + RequestContext context = new RequestContext( + this, + PublicApi.ACQUIRE_TOKEN_FOR_AGENT, + parameters); + + AcquireTokenForAgentRequest agentRequest = + new AcquireTokenForAgentRequest( + parameters, + this, + context); + + return this.executeRequest(agentRequest); + } + private ConfidentialClientApplication(Builder builder) { super(builder); sendX5c = builder.sendX5c; diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/IConfidentialClientApplication.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/IConfidentialClientApplication.java index eaa7cd2c..bfb89443 100644 --- a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/IConfidentialClientApplication.java +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/IConfidentialClientApplication.java @@ -62,4 +62,18 @@ public interface IConfidentialClientApplication extends IClientApplicationBase { * @return {@link CompletableFuture} containing an {@link IAuthenticationResult} */ CompletableFuture acquireToken(UserFederatedIdentityCredentialParameters parameters); + + /** + * Acquires a token for agent scenarios by orchestrating the full three-leg + * FMI/FIC token exchange. The developer passes scopes and an {@link AgentIdentity}; + * MSAL handles Legs 1-3 internally, including caching intermediate tokens. + *

+ * For user-scoped tokens, the agent identity must include either a UPN + * ({@link AgentIdentity#withUsername}) or an Object ID ({@link AgentIdentity#AgentIdentity(String, java.util.UUID)}). + * For app-only tokens, use {@link AgentIdentity#appOnly}. + * + * @param parameters instance of {@link AcquireTokenForAgentParameters} + * @return {@link CompletableFuture} containing an {@link IAuthenticationResult} + */ + CompletableFuture acquireTokenForAgent(AcquireTokenForAgentParameters parameters); } diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/PublicApi.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/PublicApi.java index faf19722..c30c5c98 100644 --- a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/PublicApi.java +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/PublicApi.java @@ -14,6 +14,7 @@ enum PublicApi { ACQUIRE_TOKEN_FOR_CLIENT(729), ACQUIRE_TOKEN_BY_AUTHORIZATION_CODE(831), ACQUIRE_TOKEN_BY_USER_FEDERATED_IDENTITY_CREDENTIAL(900), + ACQUIRE_TOKEN_FOR_AGENT(1020), ACQUIRE_TOKEN_SILENTLY(800), GET_ACCOUNTS(801), REMOVE_ACCOUNTS(802), diff --git a/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AcquireTokenForAgentTest.java b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AcquireTokenForAgentTest.java new file mode 100644 index 00000000..c7ad5b78 --- /dev/null +++ b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AcquireTokenForAgentTest.java @@ -0,0 +1,374 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.microsoft.aad.msal4j; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; + +import java.nio.charset.StandardCharsets; +import java.util.Base64; +import java.util.Collections; +import java.util.HashMap; +import java.util.Set; +import java.util.UUID; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; + +/** + * Tests for the composite AcquireTokenForAgent flow (§10 from AgentIDs_ComponentsReference). + * Validates the three-leg orchestration, internal CCA caching, per-user token isolation, + * and silent retrieval. + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +class AcquireTokenForAgentTest { + + private static final String BLUEPRINT_CLIENT_ID = "blueprint-client-id"; + private static final String AUTHORITY = "https://login.microsoftonline.com/tenant/"; + private static final Set CALLER_SCOPES = Collections.singleton("https://graph.microsoft.com/.default"); + + private static final String AGENT_APP_ID = "agent-app-id-abc"; + private static final String TENANT_ID = "f645ad92-e38d-4d1a-b510-d1b09a74a8ca"; + + private static final String USER1_UPN = "alice@contoso.com"; + private static final String USER1_OID = "11111111-1111-1111-1111-111111111111"; + private static final String USER2_UPN = "bob@contoso.com"; + private static final String USER2_OID = "22222222-2222-2222-2222-222222222222"; + + private ConfidentialClientApplication createBlueprintCca(DefaultHttpClient httpClientMock) throws Exception { + return ConfidentialClientApplication.builder( + BLUEPRINT_CLIENT_ID, + ClientCredentialFactory.createFromSecret("secret")) + .authority(AUTHORITY) + .instanceDiscovery(false) + .validateAuthority(false) + .httpClient(httpClientMock) + .build(); + } + + /** + * Creates a simple app-token response (for Leg 1 FMI credential or Leg 2 assertion token). + */ + private HttpResponse createAppTokenResponse(String accessToken) { + HashMap responseValues = new HashMap<>(); + responseValues.put("access_token", accessToken); + return TestHelper.expectedResponse(HttpStatus.HTTP_OK, + TestHelper.getSuccessfulTokenResponse(responseValues)); + } + + /** + * Creates a user-token response (for Leg 3) with id_token and client_info so that + * an account is properly stored in the cache. + */ + private HttpResponse createUserTokenResponse(String accessToken, String upn, String oid) { + HashMap idTokenValues = new HashMap<>(); + idTokenValues.put("oid", oid); + idTokenValues.put("preferred_username", upn); + String idToken = TestHelper.createIdToken(idTokenValues); + + String clientInfo = createClientInfo(oid); + + HashMap responseValues = new HashMap<>(); + responseValues.put("access_token", accessToken); + responseValues.put("id_token", idToken); + responseValues.put("client_info", clientInfo); + return TestHelper.expectedResponse(HttpStatus.HTTP_OK, + TestHelper.getSuccessfulTokenResponse(responseValues)); + } + + private String createClientInfo(String uid) { + String json = String.format("{\"uid\":\"%s\",\"utid\":\"%s\"}", uid, TENANT_ID); + return Base64.getEncoder().encodeToString(json.getBytes(StandardCharsets.UTF_8)); + } + + // ======================================================================== + // Core composite test: two users + caching (matches .NET TwoUpns test) + // ======================================================================== + + @Test + void acquireTokenForAgent_twoUpns_cacheReturnsCorrectUserToken() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + // Queue 4 HTTP responses in order: + // 1. Leg 1 (FMI credential for blueprint) + // 2. Leg 2 (assertion token for agent CCA) + // 3. Leg 3 (user token for alice) + // 4. Leg 3 (user token for bob — Legs 1+2 are cached) + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-credential-token"), + createAppTokenResponse("assertion-token"), + createUserTokenResponse("alice-access-token", USER1_UPN, USER1_OID), + createUserTokenResponse("bob-access-token", USER2_UPN, USER2_OID)); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + + AgentIdentity aliceAgent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + AgentIdentity bobAgent = AgentIdentity.withUsername(AGENT_APP_ID, USER2_UPN); + + // Act 1: Alice — should trigger 3 HTTP calls (Leg 1 + Leg 2 + Leg 3) + IAuthenticationResult result1 = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceAgent).build() + ).get(); + + assertEquals("alice-access-token", result1.accessToken()); + verify(httpClientMock, times(3)).send(any(HttpRequest.class)); + + // Act 2: Bob — should trigger only 1 HTTP call (Leg 3; Legs 1+2 cached) + IAuthenticationResult result2 = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, bobAgent).build() + ).get(); + + assertEquals("bob-access-token", result2.accessToken()); + verify(httpClientMock, times(4)).send(any(HttpRequest.class)); + + // Act 3: Alice again — should return from cache (0 HTTP calls) + IAuthenticationResult result3 = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceAgent).build() + ).get(); + + assertEquals("alice-access-token", result3.accessToken()); + verify(httpClientMock, times(4)).send(any(HttpRequest.class)); // still 4 total + } + + // ======================================================================== + // App-only flow + // ======================================================================== + + @Test + void acquireTokenForAgent_appOnly_acquiresTokenWithoutUserLeg() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + // App-only: Leg 1 (FMI credential) + Leg 2 (app token for caller scopes) + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-credential-token"), + createAppTokenResponse("agent-app-token")); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + + AgentIdentity appOnlyAgent = AgentIdentity.appOnly(AGENT_APP_ID); + + // Act + IAuthenticationResult result = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, appOnlyAgent).build() + ).get(); + + // Assert + assertEquals("agent-app-token", result.accessToken()); + verify(httpClientMock, times(2)).send(any(HttpRequest.class)); + } + + @Test + void acquireTokenForAgent_appOnly_secondCallReturnsCached() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-credential-token"), + createAppTokenResponse("agent-app-token")); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + AgentIdentity appOnlyAgent = AgentIdentity.appOnly(AGENT_APP_ID); + + // Act 1: first call triggers HTTP + blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, appOnlyAgent).build() + ).get(); + verify(httpClientMock, times(2)).send(any(HttpRequest.class)); + + // Act 2: second call should return from app cache (0 new HTTP calls) + IAuthenticationResult result = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, appOnlyAgent).build() + ).get(); + + assertEquals("agent-app-token", result.accessToken()); + verify(httpClientMock, times(2)).send(any(HttpRequest.class)); // still 2 total + } + + // ======================================================================== + // ForceRefresh bypasses user cache + // ======================================================================== + + @Test + void acquireTokenForAgent_forceRefresh_bypassesUserCache() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + // First call: 3 HTTP calls (Legs 1+2+3) + // Second call with forceRefresh: Leg 3 again (Legs 1+2 still cached) + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-credential-token"), + createAppTokenResponse("assertion-token"), + createUserTokenResponse("alice-token-1", USER1_UPN, USER1_OID), + createUserTokenResponse("alice-token-2", USER1_UPN, USER1_OID)); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + AgentIdentity aliceAgent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + + // Act 1: normal call + IAuthenticationResult result1 = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceAgent).build() + ).get(); + assertEquals("alice-token-1", result1.accessToken()); + verify(httpClientMock, times(3)).send(any(HttpRequest.class)); + + // Act 2: forceRefresh — should bypass user cache and execute Leg 3 again + IAuthenticationResult result2 = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceAgent) + .forceRefresh(true) + .build() + ).get(); + assertEquals("alice-token-2", result2.accessToken()); + // Leg 3 fires again (1 more HTTP call), Legs 1+2 still from cache + verify(httpClientMock, times(4)).send(any(HttpRequest.class)); + } + + // ======================================================================== + // User identity by OID + // ======================================================================== + + @Test + void acquireTokenForAgent_withOid_acquiresUserToken() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-credential-token"), + createAppTokenResponse("assertion-token"), + createUserTokenResponse("alice-token", USER1_UPN, USER1_OID)); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + + UUID aliceOid = UUID.fromString(USER1_OID); + AgentIdentity aliceByOid = new AgentIdentity(AGENT_APP_ID, aliceOid); + + // Act + IAuthenticationResult result = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceByOid).build() + ).get(); + + // Assert + assertEquals("alice-token", result.accessToken()); + verify(httpClientMock, times(3)).send(any(HttpRequest.class)); + } + + // ======================================================================== + // Agent CCA caching: same agent ID reuses CCA + // ======================================================================== + + @Test + void acquireTokenForAgent_samAgentId_reusesCca() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-token"), + createAppTokenResponse("assertion-token"), + createUserTokenResponse("user-token", USER1_UPN, USER1_OID)); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + AgentIdentity agent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + + // Act + blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, agent).build() + ).get(); + + // Assert — agent CCA cache should have exactly one entry + assertEquals(1, blueprintCca.agentCcaCache.size()); + assertTrue(blueprintCca.agentCcaCache.containsKey("agent_" + AGENT_APP_ID)); + } + + // ======================================================================== + // Leg 1 sends fmi_path in body + // ======================================================================== + + @Test + void acquireTokenForAgent_leg1_sendsFmiPathInBody() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-token"), + createAppTokenResponse("assertion-token"), + createUserTokenResponse("user-token", USER1_UPN, USER1_OID)); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + AgentIdentity agent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + + // Act + blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, agent).build() + ).get(); + + // Assert — first HTTP call (Leg 1) should include fmi_path=agentAppId + verify(httpClientMock, atLeastOnce()).send(argThat(request -> { + String body = request.body(); + return body != null && body.contains("fmi_path=" + AGENT_APP_ID); + })); + } + + // ======================================================================== + // Leg 3 sends user_fic grant type + // ======================================================================== + + @Test + void acquireTokenForAgent_leg3_sendsUserFicGrantType() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-token"), + createAppTokenResponse("assertion-token"), + createUserTokenResponse("user-token", USER1_UPN, USER1_OID)); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + AgentIdentity agent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + + // Act + blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, agent).build() + ).get(); + + // Assert — one of the HTTP calls should contain grant_type=user_fic + verify(httpClientMock, atLeastOnce()).send(argThat(request -> { + String body = request.body(); + return body != null && body.contains("grant_type=user_fic"); + })); + } + + // ======================================================================== + // Input validation + // ======================================================================== + + @Test + void acquireTokenForAgent_nullParameters_throwsException() throws Exception { + ConfidentialClientApplication cca = createBlueprintCca(mock(DefaultHttpClient.class)); + assertThrows(IllegalArgumentException.class, () -> + cca.acquireTokenForAgent(null)); + } + + @Test + void parameterBuilder_nullScopes_throwsException() { + AgentIdentity agent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + assertThrows(IllegalArgumentException.class, () -> + AcquireTokenForAgentParameters.builder(null, agent)); + } + + @Test + void parameterBuilder_nullAgentIdentity_throwsException() { + assertThrows(IllegalArgumentException.class, () -> + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, null)); + } +} diff --git a/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AgentIdentityTest.java b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AgentIdentityTest.java new file mode 100644 index 00000000..d4bf59c0 --- /dev/null +++ b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AgentIdentityTest.java @@ -0,0 +1,140 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.microsoft.aad.msal4j; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; + +import java.util.UUID; + +import static org.junit.jupiter.api.Assertions.*; + +/** + * Tests for the {@link AgentIdentity} model class (§9 from AgentIDs_ComponentsReference). + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +class AgentIdentityTest { + + private static final String AGENT_APP_ID = "agent-app-id-123"; + private static final String USER_UPN = "alice@contoso.com"; + private static final UUID USER_OID = UUID.fromString("11111111-1111-1111-1111-111111111111"); + + // ======================================================================== + // Constructor: by Object ID (recommended) + // ======================================================================== + + @Test + void constructor_withOid_setsPropertiesCorrectly() { + AgentIdentity identity = new AgentIdentity(AGENT_APP_ID, USER_OID); + + assertEquals(AGENT_APP_ID, identity.agentApplicationId()); + assertEquals(USER_OID, identity.userObjectId()); + assertNull(identity.username()); + assertTrue(identity.hasUserIdentifier()); + } + + @Test + void constructor_withOid_nullAgentAppId_throwsException() { + assertThrows(IllegalArgumentException.class, () -> + new AgentIdentity(null, USER_OID)); + } + + @Test + void constructor_withOid_emptyAgentAppId_throwsException() { + assertThrows(IllegalArgumentException.class, () -> + new AgentIdentity("", USER_OID)); + } + + @Test + void constructor_withOid_nullOid_throwsException() { + assertThrows(IllegalArgumentException.class, () -> + new AgentIdentity(AGENT_APP_ID, null)); + } + + // ======================================================================== + // Factory: withUsername (by UPN) + // ======================================================================== + + @Test + void withUsername_setsPropertiesCorrectly() { + AgentIdentity identity = AgentIdentity.withUsername(AGENT_APP_ID, USER_UPN); + + assertEquals(AGENT_APP_ID, identity.agentApplicationId()); + assertNull(identity.userObjectId()); + assertEquals(USER_UPN, identity.username()); + assertTrue(identity.hasUserIdentifier()); + } + + @Test + void withUsername_nullAgentAppId_throwsException() { + assertThrows(IllegalArgumentException.class, () -> + AgentIdentity.withUsername(null, USER_UPN)); + } + + @Test + void withUsername_emptyAgentAppId_throwsException() { + assertThrows(IllegalArgumentException.class, () -> + AgentIdentity.withUsername("", USER_UPN)); + } + + @Test + void withUsername_nullUsername_throwsException() { + assertThrows(IllegalArgumentException.class, () -> + AgentIdentity.withUsername(AGENT_APP_ID, null)); + } + + @Test + void withUsername_emptyUsername_throwsException() { + assertThrows(IllegalArgumentException.class, () -> + AgentIdentity.withUsername(AGENT_APP_ID, "")); + } + + // ======================================================================== + // Factory: appOnly (no user) + // ======================================================================== + + @Test + void appOnly_setsPropertiesCorrectly() { + AgentIdentity identity = AgentIdentity.appOnly(AGENT_APP_ID); + + assertEquals(AGENT_APP_ID, identity.agentApplicationId()); + assertNull(identity.userObjectId()); + assertNull(identity.username()); + assertFalse(identity.hasUserIdentifier()); + } + + @Test + void appOnly_nullAgentAppId_throwsException() { + assertThrows(IllegalArgumentException.class, () -> + AgentIdentity.appOnly(null)); + } + + @Test + void appOnly_emptyAgentAppId_throwsException() { + assertThrows(IllegalArgumentException.class, () -> + AgentIdentity.appOnly("")); + } + + // ======================================================================== + // hasUserIdentifier behavior + // ======================================================================== + + @Test + void hasUserIdentifier_withOid_returnsTrue() { + AgentIdentity identity = new AgentIdentity(AGENT_APP_ID, USER_OID); + assertTrue(identity.hasUserIdentifier()); + } + + @Test + void hasUserIdentifier_withUsername_returnsTrue() { + AgentIdentity identity = AgentIdentity.withUsername(AGENT_APP_ID, USER_UPN); + assertTrue(identity.hasUserIdentifier()); + } + + @Test + void hasUserIdentifier_appOnly_returnsFalse() { + AgentIdentity identity = AgentIdentity.appOnly(AGENT_APP_ID); + assertFalse(identity.hasUserIdentifier()); + } +} From 2d3c8c819b19708dc9e7429400f0511c1d970347 Mon Sep 17 00:00:00 2001 From: avdunn Date: Thu, 7 May 2026 14:22:14 -0700 Subject: [PATCH 2/2] Caching behavior fixes and improved test coverage --- .../com/microsoft/aad/msal4j/AgenticIT.java | 197 ++++++ ...cquireTokenByClientCredentialSupplier.java | 9 +- .../msal4j/AcquireTokenForAgentSupplier.java | 107 ++- .../msal4j/ClientCredentialParameters.java | 48 +- .../com/microsoft/aad/msal4j/TokenCache.java | 9 +- .../aad/msal4j/AcquireTokenForAgentTest.java | 617 ++++++++++++++++++ 6 files changed, 969 insertions(+), 18 deletions(-) diff --git a/msal4j-sdk/src/integrationtest/java/com/microsoft/aad/msal4j/AgenticIT.java b/msal4j-sdk/src/integrationtest/java/com/microsoft/aad/msal4j/AgenticIT.java index 8dd97753..1d284843 100644 --- a/msal4j-sdk/src/integrationtest/java/com/microsoft/aad/msal4j/AgenticIT.java +++ b/msal4j-sdk/src/integrationtest/java/com/microsoft/aad/msal4j/AgenticIT.java @@ -298,6 +298,203 @@ void agentCca_AppAndUserTokens_CacheIsolation() throws Exception { "Cache should have at least 2 entries (app + user)"); } + // ======================================================================== + // High-level AcquireTokenForAgent tests (composite API) + // ======================================================================== + + /** + * Tests the high-level acquireTokenForAgent API with a UPN-based AgentIdentity. + * Exercises the full 3-leg flow orchestrated internally by AcquireTokenForAgentSupplier. + */ + @Test + void acquireTokenForAgent_withUpn_fullFlow() throws Exception { + IClientCertificate clientCert = ClientCredentialFactory.createFromCertificate(privateKey, certificate); + + ConfidentialClientApplication blueprintCca = ConfidentialClientApplication.builder( + BLUEPRINT_CLIENT_ID, clientCert) + .authority(AUTHORITY) + .sendX5c(true) + .azureRegion(AZURE_REGION) + .build(); + + AgentIdentity agentId = AgentIdentity.withUsername(AGENT_APP_ID, USER_UPN); + + IAuthenticationResult result = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder( + Collections.singleton(GRAPH_SCOPE), agentId).build() + ).get(); + + assertNotNull(result, "Result should not be null"); + assertNotNull(result.accessToken(), "Access token should not be null"); + assertFalse(result.accessToken().isEmpty(), "Access token should not be empty"); + assertNotNull(result.account(), "Account should not be null for user token"); + } + + /** + * Tests the high-level acquireTokenForAgent API for app-only (no user) scenarios. + * Only Legs 1-2 are performed. + */ + @Test + void acquireTokenForAgent_appOnly() throws Exception { + IClientCertificate clientCert = ClientCredentialFactory.createFromCertificate(privateKey, certificate); + + ConfidentialClientApplication blueprintCca = ConfidentialClientApplication.builder( + BLUEPRINT_CLIENT_ID, clientCert) + .authority(AUTHORITY) + .sendX5c(true) + .azureRegion(AZURE_REGION) + .build(); + + AgentIdentity agentId = AgentIdentity.appOnly(AGENT_APP_ID); + + IAuthenticationResult result = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder( + Collections.singleton(GRAPH_SCOPE), agentId).build() + ).get(); + + assertNotNull(result, "Result should not be null"); + assertNotNull(result.accessToken(), "Access token should not be null"); + assertFalse(result.accessToken().isEmpty(), "Access token should not be empty"); + } + + /** + * Tests the high-level acquireTokenForAgent API with ForceRefresh. + * First call populates cache, second call (forceRefresh) bypasses it. + */ + @Test + void acquireTokenForAgent_forceRefresh() throws Exception { + IClientCertificate clientCert = ClientCredentialFactory.createFromCertificate(privateKey, certificate); + + ConfidentialClientApplication blueprintCca = ConfidentialClientApplication.builder( + BLUEPRINT_CLIENT_ID, clientCert) + .authority(AUTHORITY) + .sendX5c(true) + .azureRegion(AZURE_REGION) + .build(); + + AgentIdentity agentId = AgentIdentity.withUsername(AGENT_APP_ID, USER_UPN); + + // First call — populates cache + IAuthenticationResult result1 = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder( + Collections.singleton(GRAPH_SCOPE), agentId).build() + ).get(); + assertNotNull(result1.accessToken()); + + // Second call without forceRefresh — should return cached token + IAuthenticationResult result2 = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder( + Collections.singleton(GRAPH_SCOPE), agentId).build() + ).get(); + assertEquals(result1.accessToken(), result2.accessToken(), + "Second call should return cached token"); + + // Third call with forceRefresh — should get a fresh token + IAuthenticationResult result3 = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder( + Collections.singleton(GRAPH_SCOPE), agentId) + .forceRefresh(true).build() + ).get(); + assertNotNull(result3.accessToken()); + // The fresh token may be the same string (if not expired) but the flow exercised network + } + + /** + * Tests cache isolation between two blueprint CCA instances. + * Each blueprint should have its own agent CCA cache. + */ + @Test + void acquireTokenForAgent_cacheIsolation_twoBlueprintCcas() throws Exception { + IClientCertificate clientCert = ClientCredentialFactory.createFromCertificate(privateKey, certificate); + + ConfidentialClientApplication blueprint1 = ConfidentialClientApplication.builder( + BLUEPRINT_CLIENT_ID, clientCert) + .authority(AUTHORITY) + .sendX5c(true) + .azureRegion(AZURE_REGION) + .build(); + + ConfidentialClientApplication blueprint2 = ConfidentialClientApplication.builder( + BLUEPRINT_CLIENT_ID, clientCert) + .authority(AUTHORITY) + .sendX5c(true) + .azureRegion(AZURE_REGION) + .build(); + + AgentIdentity agentId = AgentIdentity.withUsername(AGENT_APP_ID, USER_UPN); + + // Acquire via blueprint1 + IAuthenticationResult result1 = blueprint1.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder( + Collections.singleton(GRAPH_SCOPE), agentId).build() + ).get(); + assertNotNull(result1.accessToken()); + + // Blueprint1 should have agent CCA cached, blueprint2 should not + assertEquals(1, blueprint1.agentCcaCache.size(), + "Blueprint1 should have one cached agent CCA"); + assertTrue(blueprint2.agentCcaCache.isEmpty(), + "Blueprint2 should have no cached agent CCAs (no bleed)"); + + // Acquire via blueprint2 + IAuthenticationResult result2 = blueprint2.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder( + Collections.singleton(GRAPH_SCOPE), agentId).build() + ).get(); + assertNotNull(result2.accessToken()); + + // Both should now have their own cache entries + assertEquals(1, blueprint1.agentCcaCache.size()); + assertEquals(1, blueprint2.agentCcaCache.size()); + } + + /** + * Tests that a UPN-based token can be found by OID lookup on the same blueprint. + * Discovers the OID via the UPN flow, then verifies OID-based call returns cached token. + */ + @Test + void acquireTokenForAgent_upnThenOid_sharesCache() throws Exception { + IClientCertificate clientCert = ClientCredentialFactory.createFromCertificate(privateKey, certificate); + + ConfidentialClientApplication blueprintCca = ConfidentialClientApplication.builder( + BLUEPRINT_CLIENT_ID, clientCert) + .authority(AUTHORITY) + .sendX5c(true) + .azureRegion(AZURE_REGION) + .build(); + + // Step 1: Acquire via UPN + AgentIdentity upnIdentity = AgentIdentity.withUsername(AGENT_APP_ID, USER_UPN); + IAuthenticationResult upnResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder( + Collections.singleton(GRAPH_SCOPE), upnIdentity).build() + ).get(); + assertNotNull(upnResult.account(), "Account should not be null"); + + // Extract OID from account's homeAccountId (format: oid.tid) + String homeAccountId = upnResult.account().homeAccountId(); + assertNotNull(homeAccountId); + String oidString = homeAccountId.contains(".") + ? homeAccountId.substring(0, homeAccountId.indexOf('.')) + : homeAccountId; + java.util.UUID userOid = java.util.UUID.fromString(oidString); + + // Step 2: Acquire via OID — should come from cache + AgentIdentity oidIdentity = new AgentIdentity(AGENT_APP_ID, userOid); + IAuthenticationResult oidResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder( + Collections.singleton(GRAPH_SCOPE), oidIdentity).build() + ).get(); + + // Should return the same cached token + assertEquals(upnResult.accessToken(), oidResult.accessToken(), + "OID-based call should return the same cached token as UPN-based call"); + } + + // ======================================================================== + // Helpers + // ======================================================================== + /** * Helper: acquires an FMI credential from the RMA (Resource Management Application). * Uses FMI_EXCHANGE_SCOPE, matching FmiIT's Flow3 pattern. diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenByClientCredentialSupplier.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenByClientCredentialSupplier.java index 8d7088b6..68871883 100644 --- a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenByClientCredentialSupplier.java +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenByClientCredentialSupplier.java @@ -40,10 +40,15 @@ AuthenticationResult execute() throws Exception { context, null); - // Propagate ext_cache_key_hash for fmi_path-based cache isolation + // Propagate ext_cache_key_hash for fmi_path/credential_fmi_path-based cache isolation + java.util.TreeMap components = new java.util.TreeMap<>(); if (!StringHelper.isBlank(this.clientCredentialRequest.parameters.fmiPath())) { - java.util.TreeMap components = new java.util.TreeMap<>(); components.put("fmi_path", this.clientCredentialRequest.parameters.fmiPath()); + } + if (!StringHelper.isBlank(this.clientCredentialRequest.parameters.credentialFmiPath())) { + components.put("credential_fmi_path", this.clientCredentialRequest.parameters.credentialFmiPath()); + } + if (!components.isEmpty()) { silentRequest.extCacheKeyHash(StringHelper.computeExtCacheKeyHash(components)); } diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentSupplier.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentSupplier.java index 47b67edf..4cc8753a 100644 --- a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentSupplier.java +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentSupplier.java @@ -8,6 +8,7 @@ import java.net.MalformedURLException; import java.util.Collections; +import java.util.Map; import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; @@ -69,7 +70,8 @@ AuthenticationResult execute() throws Exception { LOG.debug("App-only agent flow for agent app ID: {}", agentAppId); return (AuthenticationResult) joinAndUnwrap( agentCca.acquireToken( - ClientCredentialParameters.builder(callerScopes).build())); + propagateToClientCredentialParams( + ClientCredentialParameters.builder(callerScopes)).build())); } // --- User identity flow --- @@ -92,7 +94,9 @@ AuthenticationResult execute() throws Exception { LOG.debug("Executing Leg 2 (assertion token) for agent app ID: {}", agentAppId); IAuthenticationResult assertionResult = joinAndUnwrap( agentCca.acquireToken( - ClientCredentialParameters.builder(TOKEN_EXCHANGE_SCOPE).build())); + propagateToClientCredentialParams( + ClientCredentialParameters.builder(TOKEN_EXCHANGE_SCOPE) + .credentialFmiPath(agentAppId)).build())); String assertion = assertionResult.accessToken(); @@ -101,14 +105,16 @@ AuthenticationResult execute() throws Exception { LOG.debug("Executing Leg 3 (user FIC token) for agent app ID: {}", agentAppId); UserFederatedIdentityCredentialParameters ficParams; if (agentIdentity.userObjectId() != null) { - ficParams = UserFederatedIdentityCredentialParameters - .builder(callerScopes, agentIdentity.userObjectId(), assertion) - .forceRefresh(true) // always fetch from network (we already checked the cache above) + ficParams = propagateToUserFicParams( + UserFederatedIdentityCredentialParameters + .builder(callerScopes, agentIdentity.userObjectId(), assertion) + .forceRefresh(true)) // always fetch from network (we already checked the cache above) .build(); } else { - ficParams = UserFederatedIdentityCredentialParameters - .builder(callerScopes, agentIdentity.username(), assertion) - .forceRefresh(true) + ficParams = propagateToUserFicParams( + UserFederatedIdentityCredentialParameters + .builder(callerScopes, agentIdentity.username(), assertion) + .forceRefresh(true)) .build(); } @@ -131,12 +137,14 @@ private AuthenticationResult tryAcquireTokenSilent( return null; } - SilentParameters silentParams = SilentParameters - .builder(scopes, matchedAccount) - .build(); + SilentParameters.SilentParametersBuilder silentBuilder = SilentParameters + .builder(scopes, matchedAccount); + + // Propagate outer request parameters so that claims challenges cause cache bypass + propagateToSilentParams(silentBuilder); return (AuthenticationResult) joinAndUnwrap( - agentCca.acquireTokenSilently(silentParams)); + agentCca.acquireTokenSilently(silentBuilder.build())); } catch (Exception ex) { // Token expired or requires interaction — fall through to full Leg 2 + Leg 3 flow LOG.debug("Silent token acquisition failed for agent: {}", ex.getMessage()); @@ -173,6 +181,81 @@ private static String extractOid(String homeAccountId) { return dotIndex >= 0 ? homeAccountId.substring(0, dotIndex) : homeAccountId; } + // ======================================================================== + // Outer Request Parameter Propagation + // ======================================================================== + + /** + * Propagates per-request parameters from the outer AcquireTokenForAgent call to a + * ClientCredentialParameters builder (used for Legs 1-2 and app-only). + * This ensures caller-specified claims, tenant overrides, extra query parameters, + * and extra HTTP headers flow through to inner network calls. + */ + private ClientCredentialParameters.ClientCredentialParametersBuilder propagateToClientCredentialParams( + ClientCredentialParameters.ClientCredentialParametersBuilder builder) { + AcquireTokenForAgentParameters outerParams = agentRequest.parameters; + + if (outerParams.claims() != null) { + builder.claims(outerParams.claims()); + } + if (!StringHelper.isBlank(outerParams.tenant())) { + builder.tenant(outerParams.tenant()); + } + if (outerParams.extraQueryParameters() != null && !outerParams.extraQueryParameters().isEmpty()) { + builder.extraQueryParameters(outerParams.extraQueryParameters()); + } + if (outerParams.extraHttpHeaders() != null && !outerParams.extraHttpHeaders().isEmpty()) { + builder.extraHttpHeaders(outerParams.extraHttpHeaders()); + } + return builder; + } + + /** + * Propagates per-request parameters from the outer AcquireTokenForAgent call to a + * UserFederatedIdentityCredentialParameters builder (used for Leg 3). + */ + private UserFederatedIdentityCredentialParameters.UserFederatedIdentityCredentialParametersBuilder propagateToUserFicParams( + UserFederatedIdentityCredentialParameters.UserFederatedIdentityCredentialParametersBuilder builder) { + AcquireTokenForAgentParameters outerParams = agentRequest.parameters; + + if (outerParams.claims() != null) { + builder.claims(outerParams.claims()); + } + if (!StringHelper.isBlank(outerParams.tenant())) { + builder.tenant(outerParams.tenant()); + } + if (outerParams.extraQueryParameters() != null && !outerParams.extraQueryParameters().isEmpty()) { + builder.extraQueryParameters(outerParams.extraQueryParameters()); + } + if (outerParams.extraHttpHeaders() != null && !outerParams.extraHttpHeaders().isEmpty()) { + builder.extraHttpHeaders(outerParams.extraHttpHeaders()); + } + return builder; + } + + /** + * Propagates per-request parameters from the outer AcquireTokenForAgent call to a + * SilentParameters builder (used for the cache-first silent check). + * Claims propagation is important here: if a claims challenge is present, the silent + * check should recognize the cached token as insufficient and force a refresh. + */ + private void propagateToSilentParams(SilentParameters.SilentParametersBuilder builder) { + AcquireTokenForAgentParameters outerParams = agentRequest.parameters; + + if (outerParams.claims() != null) { + builder.claims(outerParams.claims()); + } + if (!StringHelper.isBlank(outerParams.tenant())) { + builder.tenant(outerParams.tenant()); + } + if (outerParams.extraQueryParameters() != null && !outerParams.extraQueryParameters().isEmpty()) { + builder.extraQueryParameters(outerParams.extraQueryParameters()); + } + if (outerParams.extraHttpHeaders() != null && !outerParams.extraHttpHeaders().isEmpty()) { + builder.extraHttpHeaders(outerParams.extraHttpHeaders()); + } + } + // ======================================================================== // Agent CCA Construction and Configuration // ======================================================================== diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ClientCredentialParameters.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ClientCredentialParameters.java index cc8ef709..c6f73c3f 100644 --- a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ClientCredentialParameters.java +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ClientCredentialParameters.java @@ -30,7 +30,9 @@ public class ClientCredentialParameters implements IAcquireTokenParameters { private String fmiPath; - private ClientCredentialParameters(Set scopes, Boolean skipCache, ClaimsRequest claims, Map extraHttpHeaders, Map extraQueryParameters, String tenant, IClientCredential clientCredential, String fmiPath) { + private String credentialFmiPath; + + private ClientCredentialParameters(Set scopes, Boolean skipCache, ClaimsRequest claims, Map extraHttpHeaders, Map extraQueryParameters, String tenant, IClientCredential clientCredential, String fmiPath, String credentialFmiPath) { this.scopes = scopes; this.skipCache = skipCache; this.claims = claims; @@ -39,6 +41,7 @@ private ClientCredentialParameters(Set scopes, Boolean skipCache, Claims this.tenant = tenant; this.clientCredential = clientCredential; this.fmiPath = fmiPath; + this.credentialFmiPath = credentialFmiPath; } private static ClientCredentialParametersBuilder builder() { @@ -101,6 +104,18 @@ public String fmiPath() { return this.fmiPath; } + /** + * Gets the credential FMI path for agent identity scenarios. + * Unlike {@link #fmiPath()}, this is cache-key-only: it contributes to the extended + * cache key hash ({@code credential_fmi_path}) but is NOT sent as an HTTP body parameter. + * This isolates assertion tokens (Leg 2) from user tokens in the agent CCA's flat cache. + * + * @return the credential FMI path, or null if not set + */ + public String credentialFmiPath() { + return this.credentialFmiPath; + } + public static class ClientCredentialParametersBuilder { private Set scopes; private Boolean skipCache = false; @@ -110,6 +125,7 @@ public static class ClientCredentialParametersBuilder { private String tenant; private IClientCredential clientCredential; private String fmiPath; + private String credentialFmiPath; ClientCredentialParametersBuilder() { } @@ -195,12 +211,38 @@ public ClientCredentialParametersBuilder fmiPath(String fmiPath) { return this; } + /** + * Sets the credential FMI path for agent identity scenarios. + * Unlike {@link #fmiPath(String)}, this is cache-key-only: it contributes to the + * extended cache key hash ({@code credential_fmi_path}) but is NOT sent as an HTTP + * body parameter. This isolates assertion tokens from user tokens in flat caches. + *

+ * Mutually exclusive with {@link #fmiPath(String)} — setting both will throw + * {@link IllegalStateException} at build time. + * + * @param credentialFmiPath the credential FMI path value (typically the agent application ID) + * @return builder that can be used to construct ClientCredentialParameters + */ + public ClientCredentialParametersBuilder credentialFmiPath(String credentialFmiPath) { + if (credentialFmiPath != null && credentialFmiPath.trim().isEmpty()) { + throw new IllegalArgumentException("credentialFmiPath cannot be empty or blank"); + } + this.credentialFmiPath = credentialFmiPath; + return this; + } + public ClientCredentialParameters build() { - return new ClientCredentialParameters(this.scopes, this.skipCache, this.claims, this.extraHttpHeaders, this.extraQueryParameters, this.tenant, this.clientCredential, this.fmiPath); + if (!StringHelper.isBlank(this.fmiPath) && !StringHelper.isBlank(this.credentialFmiPath)) { + throw new IllegalStateException( + "fmiPath and credentialFmiPath are mutually exclusive. " + + "Use fmiPath for FMI body parameter + cache key, or " + + "credentialFmiPath for cache key isolation only."); + } + return new ClientCredentialParameters(this.scopes, this.skipCache, this.claims, this.extraHttpHeaders, this.extraQueryParameters, this.tenant, this.clientCredential, this.fmiPath, this.credentialFmiPath); } public String toString() { - return "ClientCredentialParameters.ClientCredentialParametersBuilder(scopes=" + this.scopes + ", skipCache=" + this.skipCache + ", claims=" + this.claims + ", extraHttpHeaders=" + this.extraHttpHeaders + ", extraQueryParameters=" + this.extraQueryParameters + ", tenant=" + this.tenant + ", clientCredential=" + this.clientCredential + ", fmiPath=" + this.fmiPath + ")"; + return "ClientCredentialParameters.ClientCredentialParametersBuilder(scopes=" + this.scopes + ", skipCache=" + this.skipCache + ", claims=" + this.claims + ", extraHttpHeaders=" + this.extraHttpHeaders + ", extraQueryParameters=" + this.extraQueryParameters + ", tenant=" + this.tenant + ", clientCredential=" + this.clientCredential + ", fmiPath=" + this.fmiPath + ", credentialFmiPath=" + this.credentialFmiPath + ")"; } } } diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/TokenCache.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/TokenCache.java index 6f72de4d..d5a0a701 100644 --- a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/TokenCache.java +++ b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/TokenCache.java @@ -345,9 +345,16 @@ private static AccessTokenCacheEntity createAccessTokenCacheEntity(TokenRequestE private static String computeExtCacheKeyHashForRequest(MsalRequest msalRequest) { if (msalRequest instanceof ClientCredentialRequest) { ClientCredentialParameters parameters = ((ClientCredentialRequest) msalRequest).parameters; + TreeMap components = new TreeMap<>(); + if (!StringHelper.isBlank(parameters.fmiPath())) { - TreeMap components = new TreeMap<>(); components.put("fmi_path", parameters.fmiPath()); + } + if (!StringHelper.isBlank(parameters.credentialFmiPath())) { + components.put("credential_fmi_path", parameters.credentialFmiPath()); + } + + if (!components.isEmpty()) { return StringHelper.computeExtCacheKeyHash(components); } } diff --git a/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AcquireTokenForAgentTest.java b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AcquireTokenForAgentTest.java index c7ad5b78..ae2b98bc 100644 --- a/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AcquireTokenForAgentTest.java +++ b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/AcquireTokenForAgentTest.java @@ -10,11 +10,13 @@ import java.util.Base64; import java.util.Collections; import java.util.HashMap; +import java.util.Map; import java.util.Set; import java.util.UUID; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.*; /** @@ -371,4 +373,619 @@ void parameterBuilder_nullAgentIdentity_throwsException() { assertThrows(IllegalArgumentException.class, () -> AcquireTokenForAgentParameters.builder(CALLER_SCOPES, null)); } + + // ======================================================================== + // Cache isolation: two blueprint CCAs do not share agent CCA caches + // ======================================================================== + + @Test + void acquireTokenForAgent_twoBlueprintCcas_noCacheBleed() throws Exception { + // Arrange — two separate blueprint CCAs + DefaultHttpClient httpClient1 = mock(DefaultHttpClient.class); + when(httpClient1.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-token-1"), + createAppTokenResponse("assertion-token-1"), + createUserTokenResponse("alice-token-1", USER1_UPN, USER1_OID)); + + DefaultHttpClient httpClient2 = mock(DefaultHttpClient.class); + when(httpClient2.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-token-2"), + createAppTokenResponse("assertion-token-2"), + createUserTokenResponse("bob-token-2", USER2_UPN, USER2_OID)); + + ConfidentialClientApplication blueprint1 = createBlueprintCca(httpClient1); + ConfidentialClientApplication blueprint2 = createBlueprintCca(httpClient2); + + AgentIdentity aliceAgent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + AgentIdentity bobAgent = AgentIdentity.withUsername(AGENT_APP_ID, USER2_UPN); + + // Act: acquire via blueprint1 + IAuthenticationResult result1 = blueprint1.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceAgent).build() + ).get(); + assertEquals("alice-token-1", result1.accessToken()); + + // Act: acquire via blueprint2 + IAuthenticationResult result2 = blueprint2.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, bobAgent).build() + ).get(); + assertEquals("bob-token-2", result2.accessToken()); + + // Assert: each blueprint has its own agent CCA cache + assertEquals(1, blueprint1.agentCcaCache.size()); + assertEquals(1, blueprint2.agentCcaCache.size()); + + // Blueprint1's agent CCA should only have Alice's token (not Bob's) + // Verify blueprint1 still returns Alice from cache (no bleed from blueprint2) + IAuthenticationResult result1Again = blueprint1.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceAgent).build() + ).get(); + assertEquals("alice-token-1", result1Again.accessToken()); + verify(httpClient1, times(3)).send(any(HttpRequest.class)); // no new calls + } + + // ======================================================================== + // UPN → OID shared cache: same user found by either identifier + // ======================================================================== + + @Test + void acquireTokenForAgent_upnThenOid_sharesCache() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + // First call (UPN): 3 HTTP calls (Legs 1+2+3) + // Second call (OID for same user): should come from cache (0 HTTP calls) + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-token"), + createAppTokenResponse("assertion-token"), + createUserTokenResponse("alice-token", USER1_UPN, USER1_OID)); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + + AgentIdentity aliceByUpn = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + UUID aliceOid = UUID.fromString(USER1_OID); + AgentIdentity aliceByOid = new AgentIdentity(AGENT_APP_ID, aliceOid); + + // Act 1: acquire by UPN + IAuthenticationResult upnResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceByUpn).build() + ).get(); + assertEquals("alice-token", upnResult.accessToken()); + verify(httpClientMock, times(3)).send(any(HttpRequest.class)); + + // Act 2: acquire by OID for the same user — should hit cache + IAuthenticationResult oidResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceByOid).build() + ).get(); + assertEquals("alice-token", oidResult.accessToken()); + // No new HTTP calls — found via OID match in findMatchingAccount + verify(httpClientMock, times(3)).send(any(HttpRequest.class)); + } + + // ======================================================================== + // Parameter propagation: tenant override flows to inner calls + // ======================================================================== + + @Test + void acquireTokenForAgent_withTenant_propagatesToInnerCalls() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-token"), + createAppTokenResponse("assertion-token"), + createUserTokenResponse("alice-token", USER1_UPN, USER1_OID)); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + AgentIdentity agent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + + String overrideTenant = "override-tenant-id"; + + // Act + IAuthenticationResult result = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, agent) + .tenant(overrideTenant) + .build() + ).get(); + + // Assert — at least one inner HTTP call should target the override tenant + verify(httpClientMock, atLeastOnce()).send(argThat(request -> { + String url = request.url() != null ? request.url().toString() : ""; + return url.contains(overrideTenant); + })); + } + + // ======================================================================== + // Parameter propagation: extra query parameters flow to inner calls + // ======================================================================== + + @Test + void acquireTokenForAgent_withExtraQueryParams_propagatesToInnerCalls() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-token"), + createAppTokenResponse("assertion-token"), + createUserTokenResponse("alice-token", USER1_UPN, USER1_OID)); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + AgentIdentity agent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + + Map extraParams = new HashMap<>(); + extraParams.put("custom_param", "custom_value"); + + // Act + IAuthenticationResult result = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, agent) + .extraQueryParameters(extraParams) + .build() + ).get(); + + // Assert — at least one inner HTTP call should include the extra query parameter + verify(httpClientMock, atLeastOnce()).send(argThat(request -> { + String body = request.body(); + return body != null && body.contains("custom_param=custom_value"); + })); + } + + // ======================================================================== + // Parameter propagation: claims flow to inner calls + // ======================================================================== + + @Test + void acquireTokenForAgent_withClaims_propagatesToInnerCalls() throws Exception { + // Arrange + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createAppTokenResponse("fmi-token"), + createAppTokenResponse("assertion-token"), + createUserTokenResponse("alice-token", USER1_UPN, USER1_OID)); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + AgentIdentity agent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + + ClaimsRequest claims = new ClaimsRequest(); + claims.requestClaimInAccessToken("xms_cc", null); + + // Act + IAuthenticationResult result = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, agent) + .claims(claims) + .build() + ).get(); + + // Assert — at least one inner HTTP call should include the claims parameter + verify(httpClientMock, atLeastOnce()).send(argThat(request -> { + String body = request.body(); + return body != null && body.contains("claims="); + })); + } + + // ======================================================================== + // Comprehensive cache behavior test: verifies token counts, isolation + // between users, isolation between agent and non-agent flows, and + // correct silent lookups across all scenarios. + // ======================================================================== + + @Test + void acquireTokenForAgent_comprehensiveCacheBehavior() throws Exception { + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + // Queue responses in order: + // --- Agent flow for Alice (Legs 1+2+3 = 3 HTTP calls) --- + // 1. Leg 1: FMI credential (blueprint app token) + // 2. Leg 2: assertion token (agent CCA app token) + // 3. Leg 3: user token for Alice + // --- Agent flow for Bob (Leg 3 only = 1 HTTP call, Legs 1+2 cached) --- + // 4. Leg 3: user token for Bob + // --- Non-agent client_credentials for Charlie on BLUEPRINT CCA (1 HTTP call) --- + // 5. App token for Charlie scopes on the blueprint + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + // Alice (Legs 1+2+3) + createAppTokenResponse("fmi-credential-token"), + createAppTokenResponse("assertion-token"), + createUserTokenResponse("alice-agent-token", USER1_UPN, USER1_OID), + // Bob (Leg 3 only) + createUserTokenResponse("bob-agent-token", USER2_UPN, USER2_OID), + // Charlie (non-agent client_credentials on blueprint) + createAppTokenResponse("charlie-app-token")); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + + AgentIdentity aliceAgent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + AgentIdentity bobAgent = AgentIdentity.withUsername(AGENT_APP_ID, USER2_UPN); + + // ---- Step 1: Agent flow for Alice ---- + IAuthenticationResult aliceResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceAgent).build() + ).get(); + + assertEquals("alice-agent-token", aliceResult.accessToken()); + verify(httpClientMock, times(3)).send(any(HttpRequest.class)); + + // ---- Step 2: Agent flow for Bob (Legs 1+2 should be cached) ---- + IAuthenticationResult bobResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, bobAgent).build() + ).get(); + + assertEquals("bob-agent-token", bobResult.accessToken()); + verify(httpClientMock, times(4)).send(any(HttpRequest.class)); // +1 for Leg 3 only + + // ---- Step 3: Non-agent client_credentials on the BLUEPRINT CCA ---- + // This uses the blueprint's own token cache, NOT the agent CCA's cache. + IAuthenticationResult charlieResult = blueprintCca.acquireToken( + ClientCredentialParameters.builder(CALLER_SCOPES).build() + ).get(); + + assertEquals("charlie-app-token", charlieResult.accessToken()); + verify(httpClientMock, times(5)).send(any(HttpRequest.class)); + + // ---- Cache state verification ---- + + // Blueprint CCA's token cache: + // - 1 FMI credential (Leg 1, with extCacheKeyHash for fmi_path) + // - 1 non-agent app token (Charlie's client_credentials) + // Total: 2 access tokens in the blueprint's own cache + assertEquals(2, blueprintCca.tokenCache.accessTokens.size(), + "Blueprint cache should have 2 tokens: FMI credential + non-agent app token"); + + // Agent CCA cache should have exactly one entry (for AGENT_APP_ID) + assertEquals(1, blueprintCca.agentCcaCache.size(), + "Blueprint should have 1 agent CCA cached"); + + ConfidentialClientApplication agentCca = + blueprintCca.agentCcaCache.get("agent_" + AGENT_APP_ID); + assertNotNull(agentCca, "Agent CCA should exist in cache"); + + // Agent CCA's token cache: + // - 1 assertion token (Leg 2, app-level, scope=api://AzureADTokenExchange/.default) + // - 2 user tokens (Alice + Bob, scope=graph.microsoft.com/.default) + // Total: 3 access tokens in the agent CCA's cache + assertEquals(3, agentCca.tokenCache.accessTokens.size(), + "Agent CCA cache should have 3 tokens: 1 assertion + 2 user tokens"); + + // ---- Step 4: Silent retrieval for Alice (agent flow, should hit cache) ---- + IAuthenticationResult aliceSilent = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceAgent).build() + ).get(); + + assertEquals("alice-agent-token", aliceSilent.accessToken()); + verify(httpClientMock, times(5)).send(any(HttpRequest.class)); // still 5, no new calls + + // ---- Step 5: Silent retrieval for Bob (agent flow, should hit cache) ---- + IAuthenticationResult bobSilent = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, bobAgent).build() + ).get(); + + assertEquals("bob-agent-token", bobSilent.accessToken()); + verify(httpClientMock, times(5)).send(any(HttpRequest.class)); // still 5 + + // ---- Step 6: Non-agent call again (should hit blueprint's cache) ---- + IAuthenticationResult charlieAgain = blueprintCca.acquireToken( + ClientCredentialParameters.builder(CALLER_SCOPES).build() + ).get(); + + assertEquals("charlie-app-token", charlieAgain.accessToken()); + verify(httpClientMock, times(5)).send(any(HttpRequest.class)); // still 5 + + // ---- Step 7: Verify Alice by OID also hits cache (UPN→OID shared cache) ---- + UUID aliceOid = UUID.fromString(USER1_OID); + AgentIdentity aliceByOid = new AgentIdentity(AGENT_APP_ID, aliceOid); + + IAuthenticationResult aliceByOidResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceByOid).build() + ).get(); + + assertEquals("alice-agent-token", aliceByOidResult.accessToken()); + verify(httpClientMock, times(5)).send(any(HttpRequest.class)); // still 5 + + // ---- Step 8: Verify cache counts haven't changed after all silent calls ---- + assertEquals(2, blueprintCca.tokenCache.accessTokens.size(), + "Blueprint cache should still have 2 tokens after silent calls"); + assertEquals(3, agentCca.tokenCache.accessTokens.size(), + "Agent CCA cache should still have 3 tokens after silent calls"); + + // ---- Step 9: Verify cache key isolation between FMI and non-FMI tokens ---- + // The blueprint cache has tokens with different cache key structures: + // - FMI token has extCacheKeyHash (credential_type=AccessToken_Extended) + // - Non-agent token has no extCacheKeyHash (credential_type=AccessToken) + boolean hasFmiToken = blueprintCca.tokenCache.accessTokens.values().stream() + .anyMatch(at -> !StringHelper.isBlank(at.extCacheKeyHash())); + boolean hasNonFmiToken = blueprintCca.tokenCache.accessTokens.values().stream() + .anyMatch(at -> StringHelper.isBlank(at.extCacheKeyHash())); + + assertTrue(hasFmiToken, "Blueprint cache should contain an FMI token with extCacheKeyHash"); + assertTrue(hasNonFmiToken, "Blueprint cache should contain a non-FMI token without extCacheKeyHash"); + + // ---- Step 10: Verify agent CCA user tokens have distinct homeAccountIds ---- + long distinctHomeAccountIds = agentCca.tokenCache.accessTokens.values().stream() + .map(at -> at.homeAccountId) + .filter(id -> id != null && !id.isEmpty()) + .distinct() + .count(); + assertEquals(2, distinctHomeAccountIds, + "Agent CCA should have 2 distinct homeAccountIds (Alice + Bob)"); + } + + // ======================================================================== + // Leg 2 cache key isolation: verifies that credential_fmi_path produces an + // extCacheKeyHash on the assertion token, preventing collisions with user + // tokens that share the same scope. Also verifies credential_fmi_path is + // NOT sent in the HTTP request body (cache-key-only). + // ======================================================================== + + @Test + void acquireTokenForAgent_leg2CacheIsolation_credentialFmiPathPreventsCollision() throws Exception { + // Both the caller and Leg 2 use the same scope (api://AzureADTokenExchange/.default). + // Without credential_fmi_path isolation, Leg 2's cache lookup could return Alice's + // user token instead of the assertion token — an order-dependent collision. + Set exchangeScope = Collections.singleton("api://AzureADTokenExchange/.default"); + + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + // Alice: Legs 1+2+3 (3 HTTP calls) + createAppTokenResponse("fmi-credential"), + createAppTokenResponse("correct-assertion-token"), + createUserTokenResponse("alice-user-token", USER1_UPN, USER1_OID), + // Bob: Leg 3 only (1 HTTP call — Legs 1+2 cached) + createUserTokenResponse("bob-user-token", USER2_UPN, USER2_OID)); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + + AgentIdentity aliceAgent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + AgentIdentity bobAgent = AgentIdentity.withUsername(AGENT_APP_ID, USER2_UPN); + + // ---- Step 1: Alice's full agent flow with exchange scope ---- + IAuthenticationResult aliceResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(exchangeScope, aliceAgent).build() + ).get(); + + assertEquals("alice-user-token", aliceResult.accessToken()); + verify(httpClientMock, times(3)).send(any(HttpRequest.class)); + + // ---- Step 2: Verify Leg 2 token has credential_fmi_path cache isolation ---- + ConfidentialClientApplication agentCca = + blueprintCca.agentCcaCache.get("agent_" + AGENT_APP_ID); + assertNotNull(agentCca); + + // Agent CCA has 2 tokens: Leg 2 app token + Alice's user token, both with exchange scope + assertEquals(2, agentCca.tokenCache.accessTokens.size()); + + // Compute the expected hash for credential_fmi_path = agentAppId + java.util.TreeMap expectedComponents = new java.util.TreeMap<>(); + expectedComponents.put("credential_fmi_path", AGENT_APP_ID); + String expectedHash = StringHelper.computeExtCacheKeyHash(expectedComponents); + + // The Leg 2 token (app-level, empty homeAccountId) must have the correct extCacheKeyHash + AccessTokenCacheEntity leg2Token = agentCca.tokenCache.accessTokens.values().stream() + .filter(at -> StringHelper.isBlank(at.homeAccountId) || at.homeAccountId.isEmpty()) + .findFirst() + .orElseThrow(() -> new AssertionError("Leg 2 app token not found")); + + assertEquals(expectedHash, leg2Token.extCacheKeyHash(), + "Leg 2 token should have extCacheKeyHash from credential_fmi_path"); + + // Alice's user token must NOT have an extCacheKeyHash + AccessTokenCacheEntity aliceToken = agentCca.tokenCache.accessTokens.values().stream() + .filter(at -> !StringHelper.isBlank(at.homeAccountId) && !at.homeAccountId.isEmpty()) + .findFirst() + .orElseThrow(() -> new AssertionError("Alice's user token not found")); + + assertTrue(StringHelper.isBlank(aliceToken.extCacheKeyHash()), + "Alice's user token should NOT have an extCacheKeyHash"); + + // ---- Step 3: Verify credential_fmi_path is NOT sent in any HTTP body ---- + // (It's cache-key-only, not a wire parameter) + org.mockito.ArgumentCaptor requestCaptor = org.mockito.ArgumentCaptor.forClass(HttpRequest.class); + verify(httpClientMock, times(3)).send(requestCaptor.capture()); + + for (HttpRequest req : requestCaptor.getAllValues()) { + String body = req.body() != null ? req.body() : ""; + assertFalse(body.contains("credential_fmi_path"), + "credential_fmi_path should NOT appear in any HTTP request body"); + } + + // ---- Step 4: Bob's agent flow — Leg 2 from cache, Leg 3 fresh ---- + // Clear invocations to count only Bob's HTTP calls + clearInvocations(httpClientMock); + + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createUserTokenResponse("bob-user-token", USER2_UPN, USER2_OID)); + + IAuthenticationResult bobResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(exchangeScope, bobAgent).build() + ).get(); + + assertEquals("bob-user-token", bobResult.accessToken()); + + // Only 1 HTTP call: Leg 3 for Bob. If Leg 2 had a cache collision, it would + // miss the cache (wrong extCacheKeyHash) and make an extra network call. + org.mockito.ArgumentCaptor bobCaptor = org.mockito.ArgumentCaptor.forClass(HttpRequest.class); + verify(httpClientMock, times(1)).send(bobCaptor.capture()); + + // Verify Bob's Leg 3 used the correct assertion (not Alice's user token) + HttpRequest bobLeg3Request = bobCaptor.getValue(); + String bobBody = bobLeg3Request.body() != null ? bobLeg3Request.body() : ""; + assertTrue(bobBody.contains("correct-assertion-token"), + "Bob's Leg 3 should use the cached Leg 2 assertion token as client_assertion"); + assertFalse(bobBody.contains("alice-user-token"), + "Bob's Leg 3 should NOT use Alice's user token as the assertion"); + + // ---- Step 5: Final cache state ---- + assertEquals(3, agentCca.tokenCache.accessTokens.size(), + "Agent CCA should have 3 tokens: 1 Leg 2 assertion + 2 user tokens"); + } + + // ======================================================================== + // Scope collision test: caller uses api://AzureADTokenExchange/.default + // which is the same scope used internally for Leg 2 assertion tokens. + // This probes whether the flat cache can distinguish between an app-level + // assertion token and a user-level FIC token when both share the same scope. + // ======================================================================== + + @Test + void acquireTokenForAgent_callerScopeMatchesInternalAssertionScope_noCollision() throws Exception { + // The internal Leg 2 uses api://AzureADTokenExchange/.default for the assertion token. + // If an external caller also requests that scope for a user token, the agent CCA's + // flat cache will contain both an app token and a user token with the same scope. + Set tokenExchangeScope = Collections.singleton("api://AzureADTokenExchange/.default"); + + String user3Upn = "charlie@contoso.com"; + String user3Oid = "33333333-3333-3333-3333-333333333333"; + + DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class); + + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + // --- Alice: normal scopes (Legs 1+2+3 = 3 HTTP calls) --- + createAppTokenResponse("fmi-credential-token"), + createAppTokenResponse("assertion-token-leg2"), + createUserTokenResponse("alice-graph-token", USER1_UPN, USER1_OID), + + // --- Bob: api://AzureADTokenExchange scope (Leg 3 only = 1 HTTP call) --- + // Legs 1+2 are cached from Alice's flow. + // Leg 3 creates a USER token with the same scope as the Leg 2 app token. + createUserTokenResponse("bob-exchange-token", USER2_UPN, USER2_OID), + + // --- Charlie: non-agent client_credentials with same exchange scope --- + // This goes through the BLUEPRINT CCA (not the agent CCA). + createAppTokenResponse("charlie-exchange-app-token")); + + ConfidentialClientApplication blueprintCca = createBlueprintCca(httpClientMock); + + AgentIdentity aliceAgent = AgentIdentity.withUsername(AGENT_APP_ID, USER1_UPN); + AgentIdentity bobAgent = AgentIdentity.withUsername(AGENT_APP_ID, USER2_UPN); + + // ---- Step 1: Alice with normal Graph scopes ---- + IAuthenticationResult aliceResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceAgent).build() + ).get(); + + assertEquals("alice-graph-token", aliceResult.accessToken()); + verify(httpClientMock, times(3)).send(any(HttpRequest.class)); + + // ---- Step 2: Bob with api://AzureADTokenExchange/.default ---- + // This is the dangerous scenario: the agent CCA already has an app token + // (assertion-token-leg2) for this exact scope from Leg 2. Now we're asking + // for a USER token with the same scope. If the cache lookup for Leg 2 on + // future calls uses findAny() without filtering by homeAccountId, it could + // return Bob's user token instead of the assertion token. + IAuthenticationResult bobResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(tokenExchangeScope, bobAgent).build() + ).get(); + + assertEquals("bob-exchange-token", bobResult.accessToken()); + verify(httpClientMock, times(4)).send(any(HttpRequest.class)); // +1 for Bob's Leg 3 + + // ---- Step 3: Verify agent CCA cache state ---- + ConfidentialClientApplication agentCca = + blueprintCca.agentCcaCache.get("agent_" + AGENT_APP_ID); + assertNotNull(agentCca); + + // Agent CCA should have: + // - 1 Leg 2 assertion token (app token, scope=api://AzureADTokenExchange/.default, homeAccountId="") + // - 1 Alice user token (scope=graph.microsoft.com/.default, homeAccountId=alice-oid.tenant) + // - 1 Bob user token (scope=api://AzureADTokenExchange/.default, homeAccountId=bob-oid.tenant) + assertEquals(3, agentCca.tokenCache.accessTokens.size(), + "Agent CCA cache should have 3 tokens (1 assertion + 2 user)"); + + // Count how many tokens have api://AzureADTokenExchange scope in the agent CCA. + // There should be 2: the Leg 2 assertion token (app) and Bob's user token. + long exchangeScopeTokenCount = agentCca.tokenCache.accessTokens.values().stream() + .filter(at -> at.target() != null && + at.target().toLowerCase().contains("azureadtokenexchange")) + .count(); + assertEquals(2, exchangeScopeTokenCount, + "Agent CCA should have 2 tokens with the exchange scope (1 app + 1 user)"); + + // ---- Step 4: Alice again — should still return from cache (no collision) ---- + IAuthenticationResult aliceAgain = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(CALLER_SCOPES, aliceAgent).build() + ).get(); + + assertEquals("alice-graph-token", aliceAgain.accessToken()); + verify(httpClientMock, times(4)).send(any(HttpRequest.class)); // still 4 + + // ---- Step 5: Bob again — this is the critical test ---- + // When we request Bob's token again, the silent lookup should find Bob's + // USER token (not the Leg 2 assertion token) even though both share the same scope. + IAuthenticationResult bobAgain = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(tokenExchangeScope, bobAgent).build() + ).get(); + + assertEquals("bob-exchange-token", bobAgain.accessToken(), + "Bob's silent retrieval should return the user token, not the assertion token"); + verify(httpClientMock, times(4)).send(any(HttpRequest.class)); // still 4 + + // ---- Step 6: Now trigger a NEW agent flow for a different user ---- + // This is where the Leg 2 collision matters most. A new user triggers Leg 2 + // (acquireToken(ClientCredentialParameters)) which uses getApplicationAccessTokenCacheEntity. + // If that lookup returns Bob's USER token instead of the app assertion token, + // Leg 3 will use the wrong assertion and fail or return incorrect results. + // + // We queue a Leg 3 response for a hypothetical third user to test this. + // If Leg 2 correctly returns the cached assertion token, only 1 HTTP call (Leg 3) fires. + // If Leg 2 gets a collision and returns Bob's user token as the assertion, the behavior + // will be unpredictable (wrong assertion value, possibly an error, or 2+ HTTP calls). + + String user3_upn = "charlie@contoso.com"; + String user3_oid = "33333333-3333-3333-3333-333333333333"; + + // Clear invocation history so we can count only Charlie's HTTP calls + clearInvocations(httpClientMock); + + // Reset mock for the next sequence: only Leg 3 should fire (1 call). + // If Leg 2 has a cache collision, it may re-fetch (2+ calls), or if + // the wrong token is used as the assertion, it may error out. + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn( + createUserTokenResponse("charlie-exchange-token", user3_upn, user3_oid)); + + AgentIdentity charlieAgent = AgentIdentity.withUsername(AGENT_APP_ID, user3_upn); + IAuthenticationResult charlieResult = blueprintCca.acquireTokenForAgent( + AcquireTokenForAgentParameters.builder(tokenExchangeScope, charlieAgent).build() + ).get(); + + // If this assertion fails, it means Leg 2 returned Bob's user token instead of the + // cached assertion token, causing downstream problems. + assertEquals("charlie-exchange-token", charlieResult.accessToken(), + "Charlie should get a fresh Leg 3 token using the cached Leg 2 assertion"); + + // Verify only 1 new HTTP call was made (Leg 3 for Charlie). + // If 2+ new calls were made, Leg 2 had to re-fetch because of a cache collision. + verify(httpClientMock, times(1)).send(any(HttpRequest.class)); + + // ---- Step 7: Final cache state verification ---- + assertEquals(4, agentCca.tokenCache.accessTokens.size(), + "Agent CCA should now have 4 tokens: 1 assertion + 3 user tokens"); + + // ---- Step 8: Non-agent client_credentials with exchange scope on BLUEPRINT ---- + // This tests that the blueprint's own cache doesn't collide with the FMI token + // (which also targets api://AzureADTokenExchange/.default but has an extCacheKeyHash). + clearInvocations(httpClientMock); + when(httpClientMock.send(any(HttpRequest.class))) + .thenReturn(createAppTokenResponse("blueprint-exchange-app-token")); + + IAuthenticationResult blueprintExchangeResult = blueprintCca.acquireToken( + ClientCredentialParameters.builder(tokenExchangeScope).build() + ).get(); + + // The FMI token has extCacheKeyHash set (from fmi_path), so a plain + // client_credentials call without fmi_path should NOT match it. + // It should trigger a new HTTP call and store a separate cache entry. + assertEquals("blueprint-exchange-app-token", blueprintExchangeResult.accessToken(), + "Blueprint's non-FMI exchange token should not collide with the FMI token"); + } }