fix: Survive a shutdown from inside receive dispatch

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

Source/include/mafianet/peer.h

@@ -1110,6 +1110,12 @@ class RAK_DLL_EXPORT RakPeer : public RakPeerInterface, public RNS2EventHandler
 
 	SimpleMutex packetAllocationPoolMutex;
 	DataStructures::MemoryPool<Packet> 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<Packet*> packetReturnQueue;

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<FileListProgress*> &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)
 			{

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<Packet>::MemoryWithPage)*32);
+	packetsOutstanding = 0;
 	packetAllocationPoolMutex.Unlock();
+	shutdownGeneration = 0;
 
 	remoteSystemIndexPool.SetPageSize(sizeof(DataStructures::MemoryPool<RemoteSystemIndex>::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