Throw from OnConnected when using client results (#42447)

Author: BrennanConroy

src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs

@@ -90,6 +90,9 @@ public override async Task OnConnectedAsync(HubConnectionContext connection)
             {
                 await hub.OnConnectedAsync();
             }
+
+            // OnConnectedAsync is finished, allow hub methods to use client results (ISingleClientProxy.InvokeAsync)
+            connection.HubCallerClients.InvokeAllowed = true;
         }
         finally
         {

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
+    // 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 @@ public ISingleClientProxy Caller
             {
                 return new NotParallelSingleClientProxy(_hubClients.Client(_connectionId));
             }
+            if (!InvokeAllowed)
+            {
+                return new NoInvokeSingleClientProxy(_hubClients.Client(_connectionId));
+            }
             return _hubClients.Client(_connectionId);
         }
     }
@@ -47,6 +56,10 @@ public ISingleClientProxy Client(string connectionId)
         {
             return new NotParallelSingleClientProxy(_hubClients.Client(connectionId));
         }
+        if (!InvokeAllowed)
+        {
+            return new NoInvokeSingleClientProxy(_hubClients.Client(_connectionId));
+        }
         return _hubClients.Client(connectionId);
     }
 
@@ -104,4 +117,24 @@ public Task SendCoreAsync(string method, object?[] args, CancellationToken cance
             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 Hub methods are not allowed.");
+        }
+
+        public Task SendCoreAsync(string method, object?[] args, CancellationToken cancellationToken = default)
+        {
+            return _proxy.SendCoreAsync(method, args, cancellationToken);
+        }
+    }
 }

src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Hubs.cs

@@ -1254,6 +1254,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;

src/SignalR/server/SignalR/test/HubConnectionHandlerTests.ClientResult.cs

@@ -104,6 +104,66 @@ public async Task ThrowsWhenParallelHubInvokesNotEnabled()
         }
     }
 
+    [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 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<IOException>(() => connectionHandlerTask).DefaultTimeout();
+                Assert.Equal($"Connection '{client.Connection.ConnectionId}' disconnected.", ex.Message);
+            }
+        }
+
+        Assert.Single(TestSink.Writes.Where(write => write.EventId.Name == "ErrorDispatchingHubEvent"));
+    }
+
     [Fact]
     public async Task CanUseClientResultsWithIHubContext()
     {