From 6ede3cb14e9aeb1c8f11f365c76fa0c1e5c1c50a Mon Sep 17 00:00:00 2001 From: KHeartz Date: Sat, 10 Oct 2026 02:04:53 -0400 Subject: [PATCH] Networking: Defer the client peer shutdown out of MafiaNet dispatch Disconnect() called from inside a MafiaNet callback shut the peer down while FileListTransfer and RakPeer::Receive were still on the stack. The shutdown frees every FileListReceiver and the packet pool, so a join refused from OnAssetsDownloaded used the freed receiver or freed it twice. Disconnect() during Update() now resets the session at once and leaves the peer shutdown to Update() once dispatch unwinds. Packets and RPCs still queued for the dropped session are not dispatched. Fixes #315 --- .../src/networking/network_client.cpp | 31 +++++-- .../framework/src/networking/network_client.h | 10 +++ .../framework/src/networking/network_peer.cpp | 4 +- code/framework/src/networking/network_peer.h | 11 ++- code/tests/framework_ut.cpp | 2 + code/tests/modules/download_disconnect_ut.h | 90 +++++++++++++++++++ 6 files changed, 139 insertions(+), 9 deletions(-) create mode 100644 code/tests/modules/download_disconnect_ut.h diff --git a/code/framework/src/networking/network_client.cpp b/code/framework/src/networking/network_client.cpp index c1fc21a82..36fa80f81 100644 --- a/code/framework/src/networking/network_client.cpp +++ b/code/framework/src/networking/network_client.cpp @@ -121,13 +121,16 @@ namespace Framework::Networking { return {}; } - _peer->Shutdown(kShutdownBlockDurationMs, 0, MafiaNet::Priority::Immediate); + // Called from inside a MafiaNet callback (MafiaHub/Framework#315): Update() shuts the peer + // down once dispatch has unwound. + if (_dispatching) { + _shutdownPending = true; + } + else { + ShutdownPeer(); + } Logging::GetLogger(FRAMEWORK_INNER_NETWORKING)->debug("Disconnecting from the server..."); - // Peer shutdown fires TwoWayAuthentication::Clear(), wiping the registered build token; - // drop the cache so the next Connect()'s Init() re-registers instead of no-op'ing. - _registeredToken.clear(); - if (_onPlayerDisconnectedCallback) { // Locally initiated: there is no inbound packet, so pass null rather than a stale _packet. _onPlayerDisconnectedCallback(nullptr, DisconnectionReason::GRACEFUL_SHUTDOWN, ""); @@ -137,8 +140,17 @@ namespace Framework::Networking { return {}; } + void NetworkClient::ShutdownPeer() { + _peer->Shutdown(kShutdownBlockDurationMs, 0, MafiaNet::Priority::Immediate); + + // Peer shutdown fires TwoWayAuthentication::Clear(), wiping the registered build token; + // drop the cache so the next Connect()'s Init() re-registers instead of no-op'ing. + _registeredToken.clear(); + } + void NetworkClient::ResetConnectionState() { - if (_peer) { + // A pending shutdown closes every connection itself, after the plugins have unwound. + if (_peer && !_shutdownPending) { MafiaNet::SystemAddress systems[kMaxConnections]; unsigned short count = kMaxConnections; if (_peer->GetConnectionList(systems, &count)) { @@ -174,7 +186,14 @@ namespace Framework::Networking { return; } + _dispatching = true; NetworkPeer::Update(); + _dispatching = false; + + if (_shutdownPending) { + _shutdownPending = false; + ShutdownPeer(); + } } bool NetworkClient::HandlePacket(uint8_t packetID, MafiaNet::Packet *packet) { diff --git a/code/framework/src/networking/network_client.h b/code/framework/src/networking/network_client.h index 3e4021420..8ae88a420 100644 --- a/code/framework/src/networking/network_client.h +++ b/code/framework/src/networking/network_client.h @@ -48,9 +48,16 @@ namespace Framework::Networking { AssetFileTransfer _fileListTransfer; bool _initialReplicationDownloadComplete {}; + // Set while Update() dispatches: a MafiaNet plugin or RakPeer::Receive may be on the stack, + // and the peer shutdown frees the receivers and packets they still hold. + bool _dispatching {}; + bool _shutdownPending {}; + // Single teardown for every terminal disconnect: drops the connection and clears the flags. void ResetConnectionState(); + void ShutdownPeer(); + // Empties the peer's receive queue without dispatching. void DrainStalePackets(); @@ -64,6 +71,9 @@ namespace Framework::Networking { void Update() override; bool HandlePacket(uint8_t packetID, MafiaNet::Packet *packet) override; + bool IsDispatchHalted() const override { + return _shutdownPending; + } // sessionPayload rides the connection request (MafiaNet's session handshake): it is what the // server's admission gate decides on before either side reports a connection. diff --git a/code/framework/src/networking/network_peer.cpp b/code/framework/src/networking/network_peer.cpp index ecf19dd75..ebb355417 100644 --- a/code/framework/src/networking/network_peer.cpp +++ b/code/framework/src/networking/network_peer.cpp @@ -85,8 +85,8 @@ namespace Framework::Networking { // Receive() updates every plugin before each packet; the replication pass runs on the first // one only, ahead of this drain's packets, which the next pass then serializes. _replicationManager->BeginNetworkUpdate(); - for (_packet = _peer->Receive(); _packet; _peer->DeallocatePacket(_packet), _packet = _peer->Receive()) { - if (_packet->length == 0) { + for (_packet = _peer->Receive(); _packet; _peer->DeallocatePacket(_packet), _packet = IsDispatchHalted() ? nullptr : _peer->Receive()) { + if (_packet->length == 0 || IsDispatchHalted()) { continue; } const int offset = ResolvePacketDataOffset(_packet->data, _packet->length); diff --git a/code/framework/src/networking/network_peer.h b/code/framework/src/networking/network_peer.h index d39c10761..211b6272b 100644 --- a/code/framework/src/networking/network_peer.h +++ b/code/framework/src/networking/network_peer.h @@ -150,7 +150,11 @@ namespace Framework::Networking { // is stored for the peer's lifetime and reached through RPC4's per-slot context, so no // file-static handler pointers are needed. void RegisterRawRPC(const char *identifier, fu2::function handler) { - auto slot = std::make_unique(std::move(handler)); + auto slot = std::make_unique([this, cb = std::move(handler)](MafiaNet::BitStream *bs, MafiaNet::Packet *packet) mutable { + if (!IsDispatchHalted()) { + cb(bs, packet); + } + }); void *context = slot.get(); _rpcHandlers.push_back(std::move(slot)); _rpc.RegisterSlot(identifier, &NetworkPeer::DispatchRPC, context, 0); @@ -164,6 +168,11 @@ namespace Framework::Networking { void Update() override; virtual bool HandlePacket(uint8_t packetID, MafiaNet::Packet *packet) = 0; + // True once the session is being torn down mid-dispatch: the rest of the queue belongs to it. + virtual bool IsDispatchHalted() const { + return false; + } + // Byte offset of the packet id in a datagram, skipping an optional ID_TIMESTAMP + 8-byte // MafiaNet::Time prefix. Returns -1 if too short. Pure + tested so the skip width can't drift // from what RakNet writes (a wrong width reads the id mid-timestamp — phantom control packets). diff --git a/code/tests/framework_ut.cpp b/code/tests/framework_ut.cpp index ac87303c5..a0c348477 100644 --- a/code/tests/framework_ut.cpp +++ b/code/tests/framework_ut.cpp @@ -15,6 +15,7 @@ #include "modules/connection_admission_ut.h" #include "modules/connection_gate_ut.h" #include "modules/delegation_ut.h" +#include "modules/download_disconnect_ut.h" #include "modules/gui_resources_ut.h" #include "modules/input_state_ut.h" #include "modules/interest_grid_ut.h" @@ -126,6 +127,7 @@ int main() { UNIT_MODULE(timer_context); UNIT_MODULE(connection_gate); UNIT_MODULE(connection_admission); + UNIT_MODULE(download_disconnect); UNIT_MODULE(nametag); return UNIT_RUN(); diff --git a/code/tests/modules/download_disconnect_ut.h b/code/tests/modules/download_disconnect_ut.h new file mode 100644 index 000000000..177b21449 --- /dev/null +++ b/code/tests/modules/download_disconnect_ut.h @@ -0,0 +1,90 @@ +/* + * MafiaHub OSS license + * Copyright (c) 2026, MafiaHub. All rights reserved. + * + * This file comes from MafiaHub, hosted at https://github.com/MafiaHub/Framework. + * See LICENSE file in the source repository for information regarding licensing. + */ + +#pragma once + +#include "connection_admission_ut.h" + +#include +#include +#include + +// MafiaHub/Framework#315: the join refusal disconnects from inside FileListTransfer's completion +// callback, and the peer shutdown frees the receiver and packet pool that DecodeSetHeader and +// RakPeer::Receive still hold. +class DisconnectingDownloadHandler final: public MafiaNet::FileListTransferCBInterface { + public: + // Thrown when the shutdown has already freed what the callers hold: unwinding out of + // RakPeer::Receive is the only way past the freed receiver and packet without touching them. + struct TornDownUnderCaller {}; + + Framework::Networking::NetworkClient *client = nullptr; + unsigned short setID = 0; + int completions = 0; + int dereferences = 0; + bool receiverSurvived = false; + + bool OnFile(OnFileStruct *) override { + return true; + } + + void OnFileProgress(FileProgressStruct *) override {} + + bool OnDownloadComplete(DownloadCompleteStruct *) override { + ++completions; + (void)client->Disconnect(); + receiverSurvived = client->GetFileListTransfer()->IsHandlerActive(setID); + if (!receiverSurvived) { + throw TornDownUnderCaller {}; + } + // FileListTransfer then dereferences and deletes the receiver it still holds. + return false; + } + + void OnDereference() override { + ++dereferences; + } +}; + +MODULE(download_disconnect, { + using Rig = ConnectionAdmissionRig; + + IT("keeps the download receiver alive when the completion callback disconnects", { + // Outlives the rig: a receiver left behind by a failed step still points at it. + DisconnectingDownloadHandler handler; + Rig rig; + EQUALS(rig.StartServer(false), true); + EQUALS(rig.Connect(Rig::IdentityPayload("Downloader", "")), true); + EQUALS(rig.PumpUntil([&] { return rig.seen.clientConnected && rig.server->GetPlayerCount() == 1u; }), true); + + handler.client = rig.client.get(); + handler.setID = rig.client->GetFileListTransfer()->SetupReceive(&handler, false, rig.client->GetPeer()->GetSystemAddressFromIndex(0)); + EQUALS(handler.setID != static_cast(-1), true); + + // An empty set: the header alone completes it, through DecodeSetHeader, which is how a + // download ends when the client already has every file. + MafiaNet::FileList nothing; + MafiaNet::FileListTransfer sender; + sender.Send(¬hing, rig.server->GetPeer(), rig.server->GetPeer()->GetSystemAddressFromIndex(0), handler.setID, MafiaNet::Priority::High, 0); + + bool tornDown = false; + try { + rig.PumpUntil([&] { return handler.completions > 0 && rig.seen.clientDisconnected; }); + } + catch (const DisconnectingDownloadHandler::TornDownUnderCaller &) { + tornDown = true; + } + + EQUALS(handler.completions, 1); + EQUALS(tornDown, false); + EQUALS(handler.receiverSurvived, true); + // Released once, by FileListTransfer, and not again by the shutdown. + EQUALS(handler.dereferences, 1); + EQUALS(rig.seen.clientDisconnected, true); + }); +});