Amend middleware logic

Author: localden

src/ModelContextProtocol.AspNetCore/AuthorizationMiddleware.cs

@@ -47,8 +47,9 @@ public async Task InvokeAsync(
             return;
         }
 
-        // Handle the PRM document endpoint
-        if (context.Request.Path.StartsWithSegments("/.well-known/oauth-protected-resource"))
+        // Handle the PRM document endpoint if not handled by the endpoint
+        if (context.Request.Path.StartsWithSegments("/.well-known/oauth-protected-resource") &&
+            context.GetEndpoint() == null)
         {
             _logger.LogDebug("Serving Protected Resource Metadata document");
             context.Response.ContentType = "application/json";
@@ -59,40 +60,8 @@ await JsonSerializer.SerializeAsync(
             return;
         }
 
-        // Serve SSE and message endpoints with authorization
-        if (context.Request.Path.StartsWithSegments("/sse") || 
-            (context.Request.Path.Value?.EndsWith("/message") == true))
-        {
-            // Check if the Authorization header is present
-            if (!context.Request.Headers.TryGetValue("Authorization", out var authHeader) || string.IsNullOrEmpty(authHeader))
-            {
-                // No Authorization header present, return 401 Unauthorized
-                var prm = authProvider.GetProtectedResourceMetadata();
-                var prmUrl = GetPrmUrl(context, prm.Resource);
-                
-                _logger.LogDebug("Authorization required, returning 401 Unauthorized with WWW-Authenticate header");
-                context.Response.StatusCode = StatusCodes.Status401Unauthorized;
-                context.Response.Headers.Append("WWW-Authenticate", $"Bearer resource_metadata=\"{prmUrl}\"");
-                return;
-            }
-
-            // Validate the token - ensuring authHeader is a non-null string
-            string authHeaderValue = authHeader.ToString();
-            bool isValid = await authProvider.ValidateTokenAsync(authHeaderValue);
-            if (!isValid)
-            {
-                // Invalid token, return 401 Unauthorized
-                var prm = authProvider.GetProtectedResourceMetadata();
-                var prmUrl = GetPrmUrl(context, prm.Resource);
-                
-                _logger.LogDebug("Invalid authorization token, returning 401 Unauthorized");
-                context.Response.StatusCode = StatusCodes.Status401Unauthorized;
-                context.Response.Headers.Append("WWW-Authenticate", $"Bearer resource_metadata=\"{prmUrl}\"");
-                return;
-            }
-        }
-
-        // Token is valid or endpoint doesn't require authentication, proceed to the next middleware
+        // Proceed to the next middleware - authorization for SSE and message endpoints
+        // is now handled by endpoint filters
         await _next(context);
     }
 

src/ModelContextProtocol.AspNetCore/HttpMcpServerBuilderExtensions.cs

@@ -18,12 +18,13 @@ public static class HttpMcpServerBuilderExtensions
     /// <param name="configureOptions">Configures options for the Streamable HTTP transport. This allows configuring per-session
     /// <see cref="McpServerOptions"/> and running logic before and after a session.</param>
     /// <returns>The builder provided in <paramref name="builder"/>.</returns>
-    /// <exception cref="ArgumentNullException"><paramref name="builder"/> is <see langword="null"/>.</exception>
+    /// <exception cref="ArgumentNullException"><paramref name="builder"/> is <see langword="null"/>.</exception>    
     public static IMcpServerBuilder WithHttpTransport(this IMcpServerBuilder builder, Action<HttpServerTransportOptions>? configureOptions = null)
     {
         ArgumentNullException.ThrowIfNull(builder);
         builder.Services.TryAddSingleton<StreamableHttpHandler>();
         builder.Services.TryAddSingleton<SseHandler>();
+        builder.Services.TryAddSingleton<McpAuthorizationFilterFactory>();
         builder.Services.AddHostedService<IdleTrackingBackgroundService>();
 
         if (configureOptions is not null)

src/ModelContextProtocol.AspNetCore/McpAuthorizationFilter.cs

@@ -0,0 +1,79 @@
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Logging;
+using ModelContextProtocol.Protocol.Auth;
+
+namespace ModelContextProtocol.AspNetCore;
+
+/// <summary>
+/// An endpoint filter that handles authorization for MCP endpoints.
+/// </summary>
+internal class McpAuthorizationFilter : IEndpointFilter
+{
+    private readonly ILogger<McpAuthorizationFilter> _logger;
+    private readonly IServerAuthorizationProvider _authProvider;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="McpAuthorizationFilter"/> class.
+    /// </summary>
+    /// <param name="logger">The logger.</param>
+    /// <param name="authProvider">The authorization provider.</param>
+    public McpAuthorizationFilter(
+        ILogger<McpAuthorizationFilter> logger,
+        IServerAuthorizationProvider authProvider)
+    {
+        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
+        _authProvider = authProvider ?? throw new ArgumentNullException(nameof(authProvider));
+    }
+
+    /// <inheritdoc/>
+    public async ValueTask<object?> InvokeAsync(EndpointFilterInvocationContext context, EndpointFilterDelegate next)
+    {
+        var httpContext = context.HttpContext;
+
+        // Check if the Authorization header is present
+        if (!httpContext.Request.Headers.TryGetValue("Authorization", out var authHeader) || string.IsNullOrEmpty(authHeader))
+        {
+            // No Authorization header present, return 401 Unauthorized
+            var prm = _authProvider.GetProtectedResourceMetadata();
+            var prmUrl = GetPrmUrl(httpContext, prm.Resource);
+            
+            _logger.LogDebug("Authorization required, returning 401 Unauthorized with WWW-Authenticate header");
+            httpContext.Response.StatusCode = StatusCodes.Status401Unauthorized;
+            httpContext.Response.Headers.Append("WWW-Authenticate", $"Bearer resource_metadata=\"{prmUrl}\"");
+            return Results.Empty;
+        }
+
+        // Validate the token - ensuring authHeader is a non-null string
+        string authHeaderValue = authHeader.ToString();
+        bool isValid = await _authProvider.ValidateTokenAsync(authHeaderValue);
+        if (!isValid)
+        {
+            // Invalid token, return 401 Unauthorized
+            var prm = _authProvider.GetProtectedResourceMetadata();
+            var prmUrl = GetPrmUrl(httpContext, prm.Resource);
+            
+            _logger.LogDebug("Invalid authorization token, returning 401 Unauthorized");
+            httpContext.Response.StatusCode = StatusCodes.Status401Unauthorized;
+            httpContext.Response.Headers.Append("WWW-Authenticate", $"Bearer resource_metadata=\"{prmUrl}\"");
+            return Results.Empty;
+        }
+
+        // Token is valid, proceed to the next filter
+        return await next(context);
+    }
+
+    private static string GetPrmUrl(HttpContext context, string resourceUri)
+    {
+        // Use the actual resource URI from PRM if it's an absolute URL, otherwise build the URL
+        if (Uri.TryCreate(resourceUri, UriKind.Absolute, out _))
+        {
+            return $"{resourceUri.TrimEnd('/')}/.well-known/oauth-protected-resource";
+        }
+
+        // Build the URL from the current request
+        var request = context.Request;
+        var scheme = request.Scheme;
+        var host = request.Host.Value;
+        return $"{scheme}://{host}/.well-known/oauth-protected-resource";
+    }
+}

src/ModelContextProtocol.AspNetCore/McpAuthorizationFilterFactory.cs

@@ -0,0 +1,53 @@
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
+using ModelContextProtocol.Protocol.Auth;
+
+namespace ModelContextProtocol.AspNetCore;
+
+/// <summary>
+/// Factory for creating <see cref="McpAuthorizationFilter"/> instances.
+/// </summary>
+internal class McpAuthorizationFilterFactory
+{
+    private readonly IServiceProvider _serviceProvider;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="McpAuthorizationFilterFactory"/> class.
+    /// </summary>
+    /// <param name="serviceProvider">The service provider.</param>
+    public McpAuthorizationFilterFactory(IServiceProvider serviceProvider)
+    {
+        _serviceProvider = serviceProvider ?? throw new ArgumentNullException(nameof(serviceProvider));
+    }
+
+    /// <summary>
+    /// Creates an endpoint filter delegate for authorization.
+    /// </summary>
+    /// <param name="context">The endpoint filter factory context.</param>
+    /// <param name="next">The next filter delegate in the pipeline.</param>
+    /// <returns>The filter delegate.</returns>
+    public EndpointFilterDelegate Create(EndpointFilterFactoryContext context, EndpointFilterDelegate next)
+    {
+        // This factory creates a filter that checks if the current endpoint is an SSE or message endpoint
+        // and applies authorization only to those endpoints
+        return async invocationContext =>
+        {
+            var httpContext = invocationContext.HttpContext;
+            var path = httpContext.Request.Path.Value?.TrimEnd('/');
+
+            // Only apply authorization to /sse and /message endpoints
+            if (path != null && (path.EndsWith("/sse") || path.EndsWith("/message")))
+            {
+                var authProvider = _serviceProvider.GetRequiredService<IServerAuthorizationProvider>();
+                var logger = _serviceProvider.GetRequiredService<ILogger<McpAuthorizationFilter>>();
+
+                var filter = new McpAuthorizationFilter(logger, authProvider);
+                return await filter.InvokeAsync(invocationContext, next);
+            }
+
+            // For all other endpoints, just invoke the next filter
+            return await next(invocationContext);
+        };
+    }
+}

src/ModelContextProtocol.AspNetCore/McpEndpointRouteBuilderExtensions.cs

@@ -2,7 +2,9 @@
 using Microsoft.AspNetCore.Http.Metadata;
 using Microsoft.AspNetCore.Routing;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
 using ModelContextProtocol.AspNetCore;
+using ModelContextProtocol.Protocol.Auth;
 using ModelContextProtocol.Protocol.Messages;
 using System.Diagnostics.CodeAnalysis;
 
@@ -20,12 +22,36 @@ public static class McpEndpointRouteBuilderExtensions
     /// </summary>
     /// <param name="endpoints">The web application to attach MCP HTTP endpoints.</param>
     /// <param name="pattern">The route pattern prefix to map to.</param>
-    /// <returns>Returns a builder for configuring additional endpoint conventions like authorization policies.</returns>
+    /// <returns>Returns a builder for configuring additional endpoint conventions like authorization policies.</returns>    
     public static IEndpointConventionBuilder MapMcp(this IEndpointRouteBuilder endpoints, [StringSyntax("Route")] string pattern = "")
     {
         var streamableHttpHandler = endpoints.ServiceProvider.GetService<StreamableHttpHandler>() ??
             throw new InvalidOperationException("You must call WithHttpTransport(). Unable to find required services. Call builder.Services.AddMcpServer().WithHttpTransport() in application startup code.");
 
+        // Map the protected resource metadata endpoint if authorization is configured
+        var authProvider = endpoints.ServiceProvider.GetService<IServerAuthorizationProvider>();
+        if (authProvider != null)
+        {
+            // Create and register the ProtectedResourceMetadataHandler if it's not already registered
+            ProtectedResourceMetadataHandler? prmHandler = null;
+            try
+            {
+                prmHandler = endpoints.ServiceProvider.GetService<ProtectedResourceMetadataHandler>();
+            }
+            catch
+            {
+                // Ignore - we'll create it below
+            }
+
+            if (prmHandler == null)
+            {
+                var logger = endpoints.ServiceProvider.GetRequiredService<ILogger<ProtectedResourceMetadataHandler>>();
+                prmHandler = new ProtectedResourceMetadataHandler(logger, authProvider);
+            }
+
+            endpoints.MapGet("/.well-known/oauth-protected-resource", prmHandler.HandleAsync);
+        }
+
         var mcpGroup = endpoints.MapGroup(pattern);
         var streamableHttpGroup = mcpGroup.MapGroup("")
             .WithDisplayName(b => $"MCP Streamable HTTP | {b.DisplayName}")
@@ -44,6 +70,16 @@ public static IEndpointConventionBuilder MapMcp(this IEndpointRouteBuilder endpo
         var sseGroup = mcpGroup.MapGroup("")
             .WithDisplayName(b => $"MCP HTTP with SSE | {b.DisplayName}");
 
+        // Apply authorization filter to SSE endpoints if authorization is configured
+        if (authProvider != null)
+        {
+            // Create the filter factory
+            var filterFactory = endpoints.ServiceProvider.GetRequiredService<McpAuthorizationFilterFactory>();
+            
+            // Apply filter to SSE and message endpoints
+            sseGroup.AddEndpointFilterFactory(filterFactory.Create);
+        }
+
         sseGroup.MapGet("/sse", sseHandler.HandleSseRequestAsync)
             .WithMetadata(new ProducesResponseTypeMetadata(StatusCodes.Status200OK, contentTypes: ["text/event-stream"]));
         sseGroup.MapPost("/message", sseHandler.HandleMessageRequestAsync)

src/ModelContextProtocol.AspNetCore/ProtectedResourceMetadataHandler.cs

@@ -0,0 +1,44 @@
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Logging;
+using ModelContextProtocol.Protocol.Auth;
+using ModelContextProtocol.Utils.Json;
+using System.Text.Json;
+
+namespace ModelContextProtocol.AspNetCore;
+
+/// <summary>
+/// Handler for the Protected Resource Metadata document endpoint.
+/// </summary>
+internal class ProtectedResourceMetadataHandler
+{
+    private readonly ILogger<ProtectedResourceMetadataHandler> _logger;
+    private readonly IServerAuthorizationProvider _authProvider;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="ProtectedResourceMetadataHandler"/> class.
+    /// </summary>
+    /// <param name="logger">The logger.</param>
+    /// <param name="authProvider">The authorization provider.</param>
+    public ProtectedResourceMetadataHandler(
+        ILogger<ProtectedResourceMetadataHandler> logger,
+        IServerAuthorizationProvider authProvider)
+    {
+        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
+        _authProvider = authProvider ?? throw new ArgumentNullException(nameof(authProvider));
+    }
+
+    /// <summary>
+    /// Handles the request for the Protected Resource Metadata document.
+    /// </summary>
+    /// <param name="context">The HTTP context.</param>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    public async Task HandleAsync(HttpContext context)
+    {
+        _logger.LogDebug("Serving Protected Resource Metadata document");
+        context.Response.ContentType = "application/json";
+        await JsonSerializer.SerializeAsync(
+            context.Response.Body, 
+            _authProvider.GetProtectedResourceMetadata(),
+            McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ProtectedResourceMetadata)));
+    }
+}