From 2e958ee5463e605983cbd4d9023f46997a2dd163 Mon Sep 17 00:00:00 2001
From: staphen <staphen@gmail.com>
Date: Tue, 21 Nov 2023 16:39:46 -0500
Subject: [PATCH] Avoid sending PT_DISCONNECT packet to dropped player

---
 Source/dvlnet/base.cpp       |  8 ++++++--
 Source/dvlnet/tcp_client.cpp |  6 ++++++
 Source/dvlnet/tcp_client.h   |  1 +
 Source/dvlnet/tcp_server.cpp | 37 +++++++++++++++++++++++-------------
 Source/dvlnet/tcp_server.h   |  1 +
 5 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/Source/dvlnet/base.cpp b/Source/dvlnet/base.cpp
index e0249eaa2..93907be45 100644
--- a/Source/dvlnet/base.cpp
+++ b/Source/dvlnet/base.cpp
@@ -470,16 +470,20 @@ bool base::SNetLeaveGame(int type)
 
 bool base::SNetDropPlayer(int playerid, uint32_t flags)
 {
+	plr_t plr = static_cast<plr_t>(playerid);
 	tl::expected<std::unique_ptr<packet>, PacketError> pkt
 	    = pktfty->make_packet<PT_DISCONNECT>(
 	        plr_self,
 	        PLR_BROADCAST,
-	        (plr_t)playerid,
-	        (leaveinfo_t)flags);
+	        plr,
+	        static_cast<leaveinfo_t>(flags));
 	if (!pkt.has_value()) {
 		LogError("make_packet: {}", pkt.error().what());
 		return false;
 	}
+	// Disconnect at the network layer first so we
+	// don't send players their own disconnect packet
+	DisconnectNet(plr);
 	tl::expected<void, PacketError> sendResult = send(**pkt);
 	if (!sendResult.has_value()) {
 		LogError("send: {}", sendResult.error().what());
diff --git a/Source/dvlnet/tcp_client.cpp b/Source/dvlnet/tcp_client.cpp
index 546159901..25e15ae50 100644
--- a/Source/dvlnet/tcp_client.cpp
+++ b/Source/dvlnet/tcp_client.cpp
@@ -169,6 +169,12 @@ tl::expected<void, PacketError> tcp_client::send(packet &pkt)
 	return {};
 }
 
+void tcp_client::DisconnectNet(plr_t plr)
+{
+	if (local_server != nullptr)
+		local_server->DisconnectNet(plr);
+}
+
 bool tcp_client::SNetLeaveGame(int type)
 {
 	auto ret = base::SNetLeaveGame(type);
diff --git a/Source/dvlnet/tcp_client.h b/Source/dvlnet/tcp_client.h
index c4029afd3..7a68072b7 100644
--- a/Source/dvlnet/tcp_client.h
+++ b/Source/dvlnet/tcp_client.h
@@ -32,6 +32,7 @@ public:
 
 	tl::expected<void, PacketError> poll() override;
 	tl::expected<void, PacketError> send(packet &pkt) override;
+	void DisconnectNet(plr_t plr) override;
 
 	bool SNetLeaveGame(int type) override;
 
diff --git a/Source/dvlnet/tcp_server.cpp b/Source/dvlnet/tcp_server.cpp
index e95771606..66fef9728 100644
--- a/Source/dvlnet/tcp_server.cpp
+++ b/Source/dvlnet/tcp_server.cpp
@@ -257,21 +257,22 @@ void tcp_server::HandleTimeout(const scc &con, const asio::error_code &ec)
 
 void tcp_server::DropConnection(const scc &con)
 {
-	if (con->plr != PLR_BROADCAST) {
-		tl::expected<std::unique_ptr<packet>, PacketError> pkt
-		    = pktfty.make_packet<PT_DISCONNECT>(PLR_MASTER, PLR_BROADCAST,
-		        con->plr, LEAVE_DROP);
-		connections[con->plr] = nullptr;
-		if (pkt.has_value()) {
-			SendPacket(**pkt);
-		} else {
-			LogError("make_packet<PT_DISCONNECT>: {}", pkt.error().what());
-		}
-		// TODO: investigate if it is really ok for the server to
-		//       drop a client directly.
-	}
+	plr_t plr = con->plr;
 	con->timer.cancel();
 	con->socket.close();
+	if (plr == PLR_BROADCAST) {
+		return;
+	}
+	connections[plr] = nullptr;
+
+	tl::expected<std::unique_ptr<packet>, PacketError> pkt
+	    = pktfty.make_packet<PT_DISCONNECT>(PLR_MASTER, PLR_BROADCAST,
+	        plr, LEAVE_DROP);
+	if (pkt.has_value()) {
+		SendPacket(**pkt);
+	} else {
+		LogError("make_packet<PT_DISCONNECT>: {}", pkt.error().what());
+	}
 }
 
 void tcp_server::RaiseIoHandlerError(const PacketError &error)
@@ -288,6 +289,16 @@ tl::expected<void, PacketError> tcp_server::CheckIoHandlerError()
 	return packetError;
 }
 
+void tcp_server::DisconnectNet(plr_t plr)
+{
+	scc &con = connections[plr];
+	if (con == nullptr)
+		return;
+	con->timer.cancel();
+	con->socket.close();
+	con = nullptr;
+}
+
 void tcp_server::Close()
 {
 	acceptor->close();
diff --git a/Source/dvlnet/tcp_server.h b/Source/dvlnet/tcp_server.h
index c9907d72c..4827cbadb 100644
--- a/Source/dvlnet/tcp_server.h
+++ b/Source/dvlnet/tcp_server.h
@@ -39,6 +39,7 @@ public:
 	    unsigned short port, packet_factory &pktfty);
 	std::string LocalhostSelf();
 	tl::expected<void, PacketError> CheckIoHandlerError();
+	void DisconnectNet(plr_t plr);
 	void Close();
 	virtual ~tcp_server();