Introduce ServerDisconnectPacket (#774)

mibac138 · web-flow · commit 8e97cd38a3a5 · 2025-10-27T12:16:54.000-05:00
* Add PacketHandlerClassAttribute

Provides a cleaner way of dealing with common packets across different connection states

* Introduce ServerDisconnectPacket

Reduces reliance on a LiteNetLib-specific feature, unifies disconnection
process across different transports, and allows to arbitrarily evolve
the new typed packet.
diff --git a/Source/Client/Networking/NetworkingInMemory.cs b/Source/Client/Networking/NetworkingInMemory.cs
@@ -1,5 +1,5 @@
-using Multiplayer.Common;
 using System;
+using Multiplayer.Common;
 using Verse;
 
 namespace Multiplayer.Client.Networking
@@ -46,7 +46,7 @@ protected override void SendRaw(byte[] raw, bool reliable = true)
             });
         }
 
-        public override void Close(MpDisconnectReason reason, byte[] data = null)
+        protected override void OnClose()
         {
         }
 
diff --git a/Source/Client/Networking/NetworkingLiteNet.cs b/Source/Client/Networking/NetworkingLiteNet.cs
@@ -44,13 +44,14 @@ public static ClientLiteNetConnection Connect(string address, int port)
 
         public void OnDisconnect(MpDisconnectReason reason, ByteReader data)
         {
+            if (State == ConnectionStateEnum.Disconnected) return;
             ConnectionStatusListeners.TryNotifyAll_Disconnected(SessionDisconnectInfo.From(reason, data));
             Multiplayer.StopMultiplayer();
         }
 
-        public override void Close(MpDisconnectReason reason, byte[] data)
+        protected override void OnClose()
         {
-            base.Close(reason, data);
+            base.OnClose();
             netManager.Stop();
         }
 
@@ -81,7 +82,8 @@ public void OnPeerDisconnected(NetPeer peer, DisconnectInfo info)
                 MpDisconnectReason reason;
                 ByteReader reader;
 
-                if (info.AdditionalData.IsNull)
+                // Fallback: should generally be handled by ClientBaseState.HandleDisconnected.
+                if (info.AdditionalData.IsNull || info.AdditionalData.AvailableBytes == 0)
                 {
                     if (info.Reason is DisconnectReason.DisconnectPeerCalled or DisconnectReason.RemoteConnectionClose)
                         reason = MpDisconnectReason.Generic;
diff --git a/Source/Client/Networking/NetworkingSteam.cs b/Source/Client/Networking/NetworkingSteam.cs
@@ -38,18 +38,13 @@ public void SendRawSteam(byte[] raw, bool reliable)
             ServerLog.Error($"Failed to send packet (len: {raw.Length}): {hex}");
         }
 
-        public override void Close(MpDisconnectReason reason, byte[] data = null)
-        {
-            if (State != ConnectionStateEnum.ClientSteam)
-                Send(Packets.Special_Steam_Disconnect, GetDisconnectBytes(reason, data));
-        }
-
         public abstract void OnError(EP2PSessionError error);
 
-        public override string ToString()
+        protected override void OnClose()
         {
-            return $"SteamP2P ({remoteId}:{username})";
         }
+
+        public override string ToString() => $"SteamP2P ({remoteId}:{username})";
     }
 
     public class SteamClientConn(CSteamID remoteId) : SteamBaseConn(remoteId, RandomChannelId(), 0), ITickableConnection
@@ -66,30 +61,12 @@ public void Tick()
             }
         }
 
-        protected override void HandleReceiveMsg(int msgId, int fragState, ByteReader reader, bool reliable)
-        {
-            if (msgId == (int)Packets.Special_Steam_Disconnect)
-            {
-                var info = SessionDisconnectInfo.From(reader.ReadEnum<MpDisconnectReason>(), reader);
-                OnDisconnect(info);
-                return;
-            }
-
-            base.HandleReceiveMsg(msgId, fragState, reader, reliable);
-        }
-
         public override void OnError(EP2PSessionError error)
         {
             var title = error == EP2PSessionError.k_EP2PSessionErrorTimeout
                 ? "MpSteamTimedOut".Translate()
                 : "MpSteamGenericError".Translate();
-
-            OnDisconnect(new SessionDisconnectInfo { titleTranslated = title });
-        }
-
-        private void OnDisconnect(SessionDisconnectInfo info)
-        {
-            ConnectionStatusListeners.TryNotifyAll_Disconnected(info);
+            ConnectionStatusListeners.TryNotifyAll_Disconnected(new SessionDisconnectInfo { titleTranslated = title });
             Multiplayer.StopMultiplayer();
         }
     }
@@ -104,17 +81,6 @@ public override void Send(Packets id, byte[] message, bool reliable = true)
             base.Send(id, message, reliable);
         }
 
-        protected override void HandleReceiveMsg(int msgId, int fragState, ByteReader reader, bool reliable)
-        {
-            if (msgId == (int)Packets.Special_Steam_Disconnect)
-            {
-                OnDisconnect();
-                return;
-            }
-
-            base.HandleReceiveMsg(msgId, fragState, reader, reliable);
-        }
-
         public override void OnError(EP2PSessionError error)
         {
             OnDisconnect();
diff --git a/Source/Client/Networking/State/ClientBaseState.cs b/Source/Client/Networking/State/ClientBaseState.cs
@@ -1,3 +1,4 @@
+using Multiplayer.Client.Networking;
 using Multiplayer.Common;
 using Multiplayer.Common.Networking.Packet;
 
@@ -7,15 +8,17 @@ public abstract class ClientBaseState(ConnectionBase connection) : MpConnectionS
 {
     protected MultiplayerSession Session => Multiplayer.session;
 
-    protected void HandleKeepAlive(ServerKeepAlivePacket packet)
+    [TypedPacketHandler]
+    public void HandleKeepAlive(ServerKeepAlivePacket packet)
     {
         int ticksBehind = TickPatch.tickUntil - TickPatch.Timer;
 
         connection.Send(new ClientKeepAlivePacket(packet.id, ticksBehind, TickPatch.Simulating, TickPatch.workTicks),
             false);
     }
 
-    protected void HandleTimeControl(ServerTimeControlPacket packet)
+    [TypedPacketHandler]
+    public void HandleTimeControl(ServerTimeControlPacket packet)
     {
         if (Multiplayer.session.remoteTickUntil >= packet.tickUntil) return;
 
@@ -24,4 +27,12 @@ protected void HandleTimeControl(ServerTimeControlPacket packet)
         Multiplayer.session.remoteSentCmds = packet.sentCmds;
         Multiplayer.session.ProcessTimeControl();
     }
+
+    [TypedPacketHandler]
+    public void HandleDisconnected(ServerDisconnectPacket packet)
+    {
+        ConnectionStatusListeners.TryNotifyAll_Disconnected(SessionDisconnectInfo.From(packet.reason,
+            new ByteReader(packet.data)));
+        Multiplayer.StopMultiplayer();
+    }
 }
diff --git a/Source/Client/Networking/State/ClientJoiningState.cs b/Source/Client/Networking/State/ClientJoiningState.cs
@@ -9,12 +9,17 @@
 
 namespace Multiplayer.Client
 {
+
+    [PacketHandlerClass(inheritHandlers: false)]
     public class ClientJoiningState : ClientBaseState
     {
         public ClientJoiningState(ConnectionBase connection) : base(connection)
         {
         }
 
+        [TypedPacketHandler]
+        public new void HandleDisconnected(ServerDisconnectPacket packet) => base.HandleDisconnected(packet);
+
         public override void StartState()
         {
             connection.Send(ClientProtocolPacket.Current());
diff --git a/Source/Client/Networking/State/ClientLoadingState.cs b/Source/Client/Networking/State/ClientLoadingState.cs
@@ -4,7 +4,6 @@
 using Ionic.Zlib;
 using Multiplayer.Client.Saving;
 using Multiplayer.Common;
-using Multiplayer.Common.Networking.Packet;
 using Verse;
 
 namespace Multiplayer.Client;
@@ -15,6 +14,7 @@ public enum LoadingState
     Downloading
 }
 
+[PacketHandlerClass(inheritHandlers: true)]
 public class ClientLoadingState(ConnectionBase connection) : ClientBaseState(connection)
 {
     public LoadingState subState = LoadingState.Waiting;
@@ -139,10 +139,4 @@ public void HandleWorldData(ByteReader data)
         Log.Message($"Loaded game in {loadingMs}ms");
         connection.ChangeState(ConnectionStateEnum.ClientPlaying);
     }
-
-    [TypedPacketHandler]
-    public new void HandleKeepAlive(ServerKeepAlivePacket packet) => base.HandleKeepAlive(packet);
-
-    [TypedPacketHandler]
-    public new void HandleTimeControl(ServerTimeControlPacket packet) => base.HandleTimeControl(packet);
 }
diff --git a/Source/Client/Networking/State/ClientPlayingState.cs b/Source/Client/Networking/State/ClientPlayingState.cs
@@ -9,14 +9,9 @@
 
 namespace Multiplayer.Client
 {
+    [PacketHandlerClass(inheritHandlers: true)]
     public class ClientPlayingState(ConnectionBase connection) : ClientBaseState(connection)
     {
-        [TypedPacketHandler]
-        public new void HandleKeepAlive(ServerKeepAlivePacket packet) => base.HandleKeepAlive(packet);
-
-        [TypedPacketHandler]
-        public new void HandleTimeControl(ServerTimeControlPacket packet) => base.HandleTimeControl(packet);
-
         [TypedPacketHandler]
         public void HandleCommand(ServerCommandPacket packet)
         {
diff --git a/Source/Client/Saving/ReplayConnection.cs b/Source/Client/Saving/ReplayConnection.cs
@@ -30,7 +30,7 @@ public override void HandleReceiveRaw(ByteReader data, bool reliable)
     {
     }
 
-    public override void Close(MpDisconnectReason reason, byte[] data)
+    protected override void OnClose()
     {
     }
 }
diff --git a/Source/Client/Session/SessionDisconnectInfo.cs b/Source/Client/Session/SessionDisconnectInfo.cs
@@ -94,7 +94,7 @@ public static SessionDisconnectInfo From(MpDisconnectReason reason, ByteReader r
         disconnectInfo.titleTranslated ??= titleKey?.Translate();
         disconnectInfo.descTranslated ??= descKey?.Translate();
 
-        Log.Message($"Processed disconnect packet. Title: {disconnectInfo.titleTranslated} ({titleKey}), " +
+        Log.Message($"Processed disconnect packet ({reason}). Title: {disconnectInfo.titleTranslated} ({titleKey}), " +
                     $"description: {disconnectInfo.descTranslated} ({descKey})");
 
         return disconnectInfo;
diff --git a/Source/Common/Networking/ConnectionBase.cs b/Source/Common/Networking/ConnectionBase.cs
@@ -264,7 +264,15 @@ private void ExecuteMessageHandler(PacketHandlerInfo handler, Packets packet, By
             }
         }
 
-        public abstract void Close(MpDisconnectReason reason, byte[]? data = null);
+        public void Close(MpDisconnectReason reason, byte[]? data = null)
+        {
+            // State.IsServer check only used when disconnecting from a self-hosted local server
+            if (State != ConnectionStateEnum.Disconnected && State.IsServer())
+                Send(new ServerDisconnectPacket { reason = reason, data = data ?? [] });
+            OnClose();
+        }
+
+        protected abstract void OnClose();
 
         /// Invoked after a keep alive timer arrives. Only used by the server
         public virtual void OnKeepAliveArrived(bool idMatched)
diff --git a/Source/Common/Networking/ConnectionStateEnum.cs b/Source/Common/Networking/ConnectionStateEnum.cs
@@ -15,3 +15,12 @@ public enum ConnectionStateEnum : byte
     Count,
     Disconnected
 }
+
+public static class ConnectionStateEnumExt
+{
+    public static bool IsClient(this ConnectionStateEnum state) =>
+        state is >= ConnectionStateEnum.ClientJoining and <= ConnectionStateEnum.ClientSteam;
+
+    public static bool IsServer(this ConnectionStateEnum state) =>
+        state is >= ConnectionStateEnum.ServerJoining and <= ConnectionStateEnum.ServerSteam;
+}
diff --git a/Source/Common/Networking/LiteNetConnection.cs b/Source/Common/Networking/LiteNetConnection.cs
@@ -14,10 +14,10 @@ protected override void SendRaw(byte[] raw, bool reliable)
                 ServerLog.Error($"SendRaw() called with invalid connection state ({peer}): {peer.ConnectionState}");
         }
 
-        public override void Close(MpDisconnectReason reason, byte[]? data)
+        protected override void OnClose()
         {
             peer.NetManager.TriggerUpdate(); // todo: is this needed?
-            peer.NetManager.DisconnectPeer(peer, GetDisconnectBytes(reason, data));
+            peer.NetManager.DisconnectPeer(peer);
         }
 
         public override void OnKeepAliveArrived(bool idMatched)
diff --git a/Source/Common/Networking/MpConnectionState.cs b/Source/Common/Networking/MpConnectionState.cs
@@ -36,7 +36,14 @@ public static void SetImplementation(ConnectionStateEnum state, Type type)
 
             stateImpls[(int)state] = type;
 
-            foreach (var method in type.GetMethods(BindingFlags.Instance | BindingFlags.Public | BindingFlags.DeclaredOnly))
+            var typeAttr = type.GetAttribute<PacketHandlerClassAttribute>();
+            if (typeAttr == null)
+                ServerLog.Log($"Packet handler {type.FullName} does not have a PacketHandlerClass attribute");
+
+            var bindingFlags = BindingFlags.Instance | BindingFlags.Public;
+            if (typeAttr?.inheritHandlers != true) bindingFlags |= BindingFlags.DeclaredOnly;
+
+            foreach (var method in type.GetMethods(bindingFlags))
             {
                 var attr = method.GetAttribute<PacketHandlerAttribute>();
                 if (attr != null) RegisterPacketHandler(state, method, attr);
diff --git a/Source/Common/Networking/Packet/DisconnectPacket.cs b/Source/Common/Networking/Packet/DisconnectPacket.cs
@@ -0,0 +1,14 @@
+namespace Multiplayer.Common.Networking.Packet;
+
+[PacketDefinition(Packets.Server_Disconnect)]
+public record struct ServerDisconnectPacket : IPacket
+{
+    public MpDisconnectReason reason;
+    public byte[] data;
+
+    public void Bind(PacketBuffer buf)
+    {
+        buf.BindEnum(ref reason);
+        buf.BindRemaining(ref data);
+    }
+}
diff --git a/Source/Common/Networking/PacketHandlerAttribute.cs b/Source/Common/Networking/PacketHandlerAttribute.cs
@@ -16,6 +16,12 @@ public class PacketHandlerAttribute(Packets packet, bool allowFragmented = false
     [AttributeUsage(AttributeTargets.Method)]
     public class TypedPacketHandlerAttribute : Attribute;
 
+    [AttributeUsage(AttributeTargets.Class)]
+    public class PacketHandlerClassAttribute(bool inheritHandlers = false) : Attribute
+    {
+        public readonly bool inheritHandlers = inheritHandlers;
+    }
+
     [MeansImplicitUse]
     [AttributeUsage(AttributeTargets.Method)]
     public class FragmentedPacketHandlerAttribute(Packets packet) : Attribute
diff --git a/Source/Common/Networking/Packets.cs b/Source/Common/Networking/Packets.cs
@@ -62,6 +62,9 @@ public enum Packets : byte
     Server_Traces,
     Server_SetFaction,
 
+    // All states (Joining, Loading, Playing)
+    Server_Disconnect,
+
     Count,
-    Special_Steam_Disconnect = 63 // Also the max packet id
+    Max = 63 // max packet id
 }

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`		`-using Multiplayer.Common;`
`2`	`1`	`using System;`
	`2`	`+using Multiplayer.Common;`
`3`	`3`	`using Verse;`
`4`	`4`
`5`	`5`	`namespace Multiplayer.Client.Networking`
`@@ -46,7 +46,7 @@ protected override void SendRaw(byte[] raw, bool reliable = true)`
`46`	`46`	`});`
`47`	`47`	`}`
`48`	`48`
`49`		`- public override void Close(MpDisconnectReason reason, byte[] data = null)`
	`49`	`+ protected override void OnClose()`
`50`	`50`	`{`
`51`	`51`	`}`
`52`	`52`
Original file line number	Diff line number	Diff line change
`@@ -44,13 +44,14 @@ public static ClientLiteNetConnection Connect(string address, int port)`
`44`	`44`
`45`	`45`	`public void OnDisconnect(MpDisconnectReason reason, ByteReader data)`
`46`	`46`	`{`
	`47`	`+ if (State == ConnectionStateEnum.Disconnected) return;`
`47`	`48`	`ConnectionStatusListeners.TryNotifyAll_Disconnected(SessionDisconnectInfo.From(reason, data));`
`48`	`49`	`Multiplayer.StopMultiplayer();`
`49`	`50`	`}`
`50`	`51`
`51`		`- public override void Close(MpDisconnectReason reason, byte[] data)`
	`52`	`+ protected override void OnClose()`
`52`	`53`	`{`
`53`		`- base.Close(reason, data);`
	`54`	`+ base.OnClose();`
`54`	`55`	`netManager.Stop();`
`55`	`56`	`}`
`56`	`57`
`@@ -81,7 +82,8 @@ public void OnPeerDisconnected(NetPeer peer, DisconnectInfo info)`
`81`	`82`	`MpDisconnectReason reason;`
`82`	`83`	`ByteReader reader;`
`83`	`84`
`84`		`- if (info.AdditionalData.IsNull)`
	`85`	`+ // Fallback: should generally be handled by ClientBaseState.HandleDisconnected.`
	`86`	`+ if (info.AdditionalData.IsNull \|\| info.AdditionalData.AvailableBytes == 0)`
`85`	`87`	`{`
`86`	`88`	`if (info.Reason is DisconnectReason.DisconnectPeerCalled or DisconnectReason.RemoteConnectionClose)`
`87`	`89`	`reason = MpDisconnectReason.Generic;`
Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+using Multiplayer.Client.Networking;`
`1`	`2`	`using Multiplayer.Common;`
`2`	`3`	`using Multiplayer.Common.Networking.Packet;`
`3`	`4`
`@@ -7,15 +8,17 @@ public abstract class ClientBaseState(ConnectionBase connection) : MpConnectionS`
`7`	`8`	`{`
`8`	`9`	`protected MultiplayerSession Session => Multiplayer.session;`
`9`	`10`
`10`		`- protected void HandleKeepAlive(ServerKeepAlivePacket packet)`
	`11`	`+ [TypedPacketHandler]`
	`12`	`+ public void HandleKeepAlive(ServerKeepAlivePacket packet)`
`11`	`13`	`{`
`12`	`14`	`int ticksBehind = TickPatch.tickUntil - TickPatch.Timer;`
`13`	`15`
`14`	`16`	`connection.Send(new ClientKeepAlivePacket(packet.id, ticksBehind, TickPatch.Simulating, TickPatch.workTicks),`
`15`	`17`	`false);`
`16`	`18`	`}`
`17`	`19`
`18`		`- protected void HandleTimeControl(ServerTimeControlPacket packet)`
	`20`	`+ [TypedPacketHandler]`
	`21`	`+ public void HandleTimeControl(ServerTimeControlPacket packet)`
`19`	`22`	`{`
`20`	`23`	`if (Multiplayer.session.remoteTickUntil >= packet.tickUntil) return;`
`21`	`24`
`@@ -24,4 +27,12 @@ protected void HandleTimeControl(ServerTimeControlPacket packet)`
`24`	`27`	`Multiplayer.session.remoteSentCmds = packet.sentCmds;`
`25`	`28`	`Multiplayer.session.ProcessTimeControl();`
`26`	`29`	`}`
	`30`	`+`
	`31`	`+ [TypedPacketHandler]`
	`32`	`+ public void HandleDisconnected(ServerDisconnectPacket packet)`
	`33`	`+ {`
	`34`	`+ ConnectionStatusListeners.TryNotifyAll_Disconnected(SessionDisconnectInfo.From(packet.reason,`
	`35`	`+ new ByteReader(packet.data)));`
	`36`	`+ Multiplayer.StopMultiplayer();`
	`37`	`+ }`
`27`	`38`	`}`
Original file line number	Diff line number	Diff line change
`@@ -9,12 +9,17 @@`
`9`	`9`
`10`	`10`	`namespace Multiplayer.Client`
`11`	`11`	`{`
	`12`	`+`
	`13`	`+ [PacketHandlerClass(inheritHandlers: false)]`
`12`	`14`	`public class ClientJoiningState : ClientBaseState`
`13`	`15`	`{`
`14`	`16`	`public ClientJoiningState(ConnectionBase connection) : base(connection)`
`15`	`17`	`{`
`16`	`18`	`}`
`17`	`19`
	`20`	`+ [TypedPacketHandler]`
	`21`	`+ public new void HandleDisconnected(ServerDisconnectPacket packet) => base.HandleDisconnected(packet);`
	`22`	`+`
`18`	`23`	`public override void StartState()`
`19`	`24`	`{`
`20`	`25`	`connection.Send(ClientProtocolPacket.Current());`
Original file line number	Diff line number	Diff line change
`@@ -4,7 +4,6 @@`
`4`	`4`	`using Ionic.Zlib;`
`5`	`5`	`using Multiplayer.Client.Saving;`
`6`	`6`	`using Multiplayer.Common;`
`7`		`-using Multiplayer.Common.Networking.Packet;`
`8`	`7`	`using Verse;`
`9`	`8`
`10`	`9`	`namespace Multiplayer.Client;`
`@@ -15,6 +14,7 @@ public enum LoadingState`
`15`	`14`	`Downloading`
`16`	`15`	`}`
`17`	`16`
	`17`	`+[PacketHandlerClass(inheritHandlers: true)]`
`18`	`18`	`public class ClientLoadingState(ConnectionBase connection) : ClientBaseState(connection)`
`19`	`19`	`{`
`20`	`20`	`public LoadingState subState = LoadingState.Waiting;`
`@@ -139,10 +139,4 @@ public void HandleWorldData(ByteReader data)`
`139`	`139`	`Log.Message($"Loaded game in {loadingMs}ms");`
`140`	`140`	`connection.ChangeState(ConnectionStateEnum.ClientPlaying);`
`141`	`141`	`}`
`142`		`-`
`143`		`- [TypedPacketHandler]`
`144`		`- public new void HandleKeepAlive(ServerKeepAlivePacket packet) => base.HandleKeepAlive(packet);`
`145`		`-`
`146`		`- [TypedPacketHandler]`
`147`		`- public new void HandleTimeControl(ServerTimeControlPacket packet) => base.HandleTimeControl(packet);`
`148`	`142`	`}`
Original file line number	Diff line number	Diff line change
`@@ -9,14 +9,9 @@`
`9`	`9`
`10`	`10`	`namespace Multiplayer.Client`
`11`	`11`	`{`
	`12`	`+ [PacketHandlerClass(inheritHandlers: true)]`
`12`	`13`	`public class ClientPlayingState(ConnectionBase connection) : ClientBaseState(connection)`
`13`	`14`	`{`
`14`		`- [TypedPacketHandler]`
`15`		`- public new void HandleKeepAlive(ServerKeepAlivePacket packet) => base.HandleKeepAlive(packet);`
`16`		`-`
`17`		`- [TypedPacketHandler]`
`18`		`- public new void HandleTimeControl(ServerTimeControlPacket packet) => base.HandleTimeControl(packet);`
`19`		`-`
`20`	`15`	`[TypedPacketHandler]`
`21`	`16`	`public void HandleCommand(ServerCommandPacket packet)`
`22`	`17`	`{`
Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ public override void HandleReceiveRaw(ByteReader data, bool reliable)`
`30`	`30`	`{`
`31`	`31`	`}`
`32`	`32`
`33`		`- public override void Close(MpDisconnectReason reason, byte[] data)`
	`33`	`+ protected override void OnClose()`
`34`	`34`	`{`
`35`	`35`	`}`
`36`	`36`	`}`