Make auth provider token caches thread-safe (#1595) (#1708)

Co-authored-by: Tarek Mahmoud Sayed <tarekms@ntdev.microsoft.com>
This commit is contained in:
Tarek Mahmoud Sayed 2026-07-22 09:43:14 -07:00 • committed by GitHub
commit 28806b1b6e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 341 additions and 36 deletions

View file

@ -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<string> _accumulatedScopes = new(StringComparer.Ordinal);
private readonly object _scopeAccumulatorLock = new();
private bool _hasAttemptedStepUp;
/// <summary>
@ -145,7 +156,10 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
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 sealed partial class ClientOAuthProvider : McpHttpClient
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 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
JsonRpcMessage? originalJsonRpcMessage,
HttpResponseMessage response,
bool attemptedRefresh,
string? usedAccessToken,
CancellationToken cancellationToken)
{
if (response.Headers.WwwAuthenticate.Count == 0)
@ -220,7 +252,7 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
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 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
/// </summary>
/// <param name="response">The HTTP response that triggered the authentication challenge.</param>
/// <param name="attemptedRefresh">Indicates whether a token refresh has already been attempted.</param>
/// <param name="usedAccessToken">The access token that produced the challenge, or <see langword="null"/> if none was sent.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
private async Task<string> GetAccessTokenAsync(HttpResponseMessage response, bool attemptedRefresh, CancellationToken cancellationToken)
private async Task<string> 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<string> 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,10 +317,20 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
if (response.StatusCode == System.Net.HttpStatusCode.Forbidden)
{
bool introducesNewScopes = ChallengeIntroducesNewScopes(protectedResourceMetadata);
lock (_scopeAccumulatorLock)
{
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))
{
return steppedUpToken.AccessToken;
}
ThrowFailedToHandleUnauthorizedResponse(
"A repeated insufficient_scope challenge added no scope beyond those already requested, " +
"so step-up authorization cannot satisfy the request.");
@ -273,7 +338,6 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
_hasAttemptedStepUp = true;
}
}
// Convert string URIs to Uri objects for the selector
List<Uri> authServerUris = [];
@ -811,15 +875,12 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
var currentOperationScopes = GetCurrentOperationScopes(protectedResourceMetadata);
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;
}
}
// Per SEP-2350: Compute the union of previously requested scopes and newly challenged scopes
// to avoid losing permissions needed for other operations during step-up authorization.
@ -827,8 +888,6 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
// 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)
{
_accumulatedScopes.Add(scope);
@ -837,7 +896,6 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
// Sort scopes for stable, deterministic output (scopes are unordered per RFC 6749 §3.3).
return string.Join(" ", _accumulatedScopes.OrderBy(s => s, StringComparer.Ordinal));
}
}
/// <summary>
/// Determines the scopes required for the current operation, preferring the <c>WWW-Authenticate</c>
@ -881,8 +939,6 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
return false;
}
lock (_scopeAccumulatorLock)
{
foreach (var scope in currentOperationScopes)
{
if (!_accumulatedScopes.Contains(scope))
@ -890,7 +946,6 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
return true;
}
}
}
return false;
}

View file

@ -3,6 +3,11 @@ namespace ModelContextProtocol.Authentication;
/// <summary>
/// Allows the client to cache access tokens beyond the lifetime of the transport.
/// </summary>
/// <remarks>
/// Implementations must be safe for concurrent use. A single cache instance may be shared by multiple
/// in-flight requests, and <see cref="GetTokensAsync"/> in particular can be invoked concurrently
/// (it is called on the request hot path without holding the provider's token-acquisition lock).
/// </remarks>
public interface ITokenCache
{
/// <summary>

View file

@ -28,6 +28,12 @@ namespace ModelContextProtocol.Authentication;
/// via the RFC 7523 JWT Bearer grant.
/// </description></item>
/// </list>
/// <para>
/// 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.
/// </para>
/// </remarks>
/// <example>
/// <code>
@ -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);
/// <summary>
/// Initializes a new instance of the <see cref="IdentityAssertionGrantProvider"/> class.
/// </summary>
@ -105,12 +120,33 @@ public sealed class IdentityAssertionGrantProvider
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<TokenContainer> 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,10 +209,22 @@ public sealed class IdentityAssertionGrantProvider
/// <summary>
/// Clears any cached tokens, forcing a fresh token exchange on the next call to <see cref="GetAccessTokenAsync"/>.
/// </summary>
/// <remarks>
/// 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.
/// </remarks>
public void InvalidateCache()
{
_tokenAcquisitionLock.Wait();
try
{
_cachedTokens = null;
}
finally
{
_tokenAcquisitionLock.Release();
}
}
private string? _resolvedIdpTokenEndpoint;

View file

@ -11,6 +11,7 @@ using ModelContextProtocol.Authentication;
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 class AuthTests : OAuthTestBase
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<string?> 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<string>(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<string>(requestedScopes[1]!.Split(' '));
Assert.Contains("mcp:tools", stepUpScopes);
Assert.Contains("files:read", stepUpScopes);
}
[Fact]
public async Task AuthorizationFlow_StopsSteppingUpWhenChallengeAddsNoNewScope()
{

View file

@ -313,6 +313,88 @@ public sealed class IdentityAssertionGrantTests : IDisposable
_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<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
var release = new TaskCompletionSource<bool>(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