Adding MaxResponseHeadersLimit to H2 framewriter

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

@@ -5,6 +5,7 @@
 using System.Buffers.Binary;
 using System.Diagnostics;
 using System.IO.Pipelines;
+using System.Net.Http;
 using System.Net.Http.HPack;
 using System.Threading.Channels;
 using Microsoft.AspNetCore.Connections;
@@ -73,6 +74,8 @@ internal sealed class Http2FrameWriter
 
     private int _maxFrameSize = Http2PeerSettings.MinAllowedMaxFrameSize;
     private readonly ArrayBufferWriter<byte> _headerEncodingBuffer;
+    private readonly int? _maxResponseHeadersTotalSize;
+    private int _currentResponseHeadersTotalSize;
     private long _unflushedBytes;
 
     private bool _completed;
@@ -110,7 +113,7 @@ internal sealed class Http2FrameWriter
         _headerEncodingBuffer = new ArrayBufferWriter<byte>(_maxFrameSize);
 
         _scheduleInline = serviceContext.Scheduler == PipeScheduler.Inline;
-
+        _maxResponseHeadersTotalSize = _http2Connection.Limits.MaxResponseHeadersTotalSize;
         _hpackEncoder = new DynamicHPackEncoder(serviceContext.ServerOptions.AllowResponseHeaderCompression);
 
         _maximumFlowControlQueueSize = AppContextMaximumFlowControlQueueSize is null
@@ -484,25 +487,32 @@ 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);
             _headerEncodingBuffer.ResetWrittenCount();
             var buffer = _headerEncodingBuffer.GetSpan(_maxFrameSize)[0.._maxFrameSize]; // GetSpan might return more data that can result in a less deterministic behavior on the way headers are split into frames.
             var done = HPackHeaderWriter.BeginEncodeHeaders(statusCode, _hpackEncoder, _headersEnumerator, buffer, out var payloadLength);
             Debug.Assert(done != HPackHeaderWriter.HeaderWriteResult.BufferTooSmall, "Oversized frames should not be returned, beucase this always writes the status.");
+            if (_maxResponseHeadersTotalSize.HasValue && payloadLength > _maxResponseHeadersTotalSize.Value)
+            {
+                ThrowResponseHeadersLimitException();
+            }
+            _currentResponseHeadersTotalSize = payloadLength;
             if (done == HPackHeaderWriter.HeaderWriteResult.Done)
             {
-                // Fast path
+                // Fast path, only a single HEADER frame.
                 _outgoingFrame.PayloadLength = payloadLength;
                 _outgoingFrame.HeadersFlags |= Http2HeadersFrameFlags.END_HEADERS;
                 WriteHeaderUnsynchronized();
                 _outputWriter.Write(buffer[0..payloadLength]);
             }
             else
             {
-                // Slow path
+                // More headers sent in CONTINUATION frames.
                 _headerEncodingBuffer.Advance(payloadLength);
-                FinishWritingHeaders(streamId, payloadLength, done);
+                SplitHeaderFramesToOutput(streamId, done, isFramePrepared: true);
+                FinishWritingHeaders(streamId);
             }
         }
         // Any exception from the HPack encoder can leave the dynamic table in a corrupt state.
@@ -540,19 +550,46 @@ private ValueTask<FlushResult> WriteDataAndTrailersAsync(Http2Stream stream, in
 
             try
             {
+                // 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 done = HPackHeaderWriter.HeaderWriteResult.MoreHeaders;
-                int payloadLength;
+                var bufferSize = _headerEncodingBuffer.Capacity;
+                HPackHeaderWriter.HeaderWriteResult done;
                 do
                 {
                     _headersEnumerator.Initialize(headers);
                     _headerEncodingBuffer.ResetWrittenCount();
-                    var bufferSize = done == HPackHeaderWriter.HeaderWriteResult.BufferTooSmall ? _headerEncodingBuffer.Capacity * 2 : _headerEncodingBuffer.Capacity;
                     var buffer = _headerEncodingBuffer.GetSpan(bufferSize)[0..bufferSize]; // GetSpan might return more data that can result in a less deterministic behavior on the way headers are split into frames.
-                    done = HPackHeaderWriter.BeginEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, out payloadLength);
+                    done = HPackHeaderWriter.BeginEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, out var payloadLength);
+                    if (done == HPackHeaderWriter.HeaderWriteResult.Done)
+                    {
+                        if (_maxResponseHeadersTotalSize.HasValue && _currentResponseHeadersTotalSize + payloadLength > _maxResponseHeadersTotalSize.Value)
+                        {
+                            ThrowResponseHeadersLimitException();
+                        }
+                        _headerEncodingBuffer.Advance(payloadLength);
+                        SplitHeaderFramesToOutput(streamId, done, isFramePrepared: true);
+                    }
+                    else if (done == HPackHeaderWriter.HeaderWriteResult.MoreHeaders)
+                    {
+                        // More headers sent in CONTINUATION frames.
+                        _currentResponseHeadersTotalSize += payloadLength;
+                        if (_maxResponseHeadersTotalSize.HasValue && _currentResponseHeadersTotalSize > _maxResponseHeadersTotalSize.Value)
+                        {
+                            ThrowResponseHeadersLimitException();
+                        }
+                        _headerEncodingBuffer.Advance(payloadLength);
+                        SplitHeaderFramesToOutput(streamId, done, isFramePrepared: true);
+                        FinishWritingHeaders(streamId);
+                    }
+                    else
+                    {
+                        if (_maxResponseHeadersTotalSize.HasValue && _currentResponseHeadersTotalSize + bufferSize > _maxResponseHeadersTotalSize.Value)
+                        {
+                            ThrowResponseHeadersLimitException();
+                        }
+                        bufferSize *= 2;
+                    }
                 } while (done == HPackHeaderWriter.HeaderWriteResult.BufferTooSmall);
-                _headerEncodingBuffer.Advance(payloadLength);
-                FinishWritingHeaders(streamId, payloadLength, done);
             }
             // 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.
@@ -594,18 +631,36 @@ private void SplitHeaderFramesToOutput(int streamId, HPackHeaderWriter.HeaderWri
         }
     }
 
-    private void FinishWritingHeaders(int streamId, int payloadLength, HPackHeaderWriter.HeaderWriteResult done)
+    private void FinishWritingHeaders(int streamId)
     {
-        SplitHeaderFramesToOutput(streamId, done, isFramePrepared: true);
-        while (done != HPackHeaderWriter.HeaderWriteResult.Done)
+        HPackHeaderWriter.HeaderWriteResult done;
+        var bufferSize = _headerEncodingBuffer.Capacity;
+        do
         {
             _headerEncodingBuffer.ResetWrittenCount();
-            var bufferSize = done == HPackHeaderWriter.HeaderWriteResult.BufferTooSmall ? _headerEncodingBuffer.Capacity * 2 : _headerEncodingBuffer.Capacity;
             var buffer = _headerEncodingBuffer.GetSpan(bufferSize)[0..bufferSize];
-            done = HPackHeaderWriter.ContinueEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, out payloadLength);
-            _headerEncodingBuffer.Advance(payloadLength);
-            SplitHeaderFramesToOutput(streamId, done, isFramePrepared: false);
-        }
+            done = HPackHeaderWriter.ContinueEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, out var payloadLength);
+
+            if (done == HPackHeaderWriter.HeaderWriteResult.BufferTooSmall)
+            {
+                if (_maxResponseHeadersTotalSize.HasValue && _currentResponseHeadersTotalSize + bufferSize > _maxResponseHeadersTotalSize.Value)
+                {
+                    ThrowResponseHeadersLimitException();
+                }
+                bufferSize *= 2;
+            }
+            else
+            {
+                // In case of Done or MoreHeaders: write to output.
+                _currentResponseHeadersTotalSize += payloadLength;
+                if (_maxResponseHeadersTotalSize.HasValue && _currentResponseHeadersTotalSize > _maxResponseHeadersTotalSize.Value)
+                {
+                    ThrowResponseHeadersLimitException();
+                }
+                _headerEncodingBuffer.Advance(payloadLength);
+                SplitHeaderFramesToOutput(streamId, done, isFramePrepared: false);
+            }
+        } while (done != HPackHeaderWriter.HeaderWriteResult.Done);
     }
 
     /*  Padding is not implemented
@@ -1031,4 +1086,6 @@ private void EnqueueWaitingForMoreConnectionWindow(Http2OutputProducer producer)
             _http2Connection.Abort(new ConnectionAbortedException("HTTP/2 connection exceeded the outgoing flow control maximum queue size."));
         }
     }
+
+    private void ThrowResponseHeadersLimitException() => throw new HPackEncodingException(SR.Format(SR.net_http_headers_exceeded_length, _maxResponseHeadersTotalSize!));
 }

src/Servers/Kestrel/Core/src/KestrelServerLimits.cs

@@ -15,6 +15,9 @@ public class KestrelServerLimits
     // Matches the non-configurable default response buffer size for Kestrel in 1.0.0
     private long? _maxResponseBufferSize = 64 * 1024;
 
+    // Matches the HttpClientHandler.MaxResponseHeadersLength's response header size.
+    private int? _maxResponseHeadersTotalSize = 64 * 1024;
+
     // Matches the default client_max_body_size in nginx.
     // Also large enough that most requests should be under the limit.
     private long? _maxRequestBufferSize = 1024 * 1024;
@@ -256,6 +259,27 @@ public TimeSpan RequestHeadersTimeout
         }
     }
 
+    /// <summary>
+    /// Gets or sets the maximum size of the total response headers. When set to null, the response headers total size is unlimited.
+    /// Defaults to 65,536 bytes (64 KB).
+    /// </summary>
+    /// <remarks>
+    /// When set to null, the size of the response buffer is unlimited.
+    /// When set to zero, no headers are allowed to be returned.
+    /// </remarks>
+    public int? MaxResponseHeadersTotalSize
+    {
+        get => _maxResponseHeadersTotalSize;
+        set
+        {
+            if (value.HasValue && value.Value < 0)
+            {
+                throw new ArgumentOutOfRangeException(nameof(value), CoreStrings.NonNegativeNumberOrNullRequired);
+            }
+            _maxResponseHeadersTotalSize = value;
+        }
+    }
+
     internal void Serialize(Utf8JsonWriter writer)
     {
         writer.WriteString(nameof(KeepAliveTimeout), KeepAliveTimeout.ToString());
@@ -323,6 +347,16 @@ internal void Serialize(Utf8JsonWriter writer)
         writer.WriteString(nameof(MinResponseDataRate), MinResponseDataRate?.ToString());
         writer.WriteString(nameof(RequestHeadersTimeout), RequestHeadersTimeout.ToString());
 
+        writer.WritePropertyName(nameof(MaxResponseHeadersTotalSize));
+        if (MaxResponseHeadersTotalSize is null)
+        {
+            writer.WriteNullValue();
+        }
+        else
+        {
+            writer.WriteNumberValue(MaxResponseHeadersTotalSize.Value);
+        }
+
         // HTTP2
         writer.WritePropertyName(nameof(Http2));
         writer.WriteStartObject();

src/Servers/Kestrel/Core/src/PublicAPI.Unshipped.txt

@@ -1 +1,3 @@
 #nullable enable
+Microsoft.AspNetCore.Server.Kestrel.Core.KestrelServerLimits.MaxResponseHeadersTotalSize.get -> int?
+Microsoft.AspNetCore.Server.Kestrel.Core.KestrelServerLimits.MaxResponseHeadersTotalSize.set -> void

src/Servers/Kestrel/Core/test/Http2/Http2FrameWriterTests.cs

@@ -12,6 +12,9 @@
 using Microsoft.AspNetCore.InternalTesting;
 using Moq;
 using Xunit;
+using Microsoft.AspNetCore.Http.Features;
+using Castle.Core;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests;
 
@@ -56,7 +59,16 @@ public async Task WriteWindowUpdate_UnsetsReservedBit()
     private Http2FrameWriter CreateFrameWriter(Pipe pipe)
     {
         var serviceContext = TestContextFactory.CreateServiceContext(new KestrelServerOptions());
-        return new Http2FrameWriter(pipe.Writer, null, null, 1, null, null, null, _dirtyMemoryPool, serviceContext);
+        var featureCollection = new FeatureCollection();
+        featureCollection.Set<IConnectionMetricsContextFeature>(new TestConnectionMetricsContextFeature());
+        var connectionContext = TestContextFactory.CreateHttpConnectionContext(
+                   serviceContext: serviceContext,
+                   connectionContext: null,
+                   transport: new DuplexPipe(pipe.Reader, pipe.Writer),
+                   connectionFeatures: featureCollection);
+
+        var http2Connection = new Http2Connection(connectionContext);
+        return new Http2FrameWriter(pipe.Writer, null, http2Connection, 1, null, null, null, _dirtyMemoryPool, serviceContext);
     }
 
     [Fact]
@@ -92,6 +104,11 @@ public async Task WriteHeader_UnsetsReservedBit()
 
         Assert.Equal(new byte[] { 0x00, 0x00, 0x00, 0x00 }, payload.Skip(5).Take(4).ToArray());
     }
+
+    private sealed class TestConnectionMetricsContextFeature : IConnectionMetricsContextFeature
+    {
+        public ConnectionMetricsContext MetricsContext { get; }
+    }
 }
 
 public static class PipeReaderExtensions

src/Servers/Kestrel/perf/Microbenchmarks/Http2/Http2FrameWriterBenchmark.cs

@@ -11,6 +11,8 @@
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2.FlowControl;
 using Microsoft.AspNetCore.InternalTesting;
+using Microsoft.AspNetCore.Http.Features;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Microbenchmarks;
 
@@ -34,10 +36,19 @@ public void GlobalSetup()
             httpParser: new HttpParser<Http1ParsingHandler>(),
             dateHeaderValueManager: new DateHeaderValueManager(TimeProvider.System));
 
+        var featureCollection = new FeatureCollection();
+        featureCollection.Set<IConnectionMetricsContextFeature>(new TestConnectionMetricsContextFeature());
+        var connectionContext = TestContextFactory.CreateHttpConnectionContext(
+                   serviceContext: serviceContext,
+                   connectionContext: null,
+                   transport: new DuplexPipe(_pipe.Reader, _pipe.Writer),
+                   connectionFeatures: featureCollection);
+        var http2Connection = new Http2Connection(connectionContext);
+
         _frameWriter = new Http2FrameWriter(
             new NullPipeWriter(),
             connectionContext: null,
-            http2Connection: null,
+            http2Connection: http2Connection,
             maxStreamsPerConnection: 1,
             timeoutControl: null,
             minResponseDataRate: null,
@@ -63,4 +74,9 @@ public void Dispose()
         _pipe.Writer.Complete();
         _memoryPool?.Dispose();
     }
+
+    private sealed class TestConnectionMetricsContextFeature : IConnectionMetricsContextFeature
+    {
+        public ConnectionMetricsContext MetricsContext { get; }
+    }
 }

src/Servers/Kestrel/samples/Http2SampleApp/Http2SampleApp.csproj

@@ -9,6 +9,7 @@
   <ItemGroup>
     <Reference Include="Microsoft.AspNetCore.Server.Kestrel" />
     <Reference Include="Microsoft.Extensions.Logging.Console" />
+    <Reference Include="Microsoft.AspNetCore" />
   </ItemGroup>
 
 </Project>