1+ using System ;
2+ using System . Net ;
3+ using System . Text ;
4+ using System . Threading ;
5+ using System . Threading . Tasks ;
6+
7+ namespace ModelContextProtocol . Auth ;
8+
9+ /// <summary>
10+ /// Provides helper methods for handling OAuth authorization.
11+ /// </summary>
12+ public static class OAuthAuthorizationHelpers
13+ {
14+ /// <summary>
15+ /// Creates an HTTP listener callback for handling OAuth 2.0 authorization code flow.
16+ /// </summary>
17+ /// <param name="openBrowser">A function that opens a browser with the given URL.</param>
18+ /// <param name="hostname">The hostname to listen on. Defaults to "localhost".</param>
19+ /// <param name="listenPort">The port to listen on. Defaults to 8888.</param>
20+ /// <param name="redirectPath">The redirect path for the HTTP listener. Defaults to "/callback".</param>
21+ /// <returns>
22+ /// A function that takes an authorization URI and returns a task that resolves to the authorization code.
23+ /// </returns>
24+ public static Func < Uri , Task < string > > CreateHttpListenerCallback (
25+ Func < string , Task > openBrowser ,
26+ string hostname = "localhost" ,
27+ int listenPort = 8888 ,
28+ string redirectPath = "/callback" )
29+ {
30+ return async ( Uri authorizationUri ) =>
31+ {
32+ string redirectUri = $ "http://{ hostname } :{ listenPort } { redirectPath } ";
33+
34+ // Add the redirect_uri parameter to the authorization URI if it's not already present
35+ string authUrl = authorizationUri . ToString ( ) ;
36+ if ( ! authUrl . Contains ( "redirect_uri=" ) )
37+ {
38+ var separator = authUrl . Contains ( "?" ) ? "&" : "?" ;
39+ authUrl = $ "{ authUrl } { separator } redirect_uri={ WebUtility . UrlEncode ( redirectUri ) } ";
40+ }
41+
42+ var authCodeTcs = new TaskCompletionSource < string > ( ) ;
43+
44+ // Ensure the path has a trailing slash for the HttpListener prefix
45+ string listenerPrefix = $ "http://{ hostname } :{ listenPort } { redirectPath } ";
46+ if ( ! listenerPrefix . EndsWith ( "/" ) )
47+ {
48+ listenerPrefix += "/" ;
49+ }
50+
51+ using var listener = new HttpListener ( ) ;
52+ listener . Prefixes . Add ( listenerPrefix ) ;
53+
54+ // Start the listener BEFORE opening the browser
55+ try
56+ {
57+ listener . Start ( ) ;
58+ }
59+ catch ( HttpListenerException ex )
60+ {
61+ throw new InvalidOperationException ( $ "Failed to start HTTP listener on { listenerPrefix } : { ex . Message } ") ;
62+ }
63+
64+ // Create a cancellation token source with a timeout
65+ using var cts = new CancellationTokenSource ( TimeSpan . FromMinutes ( 5 ) ) ;
66+
67+ _ = Task . Run ( async ( ) =>
68+ {
69+ try
70+ {
71+ // GetContextAsync doesn't accept a cancellation token, so we need to handle cancellation manually
72+ var contextTask = listener . GetContextAsync ( ) ;
73+ var completedTask = await Task . WhenAny ( contextTask , Task . Delay ( Timeout . Infinite , cts . Token ) ) ;
74+
75+ if ( completedTask == contextTask )
76+ {
77+ var context = await contextTask ;
78+ var request = context . Request ;
79+ var response = context . Response ;
80+
81+ string ? code = request . QueryString [ "code" ] ;
82+ string ? error = request . QueryString [ "error" ] ;
83+ string html ;
84+ string ? resultCode = null ;
85+
86+ if ( ! string . IsNullOrEmpty ( error ) )
87+ {
88+ html = $ "<html><body><h1>Authorization Failed</h1><p>Error: { WebUtility . HtmlEncode ( error ) } </p></body></html>";
89+ }
90+ else if ( string . IsNullOrEmpty ( code ) )
91+ {
92+ html = "<html><body><h1>Authorization Failed</h1><p>No authorization code received.</p></body></html>" ;
93+ }
94+ else
95+ {
96+ html = "<html><body><h1>Authorization Successful</h1><p>You may now close this window.</p></body></html>" ;
97+ resultCode = code ;
98+ }
99+
100+ try
101+ {
102+ // Send response to browser
103+ byte [ ] buffer = Encoding . UTF8 . GetBytes ( html ) ;
104+ response . ContentType = "text/html" ;
105+ response . ContentLength64 = buffer . Length ;
106+ response . OutputStream . Write ( buffer , 0 , buffer . Length ) ;
107+
108+ // IMPORTANT: Explicitly close the response to ensure it's fully sent
109+ response . Close ( ) ;
110+
111+ // Now that we've finished processing the browser response,
112+ // we can safely signal completion or failure with the auth code
113+ if ( resultCode != null )
114+ {
115+ authCodeTcs . TrySetResult ( resultCode ) ;
116+ }
117+ else if ( ! string . IsNullOrEmpty ( error ) )
118+ {
119+ authCodeTcs . TrySetException ( new InvalidOperationException ( $ "Authorization failed: { error } ") ) ;
120+ }
121+ else
122+ {
123+ authCodeTcs . TrySetException ( new InvalidOperationException ( "No authorization code received" ) ) ;
124+ }
125+ }
126+ catch ( Exception ex )
127+ {
128+ authCodeTcs . TrySetException ( new InvalidOperationException ( $ "Error processing browser response: { ex . Message } ") ) ;
129+ }
130+ }
131+ }
132+ catch ( Exception ex )
133+ {
134+ authCodeTcs . TrySetException ( ex ) ;
135+ }
136+ } ) ;
137+
138+ // Now open the browser AFTER the listener is started
139+ await openBrowser ( authUrl ) ;
140+
141+ try
142+ {
143+ // Use a timeout to avoid hanging indefinitely
144+ string authCode = await authCodeTcs . Task . WaitAsync ( cts . Token ) ;
145+ return authCode ;
146+ }
147+ catch ( OperationCanceledException )
148+ {
149+ throw new InvalidOperationException ( "Authorization timed out after 5 minutes." ) ;
150+ }
151+ finally
152+ {
153+ // Ensure the listener is stopped when we're done
154+ listener . Stop ( ) ;
155+ }
156+ } ;
157+ }
158+ }
0 commit comments