From a145734d43233fc347574c7be09d52c47ad43f46 Mon Sep 17 00:00:00 2001
From: Cesar Blum Silveira <cesarbs@gmail.com>
Date: Thu, 24 Aug 2017 10:17:05 -0700
Subject: [PATCH] Fix flakiness in
 ConnectionClosedWhenResponseDoesNotSatisfyMinimumDataRate (#1955).

---
 test/Kestrel.FunctionalTests/ResponseTests.cs | 60 ++++++++++---------
 1 file changed, 31 insertions(+), 29 deletions(-)

diff --git a/test/Kestrel.FunctionalTests/ResponseTests.cs b/test/Kestrel.FunctionalTests/ResponseTests.cs
index d0a4b3deef..05328b3a2c 100644
--- a/test/Kestrel.FunctionalTests/ResponseTests.cs
+++ b/test/Kestrel.FunctionalTests/ResponseTests.cs
@@ -2444,15 +2444,15 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             }
         }
 
-        [ConditionalFact]
-        [OSSkipCondition(OperatingSystems.Linux, SkipReason = "Flaky test (https://github.com/aspnet/KestrelHttpServer/issues/1955)")]
-        public async Task ConnectionClosedWhenResponseDoesNotSatisfyMinimumDataRate()
+        [Fact]
+        public void ConnectionClosedWhenResponseDoesNotSatisfyMinimumDataRate()
         {
             var chunkSize = 64 * 1024;
             var chunks = 128;
+            var responseSize = chunks * chunkSize;
 
+            var requestAborted = new ManualResetEventSlim();
             var messageLogged = new ManualResetEventSlim();
-
             var mockKestrelTrace = new Mock<IKestrelTrace>();
             mockKestrelTrace
                 .Setup(trace => trace.ResponseMininumDataRateNotSatisfied(It.IsAny<string>(), It.IsAny<string>()))
@@ -2471,16 +2471,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
                 }
             };
 
-            var aborted = new ManualResetEventSlim();
-
             using (var server = new TestServer(async context =>
             {
-                context.RequestAborted.Register(() =>
-                {
-                    aborted.Set();
-                });
+                context.RequestAborted.Register(() => requestAborted.Set());
 
-                context.Response.ContentLength = chunks * chunkSize;
+                context.Response.ContentLength = responseSize;
 
                 for (var i = 0; i < chunks; i++)
                 {
@@ -2488,29 +2483,36 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
                 }
             }, testContext))
             {
-                using (var connection = server.CreateConnection())
+                using (var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp))
                 {
-                    await connection.Send(
-                        "GET / HTTP/1.1",
-                        "Host:",
-                        "",
-                        "");
+                    socket.ReceiveBufferSize = 1;
 
-                    Assert.True(aborted.Wait(TimeSpan.FromSeconds(60)));
+                    socket.Connect(new IPEndPoint(IPAddress.Loopback, server.Port));
+                    socket.Send(Encoding.ASCII.GetBytes("GET / HTTP/1.1\r\nHost: \r\n\r\n"));
 
-                    await connection.Receive(
-                        "HTTP/1.1 200 OK",
-                        "");
-                    await connection.ReceiveStartsWith("Date: ");
-                    await connection.Receive(
-                        $"Content-Length: {chunks * chunkSize}",
-                        "",
-                        "");
+                    Assert.True(messageLogged.Wait(TimeSpan.FromSeconds(60)), "The expected message was not logged within the timeout period.");
+                    Assert.True(requestAborted.Wait(TimeSpan.FromSeconds(60)), "The request was not aborted within the timeout period.");
 
-                    await Assert.ThrowsAsync<EqualException>(async () => await connection.ReceiveForcedEnd(
-                        new string('a', chunks * chunkSize)));
+                    var totalReceived = 0;
+                    var received = 0;
 
-                    Assert.True(messageLogged.Wait(TimeSpan.FromSeconds(10)));
+                    try
+                    {
+                        var buffer = new byte[chunkSize];
+
+                        do
+                        {
+                            received = socket.Receive(buffer);
+                            totalReceived += received;
+                        } while (received > 0 && totalReceived < responseSize);
+                    }
+                    catch (IOException)
+                    {
+                        // Socket.Receive could throw, and that is fine
+                    }
+
+                    // Since we expect writes to be cut off by the rate control, we should never see the entire response
+                    Assert.NotEqual(responseSize, totalReceived);
                 }
             }
         }