Skip to content

Commit 3a8ba3e

Browse files
Avery-DunnAvery-Dunn
committed
Caching behavior fixes and improved test coverage
1 parent f99994d commit 3a8ba3e

5 files changed

Lines changed: 965 additions & 20 deletions

File tree

‎msal4j-sdk/src/integrationtest/java/com/microsoft/aad/msal4j/AgenticIT.java‎

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -298,6 +298,203 @@ void agentCca_AppAndUserTokens_CacheIsolation() throws Exception {
298298
"Cache should have at least 2 entries (app + user)");
299299
}
300300

301+
// ========================================================================
302+
// High-level AcquireTokenForAgent tests (composite API)
303+
// ========================================================================
304+
305+
/**
306+
* Tests the high-level acquireTokenForAgent API with a UPN-based AgentIdentity.
307+
* Exercises the full 3-leg flow orchestrated internally by AcquireTokenForAgentSupplier.
308+
*/
309+
@Test
310+
void acquireTokenForAgent_withUpn_fullFlow() throws Exception {
311+
IClientCertificate clientCert = ClientCredentialFactory.createFromCertificate(privateKey, certificate);
312+
313+
ConfidentialClientApplication blueprintCca = ConfidentialClientApplication.builder(
314+
BLUEPRINT_CLIENT_ID, clientCert)
315+
.authority(AUTHORITY)
316+
.sendX5c(true)
317+
.azureRegion(AZURE_REGION)
318+
.build();
319+
320+
AgentIdentity agentId = AgentIdentity.withUsername(AGENT_APP_ID, USER_UPN);
321+
322+
IAuthenticationResult result = blueprintCca.acquireTokenForAgent(
323+
AcquireTokenForAgentParameters.builder(
324+
Collections.singleton(GRAPH_SCOPE), agentId).build()
325+
).get();
326+
327+
assertNotNull(result, "Result should not be null");
328+
assertNotNull(result.accessToken(), "Access token should not be null");
329+
assertFalse(result.accessToken().isEmpty(), "Access token should not be empty");
330+
assertNotNull(result.account(), "Account should not be null for user token");
331+
}
332+
333+
/**
334+
* Tests the high-level acquireTokenForAgent API for app-only (no user) scenarios.
335+
* Only Legs 1-2 are performed.
336+
*/
337+
@Test
338+
void acquireTokenForAgent_appOnly() throws Exception {
339+
IClientCertificate clientCert = ClientCredentialFactory.createFromCertificate(privateKey, certificate);
340+
341+
ConfidentialClientApplication blueprintCca = ConfidentialClientApplication.builder(
342+
BLUEPRINT_CLIENT_ID, clientCert)
343+
.authority(AUTHORITY)
344+
.sendX5c(true)
345+
.azureRegion(AZURE_REGION)
346+
.build();
347+
348+
AgentIdentity agentId = AgentIdentity.appOnly(AGENT_APP_ID);
349+
350+
IAuthenticationResult result = blueprintCca.acquireTokenForAgent(
351+
AcquireTokenForAgentParameters.builder(
352+
Collections.singleton(GRAPH_SCOPE), agentId).build()
353+
).get();
354+
355+
assertNotNull(result, "Result should not be null");
356+
assertNotNull(result.accessToken(), "Access token should not be null");
357+
assertFalse(result.accessToken().isEmpty(), "Access token should not be empty");
358+
}
359+
360+
/**
361+
* Tests the high-level acquireTokenForAgent API with ForceRefresh.
362+
* First call populates cache, second call (forceRefresh) bypasses it.
363+
*/
364+
@Test
365+
void acquireTokenForAgent_forceRefresh() throws Exception {
366+
IClientCertificate clientCert = ClientCredentialFactory.createFromCertificate(privateKey, certificate);
367+
368+
ConfidentialClientApplication blueprintCca = ConfidentialClientApplication.builder(
369+
BLUEPRINT_CLIENT_ID, clientCert)
370+
.authority(AUTHORITY)
371+
.sendX5c(true)
372+
.azureRegion(AZURE_REGION)
373+
.build();
374+
375+
AgentIdentity agentId = AgentIdentity.withUsername(AGENT_APP_ID, USER_UPN);
376+
377+
// First call — populates cache
378+
IAuthenticationResult result1 = blueprintCca.acquireTokenForAgent(
379+
AcquireTokenForAgentParameters.builder(
380+
Collections.singleton(GRAPH_SCOPE), agentId).build()
381+
).get();
382+
assertNotNull(result1.accessToken());
383+
384+
// Second call without forceRefresh — should return cached token
385+
IAuthenticationResult result2 = blueprintCca.acquireTokenForAgent(
386+
AcquireTokenForAgentParameters.builder(
387+
Collections.singleton(GRAPH_SCOPE), agentId).build()
388+
).get();
389+
assertEquals(result1.accessToken(), result2.accessToken(),
390+
"Second call should return cached token");
391+
392+
// Third call with forceRefresh — should get a fresh token
393+
IAuthenticationResult result3 = blueprintCca.acquireTokenForAgent(
394+
AcquireTokenForAgentParameters.builder(
395+
Collections.singleton(GRAPH_SCOPE), agentId)
396+
.forceRefresh(true).build()
397+
).get();
398+
assertNotNull(result3.accessToken());
399+
// The fresh token may be the same string (if not expired) but the flow exercised network
400+
}
401+
402+
/**
403+
* Tests cache isolation between two blueprint CCA instances.
404+
* Each blueprint should have its own agent CCA cache.
405+
*/
406+
@Test
407+
void acquireTokenForAgent_cacheIsolation_twoBlueprintCcas() throws Exception {
408+
IClientCertificate clientCert = ClientCredentialFactory.createFromCertificate(privateKey, certificate);
409+
410+
ConfidentialClientApplication blueprint1 = ConfidentialClientApplication.builder(
411+
BLUEPRINT_CLIENT_ID, clientCert)
412+
.authority(AUTHORITY)
413+
.sendX5c(true)
414+
.azureRegion(AZURE_REGION)
415+
.build();
416+
417+
ConfidentialClientApplication blueprint2 = ConfidentialClientApplication.builder(
418+
BLUEPRINT_CLIENT_ID, clientCert)
419+
.authority(AUTHORITY)
420+
.sendX5c(true)
421+
.azureRegion(AZURE_REGION)
422+
.build();
423+
424+
AgentIdentity agentId = AgentIdentity.withUsername(AGENT_APP_ID, USER_UPN);
425+
426+
// Acquire via blueprint1
427+
IAuthenticationResult result1 = blueprint1.acquireTokenForAgent(
428+
AcquireTokenForAgentParameters.builder(
429+
Collections.singleton(GRAPH_SCOPE), agentId).build()
430+
).get();
431+
assertNotNull(result1.accessToken());
432+
433+
// Blueprint1 should have agent CCA cached, blueprint2 should not
434+
assertEquals(1, blueprint1.agentCcaCache.size(),
435+
"Blueprint1 should have one cached agent CCA");
436+
assertTrue(blueprint2.agentCcaCache.isEmpty(),
437+
"Blueprint2 should have no cached agent CCAs (no bleed)");
438+
439+
// Acquire via blueprint2
440+
IAuthenticationResult result2 = blueprint2.acquireTokenForAgent(
441+
AcquireTokenForAgentParameters.builder(
442+
Collections.singleton(GRAPH_SCOPE), agentId).build()
443+
).get();
444+
assertNotNull(result2.accessToken());
445+
446+
// Both should now have their own cache entries
447+
assertEquals(1, blueprint1.agentCcaCache.size());
448+
assertEquals(1, blueprint2.agentCcaCache.size());
449+
}
450+
451+
/**
452+
* Tests that a UPN-based token can be found by OID lookup on the same blueprint.
453+
* Discovers the OID via the UPN flow, then verifies OID-based call returns cached token.
454+
*/
455+
@Test
456+
void acquireTokenForAgent_upnThenOid_sharesCache() throws Exception {
457+
IClientCertificate clientCert = ClientCredentialFactory.createFromCertificate(privateKey, certificate);
458+
459+
ConfidentialClientApplication blueprintCca = ConfidentialClientApplication.builder(
460+
BLUEPRINT_CLIENT_ID, clientCert)
461+
.authority(AUTHORITY)
462+
.sendX5c(true)
463+
.azureRegion(AZURE_REGION)
464+
.build();
465+
466+
// Step 1: Acquire via UPN
467+
AgentIdentity upnIdentity = AgentIdentity.withUsername(AGENT_APP_ID, USER_UPN);
468+
IAuthenticationResult upnResult = blueprintCca.acquireTokenForAgent(
469+
AcquireTokenForAgentParameters.builder(
470+
Collections.singleton(GRAPH_SCOPE), upnIdentity).build()
471+
).get();
472+
assertNotNull(upnResult.account(), "Account should not be null");
473+
474+
// Extract OID from account's homeAccountId (format: oid.tid)
475+
String homeAccountId = upnResult.account().homeAccountId();
476+
assertNotNull(homeAccountId);
477+
String oidString = homeAccountId.contains(".")
478+
? homeAccountId.substring(0, homeAccountId.indexOf('.'))
479+
: homeAccountId;
480+
java.util.UUID userOid = java.util.UUID.fromString(oidString);
481+
482+
// Step 2: Acquire via OID — should come from cache
483+
AgentIdentity oidIdentity = new AgentIdentity(AGENT_APP_ID, userOid);
484+
IAuthenticationResult oidResult = blueprintCca.acquireTokenForAgent(
485+
AcquireTokenForAgentParameters.builder(
486+
Collections.singleton(GRAPH_SCOPE), oidIdentity).build()
487+
).get();
488+
489+
// Should return the same cached token
490+
assertEquals(upnResult.accessToken(), oidResult.accessToken(),
491+
"OID-based call should return the same cached token as UPN-based call");
492+
}
493+
494+
// ========================================================================
495+
// Helpers
496+
// ========================================================================
497+
301498
/**
302499
* Helper: acquires an FMI credential from the RMA (Resource Management Application).
303500
* Uses FMI_EXCHANGE_SCOPE, matching FmiIT's Flow3 pattern.

‎msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenByClientCredentialSupplier.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ AuthenticationResult execute() throws Exception {
4040
context,
4141
null);
4242

43-
// Propagate ext_cache_key_hash for fmi_path-based cache isolation
43+
// Propagate ext_cache_key_hash for fmi_path/credential_fmi_path-based cache isolation
4444
String extCacheKeyHash = this.clientCredentialRequest.parameters.computeFmiCacheKeyHash();
4545
if (!StringHelper.isBlank(extCacheKeyHash)) {
4646
silentRequest.extCacheKeyHash(extCacheKeyHash);

‎msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenForAgentSupplier.java‎

Lines changed: 95 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
import java.net.MalformedURLException;
1010
import java.util.Collections;
11+
import java.util.Map;
1112
import java.util.Set;
1213
import java.util.concurrent.CompletableFuture;
1314
import java.util.concurrent.CompletionException;
@@ -69,7 +70,8 @@ AuthenticationResult execute() throws Exception {
6970
LOG.debug("App-only agent flow for agent app ID: {}", agentAppId);
7071
return (AuthenticationResult) joinAndUnwrap(
7172
agentCca.acquireToken(
72-
ClientCredentialParameters.builder(callerScopes).build()));
73+
propagateToClientCredentialParams(
74+
ClientCredentialParameters.builder(callerScopes)).build()));
7375
}
7476

7577
// --- User identity flow ---
@@ -92,7 +94,9 @@ AuthenticationResult execute() throws Exception {
9294
LOG.debug("Executing Leg 2 (assertion token) for agent app ID: {}", agentAppId);
9395
IAuthenticationResult assertionResult = joinAndUnwrap(
9496
agentCca.acquireToken(
95-
ClientCredentialParameters.builder(TOKEN_EXCHANGE_SCOPE).build()));
97+
propagateToClientCredentialParams(
98+
ClientCredentialParameters.builder(TOKEN_EXCHANGE_SCOPE)
99+
.credentialFmiPath(agentAppId)).build()));
96100

97101
String assertion = assertionResult.accessToken();
98102

@@ -101,14 +105,16 @@ AuthenticationResult execute() throws Exception {
101105
LOG.debug("Executing Leg 3 (user FIC token) for agent app ID: {}", agentAppId);
102106
UserFederatedIdentityCredentialParameters ficParams;
103107
if (agentIdentity.userObjectId() != null) {
104-
ficParams = UserFederatedIdentityCredentialParameters
105-
.builder(callerScopes, agentIdentity.userObjectId(), assertion)
106-
.forceRefresh(true) // always fetch from network (we already checked the cache above)
108+
ficParams = propagateToUserFicParams(
109+
UserFederatedIdentityCredentialParameters
110+
.builder(callerScopes, agentIdentity.userObjectId(), assertion)
111+
.forceRefresh(true)) // always fetch from network (we already checked the cache above)
107112
.build();
108113
} else {
109-
ficParams = UserFederatedIdentityCredentialParameters
110-
.builder(callerScopes, agentIdentity.username(), assertion)
111-
.forceRefresh(true)
114+
ficParams = propagateToUserFicParams(
115+
UserFederatedIdentityCredentialParameters
116+
.builder(callerScopes, agentIdentity.username(), assertion)
117+
.forceRefresh(true))
112118
.build();
113119
}
114120

@@ -131,12 +137,14 @@ private AuthenticationResult tryAcquireTokenSilent(
131137
return null;
132138
}
133139

134-
SilentParameters silentParams = SilentParameters
135-
.builder(scopes, matchedAccount)
136-
.build();
140+
SilentParameters.SilentParametersBuilder silentBuilder = SilentParameters
141+
.builder(scopes, matchedAccount);
142+
143+
// Propagate outer request parameters so that claims challenges cause cache bypass
144+
propagateToSilentParams(silentBuilder);
137145

138146
return (AuthenticationResult) joinAndUnwrap(
139-
agentCca.acquireTokenSilently(silentParams));
147+
agentCca.acquireTokenSilently(silentBuilder.build()));
140148
} catch (Exception ex) {
141149
// Token expired or requires interaction — fall through to full Leg 2 + Leg 3 flow
142150
LOG.debug("Silent token acquisition failed for agent: {}", ex.getMessage());
@@ -173,6 +181,81 @@ private static String extractOid(String homeAccountId) {
173181
return dotIndex >= 0 ? homeAccountId.substring(0, dotIndex) : homeAccountId;
174182
}
175183

184+
// ========================================================================
185+
// Outer Request Parameter Propagation
186+
// ========================================================================
187+
188+
/**
189+
* Propagates per-request parameters from the outer AcquireTokenForAgent call to a
190+
* ClientCredentialParameters builder (used for Legs 1-2 and app-only).
191+
* This ensures caller-specified claims, tenant overrides, extra query parameters,
192+
* and extra HTTP headers flow through to inner network calls.
193+
*/
194+
private ClientCredentialParameters.ClientCredentialParametersBuilder propagateToClientCredentialParams(
195+
ClientCredentialParameters.ClientCredentialParametersBuilder builder) {
196+
AcquireTokenForAgentParameters outerParams = agentRequest.parameters;
197+
198+
if (outerParams.claims() != null) {
199+
builder.claims(outerParams.claims());
200+
}
201+
if (!StringHelper.isBlank(outerParams.tenant())) {
202+
builder.tenant(outerParams.tenant());
203+
}
204+
if (outerParams.extraQueryParameters() != null && !outerParams.extraQueryParameters().isEmpty()) {
205+
builder.extraQueryParameters(outerParams.extraQueryParameters());
206+
}
207+
if (outerParams.extraHttpHeaders() != null && !outerParams.extraHttpHeaders().isEmpty()) {
208+
builder.extraHttpHeaders(outerParams.extraHttpHeaders());
209+
}
210+
return builder;
211+
}
212+
213+
/**
214+
* Propagates per-request parameters from the outer AcquireTokenForAgent call to a
215+
* UserFederatedIdentityCredentialParameters builder (used for Leg 3).
216+
*/
217+
private UserFederatedIdentityCredentialParameters.UserFederatedIdentityCredentialParametersBuilder propagateToUserFicParams(
218+
UserFederatedIdentityCredentialParameters.UserFederatedIdentityCredentialParametersBuilder builder) {
219+
AcquireTokenForAgentParameters outerParams = agentRequest.parameters;
220+
221+
if (outerParams.claims() != null) {
222+
builder.claims(outerParams.claims());
223+
}
224+
if (!StringHelper.isBlank(outerParams.tenant())) {
225+
builder.tenant(outerParams.tenant());
226+
}
227+
if (outerParams.extraQueryParameters() != null && !outerParams.extraQueryParameters().isEmpty()) {
228+
builder.extraQueryParameters(outerParams.extraQueryParameters());
229+
}
230+
if (outerParams.extraHttpHeaders() != null && !outerParams.extraHttpHeaders().isEmpty()) {
231+
builder.extraHttpHeaders(outerParams.extraHttpHeaders());
232+
}
233+
return builder;
234+
}
235+
236+
/**
237+
* Propagates per-request parameters from the outer AcquireTokenForAgent call to a
238+
* SilentParameters builder (used for the cache-first silent check).
239+
* Claims propagation is important here: if a claims challenge is present, the silent
240+
* check should recognize the cached token as insufficient and force a refresh.
241+
*/
242+
private void propagateToSilentParams(SilentParameters.SilentParametersBuilder builder) {
243+
AcquireTokenForAgentParameters outerParams = agentRequest.parameters;
244+
245+
if (outerParams.claims() != null) {
246+
builder.claims(outerParams.claims());
247+
}
248+
if (!StringHelper.isBlank(outerParams.tenant())) {
249+
builder.tenant(outerParams.tenant());
250+
}
251+
if (outerParams.extraQueryParameters() != null && !outerParams.extraQueryParameters().isEmpty()) {
252+
builder.extraQueryParameters(outerParams.extraQueryParameters());
253+
}
254+
if (outerParams.extraHttpHeaders() != null && !outerParams.extraHttpHeaders().isEmpty()) {
255+
builder.extraHttpHeaders(outerParams.extraHttpHeaders());
256+
}
257+
}
258+
176259
// ========================================================================
177260
// Agent CCA Construction and Configuration
178261
// ========================================================================

0 commit comments

Comments
 (0)