Networking: Defer the client peer shutdown out of MafiaNet dispatch

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) {

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.

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

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<void(MafiaNet::BitStream *, MafiaNet::Packet *)> handler) {
-            auto slot     = std::make_unique<RPCSlot>(std::move(handler));
+            auto slot     = std::make_unique<RPCSlot>([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).

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();

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 <mafianet/FileList.h>
+#include <mafianet/FileListTransfer.h>
+#include <mafianet/FileListTransferCBInterface.h>
+
+// 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<unsigned short>(-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(&nothing, 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);
+    });
+});