From a56eed2680d807987bd993a54bbd38e76a228ac3 Mon Sep 17 00:00:00 2001
From: js6pak <kubastaron@hotmail.com>
Date: Sun, 11 Apr 2021 14:57:29 +0200
Subject: [PATCH] Use actual disconnect message in ClientBase#DisconnectAsync

---
 src/Impostor.Api/Net/IHazelConnection.cs      |  3 +-
 .../Net/Messages/S2C/MessageDisconnect.cs     | 53 +++++++++++++++++++
 src/Impostor.Server/Net/ClientBase.cs         | 11 ++--
 .../Net/Hazel/HazelConnection.cs              |  4 +-
 .../Mocks/MockHazelConnection.cs              |  2 +-
 5 files changed, 61 insertions(+), 12 deletions(-)
 create mode 100644 src/Impostor.Api/Net/Messages/S2C/MessageDisconnect.cs

diff --git a/src/Impostor.Api/Net/IHazelConnection.cs b/src/Impostor.Api/Net/IHazelConnection.cs
index e2a9cfbf9..c3ce51e15 100644
--- a/src/Impostor.Api/Net/IHazelConnection.cs
+++ b/src/Impostor.Api/Net/IHazelConnection.cs
@@ -35,7 +35,8 @@ public interface IHazelConnection
         ///     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
index 000000000..8dd5c7f94
--- /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 9874a7171..b9488c246 100644
--- a/src/Impostor.Server/Net/ClientBase.cs
+++ b/src/Impostor.Server/Net/ClientBase.cs
@@ -52,15 +52,10 @@ public async ValueTask DisconnectAsync(DisconnectReason reason, string? message
                 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 e663dec1b..c59a91c2d 100644
--- a/src/Impostor.Server/Net/Hazel/HazelConnection.cs
+++ b/src/Impostor.Server/Net/Hazel/HazelConnection.cs
@@ -32,9 +32,9 @@ public ValueTask SendAsync(IMessageWriter writer)
             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 327296f12..c611c2949 100644
--- a/src/Impostor.Tools.ServerReplay/Mocks/MockHazelConnection.cs
+++ b/src/Impostor.Tools.ServerReplay/Mocks/MockHazelConnection.cs
@@ -23,7 +23,7 @@ public ValueTask SendAsync(IMessageWriter writer)
             return ValueTask.CompletedTask;
         }
 
-        public ValueTask DisconnectAsync(string reason)
+        public ValueTask DisconnectAsync(string reason, IMessageWriter writer = null)
         {
             return ValueTask.CompletedTask;
         }