Allow HTTP2 encoder to split headers across frames (#55322)

* Initial implementation

* Review feedback

* Updating a comment

* Minor perf opt when all headers fit in a single header frame.

* Adding MaxResponseHeadersLimit to H2 framewriter

* Protect against overflow

* Validating the header length in HPackHeaderWriter

* Removing a test file commented out and adding more tests.

* Review feedback

* Removing header limits related code because it is up to the service owner to not to ddos itself

* A test for 1MB header value

* Review comments - removing an unneeded if condition

* Sketch

* using arraypool instead of arraybufferwriter

* Re-using the rented buffers

* Preserving the size of the temporary larger buffer within the framewriter. This way if the same large header is written multiple times over the lifespan of a connection, it does not need to go through the loops of increasing the headers repetitively.

* Adjusting the logic of UpdateMaxFrameSize for setting a new _headersEncodingLargeBufferSize

* Adjusting the logic of _headersEncodingLargeBufferSize to avoid 0 values.

* Extenting Benachmark with headers sized 0, 10KB and 20KB

* Remove int.Max on header size increase and using checked arithmetic instead.

* Tests

* Update src/Servers/Kestrel/Core/src/Internal/Http2/Http2FrameWriter.cs

Co-authored-by: James Newton-King <[email protected]>

* Review feedback: HeaderWriteResult not an inner type, adding comments on size increase

* Update src/Servers/Kestrel/Core/src/Internal/Http2/Http2FrameWriter.cs

Co-authored-by: Andrew Casey <[email protected]>

* - Moving comments around for _headersEncodingLargeBufferSize
- Rename splitting method to SplitHeaderAcrossFrames

* Merging the implementation of BeginEncodeHeaders and RetryBeginEncodeHeaders

* Apply suggestions from code review

Co-authored-by: James Newton-King <[email protected]>

---------

Co-authored-by: ladeak <[email protected]>
Co-authored-by: James Newton-King <[email protected]>
Co-authored-by: Andrew Casey <[email protected]>

Author: 4 people

src/Servers/Kestrel/Core/src/Internal/Http2/Http2Connection.cs

@@ -938,7 +938,9 @@ private Task ProcessSettingsFrameAsync(in ReadOnlySequence<byte> payload)
             if (_clientSettings.MaxFrameSize != previousMaxFrameSize)
             {
                 // Don't let the client choose an arbitrarily large size, this will be used for response buffers.
-                _frameWriter.UpdateMaxFrameSize(Math.Min(_clientSettings.MaxFrameSize, _serverSettings.MaxFrameSize));
+                // Safe cast, MaxFrameSize is limited to 2^24-1 bytes by the protocol and by Http2PeerSettings.
+                // Ref: https://datatracker.ietf.org/doc/html/rfc7540#section-4.2
+                _frameWriter.UpdateMaxFrameSize((int)Math.Min(_clientSettings.MaxFrameSize, _serverSettings.MaxFrameSize));
             }
 
             // This difference can be negative.
@@ -1829,4 +1831,4 @@ private static class GracefulCloseInitiator
         public const int Server = 1;
         public const int Client = 2;
     }
-}
+}

src/Servers/Kestrel/Core/src/Internal/Http2/Http2FrameWriter.cs

@@ -29,6 +29,8 @@ internal sealed class Http2FrameWriter
     /// TODO (https://github.com/dotnet/aspnetcore/issues/51309): eliminate this limit.
     private const string MaximumFlowControlQueueSizeProperty = "Microsoft.AspNetCore.Server.Kestrel.Http2.MaxConnectionFlowControlQueueSize";
 
+    private const int HeaderBufferSizeMultiplier = 2;
+
     private static readonly int? AppContextMaximumFlowControlQueueSize = GetAppContextMaximumFlowControlQueueSize();
 
     private static int? GetAppContextMaximumFlowControlQueueSize()
@@ -71,8 +73,12 @@ internal sealed class Http2FrameWriter
     // This is only set to true by tests.
     private readonly bool _scheduleInline;
 
-    private uint _maxFrameSize = Http2PeerSettings.MinAllowedMaxFrameSize;
+    private int _maxFrameSize = Http2PeerSettings.MinAllowedMaxFrameSize;
     private byte[] _headerEncodingBuffer;
+
+    // Keep track of the high-water mark of _headerEncodingBuffer's size so we don't have to grow
+    // through intermediate sizes repeatedly.
+    private int _headersEncodingLargeBufferSize = Http2PeerSettings.MinAllowedMaxFrameSize * HeaderBufferSizeMultiplier;
     private long _unflushedBytes;
 
     private bool _completed;
@@ -110,7 +116,6 @@ public Http2FrameWriter(
         _headerEncodingBuffer = new byte[_maxFrameSize];
 
         _scheduleInline = serviceContext.Scheduler == PipeScheduler.Inline;
-
         _hpackEncoder = new DynamicHPackEncoder(serviceContext.ServerOptions.AllowResponseHeaderCompression);
 
         _maximumFlowControlQueueSize = AppContextMaximumFlowControlQueueSize is null
@@ -367,12 +372,15 @@ public void UpdateMaxHeaderTableSize(uint maxHeaderTableSize)
         }
     }
 
-    public void UpdateMaxFrameSize(uint maxFrameSize)
+    public void UpdateMaxFrameSize(int maxFrameSize)
     {
         lock (_writeLock)
         {
             if (_maxFrameSize != maxFrameSize)
             {
+                // Safe multiply, MaxFrameSize is limited to 2^24-1 bytes by the protocol and by Http2PeerSettings.
+                // Ref: https://datatracker.ietf.org/doc/html/rfc7540#section-4.2
+                _headersEncodingLargeBufferSize = int.Max(_headersEncodingLargeBufferSize, maxFrameSize * HeaderBufferSizeMultiplier);
                 _maxFrameSize = maxFrameSize;
                 _headerEncodingBuffer = new byte[_maxFrameSize];
             }
@@ -507,11 +515,12 @@ private void WriteResponseHeadersUnsynchronized(int streamId, int statusCode, Ht
     {
         try
         {
+            // In the case of the headers, there is always a status header to be returned, so BeginEncodeHeaders will not return BufferTooSmall.
             _headersEnumerator.Initialize(headers);
             _outgoingFrame.PrepareHeaders(headerFrameFlags, streamId);
-            var buffer = _headerEncodingBuffer.AsSpan();
-            var done = HPackHeaderWriter.BeginEncodeHeaders(statusCode, _hpackEncoder, _headersEnumerator, buffer, out var payloadLength);
-            FinishWritingHeadersUnsynchronized(streamId, payloadLength, done);
+            var writeResult = HPackHeaderWriter.BeginEncodeHeaders(statusCode, _hpackEncoder, _headersEnumerator, _headerEncodingBuffer, out var payloadLength);
+            Debug.Assert(writeResult != HeaderWriteResult.BufferTooSmall, "This always writes the status as the first header, and it should never be an over the buffer size.");
+            FinishWritingHeadersUnsynchronized(streamId, payloadLength, writeResult);
         }
         // Any exception from the HPack encoder can leave the dynamic table in a corrupt state.
         // Since we allow custom header encoders we don't know what type of exceptions to expect.
@@ -548,11 +557,11 @@ private ValueTask<FlushResult> WriteDataAndTrailersAsync(Http2Stream stream, in
 
             try
             {
-                _headersEnumerator.Initialize(headers);
+                // In the case of the trailers, there is no status header to be written, so even the first call to BeginEncodeHeaders can return BufferTooSmall.
                 _outgoingFrame.PrepareHeaders(Http2HeadersFrameFlags.END_STREAM, streamId);
-                var buffer = _headerEncodingBuffer.AsSpan();
-                var done = HPackHeaderWriter.BeginEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, out var payloadLength);
-                FinishWritingHeadersUnsynchronized(streamId, payloadLength, done);
+                _headersEnumerator.Initialize(headers);
+                var writeResult = HPackHeaderWriter.BeginEncodeHeaders(_hpackEncoder, _headersEnumerator, _headerEncodingBuffer, out var payloadLength);
+                FinishWritingHeadersUnsynchronized(streamId, payloadLength, writeResult);
             }
             // Any exception from the HPack encoder can leave the dynamic table in a corrupt state.
             // Since we allow custom header encoders we don't know what type of exceptions to expect.
@@ -566,32 +575,102 @@ private ValueTask<FlushResult> WriteDataAndTrailersAsync(Http2Stream stream, in
         }
     }
 
-    private void FinishWritingHeadersUnsynchronized(int streamId, int payloadLength, bool done)
+    private void SplitHeaderAcrossFrames(int streamId, ReadOnlySpan<byte> dataToFrame, bool endOfHeaders, bool isFramePrepared)
     {
-        var buffer = _headerEncodingBuffer.AsSpan();
-        _outgoingFrame.PayloadLength = payloadLength;
-        if (done)
+        var shouldPrepareFrame = !isFramePrepared;
+        while (dataToFrame.Length > 0)
         {
-            _outgoingFrame.HeadersFlags |= Http2HeadersFrameFlags.END_HEADERS;
-        }
+            if (shouldPrepareFrame)
+            {
+                _outgoingFrame.PrepareContinuation(Http2ContinuationFrameFlags.NONE, streamId);
+            }
 
-        WriteHeaderUnsynchronized();
-        _outputWriter.Write(buffer.Slice(0, payloadLength));
+            // Should prepare continuation frames.
+            shouldPrepareFrame = true;
+            var currentSize = Math.Min(dataToFrame.Length, _maxFrameSize);
+            _outgoingFrame.PayloadLength = currentSize;
+            if (endOfHeaders && dataToFrame.Length == currentSize)
+            {
+                _outgoingFrame.HeadersFlags |= Http2HeadersFrameFlags.END_HEADERS;
+            }
 
-        while (!done)
-        {
-            _outgoingFrame.PrepareContinuation(Http2ContinuationFrameFlags.NONE, streamId);
+            WriteHeaderUnsynchronized();
+            _outputWriter.Write(dataToFrame[..currentSize]);
+            dataToFrame = dataToFrame.Slice(currentSize);
+        }
+    }
 
-            done = HPackHeaderWriter.ContinueEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, out payloadLength);
+    private void FinishWritingHeadersUnsynchronized(int streamId, int payloadLength, HeaderWriteResult writeResult)
+    {
+        Debug.Assert(payloadLength <= _maxFrameSize, "The initial payload lengths is written to _headerEncodingBuffer with size of _maxFrameSize");
+        byte[]? largeHeaderBuffer = null;
+        Span<byte> buffer;
+        if (writeResult == HeaderWriteResult.Done)
+        {
+            // Fast path, only a single HEADER frame.
             _outgoingFrame.PayloadLength = payloadLength;
-
-            if (done)
+            _outgoingFrame.HeadersFlags |= Http2HeadersFrameFlags.END_HEADERS;
+            WriteHeaderUnsynchronized();
+            _outputWriter.Write(_headerEncodingBuffer.AsSpan(0, payloadLength));
+            return;
+        }
+        else if (writeResult == HeaderWriteResult.MoreHeaders)
+        {
+            _outgoingFrame.PayloadLength = payloadLength;
+            WriteHeaderUnsynchronized();
+            _outputWriter.Write(_headerEncodingBuffer.AsSpan(0, payloadLength));
+        }
+        else
+        {
+            // This may happen in case of the TRAILERS after the initial encode operation.
+            // The _maxFrameSize sized _headerEncodingBuffer was too small.
+            while (writeResult == HeaderWriteResult.BufferTooSmall)
+            {
+                Debug.Assert(payloadLength == 0, "Payload written even though buffer is too small");
+                largeHeaderBuffer = ArrayPool<byte>.Shared.Rent(_headersEncodingLargeBufferSize);
+                buffer = largeHeaderBuffer.AsSpan(0, _headersEncodingLargeBufferSize);
+                writeResult = HPackHeaderWriter.RetryBeginEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, out payloadLength);
+                if (writeResult != HeaderWriteResult.BufferTooSmall)
+                {
+                    SplitHeaderAcrossFrames(streamId, buffer[..payloadLength], endOfHeaders: writeResult == HeaderWriteResult.Done, isFramePrepared: true);
+                }
+                else
+                {
+                    _headersEncodingLargeBufferSize = checked(_headersEncodingLargeBufferSize * HeaderBufferSizeMultiplier);
+                }
+                ArrayPool<byte>.Shared.Return(largeHeaderBuffer);
+                largeHeaderBuffer = null;
+            }
+            if (writeResult == HeaderWriteResult.Done)
             {
-                _outgoingFrame.ContinuationFlags = Http2ContinuationFrameFlags.END_HEADERS;
+                return;
             }
+        }
 
-            WriteHeaderUnsynchronized();
-            _outputWriter.Write(buffer.Slice(0, payloadLength));
+        // HEADERS and zero or more CONTINUATIONS sent - all subsequent frames are (unprepared) CONTINUATIONs
+        buffer = _headerEncodingBuffer;
+        while (writeResult != HeaderWriteResult.Done)
+        {
+            writeResult = HPackHeaderWriter.ContinueEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, out payloadLength);
+            if (writeResult == HeaderWriteResult.BufferTooSmall)
+            {
+                if (largeHeaderBuffer != null)
+                {
+                    ArrayPool<byte>.Shared.Return(largeHeaderBuffer);
+                    _headersEncodingLargeBufferSize = checked(_headersEncodingLargeBufferSize * HeaderBufferSizeMultiplier);
+                }
+                largeHeaderBuffer = ArrayPool<byte>.Shared.Rent(_headersEncodingLargeBufferSize);
+                buffer = largeHeaderBuffer.AsSpan(0, _headersEncodingLargeBufferSize);
+            }
+            else
+            {
+                // In case of Done or MoreHeaders: write to output.
+                SplitHeaderAcrossFrames(streamId, buffer[..payloadLength], endOfHeaders: writeResult == HeaderWriteResult.Done, isFramePrepared: false);
+            }
+        }
+        if (largeHeaderBuffer != null)
+        {
+            ArrayPool<byte>.Shared.Return(largeHeaderBuffer);
         }
     }
 
@@ -1023,4 +1102,4 @@ private void EnqueueWaitingForMoreConnectionWindow(Http2OutputProducer producer)
             _http2Connection.Abort(new ConnectionAbortedException("HTTP/2 connection exceeded the outgoing flow control maximum queue size."));
         }
     }
-}
+}