From 5096147531efd426484d687750bbdffa501f8347 Mon Sep 17 00:00:00 2001 From: avdunn Date: Tue, 7 Jul 2026 09:51:12 -0700 Subject: [PATCH 1/3] Refactor agent CCA instance management --- .../TokenAcquisition.cs | 302 ++++++++++-------- .../TokenAcquisitionTests.cs | 115 ++++++- 2 files changed, 277 insertions(+), 140 deletions(-) diff --git a/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs b/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs index 5950fb395..f49ea3f29 100644 --- a/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs +++ b/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs @@ -58,30 +58,18 @@ class OAuthConstants /// Important: call GetOrBuildConfidentialClientApplication instead of accessing _applicationsByAuthorityClientId directly. /// Write access to this dictionary is synchronized. /// - private readonly ConcurrentDictionary _applicationsByAuthorityClientId = new(); + internal readonly ConcurrentDictionary _applicationsByAuthorityClientId = new(); private readonly ConcurrentDictionary _appSemaphores = new(); /// - /// Maximum number of agent CCA instances to keep in the dictionary before - /// clearing it as a DOS protection measure. Tokens are stored in MSAL's - /// shared static cache, so clearing the dictionary only discards lightweight + /// Maximum number of agent CCA instances to keep in the shared application dictionary + /// before clearing agent entries as a DOS protection measure. Tokens are stored in + /// MSAL's shared static cache, so clearing agent entries only discards lightweight /// CCA objects — tokens remain accessible to newly-built CCAs. + /// Agent entries are identified by the ":agent:" segment in their cache key. /// internal int AgentCcaMaxCount { get; set; } = 10000; - /// - /// Caches agent CCAs for the native User FIC flow. Each agent CCA uses an assertion - /// callback that chains to the blueprint CCA for Leg 1 (FMI token acquisition). - /// Key format: "{agentAppId}:{authenticationScheme}". - /// The authenticationScheme is included because different schemes may resolve to - /// different blueprint credentials (e.g. certificates), and the agent CCA's assertion - /// callback captures the scheme to chain to the correct blueprint CCA. - /// When the dictionary exceeds , it is cleared entirely; - /// tokens survive in MSAL's shared static cache and are found by new CCAs via silent calls. - /// - internal readonly ConcurrentDictionary _agentUserFicCcas = new(); - private readonly ConcurrentDictionary _agentCcaSemaphores = new(); - /// /// Maps (agentAppId, user identifier, tenantId) tuples to MSAL account identifiers for the /// native User FIC flow. This is needed because AcquireTokenSilent requires an IAccount, @@ -341,14 +329,26 @@ private async Task GetAuthenticationResultForUserInternalA MergedOptions mergedOptions = GetMergedOptions(authenticationScheme, tokenAcquisitionOptions); user ??= await _tokenAcquisitionHost.GetAuthenticatedUserAsync(user).ConfigureAwait(false); - var application = await GetOrBuildConfidentialClientApplicationAsync(mergedOptions, isTokenBinding: false); - if (tokenAcquisitionOptions is not null) { tokenAcquisitionOptions.ExtraParameters ??= new Dictionary(); tokenAcquisitionOptions.ExtraParameters[Constants.ExtensionOptionsServiceProviderKey] = _serviceProvider; } + // Detect agentic User FIC flow early — before building the blueprint CCA. + // Agent CCAs are built via the unified builder path with their own ClientId, + // so the blueprint CCA is only built lazily (inside the assertion callback) + // when actually needed for Leg 1 token acquisition. + var agentResult = await TryGetAuthenticationResultForAgentUserFicAsync( + tenantId, scopes, mergedOptions, tokenAcquisitionOptions).ConfigureAwait(false); + if (agentResult is not null) + { + LogAuthResult(agentResult); + return agentResult; + } + + var application = await GetOrBuildConfidentialClientApplicationAsync(mergedOptions, isTokenBinding: false); + CredentialSourceLoaderParameters loaderParameters = new CredentialSourceLoaderParameters(application.AppConfig.ClientId, application.Authority) { Protocol = ProtocolNames.Bearer, @@ -442,15 +442,6 @@ private async Task GetAuthenticationResultForUserInternalA MergedOptions mergedOptions, TokenAcquisitionOptions? tokenAcquisitionOptions) { - // Handle agentic User FIC flow (both UPN and OID) using MSAL's native UserFIC API. - // This bypasses the ROPC piggybacking approach and provides proper cache behavior. - var agentUserFicResult = await TryGetAuthenticationResultForAgentUserFicAsync( - application, tenantId, scopes, mergedOptions, tokenAcquisitionOptions).ConfigureAwait(false); - if (agentUserFicResult is not null) - { - return agentUserFicResult; - } - string? username = null; string? password = null; @@ -579,7 +570,6 @@ private async Task GetAuthenticationResultForUserInternalA /// An if this is an agentic User FIC flow /// (UPN or OID); null if not an agentic flow (regular ROPC). private async Task TryGetAuthenticationResultForAgentUserFicAsync( - IConfidentialClientApplication application, string? tenantId, IEnumerable scopes, MergedOptions mergedOptions, @@ -628,7 +618,7 @@ private async Task GetAuthenticationResultForUserInternalA Logger.AgentUserFicFlowDetected(_logger, agentAppId!, identifierType); var agentCca = await GetOrBuildAgentUserFicCcaAsync( - agentAppId!, application.Authority, authScheme).ConfigureAwait(false); + agentAppId!, authScheme, mergedOptions).ConfigureAwait(false); bool forceRefresh = tokenAcquisitionOptions?.ForceRefresh ?? false; @@ -724,97 +714,59 @@ private async Task GetAuthenticationResultForUserInternalA } /// - /// Gets or builds an agent CCA for the native User FIC flow. Each agent CCA uses an - /// assertion callback that chains back to the blueprint CCA for Leg 1 (FMI token). - /// Each agent CCA has a unique ClientId (the agent app ID), providing natural cache - /// key isolation in the shared static cache. When the dictionary exceeds - /// , it is cleared entirely as DOS protection. + /// Gets or builds an agent CCA for the native User FIC flow. Delegates to the unified + /// / + /// builder path so that agent CCAs + /// receive the same configuration as normal CCAs (logging, authority, cache initialization). + /// Each agent CCA has a unique ClientId (the agent app ID), providing natural cache key + /// isolation in both the CCA dictionary and MSAL's shared static token cache. /// private async Task GetOrBuildAgentUserFicCcaAsync( string agentAppId, - string authority, - string? authenticationScheme) + string? authenticationScheme, + MergedOptions mergedOptions) { - // Include authenticationScheme in the CCA cache key so different schemes - // (pointing to different blueprint credentials) get separate CCAs. - // Authority is intentionally excluded: within a single TokenAcquisition instance, - // the authority is derived from the one configured MergedOptions and does not vary - // per call. Multi-authority scenarios would require separate TokenAcquisition instances. - string ccaCacheKey = $"{agentAppId}:{authenticationScheme ?? string.Empty}"; - - if (_agentUserFicCcas.TryGetValue(ccaCacheKey, out var existing)) - { - return existing; - } - - var semaphore = _agentCcaSemaphores.GetOrAdd(ccaCacheKey, _ => new SemaphoreSlim(1, 1)); - await semaphore.WaitAsync().ConfigureAwait(false); - try - { - if (_agentUserFicCcas.TryGetValue(ccaCacheKey, out var entry)) + // Build the assertion callback that chains to the blueprint CCA for Leg 1. + // Capture authenticationScheme so the callback resolves the correct blueprint. + string? capturedAuthScheme = authenticationScheme; + + Func> assertionCallback = async (AssertionRequestOptions options) => + { + // Leg 1: Blueprint acquires FMI token (T1) for this agent. + // AcquireTokenForClient checks cache first — only the first call + // or an expired T1 hits the network. + MergedOptions blueprintOptions = _tokenAcquisitionHost.GetOptions(capturedAuthScheme, out _); + var blueprintCca = await GetOrBuildConfidentialClientApplicationAsync( + blueprintOptions, isTokenBinding: false).ConfigureAwait(false); + + var leg1Builder = blueprintCca + .AcquireTokenForClient(s_ficScopes) + .WithFmiPath(agentAppId) + .WithSendX5C(blueprintOptions.SendX5C); + + // Propagate tenant override to Leg 1 when the caller specifies a tenant + // (e.g., via WithTenantId on Leg 2/3). MSAL's AssertionRequestOptions + // provides the resolved TenantId directly from the runtime authority. + if (!string.IsNullOrEmpty(options.TenantId)) { - return entry; + leg1Builder.WithTenantId(options.TenantId); } - // Capture authenticationScheme for the assertion callback closure. - string? capturedAuthScheme = authenticationScheme; - - var builder = ConfidentialClientApplicationBuilder - .Create(agentAppId) - .WithClientAssertion(async (AssertionRequestOptions options) => - { - // Leg 1: Blueprint acquires FMI token (T1) for this agent. - // AcquireTokenForClient checks cache first — only the first call - // or an expired T1 hits the network. - MergedOptions blueprintOptions = _tokenAcquisitionHost.GetOptions(capturedAuthScheme, out _); - var blueprintCca = await GetOrBuildConfidentialClientApplicationAsync( - blueprintOptions, isTokenBinding: false).ConfigureAwait(false); - - var leg1Builder = blueprintCca - .AcquireTokenForClient(s_ficScopes) - .WithFmiPath(agentAppId) - .WithSendX5C(blueprintOptions.SendX5C); - - // Propagate tenant override to Leg 1 when the caller specifies a tenant - // (e.g., via WithTenantId on Leg 2/3). MSAL's AssertionRequestOptions - // provides the resolved TenantId directly from the runtime authority. - if (!string.IsNullOrEmpty(options.TenantId)) - { - leg1Builder.WithTenantId(options.TenantId); - } - - var leg1 = await leg1Builder - .ExecuteAsync(options.CancellationToken) - .ConfigureAwait(false); - - return leg1.AccessToken; - }) - .WithAuthority(authority) - .WithHttpClientFactory(_httpClientFactory) - .WithCacheOptions(CacheOptions.EnableSharedCacheOptions); + var leg1 = await leg1Builder + .ExecuteAsync(options.CancellationToken) + .ConfigureAwait(false); - var newApp = builder.Build(); + return leg1.AccessToken; + }; - _agentUserFicCcas[ccaCacheKey] = newApp; - Logger.AgentCcaCreated(_logger, ccaCacheKey); + // Delegate to the unified builder path. The agent app ID is incorporated into + // the cache key automatically, and the CCA gets all the same configuration as + // normal CCAs (logging, authority, cache initialization, etc.). + var agentCca = await GetOrBuildConfidentialClientApplicationAsync( + mergedOptions, isTokenBinding: false, agentAppId, assertionCallback).ConfigureAwait(false); - // DOS protection: if the dictionary grows beyond the threshold, clear it. - // Tokens survive in MSAL's shared static cache and will be found by new CCAs. - if (_agentUserFicCcas.Count > AgentCcaMaxCount) - { - int cleared = _agentUserFicCcas.Count; - _agentUserFicCcas.Clear(); - _agentUserFicAccountIds.Clear(); - _agentCcaSemaphores.Clear(); - Logger.AgentCcaEviction(_logger, cleared, 0); - } - - return newApp; - } - finally - { - semaphore.Release(); - } + Logger.AgentCcaCreated(_logger, agentAppId); + return agentCca; } private void LogAuthResult(AuthenticationResult? authenticationResult) @@ -1350,9 +1302,17 @@ public async Task GetConfidentialClientApplicati internal /* for testing */ async Task GetOrBuildConfidentialClientApplicationAsync( MergedOptions mergedOptions, - bool isTokenBinding) + bool isTokenBinding, + string? agentAppId = null, + Func>? clientAssertionProvider = null) { + // For agent CCAs, incorporate the agent app ID into the cache key so each agent + // gets its own CCA instance while sharing the same dictionary and lifecycle. string key = GetApplicationKey(mergedOptions, isTokenBinding); + if (agentAppId is not null) + { + key = $"{key}:agent:{agentAppId}"; + } // GetOrAddAsync based on https://github.com/dotnet/runtime/issues/83636#issuecomment-1474998680 // Fast path: check if already created @@ -1370,11 +1330,37 @@ public async Task GetConfidentialClientApplicati return app; // Build and store the application - var newApp = await BuildConfidentialClientApplicationAsync(mergedOptions, isTokenBinding); + var newApp = await BuildConfidentialClientApplicationAsync( + mergedOptions, isTokenBinding, agentAppId, clientAssertionProvider); // Recompute the key as BuildConfidentialClientApplicationAsync can cause it to change. key = GetApplicationKey(mergedOptions, isTokenBinding); + if (agentAppId is not null) + { + key = $"{key}:agent:{agentAppId}"; + } _applicationsByAuthorityClientId[key] = newApp; + + // DOS protection for agent CCAs: if too many agent entries accumulate, + // clear them. Tokens survive in MSAL's shared static cache and will be + // found by new CCAs via AcquireTokenSilent. + if (agentAppId is not null) + { + int agentCount = _applicationsByAuthorityClientId.Keys.Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0); + if (agentCount > AgentCcaMaxCount) + { + int cleared = 0; + foreach (var agentKey in _applicationsByAuthorityClientId.Keys.Where(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0).ToList()) + { + _applicationsByAuthorityClientId.TryRemove(agentKey, out _); + _appSemaphores.TryRemove(agentKey, out _); + cleared++; + } + _agentUserFicAccountIds.Clear(); + Logger.AgentCcaEviction(_logger, cleared, 0); + } + } + return newApp; } finally @@ -1386,10 +1372,23 @@ public async Task GetConfidentialClientApplicati /// /// Creates an MSAL confidential client application. /// + /// Merged configuration options. + /// Whether mTLS token binding (PoP) is requested. + /// When non-null, builds an agent CCA with this app ID as + /// the ClientId and uses for credentials + /// instead of the normal client credentials. The rest of the builder configuration + /// (logging, authority, redirect URI, cache initialization) is shared with the normal + /// CCA builder path. + /// Assertion callback for agent CCAs. Required + /// when is non-null. private async Task BuildConfidentialClientApplicationAsync( MergedOptions mergedOptions, - bool isTokenBinding) + bool isTokenBinding, + string? agentAppId = null, + Func>? clientAssertionProvider = null) { + bool isAgentCca = agentAppId is not null; + mergedOptions.PrepareAuthorityInstanceForMsal(); // Validate that we have enough configuration to build an authority @@ -1404,12 +1403,27 @@ private async Task BuildConfidentialClientApplic try { + // For agent CCAs, create a fresh ConfidentialClientApplicationOptions with + // the agent's ClientId. This avoids mutating the cached MergedOptions instance + // that the blueprint CCA depends on. + ConfidentialClientApplicationOptions ccaOptions; + if (isAgentCca) + { + ccaOptions = new ConfidentialClientApplicationOptions(); + MergedOptions.UpdateConfidentialClientApplicationOptionsFromMergedOptions(mergedOptions, ccaOptions); + ccaOptions.ClientId = agentAppId; + } + else + { + ccaOptions = mergedOptions.ConfidentialClientApplicationOptions; + } + ConfidentialClientApplicationBuilder builder = ConfidentialClientApplicationBuilder - .CreateWithApplicationOptions(mergedOptions.ConfidentialClientApplicationOptions) + .CreateWithApplicationOptions(ccaOptions) .WithHttpClientFactory(_httpClientFactory) .WithLogging( new IdentityLoggerAdapter(_logger), - enablePiiLogging: mergedOptions.ConfidentialClientApplicationOptions.EnablePiiLogging) + enablePiiLogging: ccaOptions.EnablePiiLogging) .WithExperimentalFeatures(); if (_tokenCacheProvider is MsalMemoryTokenCacheProvider) @@ -1417,10 +1431,17 @@ private async Task BuildConfidentialClientApplic builder.WithCacheOptions(CacheOptions.EnableSharedCacheOptions); } + // Agent CCAs always use the shared cache: tokens survive CCA eviction and + // are found by newly-built CCAs via AcquireTokenSilent. + if (isAgentCca) + { + builder.WithCacheOptions(CacheOptions.EnableSharedCacheOptions); + } + string? currentUri = _tokenAcquisitionHost.GetCurrentRedirectUri(mergedOptions); - // The redirect URI is not needed for OBO - if (!string.IsNullOrEmpty(currentUri)) + // The redirect URI is not needed for OBO or agent flows + if (!string.IsNullOrEmpty(currentUri) && !isAgentCca) { builder.WithRedirectUri(currentUri); } @@ -1482,29 +1503,42 @@ private async Task BuildConfidentialClientApplic builder.WithAuthority(authority); } - try + // Configure credentials: agent CCAs use an assertion callback that chains + // to the blueprint CCA for Leg 1 (FMI token), while normal CCAs use the + // standard client credentials (certificate, secret, etc.). + if (isAgentCca && clientAssertionProvider is not null) { - await builder.WithClientCredentialsAsync( - mergedOptions, - _credentialsProvider, - new CredentialSourceLoaderParameters(mergedOptions.ClientId!, authority) - { - Protocol = isTokenBinding ? ProtocolNames.MtlsPop : ProtocolNames.Bearer, - }, - isTokenBinding); + builder.WithClientAssertion(clientAssertionProvider); } - catch (ArgumentException ex) when (ex.Message == IDWebErrorMessage.ClientCertificatesHaveExpiredOrCannotBeLoaded) + else { - Logger.TokenAcquisitionError( - _logger, - IDWebErrorMessage.ClientCertificatesHaveExpiredOrCannotBeLoaded, - ex); - throw; + try + { + await builder.WithClientCredentialsAsync( + mergedOptions, + _credentialsProvider, + new CredentialSourceLoaderParameters(mergedOptions.ClientId!, authority) + { + Protocol = isTokenBinding ? ProtocolNames.MtlsPop : ProtocolNames.Bearer, + }, + isTokenBinding); + } + catch (ArgumentException ex) when (ex.Message == IDWebErrorMessage.ClientCertificatesHaveExpiredOrCannotBeLoaded) + { + Logger.TokenAcquisitionError( + _logger, + IDWebErrorMessage.ClientCertificatesHaveExpiredOrCannotBeLoaded, + ex); + throw; + } } IConfidentialClientApplication app = builder.Build(); - // Initialize token cache providers + // Initialize token cache providers. + // For in-memory caches, the shared cache options above handle caching. + // For distributed caches (Redis, SQL, etc.), the provider must be + // initialized on both app and user token caches. if (!(_tokenCacheProvider is MsalMemoryTokenCacheProvider)) { _tokenCacheProvider.Initialize(app.AppTokenCache); diff --git a/tests/Microsoft.Identity.Web.Test/TokenAcquisitionTests.cs b/tests/Microsoft.Identity.Web.Test/TokenAcquisitionTests.cs index 71eb4b665..bf571262a 100644 --- a/tests/Microsoft.Identity.Web.Test/TokenAcquisitionTests.cs +++ b/tests/Microsoft.Identity.Web.Test/TokenAcquisitionTests.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; +using System.Linq; using System.Net.Http; using System.Threading; using System.Threading.Tasks; @@ -845,9 +846,13 @@ await authProvider.CreateAuthorizationHeaderForUserAsync( authorizationHeaderProviderOptions: CreateAgentIdentityOptionsWithUpn(agent2, user2Upn), claimsPrincipal: null); - // Evict ALL agent CCAs (simulating sweep clearing everything). + // Evict ALL agent CCAs from the shared dictionary. // Keep _agentUserFicAccountIds intact — silent lookup needs them to find the account. - tokenAcquisition._agentUserFicCcas.Clear(); + foreach (var key in tokenAcquisition._applicationsByAuthorityClientId.Keys + .Where(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0).ToList()) + { + tokenAcquisition._applicationsByAuthorityClientId.TryRemove(key, out _); + } // Act — new CCAs must be built, but tokens should come from shared static cache. // Each new CCA needs a Leg 1 handler (assertion callback fires on first use). @@ -955,7 +960,7 @@ await authProvider.CreateAuthorizationHeaderForUserAsync( authorizationHeaderProviderOptions: CreateAgentIdentityOptionsWithUpn(agent2, user1Upn), claimsPrincipal: null); - Assert.Equal(2, tokenAcquisition._agentUserFicCcas.Count); + Assert.Equal(2, tokenAcquisition._applicationsByAuthorityClientId.Keys.Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0)); Assert.Equal(2, tokenAcquisition._agentUserFicAccountIds.Count); // Add 3rd agent — exceeds threshold, triggers clear @@ -965,12 +970,110 @@ await authProvider.CreateAuthorizationHeaderForUserAsync( authorizationHeaderProviderOptions: CreateAgentIdentityOptionsWithUpn(agent3, user1Upn), claimsPrincipal: null); - // Assert — CCA dictionary was cleared; only the 3rd agent's account ID remains - // (written after the clear, during Leg 3 of the triggering call) - Assert.Empty(tokenAcquisition._agentUserFicCcas); + // Assert — agent CCA entries were cleared from the shared dictionary; + // only the 3rd agent's account ID remains (written after clear, during Leg 3) + Assert.Equal(0, tokenAcquisition._applicationsByAuthorityClientId.Keys.Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0)); Assert.Single(tokenAcquisition._agentUserFicAccountIds); } + /// + /// Verifies that agent CCAs are stored in the shared _applicationsByAuthorityClientId + /// dictionary alongside normal CCAs, using the ":agent:" key segment for identification. + /// This ensures agent CCAs go through the same builder path and get identical configuration + /// (logging, authority handling, cache initialization) as normal CCAs. + /// + [Fact] + public async Task AgentCca_StoredInSharedDictionary_WithAgentKeySegment() + { + // Arrange + string agent1 = Guid.NewGuid().ToString("N"); + string user1Upn = "user1@contoso.com"; + + var factory = InitTokenAcquirerFactoryForAgent(); + IServiceProvider serviceProvider = factory.Build(); + var mockHttp = serviceProvider.GetRequiredService() as MockHttpClientFactory; + IAuthorizationHeaderProvider authProvider = + serviceProvider.GetRequiredService(); + var tokenAcquisition = (TokenAcquisition)serviceProvider.GetRequiredService(); + + AddAgentUserFicMockHandlersForUser(mockHttp!, "token-a1u1", Guid.NewGuid().ToString("N"), user1Upn); + + // Act + await authProvider.CreateAuthorizationHeaderForUserAsync( + new[] { "https://graph.microsoft.com/.default" }, + authorizationHeaderProviderOptions: CreateAgentIdentityOptionsWithUpn(agent1, user1Upn), + claimsPrincipal: null); + + // Assert — the shared dictionary has agent entries (identified by ":agent:" segment) + var agentKeys = tokenAcquisition._applicationsByAuthorityClientId.Keys + .Where(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0) + .ToList(); + Assert.Single(agentKeys); + Assert.True(agentKeys[0].IndexOf(agent1, StringComparison.Ordinal) >= 0); + + // The blueprint CCA is also in the same dictionary (created lazily by assertion callback) + var blueprintKeys = tokenAcquisition._applicationsByAuthorityClientId.Keys + .Where(k => k.IndexOf(":agent:", StringComparison.Ordinal) < 0) + .ToList(); + Assert.Single(blueprintKeys); + } + + /// + /// Verifies that evicting agent CCAs from the shared dictionary does NOT affect + /// the blueprint CCA entry. Normal (non-agent) flows continue to work after + /// agent eviction without rebuilding the blueprint. + /// + [Fact] + public async Task AgentCcaEviction_DoesNotAffect_BlueprintCca() + { + // Arrange + string agent1 = Guid.NewGuid().ToString("N"); + string agent2 = Guid.NewGuid().ToString("N"); + string agent3 = Guid.NewGuid().ToString("N"); + string user1Upn = "user1@contoso.com"; + string user1Uid = Guid.NewGuid().ToString("N"); + + var factory = InitTokenAcquirerFactoryForAgent(); + IServiceProvider serviceProvider = factory.Build(); + var mockHttp = serviceProvider.GetRequiredService() as MockHttpClientFactory; + IAuthorizationHeaderProvider authProvider = + serviceProvider.GetRequiredService(); + var tokenAcquisition = (TokenAcquisition)serviceProvider.GetRequiredService(); + tokenAcquisition.AgentCcaMaxCount = 2; + + // Populate 2 agents (at threshold) + AddAgentUserFicMockHandlersForUser(mockHttp!, "token-a1", user1Uid, user1Upn); + await authProvider.CreateAuthorizationHeaderForUserAsync( + new[] { "https://graph.microsoft.com/.default" }, + authorizationHeaderProviderOptions: CreateAgentIdentityOptionsWithUpn(agent1, user1Upn), + claimsPrincipal: null); + + AddAgentUserFicMockHandlersForUser(mockHttp!, "token-a2", user1Uid, user1Upn); + await authProvider.CreateAuthorizationHeaderForUserAsync( + new[] { "https://graph.microsoft.com/.default" }, + authorizationHeaderProviderOptions: CreateAgentIdentityOptionsWithUpn(agent2, user1Upn), + claimsPrincipal: null); + + // Verify blueprint exists + int blueprintCount = tokenAcquisition._applicationsByAuthorityClientId.Keys + .Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) < 0); + Assert.Equal(1, blueprintCount); + + // Act — 3rd agent triggers eviction of agent entries + AddAgentUserFicMockHandlersForUser(mockHttp!, "token-a3", user1Uid, user1Upn); + await authProvider.CreateAuthorizationHeaderForUserAsync( + new[] { "https://graph.microsoft.com/.default" }, + authorizationHeaderProviderOptions: CreateAgentIdentityOptionsWithUpn(agent3, user1Upn), + claimsPrincipal: null); + + // Assert — agent entries cleared, but blueprint CCA still present + Assert.Equal(0, tokenAcquisition._applicationsByAuthorityClientId.Keys + .Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0)); + int blueprintCountAfter = tokenAcquisition._applicationsByAuthorityClientId.Keys + .Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) < 0); + Assert.Equal(1, blueprintCountAfter); + } + #endregion } } From 8622c5c790b016104be0f7c85d165b55334f3cdf Mon Sep 17 00:00:00 2001 From: avdunn Date: Tue, 7 Jul 2026 14:46:46 -0700 Subject: [PATCH 2/3] PR feedback --- .../TokenAcquisition.Logger.cs | 10 +- .../TokenAcquisition.cs | 94 +++++++++++-------- .../TokenAcquisitionTests.cs | 40 ++++---- 3 files changed, 81 insertions(+), 63 deletions(-) diff --git a/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.Logger.cs b/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.Logger.cs index 37fa4c3c1..d45c4ce57 100644 --- a/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.Logger.cs +++ b/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.Logger.cs @@ -114,11 +114,11 @@ public static void TokenAcquisitionMsalAuthenticationResultTime( LoggingEventId.AgentCcaCreated, "[MsIdWeb] Created new agent CCA for cache key '{CcaCacheKey}'."); - private static readonly Action s_agentCcaEviction = - LoggerMessage.Define( + private static readonly Action s_agentCcaEviction = + LoggerMessage.Define( LogLevel.Information, LoggingEventId.AgentCcaEviction, - "[MsIdWeb] Agent CCA cache cleared {EvictedCount} entries (exceeded size threshold). Remaining: {RemainingCount}."); + "[MsIdWeb] CCA dictionary cleared ({EvictedCount} entries exceeded size threshold)."); public static void AgentUserFicFlowDetected(ILogger logger, string agentAppId, string identifierType) => s_agentUserFicFlowDetected(logger, agentAppId, identifierType, null); @@ -135,8 +135,8 @@ public static void AgentUserFicAcquisitionComplete(ILogger logger, string agentA public static void AgentCcaCreated(ILogger logger, string ccaCacheKey) => s_agentCcaCreated(logger, ccaCacheKey, null); - public static void AgentCcaEviction(ILogger logger, int evictedCount, int remainingCount) - => s_agentCcaEviction(logger, evictedCount, remainingCount, null); + public static void AgentCcaEviction(ILogger logger, int evictedCount) + => s_agentCcaEviction(logger, evictedCount, null); } } } diff --git a/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs b/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs index f49ea3f29..26220193f 100644 --- a/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs +++ b/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs @@ -62,11 +62,13 @@ class OAuthConstants private readonly ConcurrentDictionary _appSemaphores = new(); /// - /// Maximum number of agent CCA instances to keep in the shared application dictionary - /// before clearing agent entries as a DOS protection measure. Tokens are stored in - /// MSAL's shared static cache, so clearing agent entries only discards lightweight - /// CCA objects — tokens remain accessible to newly-built CCAs. - /// Agent entries are identified by the ":agent:" segment in their cache key. + /// Maximum number of CCA instances to keep in the shared application dictionary + /// before clearing it as a DOS protection measure. Token data lives in external + /// caches (MSAL's shared static cache for in-memory providers, or the distributed + /// cache provider for Redis/SQL/etc.), so clearing the dictionary only discards + /// lightweight CCA objects — tokens remain accessible to newly-built CCAs. + /// Eviction is only triggered by agent CCA creation, since normal CCAs are bounded + /// by the number of configured authentication schemes. /// internal int AgentCcaMaxCount { get; set; } = 10000; @@ -264,13 +266,24 @@ private async Task AddAccountToCacheFromAuthorizationCodeInt /// Merged configuration options. /// Whether mTLS token binding (PoP) is requested. /// Callers must pass this explicitly to avoid accidental cache collisions. - /// Concatenated string of authority, client id, azure region, credential id, and token-binding flag. - private static string GetApplicationKey(MergedOptions mergedOptions, bool isTokenBinding) + /// When non-null, appends an agent-specific segment so each + /// agent CCA gets its own entry in the shared dictionary. + /// Concatenated string of authority, client id, azure region, credential id, + /// token-binding flag, and optional agent app id. + private static string GetApplicationKey(MergedOptions mergedOptions, bool isTokenBinding, string? agentAppId = null) { string credentialId = string.Join("-", mergedOptions.ClientCredentials?.Select(c => c.Id) ?? Enumerable.Empty()); string baseKey = DefaultTokenAcquirerFactoryImplementation.GetKey(mergedOptions.Authority, mergedOptions.ClientId, mergedOptions.AzureRegion) + credentialId; - return isTokenBinding ? baseKey + "-tokenBinding" : baseKey; + if (isTokenBinding) + { + baseKey += "-tokenBinding"; + } + if (agentAppId is not null) + { + baseKey += $":agent:{agentAppId}"; + } + return baseKey; } /// @@ -726,7 +739,17 @@ private async Task GetOrBuildAgentUserFicCcaAsyn string? authenticationScheme, MergedOptions mergedOptions) { - // Build the assertion callback that chains to the blueprint CCA for Leg 1. + // Fast path: if the agent CCA is already cached, return it without + // allocating a closure for the assertion callback. The callback captures + // authenticationScheme and agentAppId, producing a heap-allocated closure + // object + delegate on every call — wasteful when the CCA already exists. + string key = GetApplicationKey(mergedOptions, isTokenBinding: false, agentAppId); + if (_applicationsByAuthorityClientId.TryGetValue(key, out var cached) && cached != null) + { + return cached; + } + + // Cache miss — build the assertion callback that chains to the blueprint CCA for Leg 1. // Capture authenticationScheme so the callback resolves the correct blueprint. string? capturedAuthScheme = authenticationScheme; @@ -1306,13 +1329,7 @@ public async Task GetConfidentialClientApplicati string? agentAppId = null, Func>? clientAssertionProvider = null) { - // For agent CCAs, incorporate the agent app ID into the cache key so each agent - // gets its own CCA instance while sharing the same dictionary and lifecycle. - string key = GetApplicationKey(mergedOptions, isTokenBinding); - if (agentAppId is not null) - { - key = $"{key}:agent:{agentAppId}"; - } + string key = GetApplicationKey(mergedOptions, isTokenBinding, agentAppId); // GetOrAddAsync based on https://github.com/dotnet/runtime/issues/83636#issuecomment-1474998680 // Fast path: check if already created @@ -1334,31 +1351,21 @@ public async Task GetConfidentialClientApplicati mergedOptions, isTokenBinding, agentAppId, clientAssertionProvider); // Recompute the key as BuildConfidentialClientApplicationAsync can cause it to change. - key = GetApplicationKey(mergedOptions, isTokenBinding); - if (agentAppId is not null) - { - key = $"{key}:agent:{agentAppId}"; - } + key = GetApplicationKey(mergedOptions, isTokenBinding, agentAppId); _applicationsByAuthorityClientId[key] = newApp; - // DOS protection for agent CCAs: if too many agent entries accumulate, - // clear them. Tokens survive in MSAL's shared static cache and will be - // found by new CCAs via AcquireTokenSilent. - if (agentAppId is not null) - { - int agentCount = _applicationsByAuthorityClientId.Keys.Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0); - if (agentCount > AgentCcaMaxCount) - { - int cleared = 0; - foreach (var agentKey in _applicationsByAuthorityClientId.Keys.Where(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0).ToList()) - { - _applicationsByAuthorityClientId.TryRemove(agentKey, out _); - _appSemaphores.TryRemove(agentKey, out _); - cleared++; - } - _agentUserFicAccountIds.Clear(); - Logger.AgentCcaEviction(_logger, cleared, 0); - } + // DOS protection: if the dictionary grows beyond the threshold, clear it. + // All token data lives in external caches (MSAL's shared static cache for + // in-memory providers, or the distributed cache provider for Redis/SQL/etc.), + // so clearing the dictionary only discards lightweight CCA objects — tokens + // remain accessible to newly-built CCAs. + if (agentAppId is not null && _applicationsByAuthorityClientId.Count > AgentCcaMaxCount) + { + int cleared = _applicationsByAuthorityClientId.Count; + _applicationsByAuthorityClientId.Clear(); + _appSemaphores.Clear(); + _agentUserFicAccountIds.Clear(); + Logger.AgentCcaEviction(_logger, cleared); } return newApp; @@ -1387,6 +1394,15 @@ private async Task BuildConfidentialClientApplic string? agentAppId = null, Func>? clientAssertionProvider = null) { + // agentAppId and clientAssertionProvider must both be null or both be non-null. + // Agent CCAs require an assertion callback for Leg 1 (FMI token), and the callback + // is only meaningful in the context of an agent CCA. + if ((agentAppId is null) != (clientAssertionProvider is null)) + { + throw new ArgumentException( + "agentAppId and clientAssertionProvider must both be provided or both be null."); + } + bool isAgentCca = agentAppId is not null; mergedOptions.PrepareAuthorityInstanceForMsal(); diff --git a/tests/Microsoft.Identity.Web.Test/TokenAcquisitionTests.cs b/tests/Microsoft.Identity.Web.Test/TokenAcquisitionTests.cs index bf571262a..d16344366 100644 --- a/tests/Microsoft.Identity.Web.Test/TokenAcquisitionTests.cs +++ b/tests/Microsoft.Identity.Web.Test/TokenAcquisitionTests.cs @@ -938,8 +938,12 @@ public async Task AgentCcaEviction_ClearsDictionaryAtThreshold() serviceProvider.GetRequiredService(); var tokenAcquisition = (TokenAcquisition)serviceProvider.GetRequiredService(); - // Set threshold to 2 so the 3rd agent triggers a clear - tokenAcquisition.AgentCcaMaxCount = 2; + // Set threshold to 4 so the 3rd agent triggers a clear: + // 1 blueprint + 2 agents = 3 entries (≤ 4), 1 blueprint + 3 agents = 4 entries (≤ 4) + // but after adding the 3rd agent the count becomes 4 which equals the threshold. + // We use 3 so that: blueprint + 2 agents = 3 ≤ 3 (no eviction), + // blueprint + 3 agents = 4 > 3 (triggers eviction). + tokenAcquisition.AgentCcaMaxCount = 3; string agent1 = Guid.NewGuid().ToString("N"); string agent2 = Guid.NewGuid().ToString("N"); @@ -970,9 +974,11 @@ await authProvider.CreateAuthorizationHeaderForUserAsync( authorizationHeaderProviderOptions: CreateAgentIdentityOptionsWithUpn(agent3, user1Upn), claimsPrincipal: null); - // Assert — agent CCA entries were cleared from the shared dictionary; - // only the 3rd agent's account ID remains (written after clear, during Leg 3) - Assert.Equal(0, tokenAcquisition._applicationsByAuthorityClientId.Keys.Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0)); + // Assert — dictionary was cleared by DOS eviction, then repopulated during + // the rest of the 3rd agent's flow (blueprint rebuilds for Leg 1). + // Previous agent CCAs are gone; only the rebuilt blueprint remains. + Assert.Equal(0, tokenAcquisition._applicationsByAuthorityClientId.Keys + .Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0)); Assert.Single(tokenAcquisition._agentUserFicAccountIds); } @@ -1019,12 +1025,12 @@ await authProvider.CreateAuthorizationHeaderForUserAsync( } /// - /// Verifies that evicting agent CCAs from the shared dictionary does NOT affect - /// the blueprint CCA entry. Normal (non-agent) flows continue to work after - /// agent eviction without rebuilding the blueprint. + /// Verifies that DOS eviction clears the entire CCA dictionary (including blueprint), + /// but tokens survive in MSAL's shared static cache. After eviction, new CCAs are + /// rebuilt lazily and can still retrieve cached tokens via AcquireTokenSilent. /// [Fact] - public async Task AgentCcaEviction_DoesNotAffect_BlueprintCca() + public async Task AgentCcaEviction_ClearsDictionary_TokensSurvive() { // Arrange string agent1 = Guid.NewGuid().ToString("N"); @@ -1039,7 +1045,7 @@ public async Task AgentCcaEviction_DoesNotAffect_BlueprintCca() IAuthorizationHeaderProvider authProvider = serviceProvider.GetRequiredService(); var tokenAcquisition = (TokenAcquisition)serviceProvider.GetRequiredService(); - tokenAcquisition.AgentCcaMaxCount = 2; + tokenAcquisition.AgentCcaMaxCount = 3; // Populate 2 agents (at threshold) AddAgentUserFicMockHandlersForUser(mockHttp!, "token-a1", user1Uid, user1Upn); @@ -1054,24 +1060,20 @@ await authProvider.CreateAuthorizationHeaderForUserAsync( authorizationHeaderProviderOptions: CreateAgentIdentityOptionsWithUpn(agent2, user1Upn), claimsPrincipal: null); - // Verify blueprint exists - int blueprintCount = tokenAcquisition._applicationsByAuthorityClientId.Keys - .Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) < 0); - Assert.Equal(1, blueprintCount); + // Verify dictionary has entries (2 agents + 1 blueprint = 3) + Assert.True(tokenAcquisition._applicationsByAuthorityClientId.Count >= 3); - // Act — 3rd agent triggers eviction of agent entries + // Act — 3rd agent triggers eviction (clears entire dictionary) AddAgentUserFicMockHandlersForUser(mockHttp!, "token-a3", user1Uid, user1Upn); await authProvider.CreateAuthorizationHeaderForUserAsync( new[] { "https://graph.microsoft.com/.default" }, authorizationHeaderProviderOptions: CreateAgentIdentityOptionsWithUpn(agent3, user1Upn), claimsPrincipal: null); - // Assert — agent entries cleared, but blueprint CCA still present + // Assert — dictionary was cleared by eviction, then blueprint was rebuilt + // during the 3rd agent's Leg 1. Previous agent CCAs are gone. Assert.Equal(0, tokenAcquisition._applicationsByAuthorityClientId.Keys .Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) >= 0)); - int blueprintCountAfter = tokenAcquisition._applicationsByAuthorityClientId.Keys - .Count(k => k.IndexOf(":agent:", StringComparison.Ordinal) < 0); - Assert.Equal(1, blueprintCountAfter); } #endregion From 274d599ed30fc26150869f6a653dbc70aa0b1223 Mon Sep 17 00:00:00 2001 From: avdunn Date: Wed, 8 Jul 2026 08:32:27 -0700 Subject: [PATCH 3/3] PR feedback --- .../TokenAcquisition.cs | 34 +++++++++++-------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs b/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs index 26220193f..e095c99bc 100644 --- a/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs +++ b/src/Microsoft.Identity.Web.TokenAcquisition/TokenAcquisition.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. +// Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. using System; @@ -11,6 +11,7 @@ using System.Net.Http; using System.Security.Claims; using System.Security.Cryptography.X509Certificates; +using System.Text; using System.Text.Json; using System.Threading; using System.Threading.Tasks; @@ -274,16 +275,19 @@ private static string GetApplicationKey(MergedOptions mergedOptions, bool isToke { string credentialId = string.Join("-", mergedOptions.ClientCredentials?.Select(c => c.Id) ?? Enumerable.Empty()); - string baseKey = DefaultTokenAcquirerFactoryImplementation.GetKey(mergedOptions.Authority, mergedOptions.ClientId, mergedOptions.AzureRegion) + credentialId; + var keyBuilder = new StringBuilder( + DefaultTokenAcquirerFactoryImplementation.GetKey(mergedOptions.Authority, mergedOptions.ClientId, mergedOptions.AzureRegion)); + keyBuilder.Append(credentialId); if (isTokenBinding) { - baseKey += "-tokenBinding"; + keyBuilder.Append("-tokenBinding"); } if (agentAppId is not null) { - baseKey += $":agent:{agentAppId}"; + keyBuilder.Append(":agent:"); + keyBuilder.Append(agentAppId); } - return baseKey; + return keyBuilder.ToString(); } /// @@ -1327,7 +1331,7 @@ public async Task GetConfidentialClientApplicati MergedOptions mergedOptions, bool isTokenBinding, string? agentAppId = null, - Func>? clientAssertionProvider = null) + Func>? agenticAssertionProvider = null) { string key = GetApplicationKey(mergedOptions, isTokenBinding, agentAppId); @@ -1348,7 +1352,7 @@ public async Task GetConfidentialClientApplicati // Build and store the application var newApp = await BuildConfidentialClientApplicationAsync( - mergedOptions, isTokenBinding, agentAppId, clientAssertionProvider); + mergedOptions, isTokenBinding, agentAppId, agenticAssertionProvider); // Recompute the key as BuildConfidentialClientApplicationAsync can cause it to change. key = GetApplicationKey(mergedOptions, isTokenBinding, agentAppId); @@ -1382,25 +1386,25 @@ public async Task GetConfidentialClientApplicati /// Merged configuration options. /// Whether mTLS token binding (PoP) is requested. /// When non-null, builds an agent CCA with this app ID as - /// the ClientId and uses for credentials + /// the ClientId and uses for credentials /// instead of the normal client credentials. The rest of the builder configuration /// (logging, authority, redirect URI, cache initialization) is shared with the normal /// CCA builder path. - /// Assertion callback for agent CCAs. Required + /// Assertion callback for agent CCAs. Required /// when is non-null. private async Task BuildConfidentialClientApplicationAsync( MergedOptions mergedOptions, bool isTokenBinding, string? agentAppId = null, - Func>? clientAssertionProvider = null) + Func>? agenticAssertionProvider = null) { - // agentAppId and clientAssertionProvider must both be null or both be non-null. + // agentAppId and agenticAssertionProvider must both be null or both be non-null. // Agent CCAs require an assertion callback for Leg 1 (FMI token), and the callback // is only meaningful in the context of an agent CCA. - if ((agentAppId is null) != (clientAssertionProvider is null)) + if ((agentAppId is null) != (agenticAssertionProvider is null)) { throw new ArgumentException( - "agentAppId and clientAssertionProvider must both be provided or both be null."); + "agentAppId and agenticAssertionProvider must both be provided or both be null."); } bool isAgentCca = agentAppId is not null; @@ -1522,9 +1526,9 @@ private async Task BuildConfidentialClientApplic // Configure credentials: agent CCAs use an assertion callback that chains // to the blueprint CCA for Leg 1 (FMI token), while normal CCAs use the // standard client credentials (certificate, secret, etc.). - if (isAgentCca && clientAssertionProvider is not null) + if (isAgentCca && agenticAssertionProvider is not null) { - builder.WithClientAssertion(clientAssertionProvider); + builder.WithClientAssertion(agenticAssertionProvider); } else {