Skip to content

Commit 8790a35

Browse files
committed
Clean up the helper
1 parent e044475 commit 8790a35

1 file changed

Lines changed: 68 additions & 102 deletions

File tree

src/ModelContextProtocol/Auth/OAuthHelpers.cs

Lines changed: 68 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -30,31 +30,14 @@ public static Func<Uri, Task<string>> CreateHttpListenerCallback(
3030
int listenPort = 8888,
3131
string redirectPath = "/callback")
3232
{
33-
return async (Uri authorizationUri) =>
33+
return async authorizationUri =>
3434
{
35-
string redirectUri = $"http://{hostname}:{listenPort}{redirectPath}";
36-
37-
// Add the redirect_uri parameter to the authorization URI if it's not already present
38-
string authUrl = authorizationUri.ToString();
39-
if (!authUrl.Contains("redirect_uri="))
40-
{
41-
var separator = authUrl.Contains("?") ? "&" : "?";
42-
authUrl = $"{authUrl}{separator}redirect_uri={WebUtility.UrlEncode(redirectUri)}";
43-
}
44-
45-
var authCodeTcs = new TaskCompletionSource<string>();
46-
4735
// Ensure the path has a trailing slash for the HttpListener prefix
48-
string listenerPrefix = $"http://{hostname}:{listenPort}{redirectPath}";
49-
if (!listenerPrefix.EndsWith("/"))
50-
{
51-
listenerPrefix += "/";
52-
}
36+
var listenerPrefix = $"http://{hostname}:{listenPort}{redirectPath.TrimEnd('/')}/";
5337

5438
using var listener = new HttpListener();
5539
listener.Prefixes.Add(listenerPrefix);
56-
57-
// Start the listener BEFORE opening the browser
40+
5841
try
5942
{
6043
listener.Start();
@@ -64,101 +47,84 @@ public static Func<Uri, Task<string>> CreateHttpListenerCallback(
6447
throw new InvalidOperationException($"Failed to start HTTP listener on {listenerPrefix}: {ex.Message}");
6548
}
6649

67-
// Create a cancellation token source with a timeout
6850
using var cts = new CancellationTokenSource(TimeSpan.FromMinutes(5));
6951

70-
_ = Task.Run(async () =>
71-
{
72-
try
73-
{
74-
// GetContextAsync doesn't accept a cancellation token, so we need to handle cancellation manually
75-
var contextTask = listener.GetContextAsync();
76-
var completedTask = await Task.WhenAny(contextTask, Task.Delay(Timeout.Infinite, cts.Token));
77-
78-
if (completedTask == contextTask)
79-
{
80-
var context = await contextTask;
81-
var request = context.Request;
82-
var response = context.Response;
83-
84-
string? code = request.QueryString["code"];
85-
string? error = request.QueryString["error"];
86-
string html;
87-
string? resultCode = null;
52+
await openBrowser(authorizationUri.ToString());
8853

89-
if (!string.IsNullOrEmpty(error))
90-
{
91-
html = $"<html><body><h1>Authorization Failed</h1><p>Error: {WebUtility.HtmlEncode(error)}</p></body></html>";
92-
}
93-
else if (string.IsNullOrEmpty(code))
94-
{
95-
html = "<html><body><h1>Authorization Failed</h1><p>No authorization code received.</p></body></html>";
96-
}
97-
else
98-
{
99-
html = "<html><body><h1>Authorization Successful</h1><p>You may now close this window.</p></body></html>";
100-
resultCode = code;
101-
}
102-
103-
try
104-
{
105-
// Send response to browser
106-
byte[] buffer = Encoding.UTF8.GetBytes(html);
107-
response.ContentType = "text/html";
108-
response.ContentLength64 = buffer.Length;
109-
response.OutputStream.Write(buffer, 0, buffer.Length);
110-
111-
// IMPORTANT: Explicitly close the response to ensure it's fully sent
112-
response.Close();
113-
114-
// Now that we've finished processing the browser response,
115-
// we can safely signal completion or failure with the auth code
116-
if (resultCode != null)
117-
{
118-
authCodeTcs.TrySetResult(resultCode);
119-
}
120-
else if (!string.IsNullOrEmpty(error))
121-
{
122-
authCodeTcs.TrySetException(new InvalidOperationException($"Authorization failed: {error}"));
123-
}
124-
else
125-
{
126-
authCodeTcs.TrySetException(new InvalidOperationException("No authorization code received"));
127-
}
128-
}
129-
catch (Exception ex)
130-
{
131-
authCodeTcs.TrySetException(new InvalidOperationException($"Error processing browser response: {ex.Message}"));
132-
}
133-
}
134-
}
135-
catch (Exception ex)
54+
try
55+
{
56+
var contextTask = listener.GetContextAsync();
57+
var completedTask = await Task.WhenAny(contextTask, Task.Delay(Timeout.Infinite, cts.Token));
58+
59+
if (completedTask != contextTask)
13660
{
137-
authCodeTcs.TrySetException(ex);
61+
throw new InvalidOperationException("Authorization timed out after 5 minutes.");
13862
}
139-
});
140-
141-
// Now open the browser AFTER the listener is started
142-
await openBrowser(authUrl);
14363

144-
try
145-
{
146-
// Use a timeout to avoid hanging indefinitely
147-
string authCode = await authCodeTcs.Task.WaitAsync(cts.Token);
148-
return authCode;
149-
}
150-
catch (OperationCanceledException)
151-
{
152-
throw new InvalidOperationException("Authorization timed out after 5 minutes.");
64+
var context = await contextTask;
65+
return ProcessCallback(context);
15366
}
15467
finally
15568
{
156-
// Ensure the listener is stopped when we're done
15769
listener.Stop();
15870
}
15971
};
16072
}
161-
73+
74+
/// <summary>
75+
/// Processes the HTTP callback and extracts the authorization code.
76+
/// </summary>
77+
private static string ProcessCallback(HttpListenerContext context)
78+
{
79+
var request = context.Request;
80+
var response = context.Response;
81+
82+
string? code = request.QueryString["code"];
83+
string? error = request.QueryString["error"];
84+
string html;
85+
86+
if (!string.IsNullOrEmpty(error))
87+
{
88+
html = $"<html><body><h1>Authorization Failed</h1><p>Error: {WebUtility.HtmlEncode(error)}</p></body></html>";
89+
SendResponse(response, html);
90+
throw new InvalidOperationException($"Authorization failed: {error}");
91+
}
92+
93+
if (string.IsNullOrEmpty(code))
94+
{
95+
html = "<html><body><h1>Authorization Failed</h1><p>No authorization code received.</p></body></html>";
96+
SendResponse(response, html);
97+
throw new InvalidOperationException("No authorization code received");
98+
}
99+
100+
html = "<html><body><h1>Authorization Successful</h1><p>You may now close this window.</p></body></html>";
101+
SendResponse(response, html);
102+
return code;
103+
}
104+
105+
/// <summary>
106+
/// Sends an HTML response to the browser.
107+
/// </summary>
108+
private static void SendResponse(HttpListenerResponse response, string html)
109+
{
110+
try
111+
{
112+
byte[] buffer = Encoding.UTF8.GetBytes(html);
113+
response.ContentType = "text/html";
114+
response.ContentLength64 = buffer.Length;
115+
response.OutputStream.Write(buffer, 0, buffer.Length);
116+
117+
// IMPORTANT: Explicitly close the response to ensure it's fully sent
118+
response.Close();
119+
}
120+
catch
121+
{
122+
// Silently handle errors - we're already in an error handling path
123+
// and can't throw further exceptions or log to the console in a library
124+
// TODO: Need a better implementation here.
125+
}
126+
}
127+
162128
/// <summary>
163129
/// Exchanges an authorization code for an OAuth token.
164130
/// </summary>

0 commit comments

Comments
 (0)