Validating the header length in HPackHeaderWriter

dotnet · May 4, 2024 · 42ee438 · 42ee438
1 parent f3a4bee
commit 42ee438
Show file tree

Hide file tree

Showing 9 changed files with 124 additions and 112 deletions.
diff --git a/src/Servers/Kestrel/Core/src/Internal/Http2/Http2FrameWriter.cs b/src/Servers/Kestrel/Core/src/Internal/Http2/Http2FrameWriter.cs
@@ -5,7 +5,6 @@
 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;
@@ -491,14 +490,10 @@ private void WriteResponseHeadersUnsynchronized(int streamId, int statusCode, Ht
             _headersEnumerator.Initialize(headers);
             _outgoingFrame.PrepareHeaders(headerFrameFlags, streamId);
             _headerEncodingBuffer.ResetWrittenCount();
+            _currentResponseHeadersTotalSize = 0;
             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);
+            var done = HPackHeaderWriter.BeginEncodeHeaders(statusCode, _hpackEncoder, _headersEnumerator, buffer, ref _currentResponseHeadersTotalSize, _maxResponseHeadersTotalSize, 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, only a single HEADER frame.
@@ -559,34 +554,21 @@ private ValueTask<FlushResult> WriteDataAndTrailersAsync(Http2Stream stream, in
                     _headersEnumerator.Initialize(headers);
                     _headerEncodingBuffer.ResetWrittenCount();
                     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 var payloadLength);
+                    done = HPackHeaderWriter.BeginEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, ref _currentResponseHeadersTotalSize, _maxResponseHeadersTotalSize, 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);
@@ -639,24 +621,14 @@ private void FinishWritingHeaders(int streamId)
         {
             _headerEncodingBuffer.ResetWrittenCount();
             var buffer = _headerEncodingBuffer.GetSpan(bufferSize)[0..bufferSize];
-            done = HPackHeaderWriter.ContinueEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, out var payloadLength);
-
+            done = HPackHeaderWriter.ContinueEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, ref _currentResponseHeadersTotalSize, _maxResponseHeadersTotalSize, 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);
             }
@@ -1086,6 +1058,4 @@ 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!));
 }
diff --git a/src/Servers/Kestrel/Core/test/Http2/Http2FrameWriterTests.cs b/src/Servers/Kestrel/Core/test/Http2/Http2FrameWriterTests.cs
@@ -1,20 +1,13 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System;
 using System.Buffers;
 using System.IO.Pipelines;
-using System.Linq;
-using System.Threading.Tasks;
-using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http.Features;
+using Microsoft.AspNetCore.InternalTesting;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
-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;