From 02676956565e85d806c85e1235cb2998fd5a5caa Mon Sep 17 00:00:00 2001
From: Pawel Kadluczka <pawelka@microsoft.com>
Date: Fri, 22 Sep 2017 12:30:41 -0700
Subject: [PATCH] Exceptions thrown during writing should close the connection

We need to close the connection if there is an exception when writing to the transport on the server side. Currently if an exception happens it leaves the connection in an unsable state - after the exception no messages from the server will be sent to the client because the writing loop is terminated. Ignoring the message could cause hangs on the client side since we can fail while writing a completion message. In this case if the client is awaiting the invocation it will hang because the task will never be completed.
---
 .../HubEndPoint.cs                            | 21 ++++++++----
 .../HubEndpointTests.cs                       | 34 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/src/Microsoft.AspNetCore.SignalR.Core/HubEndPoint.cs b/src/Microsoft.AspNetCore.SignalR.Core/HubEndPoint.cs
index 01e95ea142..20915c0ba8 100644
--- a/src/Microsoft.AspNetCore.SignalR.Core/HubEndPoint.cs
+++ b/src/Microsoft.AspNetCore.SignalR.Core/HubEndPoint.cs
@@ -76,20 +76,27 @@ namespace Microsoft.AspNetCore.SignalR
             var protocolReaderWriter = connectionContext.ProtocolReaderWriter;
             async Task WriteToTransport()
             {
-                while (await output.In.WaitToReadAsync())
+                try
                 {
-                    while (output.In.TryRead(out var hubMessage))
+                    while (await output.In.WaitToReadAsync())
                     {
-                        var buffer = protocolReaderWriter.WriteMessage(hubMessage);
-                        while (await connection.Transport.Out.WaitToWriteAsync())
+                        while (output.In.TryRead(out var hubMessage))
                         {
-                            if (connection.Transport.Out.TryWrite(buffer))
+                            var buffer = protocolReaderWriter.WriteMessage(hubMessage);
+                            while (await connection.Transport.Out.WaitToWriteAsync())
                             {
-                                break;
+                                if (connection.Transport.Out.TryWrite(buffer))
+                                {
+                                    break;
+                                }
                             }
                         }
                     }
                 }
+                catch (Exception ex)
+                {
+                    connectionContext.Abort(ex);
+                }
             }
 
             var writingOutputTask = WriteToTransport();
@@ -244,7 +251,7 @@ namespace Microsoft.AspNetCore.SignalR
 
         private async Task DispatchMessagesAsync(HubConnectionContext connection)
         {
-            // Since we dispatch multiple hub invocations in parallel, we need a way to communicate failure back to the main processing loop. 
+            // Since we dispatch multiple hub invocations in parallel, we need a way to communicate failure back to the main processing loop.
             // This is done by aborting the connection.
 
             try
diff --git a/test/Microsoft.AspNetCore.SignalR.Tests/HubEndpointTests.cs b/test/Microsoft.AspNetCore.SignalR.Tests/HubEndpointTests.cs
index 673c3bd1bf..8159f38de2 100644
--- a/test/Microsoft.AspNetCore.SignalR.Tests/HubEndpointTests.cs
+++ b/test/Microsoft.AspNetCore.SignalR.Tests/HubEndpointTests.cs
@@ -4,11 +4,13 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Runtime.Serialization;
 using System.Security.Claims;
 using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Channels;
 using Microsoft.AspNetCore.Authorization;
+using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.SignalR.Internal.Protocol;
 using Microsoft.AspNetCore.SignalR.Tests.Common;
@@ -21,7 +23,6 @@ using MsgPack.Serialization;
 using Newtonsoft.Json;
 using Newtonsoft.Json.Serialization;
 using Xunit;
-using Microsoft.AspNetCore.Hosting;
 
 namespace Microsoft.AspNetCore.SignalR.Tests
 {
@@ -1149,6 +1150,32 @@ namespace Microsoft.AspNetCore.SignalR.Tests
             }
         }
 
+        [Fact]
+        public async Task ConnectionClosedIfWritingToTransportFails()
+        {
+            // MessagePack does not support serializing objects or private types (including anonymous types)
+            // and throws. In this test we make sure that this exception closes the connection and bubbles up.
+
+            var serviceProvider = CreateServiceProvider();
+
+            var endPoint = serviceProvider.GetService<HubEndPoint<MethodHub>>();
+
+            using (var client = new TestClient(false, new MessagePackHubProtocol()))
+            {
+                var transportFeature = new Mock<IConnectionTransportFeature>();
+                transportFeature.SetupGet(f => f.TransportCapabilities).Returns(TransferMode.Binary);
+                client.Connection.Features.Set(transportFeature.Object);
+
+                var endPointLifetime = endPoint.OnConnectedAsync(client.Connection);
+
+                await client.Connected.OrTimeout();
+
+                await client.SendInvocationAsync(nameof(MethodHub.SendAnonymousObject)).OrTimeout();
+
+                await Assert.ThrowsAsync<SerializationException>(() => endPointLifetime.OrTimeout());
+            }
+        }
+
         private static void AssertHubMessage(HubMessage expected, HubMessage actual)
         {
             // We aren't testing InvocationIds here
@@ -1567,6 +1594,11 @@ namespace Microsoft.AspNetCore.SignalR.Tests
                 return $"{b}, {i}, {c}, {s}";
             }
 
+            public Task SendAnonymousObject()
+            {
+                return Clients.Client(Context.ConnectionId).InvokeAsync("Send", new { });
+            }
+
             public override Task OnDisconnectedAsync(Exception e)
             {
                 return Task.CompletedTask;