diff --git a/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs b/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs index b49bb51c9..dec0d20e5 100644 --- a/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs +++ b/src/ModelContextProtocol.Core/Authentication/ClientOAuthProvider.cs @@ -48,13 +48,24 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient private string? _tokenEndpointAuthMethod; private ITokenCache _tokenCache; private AuthorizationServerMetadata? _authServerMetadata; + + // Coalesces concurrent token acquisition so that when multiple in-flight requests observe an + // expired token (or a 401) at the same time, only the first runs the refresh/authorization flow + // while the others await its result. This also serializes all reads and writes to the mutable auth + // state below (_authServerMetadata, _clientId, _clientSecret, _tokenEndpointAuthMethod) as well as + // the accumulated scope set and step-up tracking (_accumulatedScopes, _hasAttemptedStepUp), so those + // fields need no separate lock. + // + // Intentionally not disposed: this instance is only ever used via WaitAsync/Release (never its + // AvailableWaitHandle), so SemaphoreSlim allocates no unmanaged resource and there is nothing to + // dispose. Do not access AvailableWaitHandle, or this field will need deterministic disposal. + private readonly SemaphoreSlim _tokenAcquisitionLock = new(1, 1); // The accumulated scope set lives for this provider's lifetime and is intentionally not keyed by // resource or authorization server. This is safe today because one ClientOAuthProvider is created // per HttpClientTransport, i.e. per endpoint/resource. If a provider were ever reused across // multiple resources or auth servers, accumulated scopes could be sent to a server that rejects // them (invalid_scope). Accumulation is scoped per "resource and operation" combination (SEP-2350). private readonly HashSet _accumulatedScopes = new(StringComparer.Ordinal); - private readonly object _scopeAccumulatorLock = new(); private bool _hasAttemptedStepUp; /// @@ -145,7 +156,10 @@ internal override async Task SendAsync(HttpRequestMessage r if (ShouldRetryWithNewAccessToken(response)) { - return await HandleUnauthorizedResponseAsync(request, message, response, attemptedRefresh, cancellationToken).ConfigureAwait(false); + // Capture the token that produced this challenge so the retry path can detect whether + // another concurrent caller already replaced it in the cache. + var usedAccessToken = request.Headers.Authorization?.Parameter; + return await HandleUnauthorizedResponseAsync(request, message, response, attemptedRefresh, usedAccessToken, cancellationToken).ConfigureAwait(false); } return response; @@ -161,15 +175,32 @@ internal override async Task SendAsync(HttpRequestMessage r return (tokens.AccessToken, false); } - // Try to refresh the access token if it is invalid and we have a refresh token. - if (_authServerMetadata is not null && tokens?.RefreshToken is { Length: > 0 } refreshToken) + // A refresh is only possible if we have both the auth server metadata and a refresh token. + if (_authServerMetadata is null || tokens?.RefreshToken is not { Length: > 0 }) + { + // No valid token - auth handler will trigger the 401 flow + return (null, false); + } + + // Serialize the refresh so concurrent callers that all saw the expired token don't each fire + // their own refresh. Waiters re-check the cache after acquiring the lock and reuse the token + // produced by whoever refreshed first. + using var _ = await _tokenAcquisitionLock.LockAsync(cancellationToken).ConfigureAwait(false); + + var current = await _tokenCache.GetTokensAsync(cancellationToken).ConfigureAwait(false); + if (current is not null && !current.IsExpired) + { + return (current.AccessToken, true); + } + + if (_authServerMetadata is not null && current?.RefreshToken is { Length: > 0 } refreshToken) { var accessToken = await RefreshTokensAsync(refreshToken, resourceUri.ToString(), _authServerMetadata, cancellationToken).ConfigureAwait(false); return (accessToken, true); } // No valid token - auth handler will trigger the 401 flow - return (null, false); + return (null, true); } private static bool ShouldRetryWithNewAccessToken(HttpResponseMessage response) @@ -208,6 +239,7 @@ private async Task HandleUnauthorizedResponseAsync( JsonRpcMessage? originalJsonRpcMessage, HttpResponseMessage response, bool attemptedRefresh, + string? usedAccessToken, CancellationToken cancellationToken) { if (response.Headers.WwwAuthenticate.Count == 0) @@ -220,7 +252,7 @@ private async Task HandleUnauthorizedResponseAsync( throw new McpException($"The server does not support the '{BearerScheme}' authentication scheme. Server supports: [{serverSchemes}]."); } - var accessToken = await GetAccessTokenAsync(response, attemptedRefresh, cancellationToken).ConfigureAwait(false); + var accessToken = await GetAccessTokenAsync(response, attemptedRefresh, usedAccessToken, cancellationToken).ConfigureAwait(false); using var retryRequest = new HttpRequestMessage(originalRequest.Method, originalRequest.RequestUri); @@ -241,8 +273,31 @@ private async Task HandleUnauthorizedResponseAsync( /// /// The HTTP response that triggered the authentication challenge. /// Indicates whether a token refresh has already been attempted. + /// The access token that produced the challenge, or if none was sent. /// The to monitor for cancellation requests. - private async Task GetAccessTokenAsync(HttpResponseMessage response, bool attemptedRefresh, CancellationToken cancellationToken) + private async Task GetAccessTokenAsync(HttpResponseMessage response, bool attemptedRefresh, string? usedAccessToken, CancellationToken cancellationToken) + { + // Serialize the authorization flow so concurrent 401/403 challenges don't each run a full + // refresh/registration/interactive authorization and race on the shared auth state below. + using var _ = await _tokenAcquisitionLock.LockAsync(cancellationToken).ConfigureAwait(false); + + // While we waited for the lock, another concurrent caller may have already acquired or + // refreshed the token. Reuse the cached token if it is both still valid and different from + // the one that produced this challenge (otherwise we'd just replay the rejected token). When + // no token was sent (usedAccessToken is null, e.g. concurrent cold-start requests), any valid + // cached token was obtained by another caller and is safe to reuse. This is limited to 401; a + // 403 insufficient_scope challenge must still run the step-up flow. + if (response.StatusCode == System.Net.HttpStatusCode.Unauthorized && + await _tokenCache.GetTokensAsync(cancellationToken).ConfigureAwait(false) is { IsExpired: false } cached && + !string.Equals(cached.AccessToken, usedAccessToken, StringComparison.Ordinal)) + { + return cached.AccessToken; + } + + return await GetAccessTokenCoreAsync(response, attemptedRefresh, usedAccessToken, cancellationToken).ConfigureAwait(false); + } + + private async Task GetAccessTokenCoreAsync(HttpResponseMessage response, bool attemptedRefresh, string? usedAccessToken, CancellationToken cancellationToken) { // Get available authorization servers from the 401 or 403 response var protectedResourceMetadata = await ExtractProtectedResourceMetadata(response, cancellationToken).ConfigureAwait(false); @@ -262,17 +317,26 @@ private async Task GetAccessTokenAsync(HttpResponseMessage response, boo if (response.StatusCode == System.Net.HttpStatusCode.Forbidden) { bool introducesNewScopes = ChallengeIntroducesNewScopes(protectedResourceMetadata); - lock (_scopeAccumulatorLock) + if (_hasAttemptedStepUp && !introducesNewScopes) { - if (_hasAttemptedStepUp && !introducesNewScopes) + // A step-up has already run and this challenge asks for nothing new. If that step-up + // produced a different, still-valid token (for example another concurrent caller ran + // it while this one waited on the lock), reuse that token instead of failing, since it + // already reflects the accumulated scopes. Only fail when there is no newer token to + // try, which is the genuine repeated-failure case where the stepped-up token itself + // was rejected again. + if (await _tokenCache.GetTokensAsync(cancellationToken).ConfigureAwait(false) is { IsExpired: false } steppedUpToken && + !string.Equals(steppedUpToken.AccessToken, usedAccessToken, StringComparison.Ordinal)) { - ThrowFailedToHandleUnauthorizedResponse( - "A repeated insufficient_scope challenge added no scope beyond those already requested, " + - "so step-up authorization cannot satisfy the request."); + return steppedUpToken.AccessToken; } - _hasAttemptedStepUp = true; + ThrowFailedToHandleUnauthorizedResponse( + "A repeated insufficient_scope challenge added no scope beyond those already requested, " + + "so step-up authorization cannot satisfy the request."); } + + _hasAttemptedStepUp = true; } // Convert string URIs to Uri objects for the selector @@ -812,13 +876,10 @@ private async Task PerformDynamicClientRegistrationAsync( if (currentOperationScopes.Count == 0) { - lock (_scopeAccumulatorLock) - { - // If we have previously requested scopes but nothing new, return the accumulated set. - return _accumulatedScopes.Count > 0 - ? string.Join(" ", _accumulatedScopes.OrderBy(s => s, StringComparer.Ordinal)) - : null; - } + // If we have previously requested scopes but nothing new, return the accumulated set. + return _accumulatedScopes.Count > 0 + ? string.Join(" ", _accumulatedScopes.OrderBy(s => s, StringComparer.Ordinal)) + : null; } // Per SEP-2350: Compute the union of previously requested scopes and newly challenged scopes @@ -827,16 +888,13 @@ private async Task PerformDynamicClientRegistrationAsync( // offline_access (AugmentScopeWithOfflineAccess) and any ScopeSelector are applied per request // in ComputeEffectiveScope and are intentionally not accumulated, so the selector always sees // the full union and the operation stays idempotent. - lock (_scopeAccumulatorLock) + foreach (var scope in currentOperationScopes) { - foreach (var scope in currentOperationScopes) - { - _accumulatedScopes.Add(scope); - } - - // Sort scopes for stable, deterministic output (scopes are unordered per RFC 6749 §3.3). - return string.Join(" ", _accumulatedScopes.OrderBy(s => s, StringComparer.Ordinal)); + _accumulatedScopes.Add(scope); } + + // Sort scopes for stable, deterministic output (scopes are unordered per RFC 6749 §3.3). + return string.Join(" ", _accumulatedScopes.OrderBy(s => s, StringComparer.Ordinal)); } /// @@ -881,14 +939,11 @@ private bool ChallengeIntroducesNewScopes(ProtectedResourceMetadata protectedRes return false; } - lock (_scopeAccumulatorLock) + foreach (var scope in currentOperationScopes) { - foreach (var scope in currentOperationScopes) + if (!_accumulatedScopes.Contains(scope)) { - if (!_accumulatedScopes.Contains(scope)) - { - return true; - } + return true; } } diff --git a/src/ModelContextProtocol.Core/Authentication/ITokenCache.cs b/src/ModelContextProtocol.Core/Authentication/ITokenCache.cs index 3dc6e6351..62d57b913 100644 --- a/src/ModelContextProtocol.Core/Authentication/ITokenCache.cs +++ b/src/ModelContextProtocol.Core/Authentication/ITokenCache.cs @@ -3,6 +3,11 @@ namespace ModelContextProtocol.Authentication; /// /// Allows the client to cache access tokens beyond the lifetime of the transport. /// +/// +/// Implementations must be safe for concurrent use. A single cache instance may be shared by multiple +/// in-flight requests, and in particular can be invoked concurrently +/// (it is called on the request hot path without holding the provider's token-acquisition lock). +/// public interface ITokenCache { /// diff --git a/src/ModelContextProtocol.Core/Authentication/IdentityAssertionGrantProvider.cs b/src/ModelContextProtocol.Core/Authentication/IdentityAssertionGrantProvider.cs index 359c66fce..f881b5f7c 100644 --- a/src/ModelContextProtocol.Core/Authentication/IdentityAssertionGrantProvider.cs +++ b/src/ModelContextProtocol.Core/Authentication/IdentityAssertionGrantProvider.cs @@ -28,6 +28,12 @@ namespace ModelContextProtocol.Authentication; /// via the RFC 7523 JWT Bearer grant. /// /// +/// +/// Concurrency: a single provider instance may be shared across concurrent requests. Token +/// acquisition is coalesced through an internal lock, so if several callers observe an expired or +/// absent token at the same time, only one runs the exchange flow and the others await and reuse +/// its result. The cached token is refreshed at most once per expiry. +/// /// /// /// @@ -56,6 +62,15 @@ public sealed class IdentityAssertionGrantProvider private TokenContainer? _cachedTokens; + // Coalesces concurrent token acquisition so that when multiple in-flight requests observe an + // expired/absent token at the same time, only the first runs the exchange flow while the others + // await its result. Also serializes writes to _cachedTokens and _resolvedIdpTokenEndpoint. + // + // Intentionally not disposed: this instance is only ever used via WaitAsync/Wait/Release (never + // its AvailableWaitHandle), so SemaphoreSlim allocates no unmanaged resource and there is nothing + // to dispose. Do not access AvailableWaitHandle, or this field will need deterministic disposal. + private readonly SemaphoreSlim _tokenAcquisitionLock = new(1, 1); + /// /// Initializes a new instance of the class. /// @@ -105,12 +120,33 @@ public async Task GetAccessTokenAsync( Uri authorizationServerUrl, CancellationToken cancellationToken = default) { - // Return cached token if still valid + // Return cached token if still valid. Read the field once into a local so a concurrent + // InvalidateCache (which nulls _cachedTokens) cannot turn this lock-free check into a null + // dereference or a null return between the null check and the return. + var cachedBeforeLock = _cachedTokens; + if (cachedBeforeLock is not null && !cachedBeforeLock.IsExpired) + { + return cachedBeforeLock; + } + + // Serialize the exchange so concurrent callers that all saw the expired/absent token don't + // each run the full multi-step flow. Waiters re-check the cache after acquiring the lock and + // reuse the token produced by whoever ran the exchange first. + using var _ = await _tokenAcquisitionLock.LockAsync(cancellationToken).ConfigureAwait(false); + if (_cachedTokens is not null && !_cachedTokens.IsExpired) { return _cachedTokens; } + return await AcquireAccessTokenAsync(resourceUrl, authorizationServerUrl, cancellationToken).ConfigureAwait(false); + } + + private async Task AcquireAccessTokenAsync( + Uri resourceUrl, + Uri authorizationServerUrl, + CancellationToken cancellationToken) + { _logger.LogDebug("Starting Cross-Application Access flow for resource {ResourceUrl}", resourceUrl); // Step 1: Discover MCP authorization server metadata to find the token endpoint @@ -173,9 +209,21 @@ public async Task GetAccessTokenAsync( /// /// Clears any cached tokens, forcing a fresh token exchange on the next call to . /// + /// + /// This blocks until any token acquisition that is currently in progress completes, so that the + /// invalidation is not silently overwritten by a concurrent exchange storing a freshly obtained token. + /// public void InvalidateCache() { - _cachedTokens = null; + _tokenAcquisitionLock.Wait(); + try + { + _cachedTokens = null; + } + finally + { + _tokenAcquisitionLock.Release(); + } } private string? _resolvedIdpTokenEndpoint; diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/AuthTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/AuthTests.cs index 18e1a8e34..344654162 100644 --- a/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/AuthTests.cs +++ b/tests/ModelContextProtocol.AspNetCore.Tests/OAuth/AuthTests.cs @@ -11,6 +11,7 @@ using ModelContextProtocol.Client; using ModelContextProtocol.Protocol; using ModelContextProtocol.Server; +using ModelContextProtocol.Tests.Utils; using System.Net; using System.Net.Http.Json; using System.Security.Claims; @@ -763,6 +764,120 @@ public async Task AuthorizationFlow_AccumulatesScopesAcrossMultipleStepUps() Assert.Contains("files:write", thirdScopeSet); } + [Fact] + public async Task AuthorizationFlow_ConcurrentStepUps_ReuseSteppedUpToken_WhenChallengeAddsNoNewScope() + { + // Two concurrent calls to the same tool both receive the same insufficient_scope challenge + // before either has stepped up. They serialize on the provider's token acquisition lock: the + // first runs the step-up and caches the broader token, and the second must reuse that token + // instead of failing as a "repeated" challenge. Only one interactive step-up should run. + + Builder.Services.AddMcpServer() + .WithTools([ + McpServerTool.Create([McpServerTool(Name = "read-tool")] + (ClaimsPrincipal user) => + { + return "Read tool executed."; + }), + ]); + + List requestedScopes = []; + var scopeLock = new object(); + + // Release both initial challenges only after both concurrent calls have reached the server, so + // the second caller is guaranteed to be waiting on the token lock while the first steps up. + var bothChallengesReached = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + int challengesReached = 0; + + await using var app = await StartMcpServerAsync(configureMiddleware: app => + { + app.Use(async (context, next) => + { + if (context.Request.Method == HttpMethods.Post && context.Request.Path == "/") + { + context.Request.EnableBuffering(); + + var message = await JsonSerializer.DeserializeAsync( + context.Request.Body, + McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(JsonRpcMessage)), + context.RequestAborted) as JsonRpcMessage; + + context.Request.Body.Position = 0; + + if (message is JsonRpcRequest request && request.Method == "tools/call") + { + var toolCallParams = JsonSerializer.Deserialize( + request.Params, + McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(CallToolRequestParams))) as CallToolRequestParams; + + var user = context.User; + var scopeClaim = user.FindFirst("scope")?.Value ?? ""; + var scopeSet = new HashSet(scopeClaim.Split(' ')); + + if (toolCallParams?.Name == "read-tool" && !scopeSet.Contains("files:read")) + { + if (Interlocked.Increment(ref challengesReached) == 2) + { + bothChallengesReached.TrySetResult(); + } + + await bothChallengesReached.Task.WaitAsync(TestConstants.DefaultTimeout, context.RequestAborted); + + context.Response.StatusCode = StatusCodes.Status403Forbidden; + context.Response.Headers.WWWAuthenticate = $"Bearer error=\"insufficient_scope\", resource_metadata=\"{McpServerUrl}/.well-known/oauth-protected-resource\", scope=\"files:read\""; + await context.Response.StartAsync(context.RequestAborted); + await context.Response.Body.FlushAsync(context.RequestAborted); + return; + } + } + } + + await next(context); + }); + }); + + await using var transport = new HttpClientTransport(new() + { + Endpoint = new(McpServerUrl), + OAuth = new() + { + ClientId = "demo-client", + ClientSecret = "demo-secret", + RedirectUri = new Uri("http://localhost:1179/callback"), + AuthorizationRedirectDelegate = (uri, redirect, ct) => + { + var query = QueryHelpers.ParseQuery(uri.Query); + lock (scopeLock) + { + requestedScopes.Add(query["scope"].ToString()); + } + return HandleAuthorizationUrlAsync(uri, redirect, ct); + }, + }, + }, HttpClient, LoggerFactory); + + await using var client = await McpClient.CreateAsync( + transport, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken); + + // Initial connect requests "mcp:tools" from protected resource metadata. + Assert.Single(requestedScopes); + Assert.Equal("mcp:tools", requestedScopes[0]); + + var firstCall = client.CallToolAsync("read-tool", cancellationToken: TestContext.Current.CancellationToken).AsTask(); + var secondCall = client.CallToolAsync("read-tool", cancellationToken: TestContext.Current.CancellationToken).AsTask(); + + var results = await Task.WhenAll(firstCall, secondCall); + + Assert.Equal("Read tool executed.", results[0].Content[0].ToString()); + Assert.Equal("Read tool executed.", results[1].Content[0].ToString()); + + // Only one interactive step-up should have run; the second caller reused the token from the first. + Assert.Equal(2, requestedScopes.Count); + var stepUpScopes = new HashSet(requestedScopes[1]!.Split(' ')); + Assert.Contains("mcp:tools", stepUpScopes); + Assert.Contains("files:read", stepUpScopes); + } + [Fact] public async Task AuthorizationFlow_StopsSteppingUpWhenChallengeAddsNoNewScope() { diff --git a/tests/ModelContextProtocol.Tests/IdentityAssertionGrantTests.cs b/tests/ModelContextProtocol.Tests/IdentityAssertionGrantTests.cs index 44afcceb6..cfff60589 100644 --- a/tests/ModelContextProtocol.Tests/IdentityAssertionGrantTests.cs +++ b/tests/ModelContextProtocol.Tests/IdentityAssertionGrantTests.cs @@ -313,6 +313,88 @@ public void IdentityAssertionGrantProvider_MissingIdpConfig_ThrowsArgumentExcept _httpClient)); } + [Fact] + public async Task IdentityAssertionGrantProvider_ConcurrentCallers_RunExchangeOnce() + { + // Gate the first in-flight flow so multiple callers overlap while the first holds the + // acquisition lock. Without coalescing, each concurrent caller would run its own exchange. + var firstEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + var mcpTokenCallCount = 0; + _mockHandler.AsyncHandler = async request => + { + var url = request.RequestUri!.ToString(); + if (url.Contains(".well-known")) + { + return JsonResponse(HttpStatusCode.OK, new JsonObject + { + ["authorization_endpoint"] = "https://auth.example.com/authorize", + ["token_endpoint"] = "https://auth.example.com/token", + }); + } + + if (url.Contains("idp.example.com")) + { + return JsonResponse(HttpStatusCode.OK, new JsonObject + { + ["access_token"] = "mock-jag", + ["issued_token_type"] = "urn:ietf:params:oauth:token-type:id-jag", + ["token_type"] = "N_A", + }); + } + + // MCP token endpoint: this is the exchange we expect to run exactly once. + if (Interlocked.Increment(ref mcpTokenCallCount) == 1) + { + firstEntered.TrySetResult(true); + await release.Task; + } + + return JsonResponse(HttpStatusCode.OK, new JsonObject + { + ["access_token"] = "final-access-token", + ["token_type"] = "Bearer", + ["expires_in"] = 3600, + }); + }; + + var idTokenCallCount = 0; + var provider = new IdentityAssertionGrantProvider( + new IdentityAssertionGrantProviderOptions + { + ClientId = "mcp-client-id", + IdpTokenEndpoint = "https://idp.example.com/token", + IdpClientId = "idp-client-id", + IdTokenCallback = (_, _) => + { + Interlocked.Increment(ref idTokenCallCount); + return Task.FromResult("mock-id-token"); + }, + }, + _httpClient); + + var ct = TestContext.Current.CancellationToken; + var resourceUrl = new Uri("https://resource.example.com"); + var authUrl = new Uri("https://auth.example.com"); + + var tasks = Enumerable.Range(0, 8) + .Select(_ => provider.GetAccessTokenAsync(resourceUrl, authUrl, ct)) + .ToArray(); + + // Wait until the first flow is inside the exchange (holding the lock), then let it finish. + var entered = await Task.WhenAny(firstEntered.Task, Task.Delay(TimeSpan.FromSeconds(30), ct)); + Assert.Same(firstEntered.Task, entered); + release.SetResult(true); + + var results = await Task.WhenAll(tasks); + + Assert.Equal(1, mcpTokenCallCount); + Assert.Equal(1, idTokenCallCount); + Assert.All(results, r => Assert.Same(results[0], r)); + Assert.Equal("final-access-token", results[0].AccessToken); + } + #endregion #region IdentityAssertionGrantException Tests