From 85a2fce6c12033700947ec59aa66ed0a21181a4f Mon Sep 17 00:00:00 2001
From: js6pak <kubastaron@hotmail.com>
Date: Fri, 3 Jun 2022 20:17:59 +0200
Subject: [PATCH] Fix non custom disconnect reasons

---
 .../Extensions/IHazelConnectionExtensions.cs             | 9 +++++----
 src/Impostor.Server/Net/ClientBase.cs                    | 2 +-
 src/Impostor.Server/Net/Manager/ClientManager.cs         | 6 +++---
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/src/Impostor.Server/Extensions/IHazelConnectionExtensions.cs b/src/Impostor.Server/Extensions/IHazelConnectionExtensions.cs
index d46e370..8fc5340 100644
--- a/src/Impostor.Server/Extensions/IHazelConnectionExtensions.cs
+++ b/src/Impostor.Server/Extensions/IHazelConnectionExtensions.cs
@@ -12,9 +12,10 @@ namespace Impostor.Server
         /// Disconnect a connection using a custom message.
         /// </summary>
         /// <param name="connection">The connection to disconnect.</param>
-        /// <param name="message">The message to disconnect with.</param>
+        /// <param name="reason">The reason to disconnect with.</param>
+        /// <param name="message">The custom message to disconnect with if <paramref name="reason"/> is <see cref="DisconnectReason.Custom"/>.</param>
         /// <returns>Task that should be awaited to ensure disconnection.</returns>
-        public static async ValueTask CustomDisconnectAsync(this IHazelConnection connection, string message)
+        public static async ValueTask CustomDisconnectAsync(this IHazelConnection connection, DisconnectReason reason, string? message = null)
         {
             if (!connection.IsConnected)
             {
@@ -22,9 +23,9 @@ namespace Impostor.Server
             }
 
             using var writer = MessageWriter.Get();
-            MessageDisconnect.Serialize(writer, true, DisconnectReason.Custom, message);
+            MessageDisconnect.Serialize(writer, true, reason, message);
 
-            await connection.DisconnectAsync(DisconnectReason.Custom.ToString(), writer);
+            await connection.DisconnectAsync(reason.ToString(), writer);
         }
     }
 }
diff --git a/src/Impostor.Server/Net/ClientBase.cs b/src/Impostor.Server/Net/ClientBase.cs
index 15dfe91..0a321f8 100644
--- a/src/Impostor.Server/Net/ClientBase.cs
+++ b/src/Impostor.Server/Net/ClientBase.cs
@@ -53,7 +53,7 @@ namespace Impostor.Server.Net
 
         public async ValueTask DisconnectAsync(DisconnectReason reason, string? message = null)
         {
-            await Connection.CustomDisconnectAsync(message ?? reason.ToString());
+            await Connection.CustomDisconnectAsync(reason, message);
         }
     }
 }
diff --git a/src/Impostor.Server/Net/Manager/ClientManager.cs b/src/Impostor.Server/Net/Manager/ClientManager.cs
index 0c498e6..cca6d85 100644
--- a/src/Impostor.Server/Net/Manager/ClientManager.cs
+++ b/src/Impostor.Server/Net/Manager/ClientManager.cs
@@ -85,19 +85,19 @@ namespace Impostor.Server.Net.Manager
                     _ => throw new ArgumentOutOfRangeException(),
                 };
 
-                await connection.CustomDisconnectAsync(message);
+                await connection.CustomDisconnectAsync(DisconnectReason.Custom, message);
                 return;
             }
 
             if (name.Length > 10)
             {
-                await connection.CustomDisconnectAsync(DisconnectMessages.UsernameLength);
+                await connection.CustomDisconnectAsync(DisconnectReason.Custom, DisconnectMessages.UsernameLength);
                 return;
             }
 
             if (string.IsNullOrWhiteSpace(name))
             {
-                await connection.CustomDisconnectAsync(DisconnectMessages.UsernameIllegalCharacters);
+                await connection.CustomDisconnectAsync(DisconnectReason.Custom, DisconnectMessages.UsernameIllegalCharacters);
                 return;
             }
 
-- 
2.39.5