From 3d6e1e69fd03d0aed50f713ece482e7df65434a6 Mon Sep 17 00:00:00 2001
From: BrennanConroy <brecon@microsoft.com>
Date: Tue, 31 Jul 2018 09:52:10 -0700
Subject: [PATCH] Close LongPolling connection on poll exception (#2701)

---
 .../Internal/HttpConnectionDispatcher.cs      |  7 +++
 .../Transports/LongPollingTransport.cs        |  2 +
 .../HttpConnectionDispatcherTests.cs          | 43 +++++++++++++++++++
 3 files changed, 52 insertions(+)

diff --git a/src/Microsoft.AspNetCore.Http.Connections/Internal/HttpConnectionDispatcher.cs b/src/Microsoft.AspNetCore.Http.Connections/Internal/HttpConnectionDispatcher.cs
index 3d36933d3c..3b18c53bfd 100644
--- a/src/Microsoft.AspNetCore.Http.Connections/Internal/HttpConnectionDispatcher.cs
+++ b/src/Microsoft.AspNetCore.Http.Connections/Internal/HttpConnectionDispatcher.cs
@@ -289,6 +289,13 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
                         pollAgain = false;
                     }
                 }
+                else if (resultTask.IsFaulted)
+                {
+                    // transport task was faulted, we should remove the connection
+                    await _manager.DisposeAndRemoveAsync(connection, closeGracefully: false);
+
+                    pollAgain = false;
+                }
                 else if (context.Response.StatusCode == StatusCodes.Status204NoContent)
                 {
                     // Don't poll if the transport task was canceled
diff --git a/src/Microsoft.AspNetCore.Http.Connections/Internal/Transports/LongPollingTransport.cs b/src/Microsoft.AspNetCore.Http.Connections/Internal/Transports/LongPollingTransport.cs
index cf5638d74d..f3cd87f484 100644
--- a/src/Microsoft.AspNetCore.Http.Connections/Internal/Transports/LongPollingTransport.cs
+++ b/src/Microsoft.AspNetCore.Http.Connections/Internal/Transports/LongPollingTransport.cs
@@ -90,6 +90,8 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal.Transports
             catch (Exception ex)
             {
                 Log.LongPollingTerminated(_logger, ex);
+                context.Response.ContentType = "text/plain";
+                context.Response.StatusCode = StatusCodes.Status500InternalServerError;
                 throw;
             }
         }
diff --git a/test/Microsoft.AspNetCore.Http.Connections.Tests/HttpConnectionDispatcherTests.cs b/test/Microsoft.AspNetCore.Http.Connections.Tests/HttpConnectionDispatcherTests.cs
index 0761fe7c2b..7934e98df1 100644
--- a/test/Microsoft.AspNetCore.Http.Connections.Tests/HttpConnectionDispatcherTests.cs
+++ b/test/Microsoft.AspNetCore.Http.Connections.Tests/HttpConnectionDispatcherTests.cs
@@ -18,6 +18,7 @@ using Microsoft.AspNetCore.Authorization;
 using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Connections.Features;
 using Microsoft.AspNetCore.Http.Connections.Internal;
+using Microsoft.AspNetCore.Http.Connections.Internal.Transports;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Http.Internal;
 using Microsoft.AspNetCore.SignalR.Tests;
@@ -2016,6 +2017,48 @@ namespace Microsoft.AspNetCore.Http.Connections.Tests
             }
         }
 
+        [Fact]
+        public async Task ErrorDuringPollWillCloseConnection()
+        {
+            bool ExpectedErrors(WriteContext writeContext)
+            {
+                return (writeContext.LoggerName == typeof(LongPollingTransport).FullName &&
+                       writeContext.EventId.Name == "LongPollingTerminated") ||
+                       (writeContext.LoggerName == typeof(HttpConnectionManager).FullName &&
+                       writeContext.EventId.Name == "FailedDispose");
+            }
+
+            using (StartVerifiableLog(out var loggerFactory, LogLevel.Debug, expectedErrorsFilter: ExpectedErrors))
+            {
+                var manager = CreateConnectionManager(loggerFactory);
+                var connection = manager.CreateConnection();
+                connection.TransportType = HttpTransportType.LongPolling;
+
+                var dispatcher = new HttpConnectionDispatcher(manager, loggerFactory);
+
+                var services = new ServiceCollection();
+                services.AddSingleton<TestConnectionHandler>();
+                var builder = new ConnectionBuilder(services.BuildServiceProvider());
+                builder.UseConnectionHandler<TestConnectionHandler>();
+                var app = builder.Build();
+                var options = new HttpConnectionDispatcherOptions();
+
+                var context = MakeRequest("/foo", connection);
+
+                // Initial poll will complete immediately
+                await dispatcher.ExecuteAsync(context, options, app).OrTimeout();
+
+                var pollContext = MakeRequest("/foo", connection);
+                var pollTask = dispatcher.ExecuteAsync(pollContext, options, app);
+                // fail LongPollingTransport ReadAsync
+                connection.Transport.Output.Complete(new InvalidOperationException());
+                await pollTask.OrTimeout();
+
+                Assert.Equal(StatusCodes.Status500InternalServerError, pollContext.Response.StatusCode);
+                Assert.False(manager.TryGetConnection(connection.ConnectionId, out var _));
+            }
+        }
+
         private class RejectHandler : TestAuthenticationHandler
         {
             protected override bool ShouldAccept => false;