Add missing handshake messages when upgrading to WebSocket, fixes #380

Author: bajasoft

src/SocketIO.Serializer.Core/ISerializer.cs

@@ -1,4 +1,3 @@
-using System;
 using System.Collections.Generic;
 using SocketIO.Core;
 
@@ -19,6 +18,7 @@ public interface ISerializer
         SerializedItem SerializeConnectedMessage(EngineIO eio, string ns, object auth, IEnumerable<KeyValuePair<string, string>> queries);
 
         SerializedItem SerializePingMessage();
+        SerializedItem SerializePingProbeMessage();
         SerializedItem SerializePongMessage();
         SerializedItem SerializeUpgradeMessage();
     }

src/SocketIO.Serializer.MessagePack/SocketIOMessagePackSerializer.cs

@@ -257,6 +257,14 @@ public SerializedItem SerializePingMessage()
             };
         }
 
+        public SerializedItem SerializePingProbeMessage()
+        {
+            return new SerializedItem
+            {
+                Text = "2probe"
+            };
+        }
+
         public SerializedItem SerializePongMessage()
         {
             return new SerializedItem

src/SocketIO.Serializer.NewtonsoftJson/NewtonsoftJsonSerializer.cs

@@ -220,6 +220,14 @@ public SerializedItem SerializePingMessage()
             };
         }
 
+        public SerializedItem SerializePingProbeMessage()
+        {
+            return new SerializedItem
+            {
+                Text = "2probe"
+            };
+        }
+
         public SerializedItem SerializePongMessage()
         {
             return new SerializedItem
@@ -309,6 +317,7 @@ private static void ReadMessage(IMessage message, EngineIO eio, string text)
                 case MessageType.Ping:
                     break;
                 case MessageType.Pong:
+                    ReadPongMessage(message, text);
                     break;
                 case MessageType.Connected:
                     ReadConnectedMessage(message, text, eio);
@@ -336,6 +345,11 @@ private static void ReadMessage(IMessage message, EngineIO eio, string text)
             }
         }
 
+        private static void ReadPongMessage(IMessage message, string text)
+        {
+            message.ReceivedText = text;
+        }
+
         private static void ReadOpenedMessage(IMessage message, string text)
         {
             JsonConvert.PopulateObject(text, message, new JsonSerializerSettings

src/SocketIO.Serializer.SystemTextJson/SystemTextJsonSerializer.cs

@@ -220,6 +220,14 @@ public SerializedItem SerializePingMessage()
             };
         }
 
+        public SerializedItem SerializePingProbeMessage()
+        {
+            return new SerializedItem
+            {
+                Text = "2probe"
+            };
+        }
+
         public SerializedItem SerializeUpgradeMessage()
         {
             return new SerializedItem
@@ -316,6 +324,7 @@ private static void ReadMessage(IMessage message, EngineIO eio, string text)
                 case MessageType.Ping:
                     break;
                 case MessageType.Pong:
+                    ReadPongMessage(message, text);
                     break;
                 case MessageType.Connected:
                     ReadConnectedMessage(message, text, eio);
@@ -343,6 +352,11 @@ private static void ReadMessage(IMessage message, EngineIO eio, string text)
             }
         }
 
+        private static void ReadPongMessage(IMessage message, string text)
+        {
+            message.ReceivedText = text;
+        }
+
         private static void ReadOpenedMessage(IMessage message, string text)
         {
             // TODO: Should deserializing to existing object

src/SocketIOClient/SocketIO.cs

@@ -320,16 +320,51 @@ private async Task UpgradeToWebSocket(IMessage openedMessage)
         {
             var options = NewTransportOptions();
             options.OpenedMessage = openedMessage;
+
             for (var i = 0; i < 3; i++)
             {
-                var transport = (WebSocketTransport)NewTransport(TransportProtocol.WebSocket, options);
-                using var cts = new CancellationTokenSource(Options.ConnectionTimeout);
+                WebSocketTransport transport = (WebSocketTransport)NewTransport(TransportProtocol.WebSocket, options);
+
+                TaskCompletionSource<bool> pongProbeTcs = new();
+                using CancellationTokenSource connectionTimeoutCts = new(Options.ConnectionTimeout);
+                CancellationToken connectionTimeoutToken = connectionTimeoutCts.Token;
+
+                connectionTimeoutToken.Register(() =>
+                {
+                    pongProbeTcs.TrySetException(new TimeoutException("The upgrade operation has timed out!"));
+                });
+
                 try
                 {
-                    await transport.ConnectAsync(cts.Token).ConfigureAwait(false);
-                    var message = Serializer.SerializeUpgradeMessage();
+                    await transport.ConnectAsync(connectionTimeoutToken).ConfigureAwait(false);
+
+                    void pongProbeHandler(IMessage msg)
+                    {
+                        if (msg.Type == MessageType.Pong && msg.ReceivedText == "probe")
+                        {
+                            pongProbeTcs.SetResult(true);
+                        }
+                        else
+                        {
+                            pongProbeTcs.SetException(new Exception($"Unexpected handshake response: '{msg.Type}'"));
+                        }
+                    }
+
+                    transport.OnReceived += pongProbeHandler;
+
+                    SerializedItem message = Serializer.SerializePingProbeMessage();
+
+                    await transport
+                        .SendAsync(new List<SerializedItem> { message }, connectionTimeoutToken)
+                        .ConfigureAwait(false);
+
+                    await pongProbeTcs.Task;
+
+                    transport.OnReceived -= pongProbeHandler;
+
+                    message = Serializer.SerializeUpgradeMessage();
                     await transport
-                        .SendAsync(new List<SerializedItem> { message }, cts.Token)
+                        .SendAsync(new List<SerializedItem> { message }, connectionTimeoutToken)
                         .ConfigureAwait(false);
 
                     Transport.Dispose();