Skip to content

Commit 627e1e4

Browse files
localdenlocalden
committed
Update 401 handling logic. Remove pragma warning silencing.
1 parent 45c4b30 commit 627e1e4

1 file changed

Lines changed: 96 additions & 97 deletions

File tree

‎src/ModelContextProtocol/Authentication/AuthorizationDelegatingHandler.cs‎

Lines changed: 96 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -42,73 +42,114 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
4242

4343
if (response.StatusCode == System.Net.HttpStatusCode.Unauthorized)
4444
{
45-
// Gather the schemes the server wants us to use from WWW-Authenticate headers
46-
var serverSchemes = ExtractServerSupportedSchemes(response);
45+
return await HandleUnauthorizedResponseAsync(request, response, cancellationToken).ConfigureAwait(false);
46+
}
47+
48+
return response;
49+
}
50+
51+
/// <summary>
52+
/// Handles a 401 Unauthorized response by attempting to authenticate and retry the request.
53+
/// </summary>
54+
private async Task<HttpResponseMessage> HandleUnauthorizedResponseAsync(
55+
HttpRequestMessage originalRequest,
56+
HttpResponseMessage response,
57+
CancellationToken cancellationToken)
58+
{
59+
// Gather the schemes the server wants us to use from WWW-Authenticate headers
60+
var serverSchemes = ExtractServerSupportedSchemes(response);
61+
62+
// Find the intersection between what the server supports and what our provider supports
63+
var supportedSchemes = _authorizationProvider.SupportedSchemes.ToList();
64+
string? bestSchemeMatch = null;
65+
66+
// First try to find a direct match with the current scheme if it's still valid
67+
string schemeUsed = originalRequest.Headers.Authorization?.Scheme ?? _currentScheme ?? string.Empty;
68+
if (serverSchemes.Contains(schemeUsed) && supportedSchemes.Contains(schemeUsed))
69+
{
70+
bestSchemeMatch = schemeUsed;
71+
}
72+
else
73+
{
74+
// Try to find any matching scheme between server and provider using a HashSet for O(N) time complexity
75+
// Convert the supported schemes to a HashSet for O(1) lookups
76+
var supportedSchemesSet = new HashSet<string>(supportedSchemes, StringComparer.OrdinalIgnoreCase);
4777

48-
// Find the intersection between what the server supports and what our provider supports
49-
var supportedSchemes = _authorizationProvider.SupportedSchemes.ToList();
50-
string? bestSchemeMatch = null;
78+
// Find the first server scheme that's in our supported set
79+
bestSchemeMatch = serverSchemes.FirstOrDefault(scheme => supportedSchemesSet.Contains(scheme));
5180

52-
// First try to find a direct match with the current scheme if it's still valid
53-
string schemeUsed = request.Headers.Authorization?.Scheme ?? _currentScheme ?? string.Empty;
54-
if (serverSchemes.Contains(schemeUsed) && supportedSchemes.Contains(schemeUsed))
55-
{
56-
bestSchemeMatch = schemeUsed;
57-
}
58-
else
81+
// If no match was found, either throw an exception or use default
82+
if (bestSchemeMatch is null)
5983
{
60-
// Try to find any matching scheme between server and provider using a HashSet for O(N) time complexity
61-
// Convert the supported schemes to a HashSet for O(1) lookups
62-
var supportedSchemesSet = new HashSet<string>(supportedSchemes, StringComparer.OrdinalIgnoreCase);
63-
64-
// Find the first server scheme that's in our supported set
65-
bestSchemeMatch = serverSchemes.FirstOrDefault(scheme => supportedSchemesSet.Contains(scheme));
66-
67-
// If no match was found, either throw an exception or use default
68-
if (bestSchemeMatch is null)
84+
if (serverSchemes.Count > 0)
6985
{
70-
if (serverSchemes.Count > 0)
71-
{
72-
throw new AuthenticationSchemeMismatchException(
73-
$"No matching authentication scheme found. Server supports: [{string.Join(", ", serverSchemes)}], " +
74-
$"Provider supports: [{string.Join(", ", supportedSchemes)}].",
75-
serverSchemes,
76-
supportedSchemes);
77-
}
78-
79-
// If the server didn't specify any schemes, use the provider's default
80-
bestSchemeMatch = supportedSchemes.FirstOrDefault();
86+
throw new AuthenticationSchemeMismatchException(
87+
$"No matching authentication scheme found. Server supports: [{string.Join(", ", serverSchemes)}], " +
88+
$"Provider supports: [{string.Join(", ", supportedSchemes)}].",
89+
serverSchemes,
90+
supportedSchemes);
8191
}
92+
93+
// If the server didn't specify any schemes, use the provider's default
94+
bestSchemeMatch = supportedSchemes.FirstOrDefault();
95+
}
96+
}
97+
98+
// If we have a scheme to try, use it
99+
if (bestSchemeMatch != null)
100+
{
101+
// Try to handle the 401 response with the selected scheme
102+
var (handled, recommendedScheme) = await _authorizationProvider.HandleUnauthorizedResponseAsync(
103+
response,
104+
bestSchemeMatch,
105+
cancellationToken).ConfigureAwait(false);
106+
107+
if (!handled)
108+
{
109+
throw new McpException(
110+
$"Failed to handle unauthorized response with scheme '{bestSchemeMatch}'. " +
111+
"The authentication provider was unable to process the authentication challenge.");
82112
}
83113

84-
// If we have a scheme to try, use it
85-
if (bestSchemeMatch != null)
114+
_currentScheme = recommendedScheme ?? bestSchemeMatch;
115+
116+
var retryRequest = new HttpRequestMessage(originalRequest.Method, originalRequest.RequestUri)
86117
{
87-
// Try to handle the 401 response with the selected scheme
88-
var (handled, recommendedScheme) = await _authorizationProvider.HandleUnauthorizedResponseAsync(
89-
response,
90-
bestSchemeMatch,
91-
cancellationToken).ConfigureAwait(false);
92-
93-
if (handled)
94-
{
95-
var retryRequest = await CloneHttpRequestMessageAsync(request).ConfigureAwait(false);
96-
97-
_currentScheme = recommendedScheme ?? bestSchemeMatch;
98-
99-
await AddAuthorizationHeaderAsync(retryRequest, _currentScheme, cancellationToken).ConfigureAwait(false);
100-
return await base.SendAsync(retryRequest, cancellationToken).ConfigureAwait(false);
101-
}
102-
else
118+
Version = originalRequest.Version,
119+
#if NET
120+
VersionPolicy = originalRequest.VersionPolicy,
121+
#endif
122+
Content = originalRequest.Content
123+
};
124+
125+
// Copy headers except Authorization which we'll set separately
126+
foreach (var header in originalRequest.Headers)
127+
{
128+
if (!header.Key.Equals("Authorization", StringComparison.OrdinalIgnoreCase))
103129
{
104-
throw new McpException(
105-
$"Failed to handle unauthorized response with scheme '{bestSchemeMatch}'. " +
106-
"The authentication provider was unable to process the authentication challenge.");
130+
retryRequest.Headers.TryAddWithoutValidation(header.Key, header.Value);
107131
}
108132
}
109-
}
133+
#if NET
134+
foreach (var property in originalRequest.Options)
135+
{
136+
retryRequest.Options.Set(new HttpRequestOptionsKey<object?>(property.Key), property.Value);
137+
}
138+
#else
139+
foreach (var property in originalRequest.Properties)
140+
{
141+
retryRequest.Properties.Add(property);
142+
}
143+
#endif
110144

111-
return response;
145+
// Add the new authorization header
146+
await AddAuthorizationHeaderAsync(retryRequest, _currentScheme, cancellationToken).ConfigureAwait(false);
147+
148+
// Send the retry request
149+
return await base.SendAsync(retryRequest, cancellationToken).ConfigureAwait(false);
150+
}
151+
152+
return response; // Return the original response if we couldn't handle it
112153
}
113154

114155
/// <summary>
@@ -149,46 +190,4 @@ private async Task AddAuthorizationHeaderAsync(HttpRequestMessage request, strin
149190
}
150191
}
151192
}
152-
153-
/// <summary>
154-
/// Creates a clone of the HTTP request message.
155-
/// </summary>
156-
private static async Task<HttpRequestMessage> CloneHttpRequestMessageAsync(HttpRequestMessage request)
157-
{
158-
var clone = new HttpRequestMessage(request.Method, request.RequestUri);
159-
160-
// Copy the request headers
161-
foreach (var header in request.Headers)
162-
{
163-
clone.Headers.TryAddWithoutValidation(header.Key, header.Value);
164-
}
165-
166-
// Copy the request content if present
167-
if (request.Content != null)
168-
{
169-
var contentBytes = await request.Content.ReadAsByteArrayAsync().ConfigureAwait(false);
170-
var cloneContent = new ByteArrayContent(contentBytes);
171-
172-
// Copy the content headers
173-
foreach (var header in request.Content.Headers)
174-
{
175-
cloneContent.Headers.TryAddWithoutValidation(header.Key, header.Value);
176-
}
177-
178-
clone.Content = cloneContent;
179-
}
180-
181-
// Copy the request properties
182-
#pragma warning disable CS0618 // Type or member is obsolete
183-
foreach (var property in request.Properties)
184-
{
185-
clone.Properties.Add(property);
186-
}
187-
#pragma warning restore CS0618 // Type or member is obsolete
188-
189-
// Copy the request version
190-
clone.Version = request.Version;
191-
192-
return clone;
193-
}
194193
}

0 commit comments

Comments
 (0)