Throw from OnConnected and OnDisconnected when using client results

4 files changed, 115 insertions, 0 deletions

diff --git a/src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs b/src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs
index b6f7ffd02a..24d640645a 100644
--- a/src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs
+++ b/src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs
@@ -90,6 +90,9 @@ internal sealed partial class DefaultHubDispatcher<THub> : HubDispatcher<THub> w
             {
                 await hub.OnConnectedAsync();
             }
+
+            // OnConnectedAsync is finished, allow hub methods to use client results (ISingleClientProxy.InvokeAsync)
+            connection.HubCallerClients.InvokeAllowed = true;
         }
         finally
         {
@@ -107,6 +110,9 @@ internal sealed partial class DefaultHubDispatcher<THub> : HubDispatcher<THub> w
         {
             InitializeHub(hub, connection);
 
+            // OnDisonnectedAsync is being called, we don't want to allow client results to be used (ISingleClientProxy.InvokeAsync)
+            connection.HubCallerClients.InvokeAllowed = false;
+
             if (_onDisconnectedMiddleware != null)
             {
                 var context = new HubLifetimeContext(connection.HubCallerContext, scope.ServiceProvider, hub);
diff --git a/src/SignalR/server/Core/src/Internal/HubCallerClients.cs b/src/SignalR/server/Core/src/Internal/HubCallerClients.cs
index 04ce2c784c..1d3ed7db29 100644
--- a/src/SignalR/server/Core/src/Internal/HubCallerClients.cs
+++ b/src/SignalR/server/Core/src/Internal/HubCallerClients.cs
@@ -10,6 +10,11 @@ internal sealed class HubCallerClients : IHubCallerClients
     private readonly string[] _currentConnectionId;
     private readonly bool _parallelEnabled;
 
+    // Client results don't work in OnConnectedAsync and OnDisconnectedAsync
+    // This property is set by the hub dispatcher when those methods are being called
+    // so we can prevent users from making blocking client calls by returning a custom ISingleClientProxy instance
+    internal bool InvokeAllowed { get; set; }
+
     public HubCallerClients(IHubClients hubClients, string connectionId, bool parallelEnabled)
     {
         _connectionId = connectionId;
@@ -27,6 +32,10 @@ internal sealed class HubCallerClients : IHubCallerClients
             {
                 return new NotParallelSingleClientProxy(_hubClients.Client(_connectionId));
             }
+            if (!InvokeAllowed)
+            {
+                return new NoInvokeSingleClientProxy(_hubClients.Client(_connectionId));
+            }
             return _hubClients.Client(_connectionId);
         }
     }
@@ -47,6 +56,10 @@ internal sealed class HubCallerClients : IHubCallerClients
         {
             return new NotParallelSingleClientProxy(_hubClients.Client(connectionId));
         }
+        if (!InvokeAllowed)
+        {
+            return new NoInvokeSingleClientProxy(_hubClients.Client(_connectionId));
+        }
         return _hubClients.Client(connectionId);
     }
 
@@ -104,4 +117,24 @@ internal sealed class HubCallerClients : IHubCallerClients
             return _proxy.SendCoreAsync(method, args, cancellationToken);
         }
     }
+
+    private sealed class NoInvokeSingleClientProxy : ISingleClientProxy
+    {
+        private readonly ISingleClientProxy _proxy;
+
+        public NoInvokeSingleClientProxy(ISingleClientProxy hubClients)
+        {
+            _proxy = hubClients;
+        }
+
+        public Task<T> InvokeCoreAsync<T>(string method, object?[] args, CancellationToken cancellationToken = default)
+        {
+            throw new InvalidOperationException("Client results inside OnConnectedAsync or OnDisconnectedAsync Hub methods are not allowed.");
+        }
+
+        public Task SendCoreAsync(string method, object?[] args, CancellationToken cancellationToken = default)
+        {
+            return _proxy.SendCoreAsync(method, args, cancellationToken);
+        }
+    }
 }
diff --git a/src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Hubs.cs b/src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Hubs.cs
index 7ee5efc3b1..cb2fef8244 100644
--- a/src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Hubs.cs
+++ b/src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Hubs.cs
@@ -1231,6 +1231,22 @@ public class ConnectionLifetimeState
     public Exception DisconnectedException { get; set; }
 }
 
+public class OnConnectedClientResultHub : Hub
+{
+    public override async Task OnConnectedAsync()
+    {
+        await Clients.Caller.InvokeAsync<int>("Test");
+    }
+}
+
+public class OnDisconnectedClientResultHub : Hub
+{
+    public override async Task OnDisconnectedAsync(Exception ex)
+    {
+        await Clients.Caller.InvokeAsync<int>("Test");
+    }
+}
+
 public class CallerServiceHub : Hub
 {
     private readonly CallerService _service;
diff --git a/src/SignalR/server/SignalR/test/HubConnectionHandlerTests.ClientResult.cs b/src/SignalR/server/SignalR/test/HubConnectionHandlerTests.ClientResult.cs
index c348bce508..42c523e1eb 100644
--- a/src/SignalR/server/SignalR/test/HubConnectionHandlerTests.ClientResult.cs
+++ b/src/SignalR/server/SignalR/test/HubConnectionHandlerTests.ClientResult.cs
@@ -105,6 +105,66 @@ public partial class HubConnectionHandlerTests
     }
 
     [Fact]
+    public async Task ThrowsWhenUsedInOnConnectedAsync()
+    {
+        using (StartVerifiableLog(write => write.EventId.Name == "ErrorDispatchingHubEvent"))
+        {
+            var serviceProvider = HubConnectionHandlerTestUtils.CreateServiceProvider(builder =>
+            {
+                builder.AddSignalR(o =>
+                {
+                    o.MaximumParallelInvocationsPerClient = 2;
+                    o.EnableDetailedErrors = true;
+                });
+            }, LoggerFactory);
+            var connectionHandler = serviceProvider.GetService<HubConnectionHandler<OnConnectedClientResultHub>>();
+
+            using (var client = new TestClient())
+            {
+                var connectionHandlerTask = await client.ConnectAsync(connectionHandler).DefaultTimeout();
+
+                // Hub asks client for a result, this is an invocation message with an ID
+                var closeMessage = Assert.IsType<CloseMessage>(await client.ReadAsync().DefaultTimeout());
+                Assert.Equal("Connection closed with an error. InvalidOperationException: Client results inside OnConnectedAsync or OnDisconnectedAsync Hub methods are not allowed.", closeMessage.Error);
+            }
+        }
+
+        Assert.Single(TestSink.Writes.Where(write => write.EventId.Name == "ErrorDispatchingHubEvent"));
+    }
+
+    [Fact]
+    public async Task ThrowsWhenUsedInOnDisconnectedAsync()
+    {
+        using (StartVerifiableLog(write => write.EventId.Name == "ErrorDispatchingHubEvent"))
+        {
+            var serviceProvider = HubConnectionHandlerTestUtils.CreateServiceProvider(builder =>
+            {
+                builder.AddSignalR(o =>
+                {
+                    o.MaximumParallelInvocationsPerClient = 2;
+                    o.EnableDetailedErrors = true;
+                });
+            }, LoggerFactory);
+            var connectionHandler = serviceProvider.GetService<HubConnectionHandler<OnDisconnectedClientResultHub>>();
+
+            using (var client = new TestClient())
+            {
+                var connectionHandlerTask = await client.ConnectAsync(connectionHandler).DefaultTimeout();
+                client.Connection.Abort();
+
+                // Hub asks client for a result, this is an invocation message with an ID
+                var closeMessage = Assert.IsType<CloseMessage>(await client.ReadAsync().DefaultTimeout());
+                Assert.Null(closeMessage.Error);
+
+                var ex = await Assert.ThrowsAsync<InvalidOperationException>(() => connectionHandlerTask).DefaultTimeout();
+                Assert.Equal("Client results inside OnConnectedAsync or OnDisconnectedAsync Hub methods are not allowed.", ex.Message);
+            }
+        }
+
+        Assert.Single(TestSink.Writes.Where(write => write.EventId.Name == "ErrorDispatchingHubEvent"));
+    }
+
+    [Fact]
     public async Task CanUseClientResultsWithIHubContext()
     {
         using (StartVerifiableLog())