Fix websocket connection issues

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Subscriptions/ISocketConnection.cs

@@ -29,7 +29,15 @@ public interface ISocketConnection : ISocket, IFeatureProvider, IDisposable
     /// <summary>
     /// Specifies if the connection is connected to a client.
     /// </summary>
-    bool IsConnected { get; set; }
+    bool IsConnected { get; }
+
+    /// <summary>
+    /// Specifies if a connection initialization message has been received from the client.
+    /// This is set as soon as the message arrives, before the connection is accepted, so that
+    /// the initialization timeout only reflects whether the client sent the message in time and
+    /// not how long the acceptance (for example authentication) takes.
+    /// </summary>
+    bool ConnectionInitReceived { get; }
 
     /// <summary>
     /// Tries to accept the connection and returns the accepted protocol handler.

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Subscriptions/OperationManager.cs

@@ -76,7 +76,7 @@ public bool Enqueue(string sessionId, GraphQLRequest request)
 
         if (session is not null)
         {
-            session.Completed += (_, _) => Complete(sessionId);
+            session.Completed += (_, _) => CompleteSession(sessionId, session);
             session.BeginExecute(request, _cancellationToken);
             return true;
         }
@@ -109,7 +109,7 @@ public bool EnqueueBatch(string sessionId, GraphQLRequest[] requests)
 
         if (session is not null)
         {
-            session.Completed += (_, _) => Complete(sessionId);
+            session.Completed += (_, _) => CompleteSession(sessionId, session);
             session.BeginExecuteBatch(requests, _cancellationToken);
             return true;
         }
@@ -141,6 +141,24 @@ public bool Complete(string sessionId)
         return false;
     }
 
+    private void CompleteSession(string sessionId, IOperationSession session)
+    {
+        _lock.EnterWriteLock();
+
+        try
+        {
+            if (_subs.TryGetValue(sessionId, out var current) && ReferenceEquals(current, session))
+            {
+                _subs.Remove(sessionId);
+                session.Dispose();
+            }
+        }
+        finally
+        {
+            _lock.ExitWriteLock();
+        }
+    }
+
     private OperationSession CreateSession(string sessionId)
         => new(_socketSession, _executorSession, sessionId);
 
@@ -151,7 +169,18 @@ public void Dispose()
         {
             _cts.Cancel();
             _cts.Dispose();
-            _subs.Clear();
+
+            _lock.EnterWriteLock();
+
+            try
+            {
+                _subs.Clear();
+            }
+            finally
+            {
+                _lock.ExitWriteLock();
+            }
+
             _disposed = true;
         }
     }

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Subscriptions/OperationSession.cs

@@ -40,6 +40,7 @@ private async Task SendResultsAsync(GraphQLRequest request, CancellationToken ca
         using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _ct);
         var ct = cts.Token;
         var completeTry = false;
+        var errorSent = false;
 
         try
         {
@@ -52,6 +53,9 @@ private async Task SendResultsAsync(GraphQLRequest request, CancellationToken ca
                 case OperationResult queryResult:
                     if (queryResult.Data is null && queryResult.Errors is { Count: > 0 })
                     {
+                        // an error message terminates the operation, so no complete message
+                        // is sent afterwards.
+                        errorSent = true;
                         await _session.Protocol.SendErrorMessageAsync(
                             _session,
                             Id,
@@ -85,7 +89,7 @@ await _session.Protocol.SendErrorMessageAsync(
             // message again.
             completeTry = true;
 
-            if (!ct.IsCancellationRequested)
+            if (!errorSent && !ct.IsCancellationRequested)
             {
                 await _session.Protocol.SendCompleteMessageAsync(_session, Id, ct);
             }

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Subscriptions/PingPongJob.cs

@@ -15,8 +15,10 @@ public async Task RunAsync(CancellationToken cancellationToken)
             await Task.Delay(options.ConnectionInitializationTimeout, cancellationToken);
 
             // if after the timeout no connection initialization was sent by the client, we will
-            // close the connection.
-            if (!connection.IsConnected)
+            // close the connection. We check that the init message was received rather than that
+            // the connection was fully accepted, so a slow OnConnectAsync does not trigger a
+            // spurious timeout for a client that initialized in time.
+            if (!connection.ConnectionInitReceived)
             {
                 await session.Protocol.OnConnectionInitTimeoutAsync(session, cancellationToken);
                 return;

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Subscriptions/Protocols/Apollo/ApolloSubscriptionProtocolHandler.cs

@@ -55,13 +55,6 @@ private async ValueTask OnReceiveInternalAsync(
         var connection = session.Connection;
         var connected = connection.IsConnected;
 
-        if (connected && message.IsSingleSegment
-            && message.First.Equals(Utf8MessageBodies.KeepAlive))
-        {
-            // received a simple ping, we do not need to answer to this message.
-            return;
-        }
-
         using var document = JsonDocument.Parse(message);
         var root = document.RootElement;
         JsonElement idProp;
@@ -85,6 +78,12 @@ await connection.CloseAsync(
             return;
         }
 
+        if (type.ValueEquals("ka"u8))
+        {
+            // a client keep-alive is a no-op, we must not close the connection.
+            return;
+        }
+
         if (type.ValueEquals(Utf8Messages.ConnectionInitialize))
         {
             if (connected)
@@ -96,6 +95,10 @@ await connection.CloseAsync(
                 return;
             }
 
+            // signal that the client sent the connection init in time, even if accepting the
+            // connection (for example authentication) still needs to run.
+            ((WebSocketConnection)connection).ConnectionInitReceived = true;
+
             var operationMessageObj =
                 TryGetPayload(root, out var payload)
                     ? new InitializeConnectionMessage(payload)
@@ -109,7 +112,7 @@ await _interceptor.OnConnectAsync(
 
             if (connectionStatus.Accepted)
             {
-                connection.IsConnected = true;
+                ((WebSocketConnection)connection).IsConnected = true;
                 await SendConnectionAcceptMessage(
                     session,
                     connectionStatus.Extensions,

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Subscriptions/Protocols/GraphQLOverWebSocket/CloseReasons.cs

@@ -7,4 +7,5 @@ internal static class CloseReasons
     public const int TooManyInitAttempts = 4429;
     public const int SubscriberNotUnique = 4409;
     public const int Unauthorized = 4401;
+    public const int Forbidden = 4403;
 }

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Subscriptions/Protocols/GraphQLOverWebSocket/ConnectionExtensions.cs

@@ -60,10 +60,19 @@ public static ValueTask CloseUnauthorizedAsync(
 
     public static ValueTask CloseConnectionRefusedAsync(
         this ISocketConnection connection,
+        string message,
         CancellationToken cancellationToken)
         => connection.CloseAsync(
-            "Connection refused.",
-            CloseReasons.Unauthorized,
+            message,
+            CloseReasons.Forbidden,
+            cancellationToken);
+
+    public static ValueTask CloseInvalidMessageAsync(
+        this ISocketConnection connection,
+        CancellationToken cancellationToken)
+        => connection.CloseAsync(
+            "The message must be valid JSON.",
+            CloseReasons.ProtocolError,
             cancellationToken);
 
     public static ValueTask CloseToManyInitializationsAsync(

src/HotChocolate/AspNetCore/src/AspNetCore.Pipeline/Subscriptions/Protocols/GraphQLOverWebSocket/GraphQLOverWebSocketProtocolHandler.cs

@@ -47,7 +47,20 @@ private async ValueTask OnReceiveInternalAsync(
         var connection = session.Connection;
 
         var connected = connection.IsConnected;
-        using var document = JsonDocument.Parse(message);
+
+        JsonDocument parsedDocument;
+
+        try
+        {
+            parsedDocument = JsonDocument.Parse(message);
+        }
+        catch (JsonException)
+        {
+            await connection.CloseInvalidMessageAsync(cancellationToken);
+            return;
+        }
+
+        using var document = parsedDocument;
         var root = document.RootElement;
         JsonElement idProp;
 
@@ -97,6 +110,10 @@ private async ValueTask OnReceiveInternalAsync(
                 return;
             }
 
+            // signal that the client sent the connection init in time, even if accepting the
+            // connection (for example authentication) still needs to run.
+            ((WebSocketConnection)connection).ConnectionInitReceived = true;
+
             var operationMessageObj =
                 TryGetPayload(root, out var payload)
                     ? new ConnectionInitMessage(payload)
@@ -110,15 +127,17 @@ await interceptor.OnConnectAsync(
 
             if (connectionStatus.Accepted)
             {
-                connection.IsConnected = true;
+                ((WebSocketConnection)connection).IsConnected = true;
                 await SendConnectionAcceptMessage(
                     session,
                     connectionStatus.Extensions,
                     cancellationToken);
             }
             else
             {
-                await connection.CloseConnectionRefusedAsync(cancellationToken);
+                await connection.CloseConnectionRefusedAsync(
+                    connectionStatus.Message,
+                    cancellationToken);
             }
 
             return;
@@ -134,9 +153,14 @@ await SendConnectionAcceptMessage(
 
         if (type.ValueEquals(Utf8Messages.Subscribe))
         {
+            // the parsed requests own pooled JSON buffers (variables and extensions). Ownership
+            // transfers to the operation session only when Enqueue/EnqueueBatch succeeds. On every
+            // rejection or error path we keep this reference non-null so the finally disposes them.
+            GraphQLRequest[]? requests = null;
+
             try
             {
-                if (!TryParseSubscribeMessage(root, out var subscribeId, out var requests))
+                if (!TryParseSubscribeMessage(root, out var subscribeId, out requests))
                 {
                     await connection.CloseInvalidSubscribeMessageAsync(cancellationToken);
                     return;
@@ -165,7 +189,12 @@ await SendConnectionAcceptMessage(
                     success = session.Operations.EnqueueBatch(subscribeId, requests);
                 }
 
-                if (!success)
+                if (success)
+                {
+                    // the operation session now owns the requests and disposes them.
+                    requests = null;
+                }
+                else
                 {
                     await connection.CloseSubscriptionIdNotUniqueAsync(cancellationToken);
                 }
@@ -204,6 +233,16 @@ await SendErrorMessageAsync(
                     [syntaxError],
                     cancellationToken);
             }
+            finally
+            {
+                if (requests is not null)
+                {
+                    foreach (var request in requests)
+                    {
+                        request.Dispose();
+                    }
+                }
+            }
 
             // the operation was excepted and we are done.
             return;