Skip to content

Commit c67ef6c

Browse files
committed
Tweaks to handles
1 parent 8b89c2d commit c67ef6c

4 files changed

Lines changed: 239 additions & 77 deletions

File tree

samples/SecureWeatherClient/Program.cs

Lines changed: 64 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -14,24 +14,49 @@ static async Task Main(string[] args)
1414
Console.WriteLine("==================================================");
1515
Console.WriteLine();
1616

17-
// Create the authorization config with HTTP listener
18-
var authConfig = new AuthorizationConfig
19-
{
20-
ClientId = "04f79824-ab56-4511-a7cb-d7deaea92dc0",
21-
Scopes = ["User.Read"]
22-
}.UseHttpListener(hostname: "localhost", listenPort: 1170);
17+
Console.WriteLine("Select authentication mode:");
18+
Console.WriteLine("1. Normal OAuth flow with browser");
19+
Console.WriteLine("2. Mock authentication (accepts any token for testing)");
20+
Console.Write("Enter your choice (1-2): ");
21+
var choice = Console.ReadLine()?.Trim();
2322

2423
// Create an HTTP client with OAuth handling
25-
var oauthHandler = new OAuthDelegatingHandler(
26-
redirectUri: authConfig.RedirectUri,
27-
clientId: authConfig.ClientId,
28-
clientName: authConfig.ClientName,
29-
scopes: authConfig.Scopes,
30-
authorizationHandler: authConfig.AuthorizationHandler)
24+
DelegatingHandler oauthHandler;
25+
26+
if (choice == "2")
27+
{
28+
Console.WriteLine("\nUsing mock authentication for testing (no browser will open).\n");
29+
30+
// Create a mock OAuth handler that always returns a token
31+
oauthHandler = new MockOAuthHandler()
32+
{
33+
InnerHandler = new HttpClientHandler()
34+
};
35+
}
36+
else
3137
{
32-
// The OAuth handler needs an inner handler
33-
InnerHandler = new HttpClientHandler()
34-
};
38+
Console.WriteLine("\nUsing standard OAuth flow with browser authentication.\n");
39+
40+
// Create the authorization config with HTTP listener
41+
var authConfig = new AuthorizationConfig
42+
{
43+
ClientId = "04f79824-ab56-4511-a7cb-d7deaea92dc0",
44+
ClientName = "SecureWeatherClient",
45+
Scopes = ["weather.read"]
46+
}.UseHttpListener(hostname: "localhost", listenPort: 1170);
47+
48+
// Create an HTTP client with OAuth handling
49+
oauthHandler = new OAuthDelegatingHandler(
50+
redirectUri: authConfig.RedirectUri,
51+
clientId: authConfig.ClientId,
52+
clientName: authConfig.ClientName,
53+
scopes: authConfig.Scopes,
54+
authorizationHandler: authConfig.AuthorizationHandler)
55+
{
56+
// The OAuth handler needs an inner handler
57+
InnerHandler = new HttpClientHandler()
58+
};
59+
}
3560

3661
var httpClient = new HttpClient(oauthHandler);
3762
var serverUrl = "http://localhost:7071/sse"; // Default server URL
@@ -46,8 +71,13 @@ static async Task Main(string[] args)
4671

4772
Console.WriteLine();
4873
Console.WriteLine($"Connecting to weather server at {serverUrl}...");
49-
Console.WriteLine("When prompted for authorization, a browser window will open automatically.");
50-
Console.WriteLine("Complete the authentication in the browser, and this application will continue automatically.");
74+
75+
if (choice != "2")
76+
{
77+
Console.WriteLine("When prompted for authorization, a browser window will open automatically.");
78+
Console.WriteLine("Complete the authentication in the browser, and this application will continue automatically.");
79+
}
80+
5181
Console.WriteLine();
5282

5383
try
@@ -97,4 +127,21 @@ static async Task Main(string[] args)
97127
Console.WriteLine("Press any key to exit...");
98128
Console.ReadKey();
99129
}
130+
}
131+
132+
/// <summary>
133+
/// A mock OAuth handler that always returns a predefined token without going through the OAuth flow.
134+
/// This is useful for testing without requiring a real OAuth server.
135+
/// </summary>
136+
public class MockOAuthHandler : DelegatingHandler
137+
{
138+
private readonly string _mockToken = "mock_test_token_" + Guid.NewGuid().ToString("N");
139+
140+
protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
141+
{
142+
// Always attach the mock token to outgoing requests
143+
request.Headers.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", _mockToken);
144+
145+
return base.SendAsync(request, cancellationToken);
146+
}
100147
}

src/ModelContextProtocol/Auth/McpClientExtensions.cs

Lines changed: 16 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -96,52 +96,26 @@ public static async Task<OAuthToken> HandleUnauthorizedResponseAsync(
9696
throw new InvalidOperationException("The HTTP client has not been configured for authorization handling. Call ConfigureAuthorizationHandler() first.");
9797
}
9898

99-
// Create OAuthAuthenticationService
100-
var authService = new OAuthAuthenticationService();
99+
// Create OAuthAuthenticationService - use appropriate constructor based on whether we have a handler
100+
OAuthAuthenticationService authService = config.AuthorizationHandler != null
101+
? new OAuthAuthenticationService(config.AuthorizationHandler)
102+
: new OAuthAuthenticationService();
101103

102104
// Get resource URI
103105
var resourceUri = response.RequestMessage?.RequestUri ?? throw new InvalidOperationException("Request URI is not available.");
104106

105107
// Start the authentication flow
106-
try
107-
{
108-
var tokenResponse = await authService.HandleAuthenticationAsync(
109-
resourceUri,
110-
wwwAuthenticateHeader,
111-
config.RedirectUri,
112-
config.ClientId,
113-
config.ClientName,
114-
config.Scopes);
115-
116-
// Attach the access token to future requests
117-
httpClient.AttachToken(tokenResponse.AccessToken);
118-
119-
return tokenResponse;
120-
}
121-
catch (NotImplementedException ex) when (ex.Message.Contains("Authorization requires user interaction"))
122-
{
123-
// Extract the authorization URL from the exception message
124-
var authUrlStart = ex.Message.IndexOf("http");
125-
var authUrlEnd = ex.Message.IndexOf("\n", authUrlStart);
126-
var authUrl = ex.Message.Substring(authUrlStart, authUrlEnd - authUrlStart);
127-
128-
// Check if a handler is registered
129-
if (config.AuthorizationHandler != null)
130-
{
131-
// Call the handler to get the authorization code
132-
var authCode = await config.AuthorizationHandler(new Uri(authUrl));
133-
134-
// In a real implementation, we would use the authorization code to get a token
135-
// For now, throw an exception with instructions
136-
throw new NotImplementedException(
137-
"Authorization code acquired, but token exchange is not implemented. " +
138-
"In a real implementation, this would call ExchangeAuthorizationCodeForTokenAsync.");
139-
}
140-
else
141-
{
142-
// Re-throw the original exception
143-
throw;
144-
}
145-
}
108+
var tokenResponse = await authService.HandleAuthenticationAsync(
109+
resourceUri,
110+
wwwAuthenticateHeader,
111+
config.RedirectUri,
112+
config.ClientId,
113+
config.ClientName,
114+
config.Scopes);
115+
116+
// Attach the access token to future requests
117+
httpClient.AttachToken(tokenResponse.AccessToken);
118+
119+
return tokenResponse;
146120
}
147121
}

src/ModelContextProtocol/Auth/OAuthAuthenticationService.cs

Lines changed: 50 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,23 @@ namespace ModelContextProtocol.Auth;
1313
public class OAuthAuthenticationService
1414
{
1515
private static readonly HttpClient _httpClient = new();
16+
private readonly Func<Uri, Task<string>>? _authorizationHandler;
17+
18+
/// <summary>
19+
/// Initializes a new instance of the <see cref="OAuthAuthenticationService"/> class.
20+
/// </summary>
21+
public OAuthAuthenticationService()
22+
{
23+
}
24+
25+
/// <summary>
26+
/// Initializes a new instance of the <see cref="OAuthAuthenticationService"/> class with an authorization handler.
27+
/// </summary>
28+
/// <param name="authorizationHandler">A handler to invoke when authorization is required.</param>
29+
public OAuthAuthenticationService(Func<Uri, Task<string>> authorizationHandler)
30+
{
31+
_authorizationHandler = authorizationHandler ?? throw new ArgumentNullException(nameof(authorizationHandler));
32+
}
1633

1734
/// <summary>
1835
/// Handles the OAuth authentication flow when a 401 Unauthorized response is received.
@@ -23,15 +40,20 @@ public class OAuthAuthenticationService
2340
/// <param name="clientId">The client ID to use for authentication, or null to register a new client.</param>
2441
/// <param name="clientName">The client name to use for registration.</param>
2542
/// <param name="scopes">The requested scopes.</param>
43+
/// <param name="authorizationHandler">A handler to invoke when authorization is required. If not provided, the handler from the constructor will be used.</param>
2644
/// <returns>The OAuth token response.</returns>
2745
public async Task<OAuthToken> HandleAuthenticationAsync(
2846
Uri resourceUri,
2947
string wwwAuthenticateHeader,
3048
Uri redirectUri,
3149
string? clientId = null,
3250
string? clientName = null,
33-
IEnumerable<string>? scopes = null)
51+
IEnumerable<string>? scopes = null,
52+
Func<Uri, Task<string>>? authorizationHandler = null)
3453
{
54+
// Use the provided authorization handler or fall back to the one from the constructor
55+
var effectiveAuthHandler = authorizationHandler ?? _authorizationHandler;
56+
3557
// Extract resource metadata URL from WWW-Authenticate header
3658
var resourceMetadataUri = ExtractResourceMetadataUri(wwwAuthenticateHeader);
3759
if (resourceMetadataUri == null)
@@ -88,7 +110,8 @@ public async Task<OAuthToken> HandleAuthenticationAsync(
88110
effectiveClientId, // This is now guaranteed to be non-null
89111
clientSecret,
90112
redirectUri,
91-
scopes?.ToList() ?? resourceMetadata.ScopesSupported);
113+
scopes?.ToList() ?? resourceMetadata.ScopesSupported,
114+
effectiveAuthHandler);
92115

93116
return tokenResponse;
94117
}
@@ -219,7 +242,8 @@ private async Task<OAuthToken> PerformAuthorizationCodeFlowAsync(
219242
string clientId,
220243
string? clientSecret,
221244
Uri redirectUri,
222-
IEnumerable<string> scopes)
245+
IEnumerable<string> scopes,
246+
Func<Uri, Task<string>>? authorizationHandler)
223247
{
224248
// Generate PKCE code verifier and challenge
225249
var codeVerifier = GenerateCodeVerifier();
@@ -233,26 +257,34 @@ private async Task<OAuthToken> PerformAuthorizationCodeFlowAsync(
233257
codeChallenge,
234258
scopes);
235259

236-
// At this point, in a real application, you would redirect the user to the authorizationUrl
237-
// and then handle the callback to redirectUri with the authorization code.
238-
// For this implementation, we'll assume the code is obtained externally and passed to us.
260+
// Check if an authorization handler is available
261+
if (authorizationHandler != null)
262+
{
263+
try
264+
{
265+
// Get the authorization code using the provided handler
266+
string authorizationCode = await authorizationHandler(new Uri(authorizationUrl));
267+
268+
// Exchange the authorization code for a token
269+
return await ExchangeAuthorizationCodeForTokenAsync(
270+
authServerMetadata.TokenEndpoint,
271+
clientId,
272+
clientSecret,
273+
redirectUri,
274+
authorizationCode,
275+
codeVerifier);
276+
}
277+
catch (Exception ex)
278+
{
279+
throw new InvalidOperationException($"Failed to complete OAuth authorization flow: {ex.Message}", ex);
280+
}
281+
}
239282

240-
// Since we can't actually perform the browser interaction in this service,
241-
// we'll throw with instructions
283+
// No authorization handler available, throw with instructions
242284
throw new NotImplementedException(
243285
$"Authorization requires user interaction. Please direct the user to: {authorizationUrl}\n" +
244286
$"After authorization, the user will be redirected to: {redirectUri}?code=[authorization_code]\n" +
245287
$"You need to handle this redirect and extract the authorization code to complete the flow.");
246-
247-
// In a real implementation, after getting the authorization code:
248-
// var authorizationCode = GetAuthorizationCodeFromRedirect();
249-
// return await ExchangeAuthorizationCodeForTokenAsync(
250-
// authServerMetadata.TokenEndpoint,
251-
// clientId,
252-
// clientSecret,
253-
// redirectUri,
254-
// authorizationCode,
255-
// codeVerifier);
256288
}
257289

258290
private string GenerateCodeVerifier()

0 commit comments

Comments
 (0)