From d982a9c3f3bc299bf1f8edf7cdab1fb8891c99ec Mon Sep 17 00:00:00 2001
From: Lee Smet <lee.smet@hotmail.com>
Date: Tue, 27 Feb 2024 17:01:59 +0100
Subject: [PATCH] Add method to remove existing peers (endpoints)

Signed-off-by: Lee Smet <lee.smet@hotmail.com>
---
 src/peer_manager.rs | 33 +++++++++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/peer_manager.rs b/src/peer_manager.rs
index 90b56d0..4e13599 100644
--- a/src/peer_manager.rs
+++ b/src/peer_manager.rs
@@ -95,10 +95,14 @@ pub struct PeerStats {
     pub connection_rx_bytes: u64,
 }
 
-/// Marker to indicate a [`peer`](Endpoint) is already known.
+/// Marker error to indicate a [`peer`](Endpoint) is already known.
 #[derive(Debug)]
 pub struct PeerExists;
 
+/// Marker error to indicate a [`peer`](Endpoint) is not known.
+#[derive(Debug)]
+pub struct PeerNotFound;
+
 struct Inner {
     /// Router is unfortunately wrapped in a Mutex, because router is not Sync.
     router: Mutex<Router>,
@@ -173,7 +177,7 @@ impl PeerManager {
     ///
     /// # Errors
     ///
-    /// This function returns an error if the endpoint is already known.
+    /// This function returns an error if the [`Endpoint`] is already known.
     pub fn add_peer(&self, peer: Endpoint) -> Result<(), PeerExists> {
         let mut peer_map = self.inner.peers.lock().unwrap();
         if peer_map.contains_key(&peer) {
@@ -192,6 +196,23 @@ impl PeerManager {
         Ok(())
     }
 
+    /// Delete a peer from the system.
+    ///
+    /// The peer will be disconnected if it is currently connected.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if there is no peer identified by the given [`Endpoint`].
+    pub fn delete_peer(&self, endpoint: &Endpoint) -> Result<(), PeerNotFound> {
+        let mut peer_map = self.inner.peers.lock().unwrap();
+        peer_map.remove(endpoint).ok_or(PeerNotFound).map(|pi| {
+            // Make sure we kill the peer connection if one exists
+            if let Some(peer) = pi.pr.upgrade() {
+                peer.died();
+            }
+        })
+    }
+
     /// Get a view of all known peers and their stats.
     pub fn peers(&self) -> Vec<PeerStats> {
         let peer_map = self.inner.peers.lock().unwrap();
@@ -772,3 +793,11 @@ impl fmt::Display for PeerExists {
 }
 
 impl std::error::Error for PeerExists {}
+
+impl fmt::Display for PeerNotFound {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str("Peer identified by endpoint not known")
+    }
+}
+
+impl std::error::Error for PeerNotFound {}