From b783e6bcb318ebd1fd25ca603532b0263e6f5424 Mon Sep 17 00:00:00 2001
From: AeonLucid <aeonlucid@outlook.com>
Date: Sun, 1 Nov 2020 21:08:54 +0100
Subject: [PATCH] Packet recorder add disconnect reason

---
 src/Impostor.Server/Recorder/ClientRecorder.cs | 4 ++--
 src/Impostor.Server/Recorder/PacketRecorder.cs | 3 ++-
 src/Impostor.Tools.ServerReplay/Program.cs     | 9 ++++++++-
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/src/Impostor.Server/Recorder/ClientRecorder.cs b/src/Impostor.Server/Recorder/ClientRecorder.cs
index 94b320e..03ad965 100644
--- a/src/Impostor.Server/Recorder/ClientRecorder.cs
+++ b/src/Impostor.Server/Recorder/ClientRecorder.cs
@@ -79,8 +79,8 @@ namespace Impostor.Server.Recorder
 
         public override async ValueTask HandleDisconnectAsync(string reason)
         {
-            await _recorder.WriteDisconnectAsync(this);
+            await _recorder.WriteDisconnectAsync(this, reason);
             await base.HandleDisconnectAsync(reason);
         }
     }
-}
\ No newline at end of file
+}
diff --git a/src/Impostor.Server/Recorder/PacketRecorder.cs b/src/Impostor.Server/Recorder/PacketRecorder.cs
index 652496c..f3ff48b 100644
--- a/src/Impostor.Server/Recorder/PacketRecorder.cs
+++ b/src/Impostor.Server/Recorder/PacketRecorder.cs
@@ -55,7 +55,7 @@ namespace Impostor.Server.Recorder
             }
         }
 
-        public async Task WriteDisconnectAsync(ClientRecorder client)
+        public async Task WriteDisconnectAsync(ClientRecorder client, string reason)
         {
             _logger.LogTrace("Writing Disconnect.");
 
@@ -65,6 +65,7 @@ namespace Impostor.Server.Recorder
             {
                 WriteHeader(context, RecordedPacketType.Disconnect);
                 WriteClient(context, client, false);
+                context.Writer.Write(reason);
                 WriteLength(context);
 
                 await WriteAsync(context.Stream);
diff --git a/src/Impostor.Tools.ServerReplay/Program.cs b/src/Impostor.Tools.ServerReplay/Program.cs
index af2ffd9..f656801 100644
--- a/src/Impostor.Tools.ServerReplay/Program.cs
+++ b/src/Impostor.Tools.ServerReplay/Program.cs
@@ -146,7 +146,14 @@ namespace Impostor.Tools.ServerReplay
                     break;
 
                 case RecordedPacketType.Disconnect:
-                    await Connections[clientId].Client!.HandleDisconnectAsync(null);
+                    string reason = null;
+
+                    if (reader.BaseStream.Position < reader.BaseStream.Length)
+                    {
+                        reason = reader.ReadString();
+                    }
+
+                    await Connections[clientId].Client!.HandleDisconnectAsync(reason);
                     Connections.Remove(clientId);
                     break;
 
-- 
2.39.5