From 66c71250b85407e53148a3b0f1601fe163fd9d6f Mon Sep 17 00:00:00 2001
From: GUVWAF <thijs@havinga.eu>
Date: Thu, 8 Jun 2023 20:02:36 +0200
Subject: [PATCH] Update last_sent_by_id in FloodingRouter

---
 src/mesh/FloodingRouter.cpp        | 11 +++++++----
 src/mesh/FloodingRouter.h          |  3 ++-
 src/modules/NeighborInfoModule.cpp | 20 +++++++++++++++++++-
 src/modules/NeighborInfoModule.h   |  6 ++++++
 4 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/mesh/FloodingRouter.cpp b/src/mesh/FloodingRouter.cpp
index 171199277..f412c0e1c 100644
--- a/src/mesh/FloodingRouter.cpp
+++ b/src/mesh/FloodingRouter.cpp
@@ -44,10 +44,13 @@ void FloodingRouter::sniffReceived(const meshtastic_MeshPacket *p, const meshtas
 
                 tosend->hop_limit--; // bump down the hop count
 
-                // If it is a traceRoute request, update the route that it went via me
-                if (p->which_payload_variant == meshtastic_MeshPacket_decoded_tag && traceRouteModule &&
-                    traceRouteModule->wantPacket(p)) {
-                    traceRouteModule->updateRoute(tosend);
+                if (p->which_payload_variant == meshtastic_MeshPacket_decoded_tag) {
+                    // If it is a traceRoute request, update the route that it went via me
+                    if (traceRouteModule && traceRouteModule->wantPacket(p))
+                        traceRouteModule->updateRoute(tosend);
+                    // If it is a neighborInfo packet, update last_sent_by_id
+                    if (neighborInfoModule && neighborInfoModule->wantPacket(p))
+                        neighborInfoModule->updateLastSentById(tosend);
                 }
 
                 LOG_INFO("Rebroadcasting received floodmsg to neighbors\n");
diff --git a/src/mesh/FloodingRouter.h b/src/mesh/FloodingRouter.h
index 73dbd1f3f..bf3e34c24 100644
--- a/src/mesh/FloodingRouter.h
+++ b/src/mesh/FloodingRouter.h
@@ -2,6 +2,7 @@
 
 #include "PacketHistory.h"
 #include "Router.h"
+#include "modules/NeighborInfoModule.h"
 #include "modules/TraceRouteModule.h"
 
 /**
@@ -57,4 +58,4 @@ class FloodingRouter : public Router, protected PacketHistory
      * Look for broadcasts we need to rebroadcast
      */
     virtual void sniffReceived(const meshtastic_MeshPacket *p, const meshtastic_Routing *c) override;
-};
+};
\ No newline at end of file
diff --git a/src/modules/NeighborInfoModule.cpp b/src/modules/NeighborInfoModule.cpp
index d4b53fabc..be1d70fd8 100644
--- a/src/modules/NeighborInfoModule.cpp
+++ b/src/modules/NeighborInfoModule.cpp
@@ -174,11 +174,29 @@ bool NeighborInfoModule::handleReceivedProtobuf(const meshtastic_MeshPacket &mp,
 {
     printNeighborInfo("RECIEVED", np);
     updateNeighbors(mp, np);
-    np->last_sent_by_id = nodeDB.getNodeNum();
     // Allow others to handle this packet
     return false;
 }
 
+/*
+Copy the content of a current NeighborInfo packet into a new one and update the last_sent_by_id to our NodeNum
+*/
+void NeighborInfoModule::updateLastSentById(meshtastic_MeshPacket *p)
+{
+    auto &incoming = p->decoded;
+    meshtastic_NeighborInfo scratch;
+    meshtastic_NeighborInfo *updated = NULL;
+    memset(&scratch, 0, sizeof(scratch));
+    pb_decode_from_bytes(incoming.payload.bytes, incoming.payload.size, &meshtastic_NeighborInfo_msg, &scratch);
+    updated = &scratch;
+
+    updated->last_sent_by_id = nodeDB.getNodeNum();
+
+    // Set updated last_sent_by_id to the payload of the to be flooded packet
+    p->decoded.payload.size =
+        pb_encode_to_bytes(p->decoded.payload.bytes, sizeof(p->decoded.payload.bytes), &meshtastic_NeighborInfo_msg, updated);
+}
+
 void NeighborInfoModule::resetNeighbors()
 {
     neighborState.neighbors_count = 0;
diff --git a/src/modules/NeighborInfoModule.h b/src/modules/NeighborInfoModule.h
index a6f61b025..923ac1f7e 100644
--- a/src/modules/NeighborInfoModule.h
+++ b/src/modules/NeighborInfoModule.h
@@ -20,6 +20,9 @@ class NeighborInfoModule : public ProtobufModule<meshtastic_NeighborInfo>, priva
 
     bool saveProtoForModule();
 
+    // Let FloodingRouter call updateLastSentById upon rebroadcasting a NeighborInfo packet
+    friend class FloodingRouter;
+
   protected:
     // Note: this holds our local info.
     meshtastic_NeighborInfo neighborState;
@@ -58,6 +61,9 @@ class NeighborInfoModule : public ProtobufModule<meshtastic_NeighborInfo>, priva
     /* update neighbors with subpacket sniffed from network */
     void updateNeighbors(const meshtastic_MeshPacket &mp, meshtastic_NeighborInfo *np);
 
+    /* update a NeighborInfo packet with our NodeNum as last_sent_by_id */
+    void updateLastSentById(meshtastic_MeshPacket *p);
+
     void loadProtoForModule();
 
     /* Does our periodic broadcast */