Use actual disconnect message in ClientBase#DisconnectAsync

diff --git a/src/Impostor.Api/Net/IHazelConnection.cs b/src/Impostor.Api/Net/IHazelConnection.cs
index e2a9cfbf9abecf54abeba544584479b8b48e0d68..c3ce51e159761428b781deaafafff058af1d7c5e 100644 (file)
--- a/src/Impostor.Api/Net/IHazelConnection.cs
+++ b/src/Impostor.Api/Net/IHazelConnection.cs
@@ -35,7 +35,8 @@ namespace Impostor.Api.Net
         ///     Disconnects the client and invokes the disconnect handler.
         /// </summary>
         /// <param name="reason">A reason.</param>
+        /// <param name="writer">A message to send with disconnect packet.</param>
         /// <returns>Task that must be awaited.</returns>
-        ValueTask DisconnectAsync(string? reason);
+        ValueTask DisconnectAsync(string? reason, IMessageWriter? writer = null);
     }
 }

diff --git a/src/Impostor.Api/Net/Messages/S2C/MessageDisconnect.cs b/src/Impostor.Api/Net/Messages/S2C/MessageDisconnect.cs
new file mode 100644 (file)
index 0000000..8dd5c7f
--- /dev/null
+++ b/src/Impostor.Api/Net/Messages/S2C/MessageDisconnect.cs
@@ -0,0 +1,53 @@
+using System;
+using Impostor.Api.Innersloth;
+
+namespace Impostor.Api.Net.Messages.S2C
+{
+    public class MessageDisconnect
+    {
+        public static void Serialize(IMessageWriter writer, bool hasReason, DisconnectReason? reason, string? message)
+        {
+            writer.Write(hasReason);
+
+            if (hasReason)
+            {
+                if (reason == null)
+                {
+                    throw new ArgumentNullException(nameof(reason));
+                }
+
+                writer.StartMessage(0);
+                writer.Write((byte)reason);
+
+                if (reason == DisconnectReason.Custom)
+                {
+                    if (message == null)
+                    {
+                        throw new ArgumentNullException(nameof(message));
+                    }
+
+                    writer.Write(message);
+                }
+
+                writer.EndMessage();
+            }
+        }
+
+        public static void Deserialize(IMessageReader reader, out bool hasReason, out DisconnectReason? reason, out string? message)
+        {
+            hasReason = reader.ReadBoolean();
+
+            if (hasReason)
+            {
+                var inner = reader.ReadMessage();
+                reason = (DisconnectReason)inner.ReadByte();
+                message = reason == DisconnectReason.Custom ? inner.ReadString() : null;
+            }
+            else
+            {
+                reason = null;
+                message = null;
+            }
+        }
+    }
+}

diff --git a/src/Impostor.Server/Net/ClientBase.cs b/src/Impostor.Server/Net/ClientBase.cs
index 9874a71716934b3f41fef7e5d1142c33e69e2843..b9488c2465c34075616c1f1f6a1aa7a0deb02965 100644 (file)
--- a/src/Impostor.Server/Net/ClientBase.cs
+++ b/src/Impostor.Server/Net/ClientBase.cs
@@ -52,15 +52,10 @@ namespace Impostor.Server.Net
                 return;
             }
 
-            using var packet = MessageWriter.Get(MessageType.Reliable);
-            Message01JoinGameS2C.SerializeError(packet, false, reason, message);
+            using var writer = MessageWriter.Get();
+            MessageDisconnect.Serialize(writer, true, reason, message);
 
-            await Connection.SendAsync(packet);
-
-            // Need this to show the correct message, otherwise it shows a generic disconnect message.
-            await Task.Delay(TimeSpan.FromMilliseconds(250));
-
-            await Connection.DisconnectAsync(message ?? reason.ToString());
+            await Connection.DisconnectAsync(message ?? reason.ToString(), writer);
         }
     }
 }

diff --git a/src/Impostor.Server/Net/Hazel/HazelConnection.cs b/src/Impostor.Server/Net/Hazel/HazelConnection.cs
index e663dec1bb28e6773179ff7e279c21f2a75c60d5..c59a91c2d1761104a367b214582d071f6939b924 100644 (file)
--- a/src/Impostor.Server/Net/Hazel/HazelConnection.cs
+++ b/src/Impostor.Server/Net/Hazel/HazelConnection.cs
@@ -32,9 +32,9 @@ namespace Impostor.Server.Net.Hazel
             return InnerConnection.SendAsync(writer);
         }
 
-        public ValueTask DisconnectAsync(string? reason)
+        public ValueTask DisconnectAsync(string? reason, IMessageWriter? writer = null)
         {
-            return InnerConnection.Disconnect(reason);
+            return InnerConnection.Disconnect(reason, writer as MessageWriter);
         }
 
         public void DisposeInnerConnection()

diff --git a/src/Impostor.Tools.ServerReplay/Mocks/MockHazelConnection.cs b/src/Impostor.Tools.ServerReplay/Mocks/MockHazelConnection.cs
index 327296f12c53fec4841e7196034fc31172e67cf8..c611c2949e4bec9267af3e65d3a1a5a04d7badaa 100644 (file)
--- a/src/Impostor.Tools.ServerReplay/Mocks/MockHazelConnection.cs
+++ b/src/Impostor.Tools.ServerReplay/Mocks/MockHazelConnection.cs
@@ -23,7 +23,7 @@ namespace Impostor.Tools.ServerReplay.Mocks
             return ValueTask.CompletedTask;
         }
 
-        public ValueTask DisconnectAsync(string reason)
+        public ValueTask DisconnectAsync(string reason, IMessageWriter writer = null)
         {
             return ValueTask.CompletedTask;
         }