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); + }); +});