fix: missing clear relayed path when connection closed (#25)

* fix: missing clear relayed path when connection closed * add log
8xFF · Dec 18, 2024 · 36a167d · 36a167d
1 parent 4f07c93
commit 36a167d
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 3 deletions.
diff --git a/src/ctx.rs b/src/ctx.rs
@@ -146,7 +146,9 @@ impl SharedCtx {
             }
             RouteAction::Next(next) => {
                 let source = self.router.local_id();
-                self.conn(&next).ok_or(anyhow!("peer not found"))?.try_send(PeerMessage::Unicast(source, dest, service_id, data))?;
+                self.conn(&next)
+                    .ok_or(anyhow!("peer not found {next}"))?
+                    .try_send(PeerMessage::Unicast(source, dest, service_id, data))?;
                 Ok(())
             }
         }
@@ -160,7 +162,10 @@ impl SharedCtx {
             }
             RouteAction::Next(next) => {
                 let source = self.router.local_id();
-                self.conn(&next).ok_or(anyhow!("peer not found"))?.send(PeerMessage::Unicast(source, dest, service_id, data)).await?;
+                self.conn(&next)
+                    .ok_or(anyhow!("peer not found {next}"))?
+                    .send(PeerMessage::Unicast(source, dest, service_id, data))
+                    .await?;
                 Ok(())
             }
         }
@@ -201,7 +206,7 @@ impl SharedCtx {
             }
             RouteAction::Next(next) => {
                 let source = self.router.local_id();
-                Ok(self.conn(&next).ok_or(anyhow!("peer not found"))?.open_stream(service, source, dest, meta).await?)
+                Ok(self.conn(&next).ok_or(anyhow!("peer not found {next}"))?.open_stream(service, source, dest, meta).await?)
             }
         }
     }

diff --git a/src/router.rs b/src/router.rs
@@ -128,6 +128,14 @@ impl RouterTable {
                 }
             }
         }
+
+        // we also need to remove relayed path which go over this connection
+        for (peer, memory) in self.peers.iter_mut() {
+            if memory.paths.remove(conn).is_some() {
+                Self::select_best_for(peer, memory);
+            }
+        }
+        self.peers.retain(|_k, v| v.best().is_some());
     }
 
     fn action(&self, dest: &PeerId) -> Option<RouteAction> {
@@ -331,4 +339,25 @@ mod tests {
         // we shouldn't create sync with peer2 because it over MAX_HOPS
         assert_eq!(table.create_sync(&peer3), RouterTableSync(vec![(peer1, (0, 100).into())]));
     }
+
+    #[test_log::test]
+    fn should_remove_relay_path_after_disconnect() {
+        let mut table = RouterTable::new(PeerId(0));
+
+        let peer1 = PeerId(1);
+        let conn1 = ConnectionId(1);
+
+        let peer2 = PeerId(2);
+
+        table.set_direct(conn1, peer1, 100);
+
+        table.apply_sync(conn1, RouterTableSync(vec![(peer2, (MAX_HOPS, 200).into())]));
+        assert_eq!(table.next_remote(&peer2), Some((conn1, (MAX_HOPS + 1, 300).into())));
+
+        // after disconnect peer1
+        table.del_direct(&conn1);
+
+        // we should not have peer2 anymore
+        assert_eq!(table.next_remote(&peer2), None);
+    }
 }