dont re-read the same data

Author: DeagleGross

src/Servers/Kestrel/Core/src/Middleware/TlsListenerMiddleware.cs

@@ -24,6 +24,7 @@ public TlsListenerMiddleware(ConnectionDelegate next, Action<ConnectionContext,
     internal async Task OnTlsClientHelloAsync(ConnectionContext connection)
     {
         var input = connection.Transport.Input;
+        ClientHelloParseState parseState = ClientHelloParseState.NotEnoughData;
 
         while (true)
         {
@@ -39,8 +40,7 @@ internal async Task OnTlsClientHelloAsync(ConnectionContext connection)
                     break;
                 }
 
-                var parseState = TryParseClientHello(buffer, out var clientHelloBytes);
-
+                parseState = TryParseClientHello(buffer, out var clientHelloBytes);
                 if (parseState == ClientHelloParseState.NotEnoughData)
                 {
                     // if no data will be added, and we still lack enough bytes
@@ -63,7 +63,15 @@ internal async Task OnTlsClientHelloAsync(ConnectionContext connection)
             }
             finally
             {
-                input.AdvanceTo(buffer.Start);
+                if (parseState is ClientHelloParseState.NotEnoughData)
+                {
+                    input.AdvanceTo(buffer.Start, buffer.End);
+                }
+                else
+                {
+                    // ready to continue middleware pipeline, reset the buffer to initial state
+                    input.AdvanceTo(buffer.Start);
+                }
             }
         }
 

src/Servers/Kestrel/test/InMemory.FunctionalTests/TestTransport/InMemoryTransportConnection.cs

@@ -91,7 +91,7 @@ public override async ValueTask DisposeAsync()
     // This piece of code allows us to wait until the PipeReader has been awaited on.
     // We need to wrap lots of layers (including the ValueTask) to gain visiblity into when
     // the machinery for the await happens
-    private class ObservableDuplexPipe : IDuplexPipe
+    internal class ObservableDuplexPipe : IDuplexPipe
     {
         private readonly ObservablePipeReader _reader;
 
@@ -110,11 +110,14 @@ public ObservableDuplexPipe(IDuplexPipe duplexPipe)
 
         public PipeWriter Output { get; }
 
+        public int ReadAsyncCounter => _reader.ReadAsyncCounter;
+
         private class ObservablePipeReader : PipeReader
         {
             private readonly PipeReader _reader;
             private readonly TaskCompletionSource _tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
 
+            public int ReadAsyncCounter { get; private set; } = 0;
             public Task WaitForReadTask => _tcs.Task;
 
             public ObservablePipeReader(PipeReader reader)
@@ -144,6 +147,7 @@ public override void Complete(Exception exception = null)
 
             public override ValueTask<ReadResult> ReadAsync(CancellationToken cancellationToken = default)
             {
+                ReadAsyncCounter++;
                 var task = _reader.ReadAsync(cancellationToken);
 
                 if (_tcs.Task.IsCompleted)
@@ -152,7 +156,7 @@ public override ValueTask<ReadResult> ReadAsync(CancellationToken cancellationTo
                 }
 
                 return new ValueTask<ReadResult>(new ObservableValueTask<ReadResult>(task, _tcs), 0);
-            }
+             }
 
             public override bool TryRead(out ReadResult result)
             {

src/Servers/Kestrel/test/InMemory.FunctionalTests/TlsListenerMiddlewareTests.Units.cs

@@ -25,6 +25,7 @@
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Moq;
+using static Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests.TestTransport.InMemoryTransportConnection;
 
 namespace InMemory.FunctionalTests;
 
@@ -50,6 +51,59 @@ public Task OnTlsClientHelloAsync_ValidData_MultipleSegments(int id, List<byte[]
     public Task OnTlsClientHelloAsync_InvalidData_MultipleSegments(int id, List<byte[]> packets, bool nextMiddlewareInvoked)
         => RunTlsClientHelloCallbackTest_WithMultipleSegments(id, packets, nextMiddlewareInvoked, tlsClientHelloCallbackExpected: false);
 
+    [Fact]
+    public async Task RunTlsClientHelloCallbackTest_DeterministinglyReads()
+    {
+        var serviceContext = new TestServiceContext(LoggerFactory);
+        var logger = LoggerFactory.CreateLogger<InMemoryTransportConnection>();
+        var memoryPool = serviceContext.MemoryPoolFactory();
+        var transportConnection = new InMemoryTransportConnection(memoryPool, logger);
+
+        var nextMiddlewareInvoked = false;
+        var tlsClientHelloCallbackInvoked = false;
+
+        var middleware = new TlsListenerMiddleware(
+            next: ctx =>
+            {
+                nextMiddlewareInvoked = true;
+                var readResult = ctx.Transport.Input.ReadAsync();
+                Assert.Equal(6, readResult.Result.Buffer.Length);
+
+                return Task.CompletedTask;
+            },
+            tlsClientHelloBytesCallback: (ctx, data) =>
+            {
+                tlsClientHelloCallbackInvoked = true;
+            }
+        );
+
+        await transportConnection.Input.WriteAsync(new byte[1] { 0x16 });
+        var middlewareTask = Task.Run(() => middleware.OnTlsClientHelloAsync(transportConnection));
+        await Task.Delay(TimeSpan.FromMilliseconds(25));
+
+        await transportConnection.Input.WriteAsync(new byte[2] { 0x03, 0x01 });
+        await Task.Delay(TimeSpan.FromMilliseconds(25));
+
+        await transportConnection.Input.WriteAsync(new byte[2] { 0x00, 0x20 });
+        await Task.Delay(TimeSpan.FromMilliseconds(25));
+
+        // not correct TLS client hello byte;
+        // meaning we will not invoke the callback and advance request processing
+        await transportConnection.Input.WriteAsync(new byte[1] { 0x15 });
+        await Task.Delay(TimeSpan.FromMilliseconds(25));
+
+        await transportConnection.Input.CompleteAsync();
+
+        // ensuring that we have read only 5 times (ReadAsync() is called 5 times)
+        var observableTransport = transportConnection.Transport as ObservableDuplexPipe;
+        Assert.NotNull(observableTransport);
+        Assert.Equal(5, observableTransport.ReadAsyncCounter);
+
+        await middlewareTask;
+        Assert.True(nextMiddlewareInvoked);
+        Assert.False(tlsClientHelloCallbackInvoked);
+    }
+
     private async Task RunTlsClientHelloCallbackTest_WithMultipleSegments(
         int id,
         List<byte[]> packets,