From 2b46a4d5532fb60f2025651a16c0b9104f43675c Mon Sep 17 00:00:00 2001
From: Adrian Wann <95002807+AWann2@users.noreply.github.com>
Date: Wed, 6 Nov 2024 14:45:10 +1100
Subject: [PATCH] Remove partially received packets over NUClearNet after
 timeout (#158)

Fix for memory issue in NUClearNet due to partially received messages
never being cleaned up. This resulted in the process eventually crashing
due to being out of memory.

This fix checks all the packet assemblers each time a packet chunk is
received and removes any packets that have not had a new chunk in a
while.
---
 src/extension/network/NUClearNetwork.cpp | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/extension/network/NUClearNetwork.cpp b/src/extension/network/NUClearNetwork.cpp
index 79c993ff..778adaa1 100644
--- a/src/extension/network/NUClearNetwork.cpp
+++ b/src/extension/network/NUClearNetwork.cpp
@@ -851,6 +851,15 @@ namespace extension {
                                     // We have completed this packet, discard the data
                                     assemblers.erase(assemblers.find(packet.packet_id));
                                 }
+
+                                // Check for and delete any timed out packets
+                                for (auto it = assemblers.begin(); it != assemblers.end();) {
+                                    const auto now              = std::chrono::steady_clock::now();
+                                    const auto timeout          = remote->round_trip_time * 10.0;
+                                    const auto& last_chunk_time = it->second.first;
+
+                                    it = now > last_chunk_time + timeout ? assemblers.erase(it) : std::next(it);
+                                }
                             }
                         }
                     } break;