Feedback

Author: jkotalik

src/Servers/Kestrel/Core/src/Internal/Http/Http1OutputProducer.cs

@@ -355,11 +355,7 @@ private void WriteDataWrittenBeforeHeaders(ref BufferWriter<PipeWriter> writer)
 
                 _position = 0;
 
-                if (_currentSegmentOwner != null)
-                {
-                    _currentSegmentOwner.Dispose();
-                    _currentSegmentOwner = null;
-                }
+                DisposeCurrentSegment();
             }
         }
 
@@ -382,15 +378,19 @@ public void Dispose()
                     }
                 }
 
-                if (_currentSegmentOwner != null)
-                {
-                    _currentSegmentOwner.Dispose();
-                }
+                DisposeCurrentSegment();
 
                 CompletePipe();
             }
         }
 
+        private void DisposeCurrentSegment()
+        {
+            _currentSegmentOwner?.Dispose();
+            _currentSegmentOwner = null;
+            _currentSegment = default;
+        }
+
         private void CompletePipe()
         {
             if (!_pipeWriterCompleted)
@@ -658,7 +658,7 @@ private readonly struct CompletedBuffer
             public Memory<byte> Buffer { get; }
             public int Length { get; }
 
-            public ReadOnlySpan<byte> Span => Buffer.Span;
+            public ReadOnlySpan<byte> Span => Buffer.Span.Slice(0, Length);
 
             public CompletedBuffer(IMemoryOwner<byte> owner, Memory<byte> buffer, int length)
             {

src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs

@@ -49,7 +49,7 @@ public abstract partial class HttpProtocol : IDefaultHttpContextContainer, IHttp
         // If we are writing with GetMemory/Advance before calling StartAsync, assume we can write and throw away contents if we can't.
         private bool _canWriteResponseBody = true;
         private bool _hasAdvanced;
-        private bool _requireGetMemoryNextCall;
+        private bool _isLeasedMemoryInvalid = true;
         private bool _autoChunk;
         protected Exception _applicationException;
         private BadHttpRequestException _requestRejectedException;
@@ -925,7 +925,7 @@ private void ProduceStart(bool appCompleted)
                 return;
             }
 
-            _requireGetMemoryNextCall = true;
+            _isLeasedMemoryInvalid = true;
 
             _requestProcessingStatus = RequestProcessingStatus.HeadersCommitted;
 
@@ -1283,7 +1283,7 @@ public void Advance(int bytes)
                 _hasAdvanced = true;
             }
 
-            if (_requireGetMemoryNextCall)
+            if (_isLeasedMemoryInvalid)
             {
                 throw new InvalidOperationException("Invalid ordering of calling StartAsync and Advance. " +
                     "Call StartAsync before calling GetMemory/GetSpan and Advance.");
@@ -1305,13 +1305,13 @@ public void Advance(int bytes)
 
         public Memory<byte> GetMemory(int sizeHint = 0)
         {
-            _requireGetMemoryNextCall = false;
+            _isLeasedMemoryInvalid = false;
             return Output.GetMemory(sizeHint);
         }
 
         public Span<byte> GetSpan(int sizeHint = 0)
         {
-            _requireGetMemoryNextCall = false;
+            _isLeasedMemoryInvalid = false;
             return Output.GetSpan(sizeHint);
         }
 

src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2StreamTests.cs

@@ -2625,6 +2625,52 @@ await ExpectAsync(Http2FrameType.SETTINGS,
 
         [Fact]
         public async Task GetMemoryAdvance_Works()
+        {
+            var headers = new[]
+             {
+                new KeyValuePair<string, string>(HeaderNames.Method, "GET"),
+                new KeyValuePair<string, string>(HeaderNames.Path, "/"),
+                new KeyValuePair<string, string>(HeaderNames.Scheme, "http"),
+            };
+            await InitializeConnectionAsync(httpContext =>
+            {
+                var response = httpContext.Response;
+                var memory = response.BodyPipe.GetMemory();
+                var fisrtPartOfResponse = Encoding.ASCII.GetBytes("hello,");
+                fisrtPartOfResponse.CopyTo(memory);
+                response.BodyPipe.Advance(6);
+
+                memory = response.BodyPipe.GetMemory();
+                var secondPartOfResponse = Encoding.ASCII.GetBytes(" world");
+                secondPartOfResponse.CopyTo(memory);
+                response.BodyPipe.Advance(6);
+                return Task.CompletedTask;
+            });
+
+            await StartStreamAsync(1, headers, endStream: true);
+
+            var headersFrame = await ExpectAsync(Http2FrameType.HEADERS,
+                withLength: 37,
+                withFlags: (byte)Http2HeadersFrameFlags.END_HEADERS,
+                withStreamId: 1);
+            var dataFrame = await ExpectAsync(Http2FrameType.DATA,
+                withLength: 12,
+                withFlags: (byte)Http2DataFrameFlags.END_STREAM,
+                withStreamId: 1);
+
+            await StopConnectionAsync(expectedLastStreamId: 1, ignoreNonGoAwayFrames: false);
+
+            _hpackDecoder.Decode(headersFrame.PayloadSequence, endHeaders: false, handler: this);
+
+            Assert.Equal(2, _decodedHeaders.Count);
+            Assert.Contains("date", _decodedHeaders.Keys, StringComparer.OrdinalIgnoreCase);
+            Assert.Equal("200", _decodedHeaders[HeaderNames.Status]);
+
+            Assert.True(_helloWorldBytes.AsSpan().SequenceEqual(dataFrame.PayloadSequence.ToArray()));
+        }
+
+        [Fact]
+        public async Task GetMemoryAdvance_WithStartAsync_Works()
         {
             var headers = new[]
              {
@@ -2634,7 +2680,7 @@ public async Task GetMemoryAdvance_Works()
             };
             await InitializeConnectionAsync(async httpContext =>
             {
-                await Task.CompletedTask;
+                await httpContext.Response.StartAsync();
                 var response = httpContext.Response;
                 var memory = response.BodyPipe.GetMemory();
                 var fisrtPartOfResponse = Encoding.ASCII.GetBytes("hello,");
@@ -3214,7 +3260,6 @@ await InitializeConnectionAsync(async httpContext =>
             {
                 var response = httpContext.Response;
                 response.ContentLength = 54;
-                await Task.CompletedTask;
                 var memory = response.BodyPipe.GetMemory(4096);
                 var fisrtPartOfResponse = Encoding.ASCII.GetBytes("hello,");
                 fisrtPartOfResponse.CopyTo(memory);

src/Servers/Kestrel/test/InMemory.FunctionalTests/ResponseTests.cs

@@ -3286,7 +3286,7 @@ public async Task AdvanceWithTooLargeOfAValueThrowInvalidOperationException(bool
                 {
                     await response.StartAsync();
                 }
-                Assert.Throws<ArgumentOutOfRangeException>(() => response.BodyPipe.Advance(1));
+                Assert.Throws<InvalidOperationException>(() => response.BodyPipe.Advance(1));
             }, testContext))
             {
                 using (var connection = server.CreateConnection())