Support API Gateway websocket api via LambdaServer

.autover/changes/apigateway-websocket-aspnetcore.json

@@ -0,0 +1,19 @@
+{
+  "Projects": [
+    {
+      "Name": "Amazon.Lambda.AspNetCoreServer",
+      "Type": "Minor",
+      "ChangelogMessages": [
+        "Add APIGatewayWebsocketApiProxyFunction (and TStartup variant) so API Gateway WebSocket APIs can be hosted via LambdaServer (DI, controllers). WebSocket events are dispatched as POST requests whose path is the RouteKey, allowing controller actions like [HttpPost(\"$default\")] to handle them.",
+        "Expose ParseHttpPath, ParseHttpMethod, and AddMissingRequestHeaders as protected virtual hooks on APIGatewayProxyFunction so subclasses can customize how the API Gateway request is mapped onto the ASP.NET Core request feature."
+      ]
+    },
+    {
+      "Name": "Amazon.Lambda.AspNetCoreServer.Hosting",
+      "Type": "Minor",
+      "ChangelogMessages": [
+        "Add LambdaEventSource.WebsocketApi so AddAWSLambdaHosting can wire ASP.NET Core minimal APIs and controllers up to API Gateway WebSocket events."
+      ]
+    }
+  ]
+}

Libraries/src/Amazon.Lambda.AspNetCoreServer.Hosting/Internal/LambdaRuntimeSupportServer.cs

@@ -464,4 +464,70 @@ protected override void PostMarshallItemsFeatureFeature(IItemsFeature aspNetCore
             }
         }
     }
+
+    /// <summary>
+    /// IServer for handling Lambda events from an API Gateway Websocket API.
+    /// </summary>
+    public class APIGatewayWebsocketApiLambdaRuntimeSupportServer : APIGatewayRestApiLambdaRuntimeSupportServer
+    {
+        /// <summary>
+        /// Create instances
+        /// </summary>
+        /// <param name="serviceProvider">The IServiceProvider created for the ASP.NET Core application</param>
+        public APIGatewayWebsocketApiLambdaRuntimeSupportServer(IServiceProvider serviceProvider)
+            : base(serviceProvider)
+        {
+        }
+
+        /// <summary>
+        /// Creates HandlerWrapper for processing events from API Gateway Websocket API
+        /// </summary>
+        /// <param name="serviceProvider"></param>
+        /// <returns></returns>
+        protected override HandlerWrapper CreateHandlerWrapper(IServiceProvider serviceProvider)
+        {
+            var handler = new APIGatewayWebsocketMinimalApi(serviceProvider);
+#pragma warning disable CA2252
+            var hostingOptions = serviceProvider.GetService<HostingOptions>();
+            handler.EnableResponseStreaming = hostingOptions?.EnableResponseStreaming ?? false;
+#pragma warning restore CA2252
+            Func<APIGatewayEvents.APIGatewayProxyRequest, ILambdaContext, Task<APIGatewayEvents.APIGatewayProxyResponse>> bufferedHandler = handler.FunctionHandlerAsync;
+            return HandlerWrapper.GetHandlerWrapper(bufferedHandler, Serializer);
+        }
+
+        /// <summary>
+        /// MinimalApi variant of APIGatewayWebsocketApiProxyFunction. Reuses the REST API MinimalApi plumbing
+        /// (snapshot collectors, hosting options, PostMarshall callbacks) and applies the websocket-specific
+        /// request transformations on top.
+        /// </summary>
+        public class APIGatewayWebsocketMinimalApi : APIGatewayRestApiLambdaRuntimeSupportServer.APIGatewayRestApiMinimalApi
+        {
+            /// <summary>
+            /// Create instances
+            /// </summary>
+            /// <param name="serviceProvider">The IServiceProvider created for the ASP.NET Core application</param>
+            public APIGatewayWebsocketMinimalApi(IServiceProvider serviceProvider)
+                : base(serviceProvider)
+            {
+            }
+
+            /// <inheritdoc/>
+            protected override string ParseHttpPath(APIGatewayEvents.APIGatewayProxyRequest apiGatewayRequest)
+                => "/" + System.Net.WebUtility.UrlDecode(apiGatewayRequest.RequestContext.RouteKey ?? string.Empty);
+
+            /// <inheritdoc/>
+            protected override string ParseHttpMethod(APIGatewayEvents.APIGatewayProxyRequest apiGatewayRequest) => "POST";
+
+            /// <inheritdoc/>
+            protected override Microsoft.AspNetCore.Http.IHeaderDictionary AddMissingRequestHeaders(APIGatewayEvents.APIGatewayProxyRequest apiGatewayRequest, Microsoft.AspNetCore.Http.IHeaderDictionary headers)
+            {
+                headers = base.AddMissingRequestHeaders(apiGatewayRequest, headers);
+                if (!headers.ContainsKey("Content-Type"))
+                {
+                    headers["Content-Type"] = "application/json";
+                }
+                return headers;
+            }
+        }
+    }
 }

Libraries/src/Amazon.Lambda.AspNetCoreServer.Hosting/ServiceCollectionExtensions.cs

@@ -27,7 +27,12 @@ public enum LambdaEventSource
         /// <summary>
         /// ELB Application Load Balancer
         /// </summary>
-        ApplicationLoadBalancer
+        ApplicationLoadBalancer,
+
+        /// <summary>
+        /// API Gateway WebSocket API
+        /// </summary>
+        WebsocketApi
     }
 
     /// <summary>
@@ -167,6 +172,7 @@ private static bool TryLambdaSetup(IServiceCollection services, LambdaEventSourc
                 LambdaEventSource.HttpApi => typeof(APIGatewayHttpApiV2LambdaRuntimeSupportServer),
                 LambdaEventSource.RestApi => typeof(APIGatewayRestApiLambdaRuntimeSupportServer),
                 LambdaEventSource.ApplicationLoadBalancer => typeof(ApplicationLoadBalancerLambdaRuntimeSupportServer),
+                LambdaEventSource.WebsocketApi => typeof(APIGatewayWebsocketApiLambdaRuntimeSupportServer),
                 _ => throw new ArgumentException($"Event source type {eventSource} unknown")
             };
 

Libraries/src/Amazon.Lambda.AspNetCoreServer/APIGatewayProxyFunction.cs

@@ -5,6 +5,7 @@
 using System.Security.Claims;
 using System.Text;
 
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Primitives;
@@ -194,33 +195,9 @@ protected override void MarshallRequest(InvokeFeatures features, APIGatewayProxy
             {
                 var requestFeatures = (IHttpRequestFeature)features;
                 requestFeatures.Scheme = "https";
-                requestFeatures.Method = apiGatewayRequest.HttpMethod;
+                requestFeatures.Method = this.ParseHttpMethod(apiGatewayRequest);
 
-                string path = null;
-
-                // Replaces {proxy+} in path, if exists
-                if (apiGatewayRequest.PathParameters != null && apiGatewayRequest.PathParameters.TryGetValue("proxy", out var proxy) &&
-                    !string.IsNullOrEmpty(apiGatewayRequest.Resource))
-                {
-                    var proxyPath = proxy;
-                    path = apiGatewayRequest.Resource.Replace("{proxy+}", proxyPath);
-
-                    // Adds all the rest of non greedy parameters in apiGateway.Resource to the path
-                    foreach (var pathParameter in apiGatewayRequest.PathParameters.Where(pp => pp.Key != "proxy"))
-                    {
-                        path = path.Replace($"{{{pathParameter.Key}}}", pathParameter.Value);
-                    }
-                }
-
-                if (string.IsNullOrEmpty(path))
-                {
-                    path = apiGatewayRequest.Path;
-                }
-
-                if (!path.StartsWith("/"))
-                {
-                    path = "/" + path;
-                }
+                string path = this.ParseHttpPath(apiGatewayRequest);
 
                 var rawQueryString = Utilities.CreateQueryStringParameters(
                     apiGatewayRequest.QueryStringParameters, apiGatewayRequest.MultiValueQueryStringParameters, true);
@@ -257,13 +234,7 @@ protected override void MarshallRequest(InvokeFeatures features, APIGatewayProxy
 
                 Utilities.SetHeadersCollection(requestFeatures.Headers, apiGatewayRequest.Headers, apiGatewayRequest.MultiValueHeaders);
 
-                if (!requestFeatures.Headers.ContainsKey("Host"))
-                {
-                    var apiId = apiGatewayRequest?.RequestContext?.ApiId ?? "";
-                    var stage = apiGatewayRequest?.RequestContext?.Stage ?? "";
-
-                    requestFeatures.Headers["Host"] = $"apigateway-{apiId}-{stage}";
-                }
+                requestFeatures.Headers = this.AddMissingRequestHeaders(apiGatewayRequest, requestFeatures.Headers);
 
 
                 if (!string.IsNullOrEmpty(apiGatewayRequest.Body))
@@ -299,6 +270,7 @@ protected override void MarshallRequest(InvokeFeatures features, APIGatewayProxy
                 {
                     connectionFeatures.RemotePort = int.Parse(forwardedPort, CultureInfo.InvariantCulture);
                 }
+                connectionFeatures.ConnectionId = apiGatewayRequest.RequestContext?.ConnectionId;
 
                 // Call consumers customize method in case they want to change how API Gateway's request
                 // was marshalled into ASP.NET Core request.
@@ -380,5 +352,67 @@ protected override APIGatewayProxyResponse MarshallResponse(IHttpResponseFeature
 
             return response;
         }
+
+        /// <summary>
+        /// Determines the path that should be assigned to <see cref="IHttpRequestFeature.Path"/> for this request.
+        /// The default implementation honors <c>{proxy+}</c> resource templates and falls back to <see cref="APIGatewayProxyRequest.Path"/>.
+        /// Subclasses can override to derive the path from a different source (e.g. websocket route keys).
+        /// </summary>
+        protected virtual string ParseHttpPath(APIGatewayProxyRequest apiGatewayRequest)
+        {
+            string path = null;
+
+            // Replaces {proxy+} in path, if exists
+            if (apiGatewayRequest.PathParameters != null && apiGatewayRequest.PathParameters.TryGetValue("proxy", out var proxy) &&
+                !string.IsNullOrEmpty(apiGatewayRequest.Resource))
+            {
+                var proxyPath = proxy;
+                path = apiGatewayRequest.Resource.Replace("{proxy+}", proxyPath);
+
+                // Adds all the rest of non greedy parameters in apiGateway.Resource to the path
+                foreach (var pathParameter in apiGatewayRequest.PathParameters.Where(pp => pp.Key != "proxy"))
+                {
+                    path = path.Replace($"{{{pathParameter.Key}}}", pathParameter.Value);
+                }
+            }
+
+            if (string.IsNullOrEmpty(path))
+            {
+                path = apiGatewayRequest.Path;
+            }
+
+            if (!path.StartsWith("/"))
+            {
+                path = "/" + path;
+            }
+            return path;
+        }
+
+        /// <summary>
+        /// Determines the HTTP method that should be assigned to <see cref="IHttpRequestFeature.Method"/>.
+        /// The default returns <see cref="APIGatewayProxyRequest.HttpMethod"/>; subclasses can override to force
+        /// a fixed method (e.g. websocket events are always exposed as POST).
+        /// </summary>
+        protected virtual string ParseHttpMethod(APIGatewayProxyRequest apiGatewayRequest)
+        {
+            return apiGatewayRequest.HttpMethod;
+        }
+
+        /// <summary>
+        /// Adds any headers that API Gateway did not include but ASP.NET Core needs in order to route or bind the request.
+        /// The default implementation adds a synthesized <c>Host</c> header. Subclasses can override to add additional defaults
+        /// (e.g. a default <c>Content-Type</c> for websocket payloads).
+        /// </summary>
+        protected virtual IHeaderDictionary AddMissingRequestHeaders(APIGatewayProxyRequest apiGatewayRequest, IHeaderDictionary headers)
+        {
+            if (!headers.ContainsKey("Host"))
+            {
+                var apiId = apiGatewayRequest?.RequestContext?.ApiId ?? "";
+                var stage = apiGatewayRequest?.RequestContext?.Stage ?? "";
+
+                headers["Host"] = $"apigateway-{apiId}-{stage}";
+            }
+            return headers;
+        }
     }
 }

Libraries/src/Amazon.Lambda.AspNetCoreServer/APIGatewayWebsocketApiProxyFunction.cs

@@ -0,0 +1,78 @@
+using System;
+
+using Microsoft.AspNetCore.Http;
+
+using Amazon.Lambda.APIGatewayEvents;
+using Amazon.Lambda.AspNetCoreServer.Internal;
+
+namespace Amazon.Lambda.AspNetCoreServer
+{
+    /// <summary>
+    /// Base class for ASP.NET Core Lambda functions invoked by API Gateway Websocket APIs.
+    ///
+    /// Websocket events are surfaced as POST requests whose path is the RouteKey, so the same Lambda
+    /// should be referenced by every websocket route that has a matching ASP.NET Core controller route
+    /// (e.g. <c>[HttpPost("$default")]</c>, <c>[HttpPost("$connect")]</c>) for the ASP.NET Core IServer
+    /// to successfully dispatch requests.
+    /// </summary>
+    public abstract class APIGatewayWebsocketApiProxyFunction : APIGatewayProxyFunction
+    {
+        /// <summary>
+        /// Default constructor. The ASP.NET Core framework is initialized as part of construction.
+        /// </summary>
+        protected APIGatewayWebsocketApiProxyFunction()
+            : base()
+        {
+        }
+
+        /// <summary>
+        /// Constructor that lets the caller defer ASP.NET Core framework initialization until the first request.
+        /// </summary>
+        /// <param name="startupMode">Configures when the ASP.NET Core framework is initialized.</param>
+        protected APIGatewayWebsocketApiProxyFunction(StartupMode startupMode)
+            : base(startupMode)
+        {
+        }
+
+        /// <summary>
+        /// Constructor used by Amazon.Lambda.AspNetCoreServer.Hosting to support ASP.NET Core projects using the Minimal API style.
+        /// </summary>
+        /// <param name="hostedServices">The service provider built by the ASP.NET Core host.</param>
+        protected APIGatewayWebsocketApiProxyFunction(IServiceProvider hostedServices)
+            : base(hostedServices)
+        {
+        }
+
+        /// <summary>
+        /// Maps the websocket event to a request path of <c>/{RouteKey}</c> so ASP.NET Core can dispatch to a controller
+        /// route declared with <c>[HttpPost("{RouteKey}")]</c>.
+        /// </summary>
+        protected override string ParseHttpPath(APIGatewayProxyRequest apiGatewayRequest)
+        {
+            return "/" + Utilities.DecodeResourcePath(apiGatewayRequest.RequestContext.RouteKey);
+        }
+
+        /// <summary>
+        /// Always returns <c>POST</c> for websocket events. Combined with <see cref="ParseHttpPath"/>, this lets the same
+        /// Lambda route every websocket event into an ASP.NET Core controller action.
+        /// </summary>
+        protected override string ParseHttpMethod(APIGatewayProxyRequest apiGatewayRequest)
+        {
+            return "POST";
+        }
+
+        /// <summary>
+        /// Adds a default <c>Content-Type</c> of <c>application/json</c> when API Gateway did not supply one.
+        /// Websocket message payloads are typically JSON, but the gateway does not set headers automatically.
+        /// </summary>
+        protected override IHeaderDictionary AddMissingRequestHeaders(APIGatewayProxyRequest apiGatewayRequest, IHeaderDictionary headers)
+        {
+            headers = base.AddMissingRequestHeaders(apiGatewayRequest, headers);
+            if (!headers.ContainsKey("Content-Type"))
+            {
+                headers["Content-Type"] = "application/json";
+            }
+            return headers;
+        }
+    }
+}

Libraries/src/Amazon.Lambda.AspNetCoreServer/APIGatewayWebsocketApiProxyFunction{TStartup}.cs

@@ -0,0 +1,36 @@
+using Microsoft.AspNetCore.Hosting;
+using System.Diagnostics.CodeAnalysis;
+
+namespace Amazon.Lambda.AspNetCoreServer
+{
+    /// <summary>
+    /// Strongly-typed variant of <see cref="APIGatewayWebsocketApiProxyFunction"/> that wires up an ASP.NET Core Startup class.
+    /// The Lambda function handler should point at the inherited <c>FunctionHandlerAsync</c> method.
+    /// </summary>
+    /// <typeparam name ="TStartup">The type containing the startup methods for the application.</typeparam>
+    public abstract class APIGatewayWebsocketApiProxyFunction<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicMethods | DynamicallyAccessedMemberTypes.PublicConstructors)] TStartup> : APIGatewayWebsocketApiProxyFunction where TStartup : class
+    {
+        /// <summary>
+        /// Default constructor. The ASP.NET Core framework is initialized as part of construction.
+        /// </summary>
+        protected APIGatewayWebsocketApiProxyFunction()
+            : base()
+        {
+        }
+
+        /// <summary>
+        /// Constructor that lets the caller defer ASP.NET Core framework initialization until the first request.
+        /// </summary>
+        /// <param name="startupMode">Configures when the ASP.NET Core framework is initialized.</param>
+        protected APIGatewayWebsocketApiProxyFunction(StartupMode startupMode)
+            : base(startupMode)
+        {
+        }
+
+        /// <inheritdoc/>
+        protected override void Init(IWebHostBuilder builder)
+        {
+            builder.UseStartup<TStartup>();
+        }
+    }
+}