Co-authored-by: Tarek Mahmoud Sayed <tarekms@ntdev.microsoft.com>
This commit is contained in:
parent
7369a265a8
commit
28806b1b6e
3 changed files with 341 additions and 36 deletions
|
|
@ -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,17 +317,26 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
|
|||
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 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
|
|||
|
||||
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 @@ 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)
|
||||
{
|
||||
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));
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
|
|
@ -881,14 +939,11 @@ internal sealed partial class ClientOAuthProvider : McpHttpClient
|
|||
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;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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,9 +209,21 @@ 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()
|
||||
{
|
||||
_cachedTokens = null;
|
||||
_tokenAcquisitionLock.Wait();
|
||||
try
|
||||
{
|
||||
_cachedTokens = null;
|
||||
}
|
||||
finally
|
||||
{
|
||||
_tokenAcquisitionLock.Release();
|
||||
}
|
||||
}
|
||||
|
||||
private string? _resolvedIdpTokenEndpoint;
|
||||
|
|
|
|||
Loading…
Reference in a new issue