diff --git a/Source/include/mafianet/file_list_transfer.h b/Source/include/mafianet/file_list_transfer.h index 5bd679450..5ab234872 100644 --- a/Source/include/mafianet/file_list_transfer.h +++ b/Source/include/mafianet/file_list_transfer.h @@ -146,6 +146,7 @@ class RAK_DLL_EXPORT FileListTransfer : public PluginInterface2 bool DecodeFile(Packet *packet, bool fullFile); void Clear(void); + bool IsReceiverLive(unsigned short inSetId, const FileListReceiver *receiver); void OnReferencePush(Packet *packet, bool fullFile); void OnReferencePushAck(Packet *packet); diff --git a/Source/include/mafianet/peer.h b/Source/include/mafianet/peer.h index fe31cb18f..fb4322882 100644 --- a/Source/include/mafianet/peer.h +++ b/Source/include/mafianet/peer.h @@ -1110,6 +1110,12 @@ class RAK_DLL_EXPORT RakPeer : public RakPeerInterface, public RNS2EventHandler SimpleMutex packetAllocationPoolMutex; DataStructures::MemoryPool packetAllocationPool; + // Pool packets not yet given back through DeallocatePacket, under packetAllocationPoolMutex. + // Shutdown() keeps the pool while any are out: Receive may be dispatching one, and the + // application may still hold one it is handling. + unsigned int packetsOutstanding; + // Bumped by every Shutdown(), so Receive can tell its dispatch was cut short by one. + unsigned int shutdownGeneration; SimpleMutex packetReturnMutex; DataStructures::Queue packetReturnQueue; diff --git a/Source/src/file_list_transfer.cpp b/Source/src/file_list_transfer.cpp index b6fb94bcb..2e1306fd7 100644 --- a/Source/src/file_list_transfer.cpp +++ b/Source/src/file_list_transfer.cpp @@ -342,7 +342,7 @@ bool FileListTransfer::DecodeSetHeader(Packet *packet) dcs.senderSystemAddress=packet->systemAddress; dcs.senderGuid=packet->guid; - if (fileListReceiver->downloadHandler->OnDownloadComplete(&dcs)==false) + if (fileListReceiver->downloadHandler->OnDownloadComplete(&dcs)==false && IsReceiverLive(setID, fileListReceiver)) { fileListReceiver->downloadHandler->OnDereference(); fileListReceivers.Delete(setID); @@ -445,12 +445,19 @@ bool FileListTransfer::DecodeFile(Packet *packet, bool isTheFullFile) fps.senderSystemAddress=packet->systemAddress; fps.senderGuid=packet->guid; fileListReceiver->downloadHandler->OnFileProgress(&fps); + if (!IsReceiverLive(onFileStruct.setID, fileListReceiver)) + { + rakFree_Ex(onFileStruct.fileData, _FILE_AND_LINE_ ); + return true; + } // Got a complete file // Either we are using IncrementalReadInterface and it was a small file or // We are not using IncrementalReadInterface if (fileListReceiver->downloadHandler->OnFile(&onFileStruct)) rakFree_Ex(onFileStruct.fileData, _FILE_AND_LINE_ ); + if (!IsReceiverLive(onFileStruct.setID, fileListReceiver)) + return true; fileListReceiver->filesReceived++; @@ -464,7 +471,7 @@ bool FileListTransfer::DecodeFile(Packet *packet, bool isTheFullFile) dcs.senderSystemAddress=packet->systemAddress; dcs.senderGuid=packet->guid; - if (fileListReceiver->downloadHandler->OnDownloadComplete(&dcs)==false) + if (fileListReceiver->downloadHandler->OnDownloadComplete(&dcs)==false && IsReceiverLive(onFileStruct.setID, fileListReceiver)) { fileListReceiver->downloadHandler->OnDereference(); if (fileListReceiver->deleteDownloadHandler) @@ -672,13 +679,24 @@ void FileListTransfer::GetCallbacks(DataStructures::List &cal callbacks = fileListProgressCallbacks; } +bool FileListTransfer::IsReceiverLive(unsigned short inSetId, const FileListReceiver *receiver) +{ + // A download callback may remove its own receiver -- through CancelReceive, RemoveReceiver, + // or a Shutdown() that clears every one -- so the caller re-checks before touching it again. + return fileListReceivers.Has(inSetId) && fileListReceivers.Get(inSetId)==receiver; +} void FileListTransfer::Update(void) { unsigned i; i=0; while (i < fileListReceivers.Size()) { - if (fileListReceivers[i]->downloadHandler->Update()==false) + FileListReceiver *receiver = fileListReceivers[i]; + const bool keep = receiver->downloadHandler->Update(); + // Update() may have dropped receivers, this one included + if (i >= fileListReceivers.Size() || fileListReceivers[i] != receiver) + continue; + if (keep==false) { fileListReceivers[i]->downloadHandler->OnDereference(); if (fileListReceivers[i]->deleteDownloadHandler) @@ -861,11 +879,15 @@ void FileListTransfer::OnReferencePush(Packet *packet, bool isTheFullFile) if (fps.partTotal==1) fps.onFileStruct->fileData=fps.iriDataChunk; fileListReceiver->downloadHandler->OnFileProgress(&fps); + if (!IsReceiverLive(onFileStruct.setID, fileListReceiver)) + return; // Incremental read interface sent us a file chunk // This is the last file chunk we were waiting for to consider the file done if (fileListReceiver->downloadHandler->OnFile(&onFileStruct)) rakFree_Ex(oldFileData, _FILE_AND_LINE_ ); + if (!IsReceiverLive(onFileStruct.setID, fileListReceiver)) + return; fileListReceiver->pushedFiles.Delete(onFileStruct.fileIndex); fileListReceiver->filesReceived++; @@ -880,7 +902,7 @@ void FileListTransfer::OnReferencePush(Packet *packet, bool isTheFullFile) dcs.senderSystemAddress=packet->systemAddress; dcs.senderGuid=packet->guid; - if (fileListReceiver->downloadHandler->OnDownloadComplete(&dcs)==false) + if (fileListReceiver->downloadHandler->OnDownloadComplete(&dcs)==false && IsReceiverLive(onFileStruct.setID, fileListReceiver)) { fileListReceiver->downloadHandler->OnDereference(); fileListReceivers.Delete(onFileStruct.setID); @@ -896,6 +918,8 @@ void FileListTransfer::OnReferencePush(Packet *packet, bool isTheFullFile) { // 12/23/09 Don't use OnReferencePush anymore, just use OnFileProgress fileListReceiver->downloadHandler->OnFileProgress(&fps); + if (!IsReceiverLive(onFileStruct.setID, fileListReceiver)) + return; if (fps.allocateIrIDataChunkAutomatically==false) { diff --git a/Source/src/peer.cpp b/Source/src/peer.cpp index 68456a88a..dd6df7ecd 100644 --- a/Source/src/peer.cpp +++ b/Source/src/peer.cpp @@ -184,6 +184,7 @@ Packet *RakPeer::AllocPacket(unsigned dataSize, const char *file, unsigned int l MafiaNet::Packet *p; packetAllocationPoolMutex.Lock(); p = packetAllocationPool.Allocate(file,line); + ++packetsOutstanding; packetAllocationPoolMutex.Unlock(); p = new ((void*)p) Packet; p->data=(unsigned char*) rakMalloc_Ex(dataSize,file,line); @@ -201,6 +202,7 @@ Packet *RakPeer::AllocPacket(unsigned dataSize, unsigned char *data, const char MafiaNet::Packet *p; packetAllocationPoolMutex.Lock(); p = packetAllocationPool.Allocate(file,line); + ++packetsOutstanding; packetAllocationPoolMutex.Unlock(); p = new ((void*)p) Packet; RakAssert(p); @@ -293,7 +295,9 @@ RakPeer::RakPeer() packetAllocationPoolMutex.Lock(); packetAllocationPool.SetPageSize(sizeof(DataStructures::MemoryPool::MemoryWithPage)*32); + packetsOutstanding = 0; packetAllocationPoolMutex.Unlock(); + shutdownGeneration = 0; remoteSystemIndexPool.SetPageSize(sizeof(DataStructures::MemoryPool::MemoryWithPage)*32); @@ -998,6 +1002,8 @@ void RakPeer::Shutdown( unsigned int blockDuration, unsigned char orderingChanne //unsigned short systemListSize = remoteSystemListSize; // This is done for threading reasons unsigned int systemListSize = maximumNumberOfPeers; + ++shutdownGeneration; + if ( blockDuration > 0 ) { for ( i = 0; i < systemListSize; i++ ) @@ -1126,8 +1132,12 @@ void RakPeer::Shutdown( unsigned int blockDuration, unsigned char orderingChanne DeallocatePacket(packetReturnQueue[i]); packetReturnQueue.Clear(_FILE_AND_LINE_); packetReturnMutex.Unlock(); + // Packets still out live in the pool's pages: a Shutdown() from inside Receive's dispatch, + // or from code handling a packet Receive returned, must not free the packet under it. The + // pages are reused by the next session and freed by a later Shutdown() or the destructor. packetAllocationPoolMutex.Lock(); - packetAllocationPool.Clear(_FILE_AND_LINE_); + if (packetsOutstanding == 0) + packetAllocationPool.Clear(_FILE_AND_LINE_); packetAllocationPoolMutex.Unlock(); /* @@ -1532,12 +1542,27 @@ Packet* RakPeer::Receive( void ) // return packet; + // A callback may shut the peer down. The plugins have been reset by then and the packet + // belongs to a closed session, so it goes no further and Receive returns nothing. + const unsigned int generation = shutdownGeneration; + CallPluginCallbacks(pluginListTS, packet); CallPluginCallbacks(pluginListNTS, packet); + if (shutdownGeneration != generation) + { + DeallocatePacket( packet ); + return 0; + } for (i=0; i < pluginListTS.Size(); i++) { pluginResult=pluginListTS[i]->OnReceive(packet); + if (shutdownGeneration != generation) + { + if (pluginResult!=RR_STOP_PROCESSING) + DeallocatePacket( packet ); + return 0; + } if (pluginResult==RR_STOP_PROCESSING_AND_DEALLOCATE) { DeallocatePacket( packet ); @@ -1554,6 +1579,12 @@ Packet* RakPeer::Receive( void ) for (i=0; i < pluginListNTS.Size(); i++) { pluginResult=pluginListNTS[i]->OnReceive(packet); + if (shutdownGeneration != generation) + { + if (pluginResult!=RR_STOP_PROCESSING) + DeallocatePacket( packet ); + return 0; + } if (pluginResult==RR_STOP_PROCESSING_AND_DEALLOCATE) { DeallocatePacket( packet ); @@ -1591,6 +1622,7 @@ void RakPeer::DeallocatePacket( Packet *packet ) packet->~Packet(); packetAllocationPoolMutex.Lock(); packetAllocationPool.Release(packet,_FILE_AND_LINE_); + --packetsOutstanding; packetAllocationPoolMutex.Unlock(); } else diff --git a/Tests/Integration/ShutdownDuringDispatchLiveTests.cpp b/Tests/Integration/ShutdownDuringDispatchLiveTests.cpp new file mode 100644 index 000000000..0d49ec2d7 --- /dev/null +++ b/Tests/Integration/ShutdownDuringDispatchLiveTests.cpp @@ -0,0 +1,344 @@ +/* + * Copyright (c) 2026, MafiaHub + * + * This source code is licensed under the MIT-style license found in the + * license.txt file in the root directory of this source tree. + */ + +#include + +#include + +#include "mafianet/peer.h" +#include "mafianet/peer_interface.h" +#include "mafianet/message_identifiers.h" +#include "mafianet/plugin_interface2.h" +#include "mafianet/file_list.h" +#include "mafianet/file_list_transfer.h" +#include "mafianet/file_list_transfer_cb_interface.h" +#include "mafianet/bit_stream.h" +#include "mafianet/sleep.h" +#include "mafianet/get_time.h" + +using namespace MafiaNet; + +/* +Description: +An application may shut its peer down, or drop a transfer, from inside a callback that RakPeer::Receive +is dispatching: a plugin's OnReceive, a FileListTransfer download callback, or code handling a packet +Receive returned. Shutdown() runs every plugin's OnRakPeerShutdown and releases the packet pool while +those callers are still on the stack, so each of them must stop touching what the shutdown released. + +Success conditions: +- A FileListTransfer receiver whose callback shuts the peer down, or removes the receiver, is + dereferenced exactly once, whether the set completes in the header (empty set) or on its last file. +- A plugin that shuts the peer down in OnReceive stops the packet there: later plugins never see it, + and Receive returns nothing for it. +- A packet Receive returned stays intact across Shutdown() and can still be deallocated after it, and + the peer starts and exchanges packets again afterwards. + +Failure conditions: a second OnDereference, a later plugin seeing the packet, a packet whose fields +change across Shutdown(), or a crash. +*/ + +namespace +{ + const int kConnectTimeoutMs = 15000; + const int kWaitTimeoutMs = 15000; + + struct PeerPair + { + RakPeerInterface *server = nullptr; + RakPeerInterface *client = nullptr; + SystemAddress serverAddress = UNASSIGNED_SYSTEM_ADDRESS; + SystemAddress clientAddress = UNASSIGNED_SYSTEM_ADDRESS; + + PeerPair() + { + server = RakPeerInterface::GetInstance(); + client = RakPeerInterface::GetInstance(); + } + + ~PeerPair() + { + client->Shutdown(0); + server->Shutdown(0); + RakPeerInterface::DestroyInstance(client); + RakPeerInterface::DestroyInstance(server); + } + + // Both peers bind ephemeral ports, so tests can run side by side. + bool Connect() + { + SocketDescriptor serverSocket(0, "127.0.0.1"); + SocketDescriptor clientSocket(0, "127.0.0.1"); + if (server->Startup(4, &serverSocket, 1) != RAKNET_STARTED || client->Startup(1, &clientSocket, 1) != RAKNET_STARTED) + return false; + server->SetMaximumIncomingConnections(4); + const unsigned short port = server->GetMyBoundAddress().GetPort(); + if (client->Connect("127.0.0.1", port, 0, 0) != CONNECTION_ATTEMPT_STARTED) + return false; + + const TimeMS entry = GetTimeMS(); + while ((clientAddress == UNASSIGNED_SYSTEM_ADDRESS || serverAddress == UNASSIGNED_SYSTEM_ADDRESS) && GetTimeMS() - entry < (TimeMS)kConnectTimeoutMs) + { + for (Packet *p = server->Receive(); p; server->DeallocatePacket(p), p = server->Receive()) + if (p->data[0] == ID_NEW_INCOMING_CONNECTION) + clientAddress = p->systemAddress; + for (Packet *p = client->Receive(); p; client->DeallocatePacket(p), p = client->Receive()) + if (p->data[0] == ID_CONNECTION_REQUEST_ACCEPTED) + serverAddress = p->systemAddress; + RakSleep(5); + } + return clientAddress != UNASSIGNED_SYSTEM_ADDRESS && serverAddress != UNASSIGNED_SYSTEM_ADDRESS; + } + + // Pumps both peers until done() holds or the wait runs out. + template + bool PumpUntil(Done done) + { + const TimeMS entry = GetTimeMS(); + while (!done() && GetTimeMS() - entry < (TimeMS)kWaitTimeoutMs) + { + for (Packet *p = server->Receive(); p; server->DeallocatePacket(p), p = server->Receive()) + ; + for (Packet *p = client->Receive(); p; client->DeallocatePacket(p), p = client->Receive()) + ; + RakSleep(1); + } + return done(); + } + + // Pumps both peers for a fixed time, for traffic that must not arrive. + void Drain(int ms) + { + const TimeMS entry = GetTimeMS(); + PumpUntil([&] { return GetTimeMS() - entry >= (TimeMS)ms; }); + } + }; + + enum class Reaction + { + ShutdownPeer, + RemoveReceiver, + }; + + // Reacts from the callback that ends the set, as a client refusing a join does. + class ReactingReceiver : public FileListTransferCBInterface + { + public: + RakPeerInterface *peer = nullptr; + FileListTransfer *transfer = nullptr; + SystemAddress sender = UNASSIGNED_SYSTEM_ADDRESS; + Reaction reaction = Reaction::ShutdownPeer; + int completions = 0; + int dereferences = 0; + + bool OnFile(OnFileStruct *) override + { + return true; + } + + void OnFileProgress(FileProgressStruct *) override {} + + bool OnDownloadComplete(DownloadCompleteStruct *) override + { + ++completions; + if (reaction == Reaction::ShutdownPeer) + peer->Shutdown(0); + else + transfer->RemoveReceiver(sender); + return false; + } + + void OnDereference(void) override + { + ++dereferences; + } + }; + + void RunReactingTransfer(bool withFile, Reaction reaction) + { + PeerPair peers; + FileListTransfer serverTransfer; + FileListTransfer clientTransfer; + peers.server->AttachPlugin(&serverTransfer); + peers.client->AttachPlugin(&clientTransfer); + ASSERT_TRUE(peers.Connect()); + + ReactingReceiver receiver; + receiver.peer = peers.client; + receiver.transfer = &clientTransfer; + receiver.sender = peers.serverAddress; + receiver.reaction = reaction; + const unsigned short setId = clientTransfer.SetupReceive(&receiver, false, peers.serverAddress); + + FileList list; + if (withFile) + { + const std::vector bytes(900, 'x'); + list.AddFile("small.bin", "small.bin", bytes.data(), (unsigned)bytes.size(), (unsigned)bytes.size(), FileListNodeContext(0, 0, 0, 0)); + } + serverTransfer.Send(&list, peers.server, peers.clientAddress, setId, Priority::Medium, 0); + + EXPECT_TRUE(peers.PumpUntil([&] { return receiver.completions > 0; })); + // Anything still queued must not reach the dropped receiver either. + peers.Drain(500); + + EXPECT_EQ(receiver.completions, 1); + EXPECT_EQ(receiver.dereferences, 1); + + peers.server->DetachPlugin(&serverTransfer); + peers.client->DetachPlugin(&clientTransfer); + } + + const MessageID kProbe = ID_USER_PACKET_ENUM + 7; + + // Shuts the peer down on the probe, from inside Receive's plugin dispatch. + class ShuttingPlugin : public PluginInterface2 + { + public: + int probes = 0; + + PluginReceiveResult OnReceive(Packet *packet) override + { + if (packet->data[0] != kProbe) + return RR_CONTINUE_PROCESSING; + ++probes; + GetRakPeerInterface()->Shutdown(0); + return RR_CONTINUE_PROCESSING; + } + }; + + // Attached after the shutting plugin; must never see the probe. + class LaterPlugin : public PluginInterface2 + { + public: + int probes = 0; + int shutdowns = 0; + + PluginReceiveResult OnReceive(Packet *packet) override + { + if (packet->data[0] == kProbe) + ++probes; + return RR_CONTINUE_PROCESSING; + } + + void OnRakPeerShutdown(void) override + { + ++shutdowns; + } + }; + + void SendProbe(PeerPair &peers) + { + BitStream out; + out.Write(kProbe); + out.Write((uint32_t)0x5AFEC0DE); + peers.client->Send(&out, Priority::High, Reliability::ReliableOrdered, 0, peers.serverAddress, false); + } +} // namespace + +TEST(ShutdownDuringDispatch, EmptySetReceiverShutsThePeerDown) +{ + RunReactingTransfer(false, Reaction::ShutdownPeer); +} + +TEST(ShutdownDuringDispatch, LastFileReceiverShutsThePeerDown) +{ + RunReactingTransfer(true, Reaction::ShutdownPeer); +} + +TEST(ShutdownDuringDispatch, EmptySetReceiverRemovesItself) +{ + RunReactingTransfer(false, Reaction::RemoveReceiver); +} + +TEST(ShutdownDuringDispatch, LastFileReceiverRemovesItself) +{ + RunReactingTransfer(true, Reaction::RemoveReceiver); +} + +TEST(ShutdownDuringDispatch, PluginShutdownStopsThePacket) +{ + PeerPair peers; + ShuttingPlugin shutting; + LaterPlugin later; + peers.server->AttachPlugin(&shutting); + peers.server->AttachPlugin(&later); + ASSERT_TRUE(peers.Connect()); + + SendProbe(peers); + bool probeReturned = false; + const TimeMS entry = GetTimeMS(); + while (shutting.probes == 0 && GetTimeMS() - entry < (TimeMS)kWaitTimeoutMs) + { + for (Packet *p = peers.server->Receive(); p; peers.server->DeallocatePacket(p), p = peers.server->Receive()) + probeReturned |= p->data[0] == kProbe; + for (Packet *p = peers.client->Receive(); p; peers.client->DeallocatePacket(p), p = peers.client->Receive()) + ; + RakSleep(1); + } + + EXPECT_EQ(shutting.probes, 1); + EXPECT_EQ(later.probes, 0); + EXPECT_EQ(later.shutdowns, 1); + EXPECT_FALSE(probeReturned); + EXPECT_FALSE(peers.server->IsActive()); + + peers.server->DetachPlugin(&later); + peers.server->DetachPlugin(&shutting); +} + +TEST(ShutdownDuringDispatch, ReturnedPacketOutlivesShutdown) +{ + PeerPair peers; + ASSERT_TRUE(peers.Connect()); + + SendProbe(peers); + Packet *probe = nullptr; + const TimeMS entry = GetTimeMS(); + while (probe == nullptr && GetTimeMS() - entry < (TimeMS)kWaitTimeoutMs) + { + for (Packet *p = peers.server->Receive(); p; p = peers.server->Receive()) + { + if (p->data[0] == kProbe && probe == nullptr) + { + probe = p; + continue; + } + peers.server->DeallocatePacket(p); + } + RakSleep(1); + } + ASSERT_NE(probe, nullptr); + + const unsigned int length = probe->length; + const RakNetGUID guid = probe->guid; + const std::vector data(probe->data, probe->data + probe->length); + + // As an application does when a packet it is handling makes it disconnect. + peers.server->Shutdown(0); + + EXPECT_EQ(probe->length, length); + EXPECT_EQ(probe->guid, guid); + EXPECT_EQ(std::vector(probe->data, probe->data + probe->length), data); + peers.server->DeallocatePacket(probe); + + // The pool the packet went back to still serves a fresh session. + peers.client->Shutdown(0); + peers.serverAddress = peers.clientAddress = UNASSIGNED_SYSTEM_ADDRESS; + ASSERT_TRUE(peers.Connect()); + SendProbe(peers); + bool received = false; + const TimeMS again = GetTimeMS(); + while (!received && GetTimeMS() - again < (TimeMS)kWaitTimeoutMs) + { + for (Packet *p = peers.server->Receive(); p; peers.server->DeallocatePacket(p), p = peers.server->Receive()) + received |= p->data[0] == kProbe; + for (Packet *p = peers.client->Receive(); p; peers.client->DeallocatePacket(p), p = peers.client->Receive()) + ; + RakSleep(1); + } + EXPECT_TRUE(received); +}