Update with proper token logic

localden · localden · commit 2f44765b2677 · 2025-04-24T00:02:24.000-07:00
diff --git a/samples/AuthorizationExample/Program.cs b/samples/AuthorizationExample/Program.cs
@@ -86,6 +86,8 @@ public static async Task Main(string[] args)
             {
                 Console.WriteLine($"Inner Error: {ex.InnerException.Message}");
             }
+            // Print the stack trace for debugging
+            Console.WriteLine($"Stack Trace:\n{ex.StackTrace}");
         }
     }
 }
diff --git a/samples/AuthorizationServerExample/Program.cs b/samples/AuthorizationServerExample/Program.cs
@@ -35,11 +35,8 @@ public static async Task Main(string[] args)
         // In a real application, this would verify the token with your identity provider
         async Task<bool> ValidateToken(string token)
         {
-            // For demo purposes, we'll accept any token that starts with "valid_"
-            // In production, you would validate the token with your identity provider
-            var isValid = token.StartsWith("valid_", StringComparison.OrdinalIgnoreCase);
-            Console.WriteLine($"Token validation result: {(isValid ? "Valid" : "Invalid")}");
-            return isValid;
+            // For demo purposes, we'll accept any token.
+            return true;
         }
 
         // 3. Create an authorization provider with the PRM and token validator
diff --git a/src/ModelContextProtocol.AspNetCore/AuthorizationMiddleware.cs b/src/ModelContextProtocol.AspNetCore/AuthorizationMiddleware.cs
@@ -32,13 +32,14 @@ public AuthorizationMiddleware(RequestDelegate next, ILogger<AuthorizationMiddle
     /// </summary>
     /// <param name="context">The HTTP context.</param>
     /// <param name="serverOptions">The MCP server options.</param>
+    /// <param name="authProvider">The authorization provider.</param>
     /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
-    public async Task InvokeAsync(HttpContext context, IOptions<McpServerOptions> serverOptions)
+    public async Task InvokeAsync(
+        HttpContext context, 
+        IOptions<McpServerOptions> serverOptions,
+        IServerAuthorizationProvider? authProvider = null)
     {
         // Check if authorization is configured
-        var authCapability = serverOptions.Value.Capabilities?.Authorization;
-        var authProvider = authCapability?.AuthorizationProvider;
-
         if (authProvider == null)
         {
             // Authorization is not configured, proceed to the next middleware
diff --git a/src/ModelContextProtocol/Configuration/McpServerAuthorizationExtensions.cs b/src/ModelContextProtocol/Configuration/McpServerAuthorizationExtensions.cs
@@ -24,13 +24,12 @@ public static IMcpServerBuilder WithAuthorization(
         Throw.IfNull(builder);
         Throw.IfNull(authorizationProvider);
 
+        // Register the authorization provider in the DI container
+        builder.Services.AddSingleton(authorizationProvider);
+
         builder.Services.Configure<McpServerOptions>(options =>
         {
             options.Capabilities ??= new ServerCapabilities();
-            options.Capabilities.Authorization = new AuthorizationCapability
-            {
-                AuthorizationProvider = authorizationProvider
-            };
         });
 
         return builder;
diff --git a/src/ModelContextProtocol/Protocol/Auth/DefaultAuthorizationHandler.cs b/src/ModelContextProtocol/Protocol/Auth/DefaultAuthorizationHandler.cs
@@ -1,9 +1,8 @@
-using System.Diagnostics;
-using System.Net;
-using System.Net.Http.Headers;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 using ModelContextProtocol.Utils;
+using System.Net;
+using System.Net.Http.Headers;
 
 namespace ModelContextProtocol.Protocol.Auth;
 
diff --git a/src/ModelContextProtocol/Protocol/Transport/SseClientSessionTransport.cs b/src/ModelContextProtocol/Protocol/Transport/SseClientSessionTransport.cs
@@ -89,29 +89,12 @@ public override async Task SendMessageAsync(
         if (_messageEndpoint == null)
             throw new InvalidOperationException("Transport not connected");
 
-        using var content = new StringContent(
-            JsonSerializer.Serialize(message, McpJsonUtilities.JsonContext.Default.JsonRpcMessage),
-            Encoding.UTF8,
-            "application/json"
-        );
-
         string messageId = "(no id)";
 
         if (message is JsonRpcMessageWithId messageWithId)
         {
             messageId = messageWithId.Id.ToString();
         }
-
-        using var httpRequestMessage = new HttpRequestMessage(HttpMethod.Post, _messageEndpoint)
-        {
-            Content = content,
-        };
-        
-        // Add authorization headers if needed
-        await _authorizationHandler.AuthenticateRequestAsync(httpRequestMessage).ConfigureAwait(false);
-        
-        // Copy additional headers
-        CopyAdditionalHeaders(httpRequestMessage.Headers);
         
         // Send the request, handling potential auth challenges
         HttpResponseMessage? response = null;
@@ -120,37 +103,32 @@ public override async Task SendMessageAsync(
         do
         {
             authRetry = false;
-            response = await _httpClient.SendAsync(httpRequestMessage, cancellationToken).ConfigureAwait(false);
             
-            // Handle 401 Unauthorized response
+            // Create a new request for each attempt
+            using var currentRequest = new HttpRequestMessage(HttpMethod.Post, _messageEndpoint);
+            currentRequest.Content = new StringContent(
+                JsonSerializer.Serialize(message, McpJsonUtilities.JsonContext.Default.JsonRpcMessage),
+                Encoding.UTF8,
+                "application/json"
+            );
+            
+            // Add authorization headers if needed - the handler will only add headers if auth is required
+            await _authorizationHandler.AuthenticateRequestAsync(currentRequest).ConfigureAwait(false);
+            
+            // Copy additional headers
+            CopyAdditionalHeaders(currentRequest.Headers);
+            
+            // Dispose previous response before making a new request
+            response?.Dispose();
+            
+            response = await _httpClient.SendAsync(currentRequest, cancellationToken).ConfigureAwait(false);
+            
+            // Handle 401 Unauthorized response - this will only execute if the server requires auth
             if (response.StatusCode == HttpStatusCode.Unauthorized)
             {
                 // Try to handle the unauthorized response
                 authRetry = await _authorizationHandler.HandleUnauthorizedResponseAsync(
                     response, _messageEndpoint).ConfigureAwait(false);
-                
-                if (authRetry)
-                {
-                    // Create a new request (we can't reuse the previous one)
-                    using var newRequest = new HttpRequestMessage(HttpMethod.Post, _messageEndpoint)
-                    {
-                        Content = new StringContent(
-                            JsonSerializer.Serialize(message, McpJsonUtilities.JsonContext.Default.JsonRpcMessage),
-                            Encoding.UTF8,
-                            "application/json"
-                        )
-                    };
-                    
-                    // Add authorization headers for the new request
-                    await _authorizationHandler.AuthenticateRequestAsync(newRequest).ConfigureAwait(false);
-                    CopyAdditionalHeaders(newRequest.Headers);
-                    
-                    // Dispose the previous response
-                    response.Dispose();
-                    
-                    // Send the new request
-                    response = await _httpClient.SendAsync(newRequest, cancellationToken).ConfigureAwait(false);
-                }
             }
         } while (authRetry);
 
@@ -252,55 +230,39 @@ private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
     {
         try
         {
-            using var request = new HttpRequestMessage(HttpMethod.Get, _sseEndpoint);
-            request.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("text/event-stream"));
-            
-            // Add authorization headers if needed
-            await _authorizationHandler.AuthenticateRequestAsync(request).ConfigureAwait(false);
-            
-            // Copy additional headers
-            CopyAdditionalHeaders(request.Headers);
-
             // Send the request, handling potential auth challenges
             HttpResponseMessage? response = null;
             bool authRetry = false;
             
             do
             {
                 authRetry = false;
+                
+                // Create a new request for each attempt
+                using var currentRequest = new HttpRequestMessage(HttpMethod.Get, _sseEndpoint);
+                currentRequest.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("text/event-stream"));
+                
+                // Add authorization headers if needed - the handler will only add headers if auth is required
+                await _authorizationHandler.AuthenticateRequestAsync(currentRequest).ConfigureAwait(false);
+                
+                // Copy additional headers
+                CopyAdditionalHeaders(currentRequest.Headers);
+                
+                // Dispose previous response before making a new request
+                response?.Dispose();
+                
                 response = await _httpClient.SendAsync(
-                    request,
+                    currentRequest,
                     HttpCompletionOption.ResponseHeadersRead,
                     cancellationToken
                 ).ConfigureAwait(false);
                 
-                // Handle 401 Unauthorized response
+                // Handle 401 Unauthorized response - this will only execute if the server requires auth
                 if (response.StatusCode == HttpStatusCode.Unauthorized)
                 {
                     // Try to handle the unauthorized response
                     authRetry = await _authorizationHandler.HandleUnauthorizedResponseAsync(
                         response, _sseEndpoint).ConfigureAwait(false);
-                    
-                    if (authRetry)
-                    {
-                        // Create a new request (we can't reuse the previous one)
-                        using var newRequest = new HttpRequestMessage(HttpMethod.Get, _sseEndpoint);
-                        newRequest.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("text/event-stream"));
-                        
-                        // Add authorization headers for the new request
-                        await _authorizationHandler.AuthenticateRequestAsync(newRequest).ConfigureAwait(false);
-                        CopyAdditionalHeaders(newRequest.Headers);
-                        
-                        // Dispose the previous response
-                        response.Dispose();
-                        
-                        // Send the new request
-                        response = await _httpClient.SendAsync(
-                            newRequest,
-                            HttpCompletionOption.ResponseHeadersRead,
-                            cancellationToken
-                        ).ConfigureAwait(false);
-                    }
                 }
             } while (authRetry);
 
diff --git a/src/ModelContextProtocol/Protocol/Types/ServerCapabilities.cs b/src/ModelContextProtocol/Protocol/Types/ServerCapabilities.cs
@@ -35,12 +35,6 @@ public class ServerCapabilities
     [JsonPropertyName("experimental")]
     public Dictionary<string, object>? Experimental { get; set; }
 
-    /// <summary>
-    /// Gets or sets a server's authorization capability, supporting OAuth 2.0 authorization flows.
-    /// </summary>
-    [JsonPropertyName("authorization")]
-    public AuthorizationCapability? Authorization { get; set; }
-
     /// <summary>
     /// Gets or sets a server's logging capability, supporting sending log messages to the client.
     /// </summary>
diff --git a/src/ModelContextProtocol/Server/McpServer.cs b/src/ModelContextProtocol/Server/McpServer.cs
@@ -328,7 +328,6 @@ await originalListPromptsHandler(request, cancellationToken).ConfigureAwait(fals
             ServerCapabilities = new()
             {
                 Experimental = options.Capabilities?.Experimental,
-                Authorization = options.Capabilities?.Authorization,
                 Logging = options.Capabilities?.Logging,
                 Tools = options.Capabilities?.Tools,
                 Resources = options.Capabilities?.Resources,
@@ -427,7 +426,6 @@ await originalListToolsHandler(request, cancellationToken).ConfigureAwait(false)
             ServerCapabilities = new()
             {
                 Experimental = options.Capabilities?.Experimental,
-                Authorization = options.Capabilities?.Authorization,
                 Logging = options.Capabilities?.Logging,
                 Prompts = options.Capabilities?.Prompts,
                 Resources = options.Capabilities?.Resources,

Original file line number	Diff line number	Diff line change
`@@ -86,6 +86,8 @@ public static async Task Main(string[] args)`
`86`	`86`	`{`
`87`	`87`	`Console.WriteLine($"Inner Error: {ex.InnerException.Message}");`
`88`	`88`	`}`
	`89`	`+ // Print the stack trace for debugging`
	`90`	`+ Console.WriteLine($"Stack Trace:\n{ex.StackTrace}");`
`89`	`91`	`}`
`90`	`92`	`}`
`91`	`93`	`}`
Original file line number	Diff line number	Diff line change
`@@ -35,11 +35,8 @@ public static async Task Main(string[] args)`
`35`	`35`	`// In a real application, this would verify the token with your identity provider`
`36`	`36`	`async Task<bool> ValidateToken(string token)`
`37`	`37`	`{`
`38`		`- // For demo purposes, we'll accept any token that starts with "valid_"`
`39`		`- // In production, you would validate the token with your identity provider`
`40`		`- var isValid = token.StartsWith("valid_", StringComparison.OrdinalIgnoreCase);`
`41`		`- Console.WriteLine($"Token validation result: {(isValid ? "Valid" : "Invalid")}");`
`42`		`- return isValid;`
	`38`	`+ // For demo purposes, we'll accept any token.`
	`39`	`+ return true;`
`43`	`40`	`}`
`44`	`41`
`45`	`42`	`// 3. Create an authorization provider with the PRM and token validator`