Rework encoding and add support for using API Gateway

Author: normj

Libraries/src/Amazon.Lambda.Core/ResponseStreaming/HttpResponseStreamPrelude.cs

@@ -0,0 +1,89 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+#if NET8_0_OR_GREATER
+using System.Collections.Generic;
+using System.Net;
+using System.Runtime.Versioning;
+using System.Text.Json;
+
+namespace Amazon.Lambda.Core.ResponseStreaming
+{
+    /// <summary>
+    /// The HTTP response prelude to be sent as the first chunk of a streaming response when using <see cref="LambdaResponseStreamFactory.CreateHttpStream"/>.
+    /// </summary>
+    [RequiresPreviewFeatures(LambdaResponseStreamFactory.ParameterizedPreviewMessage)]
+    public class HttpResponseStreamPrelude
+    {
+        /// <summary>
+        /// The Http status code to include in the response prelude.
+        /// </summary>
+        public HttpStatusCode? StatusCode { get; set; }
+
+        /// <summary>
+        /// The response headers to include in the response prelude. This collection supports setting single value for the same headers.
+        /// </summary>
+        public IDictionary<string, string> Headers { get; set; } = new Dictionary<string, string>();
+
+        /// <summary>
+        /// The response headers to include in the response prelude. This collection supports setting multiple values for the same headers.
+        /// </summary>
+        public IDictionary<string, IList<string>> MultiValueHeaders { get; set; } = new Dictionary<string, IList<string>>();
+
+        /// <summary>
+        /// The list of cookies to include in the response prelude. This is used for Lambda Function URL responses, which support a separate "cookies" field in the response JSON for setting cookies, rather than requiring cookies to be set via the "Set-Cookie" header.
+        /// </summary>
+        public IList<string> Cookies { get; set; } = new List<string>();
+
+        internal byte[] ToByteArray()
+        {
+            var bufferWriter = new System.Buffers.ArrayBufferWriter<byte>();
+            using (var writer = new Utf8JsonWriter(bufferWriter))
+            {
+                writer.WriteStartObject();
+
+                if (StatusCode.HasValue)
+                    writer.WriteNumber("statusCode", (int)StatusCode);
+
+                if (Headers?.Count > 0)
+                {
+                    writer.WriteStartObject("headers");
+                    foreach (var header in Headers)
+                    {
+                        writer.WriteString(header.Key, header.Value);
+                    }
+                    writer.WriteEndObject();
+                }
+
+                if (MultiValueHeaders?.Count > 0)
+                {
+                    writer.WriteStartObject("multiValueHeaders");
+                    foreach (var header in MultiValueHeaders)
+                    {
+                        writer.WriteStartArray(header.Key);
+                        foreach (var value in header.Value)
+                        {
+                            writer.WriteStringValue(value);
+                        }
+                        writer.WriteEndArray();
+                    }
+                    writer.WriteEndObject();
+                }
+
+                if (Cookies?.Count > 0)
+                {
+                    writer.WriteStartArray("cookies");
+                    foreach (var cookie in Cookies)
+                    {
+                        writer.WriteStringValue(cookie);
+                    }
+                    writer.WriteEndArray();
+                }
+
+                writer.WriteEndObject();
+            }
+
+            return bufferWriter.WrittenSpan.ToArray();
+        }
+    }
+}
+#endif

Libraries/src/Amazon.Lambda.Core/ResponseStreaming/ILambdaResponseStream.cs

@@ -0,0 +1,40 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+#if NET8_0_OR_GREATER
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Amazon.Lambda.Core.ResponseStreaming
+{
+    /// <summary>
+    /// Interface for writing streaming responses in AWS Lambda functions.
+    /// Obtained by calling <see cref="LambdaResponseStreamFactory.CreateStream"/> within a handler.
+    /// </summary>
+    internal interface ILambdaResponseStream : IDisposable
+    {
+        /// <summary>
+        /// Asynchronously writes a portion of a byte array to the response stream.
+        /// </summary>
+        /// <param name="buffer">The byte array containing data to write.</param>
+        /// <param name="offset">The zero-based byte offset in buffer at which to begin copying bytes.</param>
+        /// <param name="count">The number of bytes to write.</param>
+        /// <param name="cancellationToken">Optional cancellation token.</param>
+        /// <returns>A task representing the asynchronous operation.</returns>
+        /// <exception cref="InvalidOperationException">Thrown if the stream is already completed or an error has been reported.</exception>
+        Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken = default);
+
+
+        /// <summary>
+        /// Gets the total number of bytes written to the stream so far.
+        /// </summary>
+        long BytesWritten { get; }
+
+
+        /// <summary>
+        /// Gets whether an error has been reported.
+        /// </summary>
+        bool HasError { get; }
+    }
+}
+#endif

…mbda.Core/LambdaResponseStreamFactory.cs …esponseStreaming/LambdaResponseStream.csLibraries/src/Amazon.Lambda.Core/LambdaResponseStreamFactory.cs renamed to Libraries/src/Amazon.Lambda.Core/ResponseStreaming/LambdaResponseStream.cs Libraries/src/Amazon.Lambda.Core/LambdaResponseStreamFactory.cs renamed to Libraries/src/Amazon.Lambda.Core/ResponseStreaming/LambdaResponseStream.cs

@@ -1,76 +1,15 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 #if NET8_0_OR_GREATER
+
 using System;
 using System.IO;
 using System.Runtime.Versioning;
 using System.Threading;
 using System.Threading.Tasks;
 
-namespace Amazon.Lambda.Core
+namespace Amazon.Lambda.Core.ResponseStreaming
 {
-    /// <summary>
-    /// Factory to create Lambda response streams for writing streaming responses in AWS Lambda functions. The created streams are write-only and non-seekable.
-    /// </summary>
-    [RequiresPreviewFeatures(LambdaResponseStreamFactory.ParameterizedPreviewMessage)]
-    public class LambdaResponseStreamFactory
-    {
-        internal const string ParameterizedPreviewMessage =
-            "Response streaming is in preview till a new version of .NET Lambda runtime client that supports response streaming " +
-            "has been deployed to the .NET Lambda managed runtime. Till deployment has been made the feature can be used by deploying as an " +
-            "executable including the latest version of Amazon.Lambda.RuntimeSupport and setting the \"EnablePreviewFeatures\" in the Lambda " +
-            "project file to \"true\"";
-
-        private static Func<byte[], ILambdaResponseStream> _streamFactory;
-
-        internal static void SetLambdaResponseStream(Func<byte[], ILambdaResponseStream> streamFactory)
-        {
-            _streamFactory = streamFactory ?? throw new ArgumentNullException(nameof(streamFactory));
-        }
-
-        /// <summary>
-        /// Creates a <see cref="Stream"/> that can be used to write streaming responses back to callers of the Lambda function. Once
-        /// A Lambda function creates a response stream all output must be returned by writing to the stream; the Lambda function's handler
-        /// return value will be ignored. The stream is write-only and non-seekable.
-        /// </summary>
-        /// <returns></returns>
-        public static Stream CreateStream()
-        {
-            var runtimeResponseStream = _streamFactory(Array.Empty<byte>());
-            return new LambdaResponseStream(runtimeResponseStream);
-        }
-    }
-
-    /// <summary>
-    /// Interface for writing streaming responses in AWS Lambda functions.
-    /// Obtained by calling <see cref="LambdaResponseStreamFactory.CreateStream"/> within a handler.
-    /// </summary>
-    internal interface ILambdaResponseStream : IDisposable
-    {
-        /// <summary>
-        /// Asynchronously writes a portion of a byte array to the response stream.
-        /// </summary>
-        /// <param name="buffer">The byte array containing data to write.</param>
-        /// <param name="offset">The zero-based byte offset in buffer at which to begin copying bytes.</param>
-        /// <param name="count">The number of bytes to write.</param>
-        /// <param name="cancellationToken">Optional cancellation token.</param>
-        /// <returns>A task representing the asynchronous operation.</returns>
-        /// <exception cref="InvalidOperationException">Thrown if the stream is already completed or an error has been reported.</exception>
-        Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken = default);
-
-
-        /// <summary>
-        /// Gets the total number of bytes written to the stream so far.
-        /// </summary>
-        long BytesWritten { get; }
-
-
-        /// <summary>
-        /// Gets whether an error has been reported.
-        /// </summary>
-        bool HasError { get; }
-    }
-
     /// <summary>
     /// A write-only, non-seekable <see cref="Stream"/> subclass that streams response data
     /// to the Lambda Runtime API. Returned by <see cref="LambdaResponseStreamFactory.CreateStream"/>.

Libraries/src/Amazon.Lambda.Core/ResponseStreaming/LambdaResponseStreamFactory.cs

@@ -0,0 +1,59 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+#if NET8_0_OR_GREATER
+using System;
+using System.IO;
+using System.Runtime.Versioning;
+
+namespace Amazon.Lambda.Core.ResponseStreaming
+{
+    /// <summary>
+    /// Factory to create Lambda response streams for writing streaming responses in AWS Lambda functions. The created streams are write-only and non-seekable.
+    /// </summary>
+    [RequiresPreviewFeatures(LambdaResponseStreamFactory.ParameterizedPreviewMessage)]
+    public class LambdaResponseStreamFactory
+    {
+        internal const string ParameterizedPreviewMessage =
+            "Response streaming is in preview till a new version of .NET Lambda runtime client that supports response streaming " +
+            "has been deployed to the .NET Lambda managed runtime. Till deployment has been made the feature can be used by deploying as an " +
+            "executable including the latest version of Amazon.Lambda.RuntimeSupport and setting the \"EnablePreviewFeatures\" in the Lambda " +
+            "project file to \"true\"";
+
+        private static Func<byte[], ILambdaResponseStream> _streamFactory;
+
+        internal static void SetLambdaResponseStream(Func<byte[], ILambdaResponseStream> streamFactory)
+        {
+            _streamFactory = streamFactory ?? throw new ArgumentNullException(nameof(streamFactory));
+        }
+
+        /// <summary>
+        /// Creates a <see cref="Stream"/> that can be used to write streaming responses back to callers of the Lambda function. Once
+        /// a Lambda function creates a response stream all output must be returned by writing to the stream; the Lambda function's handler
+        /// return value will be ignored. The stream is write-only and non-seekable.
+        /// </summary>
+        /// <returns></returns>
+        public static Stream CreateStream()
+        {
+            var runtimeResponseStream = _streamFactory(Array.Empty<byte>());
+            return new LambdaResponseStream(runtimeResponseStream);
+        }
+
+        /// <summary>
+        /// Create a <see cref="Stream"/> for writing streaming responses, with an HTTP response prelude containing status code and headers. This should be used for
+        /// Lambda functions using response streaming that are invoked via the Lambda Function URLs or API Gateway HTTP APIs, where the response format is expected to be an HTTP response.
+        /// The prelude will be serialized and sent as the first chunk of the response stream, and should contain any necessary HTTP status code and headers.
+        /// <para>
+        /// Once a Lambda function creates a response stream all output must be returned by writing to the stream; the Lambda function's handler
+        /// return value will be ignored. The stream is write-only and non-seekable.
+        /// </para>
+        /// </summary>
+        /// <param name="prelude">The HTTP response prelude including status code and headers.</param>
+        /// <returns></returns>
+        public static Stream CreateHttpStream(HttpResponseStreamPrelude prelude)
+        {
+            var runtimeResponseStream = _streamFactory(prelude.ToByteArray());
+            return new LambdaResponseStream(runtimeResponseStream);
+        }
+    }
+}
+#endif

Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/ResponseStreaming/ResponseStream.cs

@@ -27,8 +27,6 @@ namespace Amazon.Lambda.RuntimeSupport.Client.ResponseStreaming
     /// </summary>
     internal class ResponseStream
     {
-        private static readonly byte[] CrlfBytes = Encoding.ASCII.GetBytes("\r\n");
-
         private long _bytesWritten;
         private bool _isCompleted;
         private bool _hasError;
@@ -41,6 +39,8 @@ internal class ResponseStream
         private readonly SemaphoreSlim _httpStreamReady = new SemaphoreSlim(0, 1);
         private readonly SemaphoreSlim _completionSignal = new SemaphoreSlim(0, 1);
 
+        private static readonly byte[] PreludeDelimiter = new byte[8];
+
         /// <summary>
         /// The number of bytes written to the Lambda response stream so far.
         /// </summary>
@@ -54,10 +54,14 @@ internal class ResponseStream
         private readonly byte[] _prelude;
 
 
+        private readonly InternalLogger _logger;
+
+
         internal Exception ReportedError => _reportedError;
 
         internal ResponseStream(byte[] prelude)
         {
+            _logger = InternalLogger.GetDefaultLogger();
             _prelude = prelude;
         }
 
@@ -69,14 +73,14 @@ internal async Task SetHttpOutputStreamAsync(Stream httpOutputStream, Cancellati
             _httpOutputStream = httpOutputStream;
             _httpStreamReady.Release();
 
-            InternalLogger.GetDefaultLogger().LogDebug($"Writing prelude of {_prelude.Length} bytes to HTTP stream.");
             await WritePreludeAsync(cancellationToken);
         }
 
         private async Task WritePreludeAsync(CancellationToken cancellationToken = default)
         {
             if (_prelude?.Length > 0)
             {
+                _logger.LogDebug($"Writing prelude of {_prelude.Length} bytes to HTTP stream.");
                 await _httpStreamReady.WaitAsync(cancellationToken);
                 try
                 {
@@ -85,22 +89,8 @@ private async Task WritePreludeAsync(CancellationToken cancellationToken = defau
                         ThrowIfCompletedOrError();
                     }
 
-                    // Write prelude JSON chunk
-                    var chunkSizeHex = _prelude.Length.ToString("X");
-                    var chunkSizeBytes = Encoding.ASCII.GetBytes(chunkSizeHex);
-                    await _httpOutputStream.WriteAsync(chunkSizeBytes, 0, chunkSizeBytes.Length, cancellationToken);
-                    await _httpOutputStream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length, cancellationToken);
                     await _httpOutputStream.WriteAsync(_prelude, 0, _prelude.Length, cancellationToken);
-                    await _httpOutputStream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length, cancellationToken);
-
-                    // Write 8 null bytes delimiter chunk
-                    var delimiterBytes = new byte[8];
-                    chunkSizeHex = "8";
-                    chunkSizeBytes = Encoding.ASCII.GetBytes(chunkSizeHex);
-                    await _httpOutputStream.WriteAsync(chunkSizeBytes, 0, chunkSizeBytes.Length, cancellationToken);
-                    await _httpOutputStream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length, cancellationToken);
-                    await _httpOutputStream.WriteAsync(delimiterBytes, 0, delimiterBytes.Length, cancellationToken);
-                    await _httpOutputStream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length, cancellationToken);
+                    await _httpOutputStream.WriteAsync(PreludeDelimiter, 0, PreludeDelimiter.Length, cancellationToken);
 
                     await _httpOutputStream.FlushAsync(cancellationToken);
                 }
@@ -149,19 +139,15 @@ public async Task WriteAsync(byte[] buffer, int offset, int count, CancellationT
             await _httpStreamReady.WaitAsync(cancellationToken);
             try
             {
+                _logger.LogDebug($"Writing chuck of {count} bytes to HTTP stream.");
+
                 lock (_lock)
                 {
                     ThrowIfCompletedOrError();
                     _bytesWritten += count;
                 }
 
-                // Write chunk directly to the HTTP stream: size(hex) + CRLF + data + CRLF
-                var chunkSizeHex = count.ToString("X");
-                var chunkSizeBytes = Encoding.ASCII.GetBytes(chunkSizeHex);
-                await _httpOutputStream.WriteAsync(chunkSizeBytes, 0, chunkSizeBytes.Length, cancellationToken);
-                await _httpOutputStream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length, cancellationToken);
                 await _httpOutputStream.WriteAsync(buffer, offset, count, cancellationToken);
-                await _httpOutputStream.WriteAsync(CrlfBytes, 0, CrlfBytes.Length, cancellationToken);
                 await _httpOutputStream.FlushAsync(cancellationToken);
             }
             finally

Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/ResponseStreaming/ResponseStreamLambdaCoreInitializerIsolated.cs

@@ -5,7 +5,7 @@
 using System;
 using System.Threading;
 using System.Threading.Tasks;
-using Amazon.Lambda.Core;
+using Amazon.Lambda.Core.ResponseStreaming;
 using Amazon.Lambda.RuntimeSupport.Client.ResponseStreaming;
 #pragma warning disable CA2252
 namespace Amazon.Lambda.RuntimeSupport