add connection id in handshake response and add echo in ping. (#1821)

Azure · Aug 24, 2023 · d050dd2 · d050dd2
1 parent 06c5071
commit d050dd2
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 4 deletions.
diff --git a/src/Microsoft.Azure.SignalR.Common/ServiceConnections/ServiceConnectionBase.cs b/src/Microsoft.Azure.SignalR.Common/ServiceConnections/ServiceConnectionBase.cs
@@ -299,6 +299,10 @@ protected Task OnServiceErrorAsync(ServiceErrorMessage serviceErrorMessage)
 
         protected virtual Task OnPingMessageAsync(PingMessage pingMessage)
         {
+            if (RuntimeServicePingMessage.IsEchoMessage(pingMessage))
+            {
+                return WriteAsync(pingMessage);
+            }
             if (RuntimeServicePingMessage.TryGetOffline(pingMessage, out var instanceId))
             {
                 Log.ReceivedInstanceOfflinePing(Logger, instanceId);

diff --git a/src/Microsoft.Azure.SignalR.Common/ServiceMessages/RuntimeServicePingMessages.cs b/src/Microsoft.Azure.SignalR.Common/ServiceMessages/RuntimeServicePingMessages.cs
@@ -9,6 +9,7 @@ namespace Microsoft.Azure.SignalR
 {
     internal static class RuntimeServicePingMessage
     {
+        private const string EchoKey = "echo";
         private const string OfflineKey = "offline";
         private const string TargetKey = "target";
         private const string StatusKey = "status";
@@ -52,6 +53,11 @@ internal static class RuntimeServicePingMessage
         private static readonly ServicePingMessage ServersTag =
             new ServicePingMessage { Messages = new[] { ServersKey, string.Empty } };
 
+        public static bool IsEchoMessage(this ServicePingMessage ping)
+        {
+            return TryGetValue(ping, EchoKey, out _);
+        }
+
         public static bool TryGetMessageLogEnableFlag(this ServicePingMessage ping, out bool enableMessageLog)
         {
             if (TryGetValue(ping, DiagnosticLogsMessagingTypeKey, out var value))

diff --git a/src/Microsoft.Azure.SignalR.Protocols/ServiceMessage.cs b/src/Microsoft.Azure.SignalR.Protocols/ServiceMessage.cs
@@ -285,6 +285,11 @@ public class HandshakeResponseMessage : ExtensibleServiceMessage
         /// </summary>
         public string ErrorMessage { get; set; }
 
+        /// <summary>
+        /// Gets or sets the id of this connection.
+        /// </summary>
+        public string ConnectionId { get; set; }
+
         /// <summary>
         /// Initializes a new instance of the <see cref="HandshakeResponseMessage"/> class.
         /// </summary>

diff --git a/src/Microsoft.Azure.SignalR.Protocols/ServiceProtocol.cs b/src/Microsoft.Azure.SignalR.Protocols/ServiceProtocol.cs
@@ -319,10 +319,11 @@ private static void WriteHandshakeRequestMessage(ref MessagePackWriter writer, H
 
         private static void WriteHandshakeResponseMessage(ref MessagePackWriter writer, HandshakeResponseMessage message)
         {
-            writer.WriteArrayHeader(3);
+            writer.WriteArrayHeader(4);
             writer.Write(ServiceProtocolConstants.HandshakeResponseType);
             writer.Write(message.ErrorMessage);
             message.WriteExtensionMembers(ref writer);
+            writer.Write(message.ConnectionId);
         }
 
         private static void WriteAccessKeyRequestMessage(ref MessagePackWriter writer, AccessKeyRequestMessage message)
@@ -800,6 +801,10 @@ private static HandshakeResponseMessage CreateHandshakeResponseMessage(ref Messa
             {
                 result.ReadExtensionMembers(ref reader);
             }
+            if (arrayLength >= 4)
+            {
+                result.ConnectionId = ReadString(ref reader, "connectionId");
+            }
             return result;
         }
 

diff --git a/test/Microsoft.Azure.SignalR.Protocols.Tests/ServiceMessageEqualityComparer.cs b/test/Microsoft.Azure.SignalR.Protocols.Tests/ServiceMessageEqualityComparer.cs
@@ -118,7 +118,8 @@ private bool HandshakeRequestMessagesEqual(HandshakeRequestMessage x, HandshakeR
 
         private bool HandshakeResponseMessagesEqual(HandshakeResponseMessage x, HandshakeResponseMessage y)
         {
-            return StringEqual(x.ErrorMessage, y.ErrorMessage);
+            return StringEqual(x.ErrorMessage, y.ErrorMessage) &&
+                StringEqual(x.ConnectionId, y.ConnectionId);
         }
 
         private bool AccessKeyRequestMessageEqual(AccessKeyRequestMessage x, AccessKeyRequestMessage y)

diff --git a/test/Microsoft.Azure.SignalR.Protocols.Tests/ServiceProtocolFacts.cs b/test/Microsoft.Azure.SignalR.Protocols.Tests/ServiceProtocolFacts.cs
@@ -90,10 +90,18 @@ public static IEnumerable<object[]> TestParseOldData
                 name: "HandshakeResponse_NoOptionalField",
                 message: new HandshakeResponseMessage(),
                 binary: "kgKg"),
+            new ProtocolTestData(
+                name: "HandshakeResponse_NoConnectionId",
+                message: new HandshakeResponseMessage(),
+                binary: "kwKggA=="),
             new ProtocolTestData(
                 name: "HandshakeResponseWithError_NoOptionalField",
                 message: new HandshakeResponseMessage("Version mismatch."),
                 binary: "kgKxVmVyc2lvbiBtaXNtYXRjaC4="),
+            new ProtocolTestData(
+                name: "HandshakeResponseWithError_NoConnectionId",
+                message: new HandshakeResponseMessage("Version mismatch."),
+                binary: "kwKxVmVyc2lvbiBtaXNtYXRjaC6A"),
             new ProtocolTestData(
                 name: "OpenConnection_NoOptionalField",
                 message: new OpenConnectionMessage("conn1", null),
@@ -280,11 +288,15 @@ public static IEnumerable<object[]> TestParseOldData
             new ProtocolTestData(
                 name: "HandshakeResponse",
                 message: new HandshakeResponseMessage(),
-                binary: "kwKggA=="),
+                binary: "lAKggMA="),
             new ProtocolTestData(
                 name: "HandshakeResponseWithError",
                 message: new HandshakeResponseMessage("Version mismatch."),
-                binary: "kwKxVmVyc2lvbiBtaXNtYXRjaC6A"),
+                binary: "lAKxVmVyc2lvbiBtaXNtYXRjaC6AwA=="),
+            new ProtocolTestData(
+                name: "HandshakeResponseWithConnectionId",
+                message: new HandshakeResponseMessage() { ConnectionId = "abc" },
+                binary: "lAKggKNhYmM="),
             new ProtocolTestData(
                 name: "AccessKeyRequestMessage",
                 message: new AccessKeyRequestMessage("token"),