Skip to content

Commit 2f84981

Browse files
localdenlocalden
committed
Update token logic
1 parent 10cf1f7 commit 2f84981

2 files changed

Lines changed: 101 additions & 25 deletions

File tree

‎samples/ProtectedMCPClient/BasicOAuthAuthorizationProvider.cs‎

Lines changed: 83 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -29,33 +29,48 @@ public partial class BasicOAuthAuthorizationProvider(
2929
Uri? redirectUri = null,
3030
IEnumerable<string>? scopes = null) : IMcpAuthorizationProvider
3131
{
32+
// Cache for tokens, keyed by the canonical resource URI from resource metadata
3233
private readonly ConcurrentDictionary<string, TokenContainer> _tokenCache = new();
34+
35+
// Cache for resource metadata, keyed by the request URI
36+
private readonly ConcurrentDictionary<string, ProtectedResourceMetadata> _resourceMetadataCache = new();
37+
3338
private readonly Uri _serverUrl = serverUrl ?? throw new ArgumentNullException(nameof(serverUrl));
3439
private readonly Uri _redirectUri = redirectUri ?? new Uri("http://localhost:8080/callback");
3540
private readonly IEnumerable<string> _scopes = scopes ?? Array.Empty<string>();
3641

3742
/// <inheritdoc />
3843
public async Task<string?> GetCredentialAsync(Uri resourceUri, CancellationToken cancellationToken = default)
3944
{
40-
Console.WriteLine($"Getting authentication token for {resourceUri}");
45+
Console.WriteLine($"Getting credential for resource URI: {resourceUri}");
46+
47+
// First, get the resource metadata to determine the canonical resource URI
48+
var resourceMetadata = await GetCachedResourceMetadataAsync(resourceUri, cancellationToken);
49+
if (resourceMetadata == null)
50+
{
51+
Console.WriteLine("Failed to get resource metadata, cannot authenticate");
52+
return null;
53+
}
54+
55+
// Use the canonical resource URI from the metadata as the cache key
56+
string resourceKey = resourceMetadata.Resource.ToString();
57+
Console.WriteLine($"Using canonical resource key: {resourceKey}");
4158

4259
// Check if we have a valid cached token
43-
string resourceKey = resourceUri.ToString();
4460
if (_tokenCache.TryGetValue(resourceKey, out var tokenInfo))
4561
{
4662
// Check if the token is still valid or needs to be refreshed
4763
if (tokenInfo.ExpiresAt > DateTimeOffset.UtcNow.AddMinutes(5)) // 5-minute buffer
4864
{
4965
Console.WriteLine("Using cached token");
66+
// Return just the access token, not the token type + access token
5067
return tokenInfo.AccessToken;
5168
}
5269
else if (!string.IsNullOrEmpty(tokenInfo.RefreshToken))
5370
{
5471
Console.WriteLine("Token expired, attempting to refresh");
5572

56-
// Get the authorization server metadata for the resource
57-
var resourceMetadata = await GetResourceMetadataAsync(resourceUri, cancellationToken);
58-
if (resourceMetadata?.AuthorizationServers?.Count > 0)
73+
if (resourceMetadata.AuthorizationServers?.Count > 0)
5974
{
6075
var authServerUrl = resourceMetadata.AuthorizationServers[0];
6176
var authServerMetadata = await AuthorizationServerUtils.FetchAuthorizationServerMetadataAsync(
@@ -73,6 +88,7 @@ public partial class BasicOAuthAuthorizationProvider(
7388
{
7489
_tokenCache[resourceKey] = refreshedToken;
7590
Console.WriteLine("Token refreshed successfully");
91+
// Return just the access token, not the token type + access token
7692
return refreshedToken.AccessToken;
7793
}
7894
else
@@ -91,11 +107,47 @@ public partial class BasicOAuthAuthorizationProvider(
91107
_tokenCache.TryRemove(resourceKey, out _);
92108
}
93109

94-
// We don't have a valid token and need to get a new one
95-
Console.WriteLine("No valid token available");
110+
// We don't have a valid token - let the 401 handler trigger the auth flow
111+
Console.WriteLine("No valid token available for: " + resourceKey);
96112
return null;
97113
}
98-
114+
115+
/// <summary>
116+
/// Gets resource metadata, using the cache if available.
117+
/// </summary>
118+
/// <param name="resourceUri">The URI of the protected resource.</param>
119+
/// <param name="cancellationToken">A token to cancel the operation.</param>
120+
/// <returns>The protected resource metadata.</returns>
121+
private async Task<ProtectedResourceMetadata?> GetCachedResourceMetadataAsync(Uri resourceUri, CancellationToken cancellationToken)
122+
{
123+
string requestUriKey = resourceUri.ToString();
124+
125+
// Check if we already have cached metadata for this URI
126+
if (_resourceMetadataCache.TryGetValue(requestUriKey, out var cachedMetadata))
127+
{
128+
Console.WriteLine($"Using cached resource metadata for: {requestUriKey}");
129+
return cachedMetadata;
130+
}
131+
132+
// If not in cache, fetch the metadata
133+
var metadata = await GetResourceMetadataAsync(resourceUri, cancellationToken);
134+
if (metadata != null)
135+
{
136+
// Cache the metadata using the request URI as the key
137+
_resourceMetadataCache[requestUriKey] = metadata;
138+
139+
// Also cache using any alternate forms of the URI that might be used
140+
if (resourceUri.ToString() != metadata.Resource.ToString())
141+
{
142+
_resourceMetadataCache[metadata.Resource.ToString()] = metadata;
143+
}
144+
145+
Console.WriteLine($"Cached resource metadata for {requestUriKey} -> {metadata.Resource}");
146+
}
147+
148+
return metadata;
149+
}
150+
99151
/// <summary>
100152
/// Refreshes an OAuth token using the refresh token.
101153
/// </summary>
@@ -145,15 +197,9 @@ public partial class BasicOAuthAuthorizationProvider(
145197

146198
if (tokenResponse != null)
147199
{
148-
// Set the time when the token was obtained
200+
// Set the time when the token was obtained - ExpiresAt will be calculated automatically
149201
tokenResponse.ObtainedAt = DateTimeOffset.UtcNow;
150202

151-
// Calculate expiration time if not set
152-
if (tokenResponse.ExpiresIn > 0 && tokenResponse.ExpiresAt == default)
153-
{
154-
tokenResponse.ExpiresAt = tokenResponse.ObtainedAt.AddSeconds(tokenResponse.ExpiresIn);
155-
}
156-
157203
// Preserve the refresh token if the response doesn't include a new one
158204
if (string.IsNullOrEmpty(tokenResponse.RefreshToken))
159205
{
@@ -226,6 +272,24 @@ public async Task<bool> HandleUnauthorizedResponseAsync(HttpResponseMessage resp
226272
_serverUrl,
227273
cancellationToken);
228274

275+
if (resourceMetadata == null)
276+
{
277+
Console.WriteLine("Failed to extract resource metadata from response");
278+
return false;
279+
}
280+
281+
// Cache the resource metadata for future use
282+
string requestUriKey = response.RequestMessage?.RequestUri?.ToString() ?? string.Empty;
283+
if (!string.IsNullOrEmpty(requestUriKey))
284+
{
285+
_resourceMetadataCache[requestUriKey] = resourceMetadata;
286+
287+
// Also cache with the canonical resource URI
288+
_resourceMetadataCache[resourceMetadata.Resource.ToString()] = resourceMetadata;
289+
290+
Console.WriteLine($"Cached resource metadata for {requestUriKey} -> {resourceMetadata.Resource}");
291+
}
292+
229293
// If we get here, the resource metadata is valid and matches our server
230294
Console.WriteLine($"Successfully validated resource metadata for: {resourceMetadata.Resource}");
231295

@@ -247,9 +311,11 @@ public async Task<bool> HandleUnauthorizedResponseAsync(HttpResponseMessage resp
247311

248312
if (token != null)
249313
{
250-
// Store the token in the cache
314+
// Store the token in the cache using the canonical resource URI from metadata
251315
string resourceKey = resourceMetadata.Resource.ToString();
316+
Console.WriteLine($"Storing token with canonical resource key: {resourceKey}");
252317
_tokenCache[resourceKey] = token;
318+
253319
Console.WriteLine("Successfully obtained a new token");
254320
return true;
255321
}
@@ -543,15 +609,9 @@ private Uri BuildAuthorizationUrl(
543609

544610
if (tokenResponse != null)
545611
{
546-
// Set the time when the token was obtained
612+
// Set the time when the token was obtained - ExpiresAt will be calculated automatically
547613
tokenResponse.ObtainedAt = DateTimeOffset.UtcNow;
548614

549-
// Calculate expiration time if not set
550-
if (tokenResponse.ExpiresIn > 0 && tokenResponse.ExpiresAt == default)
551-
{
552-
tokenResponse.ExpiresAt = tokenResponse.ObtainedAt.AddSeconds(tokenResponse.ExpiresIn);
553-
}
554-
555615
Console.WriteLine("Token exchange successful");
556616
return tokenResponse;
557617
}

‎samples/ProtectedMCPClient/Types/TokenContainer.cs‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,38 @@
1+
using System.Text.Json.Serialization;
2+
13
namespace ProtectedMCPClient.Types;
24

35
/// <summary>
46
/// Represents a token response from the OAuth server.
57
/// </summary>
68
internal class TokenContainer
79
{
10+
[JsonPropertyName("access_token")]
811
public string AccessToken { get; set; } = string.Empty;
12+
13+
[JsonPropertyName("refresh_token")]
914
public string? RefreshToken { get; set; }
15+
16+
[JsonPropertyName("expires_in")]
1017
public int ExpiresIn { get; set; }
18+
19+
[JsonPropertyName("ext_expires_in")]
20+
public int ExtExpiresIn { get; set; }
21+
22+
[JsonPropertyName("token_type")]
1123
public string TokenType { get; set; } = string.Empty;
1224

25+
[JsonPropertyName("scope")]
26+
public string Scope { get; set; } = string.Empty;
27+
1328
/// <summary>
1429
/// Gets or sets the timestamp when the token was obtained.
1530
/// </summary>
1631
public DateTimeOffset ObtainedAt { get; set; }
1732

1833
/// <summary>
19-
/// Gets or sets the timestamp when the token expires.
34+
/// Gets the timestamp when the token expires, calculated from ObtainedAt and ExpiresIn.
2035
/// </summary>
21-
public DateTimeOffset ExpiresAt { get; set; }
36+
[JsonIgnore]
37+
public DateTimeOffset ExpiresAt => ObtainedAt.AddSeconds(ExpiresIn);
2238
}

0 commit comments

Comments
 (0)