client

Author: BrennanConroy

src/SignalR/clients/csharp/Client.Core/src/HubConnection.cs

@@ -576,6 +576,11 @@ private async Task StopAsyncCore(bool disposing)
                         TaskContinuationOptions.ExecuteSynchronously | TaskContinuationOptions.OnlyOnFaulted,
                         TaskScheduler.Default);
                 }
+
+                if (connectionState.Connection.Features.Get<IReconnectFeature>() is IReconnectFeature feature)
+                {
+                    feature.DisableReconnect();
+                }
             }
             else
             {
@@ -1085,6 +1090,12 @@ private async Task SendWithLock(ConnectionState expectedConnectionState, HubMess
                 {
                     Log.ReceivedCloseWithError(_logger, close.Error);
                 }
+
+                if (connectionState.Connection.Features.Get<IReconnectFeature>() is IReconnectFeature feature)
+                {
+                    feature.DisableReconnect();
+                }
+
                 return close;
             case PingMessage _:
                 Log.ReceivedPing(_logger);

src/SignalR/clients/csharp/Client/test/UnitTests/HubConnectionTests.cs

@@ -7,6 +7,7 @@
 using System.Net.WebSockets;
 using System.Threading.Channels;
 using Microsoft.AspNetCore.Connections;
+using Microsoft.AspNetCore.Connections.Abstractions;
 using Microsoft.AspNetCore.Connections.Features;
 using Microsoft.AspNetCore.Http.Connections.Client;
 using Microsoft.AspNetCore.SignalR.Protocol;
@@ -878,6 +879,66 @@ public async Task HubConnectionIsMockable()
         mockConnection.Verify(c => c.StopAsync(It.IsAny<CancellationToken>()), Times.Once);
     }
 
+    [Fact]
+    public async Task DisableReconnectCalledWhenCloseMessageReceived()
+    {
+        var builder = new HubConnectionBuilder().WithUrl("http://example.com");
+        var innerConnection = new TestConnection();
+        var reconnectFeature = new TestReconnectFeature();
+        innerConnection.Features.Set<IReconnectFeature>(reconnectFeature);
+
+        var delegateConnectionFactory = new DelegateConnectionFactory(
+            endPoint => innerConnection.StartAsync());
+        builder.Services.AddSingleton<IConnectionFactory>(delegateConnectionFactory);
+
+        var hubConnection = builder.Build();
+        var closedEventTcs = new TaskCompletionSource<Exception>();
+        hubConnection.Closed += e =>
+        {
+            closedEventTcs.SetResult(e);
+            return Task.CompletedTask;
+        };
+
+        await hubConnection.StartAsync().DefaultTimeout();
+
+        await innerConnection.ReceiveJsonMessage(new { type = HubProtocolConstants.CloseMessageType });
+
+        var exception = await closedEventTcs.Task.DefaultTimeout();
+        Assert.Null(exception);
+
+        await reconnectFeature.DisableReconnectCalled.DefaultTimeout();
+    }
+
+    [Fact]
+    public async Task DisableReconnectCalledWhenSendingCloseMessage()
+    {
+        var builder = new HubConnectionBuilder().WithUrl("http://example.com");
+        var innerConnection = new TestConnection();
+        var reconnectFeature = new TestReconnectFeature();
+        innerConnection.Features.Set<IReconnectFeature>(reconnectFeature);
+
+        var delegateConnectionFactory = new DelegateConnectionFactory(
+            endPoint => innerConnection.StartAsync());
+        builder.Services.AddSingleton<IConnectionFactory>(delegateConnectionFactory);
+
+        var hubConnection = builder.Build();
+        var closedEventTcs = new TaskCompletionSource<Exception>();
+        hubConnection.Closed += e =>
+        {
+            closedEventTcs.SetResult(e);
+            return Task.CompletedTask;
+        };
+
+        await hubConnection.StartAsync().DefaultTimeout();
+
+        await hubConnection.StopAsync().DefaultTimeout();
+
+        var exception = await closedEventTcs.Task.DefaultTimeout();
+        Assert.Null(exception);
+
+        await reconnectFeature.DisableReconnectCalled.DefaultTimeout();
+    }
+
     private class SampleObject
     {
         public SampleObject(string foo, int bar)
@@ -962,4 +1023,18 @@ public ReadOnlyMemory<byte> GetMessageBytes(HubMessage message)
             return HubProtocolExtensions.GetMessageBytes(this, message);
         }
     }
+
+    private sealed class TestReconnectFeature : IReconnectFeature
+    {
+        private TaskCompletionSource _disableReconnect = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        public Task DisableReconnectCalled => _disableReconnect.Task;
+
+        public Action<PipeWriter> NotifyOnReconnect { get; set; }
+
+        public void DisableReconnect()
+        {
+            _disableReconnect.TrySetResult();
+        }
+    }
 }

src/SignalR/common/Http.Connections/test/HttpConnectionDispatcherTests.cs

@@ -2414,6 +2414,77 @@ public async Task ReconnectStopsPreviousConnection()
         }
     }
 
+    [Fact]
+    public async Task DisableReconnectDisallowsReplacementConnection()
+    {
+        using (StartVerifiableLog())
+        {
+            var manager = CreateConnectionManager(LoggerFactory);
+            var options = new HttpConnectionDispatcherOptions() { AllowAcks = true };
+            options.WebSockets.CloseTimeout = TimeSpan.FromMilliseconds(1);
+            // pretend negotiate occurred
+            var connection = manager.CreateConnection(options, negotiateVersion: 1, useAck: true);
+            connection.TransportType = HttpTransportType.WebSockets;
+
+            var dispatcher = CreateDispatcher(manager, LoggerFactory);
+            var services = new ServiceCollection();
+
+            var context = MakeRequest("/foo", connection, services);
+            SetTransport(context, HttpTransportType.WebSockets);
+
+            var builder = new ConnectionBuilder(services.BuildServiceProvider());
+            builder.UseConnectionHandler<ReconnectConnectionHandler>();
+            var app = builder.Build();
+
+            var initialWebSocketTask = dispatcher.ExecuteAsync(context, options, app);
+
+            var reconnectFeature = connection.Features.Get<IReconnectFeature>();
+            Assert.NotNull(reconnectFeature);
+
+            var firstMsg = new byte[] { 1, 4, 8, 9 };
+            await connection.Application.Output.WriteAsync(firstMsg);
+
+            var websocketFeature = (TestWebSocketConnectionFeature)context.Features.Get<IHttpWebSocketFeature>();
+            await websocketFeature.Accepted.DefaultTimeout();
+            // Run the client socket
+            var webSocketMessage = await websocketFeature.Client.GetNextMessageAsync().DefaultTimeout();
+
+            Assert.Equal(firstMsg, webSocketMessage.Buffer);
+
+            var called = false;
+            var reconnectCallback = reconnectFeature.NotifyOnReconnect;
+            reconnectFeature.NotifyOnReconnect = (writer) =>
+            {
+                called = true;
+                reconnectCallback(writer);
+            };
+
+            // Disable will not allow new connection to override existing
+            reconnectFeature.DisableReconnect();
+
+            // New websocket connection with previous connection token
+            context = MakeRequest("/foo", connection, services);
+            SetTransport(context, HttpTransportType.WebSockets);
+
+            await dispatcher.ExecuteAsync(context, options, app).DefaultTimeout();
+
+            Assert.Equal(409, context.Response.StatusCode);
+
+            Assert.False(called);
+
+            // Connection still works
+            var secondMsg = new byte[] { 7, 6, 3, 2 };
+            await connection.Application.Output.WriteAsync(secondMsg);
+
+            webSocketMessage = await websocketFeature.Client.GetNextMessageAsync().DefaultTimeout();
+            Assert.Equal(secondMsg, webSocketMessage.Buffer);
+
+            connection.Abort();
+
+            await initialWebSocketTask.DefaultTimeout();
+        }
+    }
+
     private class ControllableMemoryStream : MemoryStream
     {
         private readonly SyncPoint _syncPoint;

src/SignalR/server/SignalR/test/Internal/MessageBufferTests.cs

@@ -283,7 +283,6 @@ public DuplexPipePair(IDuplexPipe transport, IDuplexPipe application)
 
     public static void UpdateConnectionPair(ref DuplexPipePair duplexPipePair, ConnectionContext connection)
     {
-        var prevPipe = duplexPipePair.Application.Input;
         var input = new Pipe();
 
         // Add new pipe for reading from and writing to transport from app code