From 77de9e8f6615f557d943f410c79da623b9a30f8e 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       |  6 +++++-
 Source/dvlnet/tcp_client.cpp |  6 ++++++
 Source/dvlnet/tcp_client.h   |  1 +
 Source/dvlnet/tcp_server.cpp | 25 ++++++++++++++++++-------
 Source/dvlnet/tcp_server.h   |  1 +
 5 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/Source/dvlnet/base.cpp b/Source/dvlnet/base.cpp
index 964bcbceb..871eb3ff7 100644
--- a/Source/dvlnet/base.cpp
+++ b/Source/dvlnet/base.cpp
@@ -409,10 +409,14 @@ bool base::SNetLeaveGame(int type)
 
 bool base::SNetDropPlayer(int playerid, uint32_t flags)
 {
+	const plr_t plr = static_cast<plr_t>(playerid);
 	auto pkt = pktfty->make_packet<PT_DISCONNECT>(plr_self,
 	    PLR_BROADCAST,
-	    static_cast<plr_t>(playerid),
+	    plr,
 	    static_cast<leaveinfo_t>(flags));
+	// Disconnect at the network layer first so we
+	// don't send players their own disconnect packet
+	DisconnectNet(plr);
 	send(*pkt);
 	RecvLocal(*pkt);
 	return true;
diff --git a/Source/dvlnet/tcp_client.cpp b/Source/dvlnet/tcp_client.cpp
index 7ad906d93..254bf8ed8 100644
--- a/Source/dvlnet/tcp_client.cpp
+++ b/Source/dvlnet/tcp_client.cpp
@@ -123,6 +123,12 @@ void tcp_client::send(packet &pkt)
 	});
 }
 
+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 8a06213e4..1e5767209 100644
--- a/Source/dvlnet/tcp_client.h
+++ b/Source/dvlnet/tcp_client.h
@@ -23,6 +23,7 @@ public:
 
 	void poll() override;
 	void 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 65277c4a9..66e2ad809 100644
--- a/Source/dvlnet/tcp_server.cpp
+++ b/Source/dvlnet/tcp_server.cpp
@@ -212,16 +212,27 @@ void tcp_server::HandleTimeout(const scc &con, const asio::error_code &ec)
 
 void tcp_server::DropConnection(const scc &con)
 {
-	if (con->plr != PLR_BROADCAST) {
-		auto pkt = pktfty.make_packet<PT_DISCONNECT>(PLR_MASTER, PLR_BROADCAST,
-		    con->plr, static_cast<leaveinfo_t>(LEAVE_DROP));
-		connections[con->plr] = nullptr;
-		SendPacket(*pkt);
-		// 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;
+
+	auto pkt = pktfty.make_packet<PT_DISCONNECT>(PLR_MASTER, PLR_BROADCAST,
+	    con->plr, static_cast<leaveinfo_t>(LEAVE_DROP));
+	SendPacket(*pkt);
+}
+
+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()
diff --git a/Source/dvlnet/tcp_server.h b/Source/dvlnet/tcp_server.h
index d56fdce8c..d6c379e7c 100644
--- a/Source/dvlnet/tcp_server.h
+++ b/Source/dvlnet/tcp_server.h
@@ -30,6 +30,7 @@ public:
 	tcp_server(asio::io_context &ioc, const std::string &bindaddr,
 	    unsigned short port, packet_factory &pktfty);
 	std::string LocalhostSelf();
+	void DisconnectNet(plr_t plr);
 	void Close();
 	virtual ~tcp_server();