Apply HeadersLengthLimit to preamble in MultipartReader (#59646)

Author: BrennanConroy

src/Http/WebUtilities/src/MultipartBoundary.cs

@@ -10,11 +10,11 @@ internal sealed class MultipartBoundary
     private readonly byte[] _boundaryBytes;
     private bool _expectLeadingCrlf;
 
-    public MultipartBoundary(string boundary, bool expectLeadingCrlf = true)
+    public MultipartBoundary(string boundary)
     {
         ArgumentNullException.ThrowIfNull(boundary);
 
-        _expectLeadingCrlf = expectLeadingCrlf;
+        _expectLeadingCrlf = false;
         _boundaryBytes = Encoding.UTF8.GetBytes("\r\n--" + boundary);
 
         FinalBoundaryLength = BoundaryBytes.Length + 2; // Include the final '--' terminator.
@@ -25,6 +25,12 @@ public void ExpectLeadingCrlf()
         _expectLeadingCrlf = true;
     }
 
+    // Lets us throw a more specific error from MultipartReaderStream when reading any preamble data.
+    public bool BeforeFirstBoundary()
+    {
+        return !_expectLeadingCrlf;
+    }
+
     // Return either "--{boundary}" or "\r\n--{boundary}" depending on if we're looking for the end of a section
     public ReadOnlySpan<byte> BoundaryBytes => _boundaryBytes.AsSpan(_expectLeadingCrlf ? 0 : 2);
 

src/Http/WebUtilities/src/MultipartReader.cs

@@ -27,7 +27,7 @@ public class MultipartReader
 
     private readonly BufferedReadStream _stream;
     private readonly MultipartBoundary _boundary;
-    private MultipartReaderStream _currentStream;
+    private MultipartReaderStream? _currentStream;
 
     /// <summary>
     /// Initializes a new instance of <see cref="MultipartReader"/>.
@@ -56,10 +56,7 @@ public MultipartReader(string boundary, Stream stream, int bufferSize)
         }
         _stream = new BufferedReadStream(stream, bufferSize);
         boundary = HeaderUtilities.RemoveQuotes(new StringSegment(boundary)).ToString();
-        _boundary = new MultipartBoundary(boundary, false);
-        // This stream will drain any preamble data and remove the first boundary marker.
-        // TODO: HeadersLengthLimit can't be modified until after the constructor.
-        _currentStream = new MultipartReaderStream(_stream, _boundary) { LengthLimit = HeadersLengthLimit };
+        _boundary = new MultipartBoundary(boundary);
     }
 
     /// <summary>
@@ -86,6 +83,10 @@ public MultipartReader(string boundary, Stream stream, int bufferSize)
     /// <returns></returns>
     public async Task<MultipartSection?> ReadNextSectionAsync(CancellationToken cancellationToken = new CancellationToken())
     {
+        // Only occurs on first call
+        // This stream will drain any preamble data and remove the first boundary marker.
+        _currentStream ??= new MultipartReaderStream(_stream, _boundary) { LengthLimit = HeadersLengthLimit };
+
         // Drain the prior section.
         await _currentStream.DrainAsync(cancellationToken);
         // If we're at the end return null

src/Http/WebUtilities/src/MultipartReaderStream.cs

@@ -145,9 +145,19 @@ private int UpdatePosition(int read)
         if (_observedLength < _position)
         {
             _observedLength = _position;
-            if (LengthLimit.HasValue && _observedLength > LengthLimit.GetValueOrDefault())
+            if (LengthLimit.HasValue &&
+                LengthLimit.GetValueOrDefault() is var lengthLimit &&
+                _observedLength > lengthLimit)
             {
-                throw new InvalidDataException($"Multipart body length limit {LengthLimit.GetValueOrDefault()} exceeded.");
+                // If we hit the limit before the first boundary then we're using the header length limit, not the body length limit.
+                if (_boundary.BeforeFirstBoundary())
+                {
+                    throw new InvalidDataException($"Multipart header length limit {lengthLimit} exceeded. Too much data before the first boundary.");
+                }
+                else
+                {
+                    throw new InvalidDataException($"Multipart body length limit {lengthLimit} exceeded.");
+                }
             }
         }
         return read;

src/Http/WebUtilities/test/MultipartReaderTests.cs

@@ -147,6 +147,43 @@ public async Task MultipartReader_HeadersLengthExceeded_Throws()
         Assert.Equal("Line length limit 17 exceeded.", exception.Message);
     }
 
+    [Fact]
+    public async Task MultipartReader_HeadersLengthExceeded_LargePreamble()
+    {
+        var body = $"preamble {new string('a', 17000)}\r\n" +
+            "--9051914041544843365972754266\r\n" +
+"\r\n" +
+"text default\r\n" +
+"--9051914041544843365972754266--\r\n";
+        var stream = MakeStream(body);
+        var reader = new MultipartReader(Boundary, stream);
+
+        var exception = await Assert.ThrowsAsync<InvalidDataException>(() => reader.ReadNextSectionAsync());
+        Assert.Equal("Multipart header length limit 16384 exceeded. Too much data before the first boundary.", exception.Message);
+    }
+
+    [Fact]
+    public async Task MultipartReader_HeadersLengthLimitSettable_LargePreamblePasses()
+    {
+        var body = $"preamble {new string('a', 100_000)}\r\n" +
+            "--9051914041544843365972754266\r\n" +
+"\r\n" +
+"text default\r\n" +
+"--9051914041544843365972754266--\r\n";
+        var stream = MakeStream(body);
+        var reader = new MultipartReader(Boundary, stream)
+        {
+            HeadersLengthLimit = 200_000,
+        };
+
+        var section = await reader.ReadNextSectionAsync();
+        Assert.NotNull(section);
+
+        var buffer = new MemoryStream();
+        await section.Body.CopyToAsync(buffer);
+        Assert.Equal("text default", Encoding.ASCII.GetString(buffer.ToArray()));
+    }
+
     [Fact]
     public async Task MultipartReader_ReadSinglePartBodyWithTrailingWhitespace_Success()
     {