From 4ade07f9a7efff6ec7e2a5ad02248e50f6b96d02 Mon Sep 17 00:00:00 2001
From: Lee Smet <lee.smet@hotmail.com>
Date: Tue, 23 Jan 2024 15:32:54 +0100
Subject: [PATCH] Fix #114: Prevent zombie quic connections

Signed-off-by: Lee Smet <lee.smet@hotmail.com>
---
 CHANGELOG.md        |  5 +++++
 src/peer.rs         |  5 +++++
 src/peer_manager.rs | 33 ++++++++++++++++++++++++++++++++-
 src/router.rs       |  2 +-
 4 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 99eee4b..a9d69b4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - If the router notices a Peer is dead, the connection is now forcibly terminated.
 
+### Fixed 
+
+- Inbound peers now replace existing peers in the peer manager. This should fix
+  an issue where Quic could leave zombie connections.
+
 ## [0.3.1] - 2024-01-23
 
 ### Added
diff --git a/src/peer.rs b/src/peer.rs
index c8036b8..d31f02b 100644
--- a/src/peer.rs
+++ b/src/peer.rs
@@ -286,6 +286,11 @@ impl PeerRef {
             false
         }
     }
+
+    /// Attempts to convert this `PeerRef` into a full [`Peer`].
+    pub fn upgrade(&self) -> Option<Peer> {
+        self.inner.upgrade().map(|inner| Peer { inner })
+    }
 }
 
 impl Default for PeerRef {
diff --git a/src/peer_manager.rs b/src/peer_manager.rs
index 6992494..b78d079 100644
--- a/src/peer_manager.rs
+++ b/src/peer_manager.rs
@@ -393,7 +393,7 @@ impl Inner {
                 Endpoint::new(Protocol::Quic, con.remote_address()),
                 PeerType::Inbound,
                 Some(new_peer),
-            );
+            )
         }
     }
 
@@ -416,6 +416,37 @@ impl Inner {
                 self.router.lock().unwrap().add_peer_interface(p);
             }
             info!("Added new peer {endpoint}");
+        } else if discovery_type == PeerType::Inbound {
+            // We got an inbound peer with a duplicate entry. This is possible if the sending port
+            // is the same as the previous one, which generally happens with our Quic setup. In
+            // this case, the old connection needs to be replaced.
+            let old_peer_info = peers.insert(
+                endpoint,
+                PeerInfo {
+                    pt: discovery_type,
+                    connecting: false,
+                    pr: if let Some(p) = &peer {
+                        p.refer()
+                    } else {
+                        PeerRef::new()
+                    },
+                    connection_attempts: 0,
+                },
+            );
+            // If we have a new peer notify insert the new one in the router, then notify it that
+            // the old one is dead.
+            if let Some(p) = peer {
+                let router = self.router.lock().unwrap();
+                router.add_peer_interface(p);
+                if let Some(old_peer) = old_peer_info
+                    .expect("We already checked the entry was occupied so this is always Some; qed")
+                    .pr
+                    .upgrade()
+                {
+                    router.handle_dead_peer(old_peer);
+                }
+            }
+            info!("Replaced existing inbound peer {endpoint}");
         } else {
             debug!("Ignoring request to add {endpoint} as it already exists");
         }
diff --git a/src/router.rs b/src/router.rs
index ae24326..4c0fd24 100644
--- a/src/router.rs
+++ b/src/router.rs
@@ -353,7 +353,7 @@ impl Router {
     }
 
     /// Remove a dead peer from the router.
-    fn handle_dead_peer(&self, dead_peer: Peer) {
+    pub fn handle_dead_peer(&self, dead_peer: Peer) {
         debug!(
             "Cleaning up peer {} which is reportedly dead",
             dead_peer.connection_identifier()