From 57dafaf8ab91de3f302376b73224cad67c65fac3 Mon Sep 17 00:00:00 2001
From: Michel de Graaf <michel@re-invention.nl>
Date: Sat, 5 Jul 2025 22:54:28 +0200
Subject: [PATCH] feat: remove peers

---
 src/client.rs            | 15 +++++++++++++++
 src/peer/default_peer.rs |  6 +++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/client.rs b/src/client.rs
index efbcb18..1fd92c1 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -21,6 +21,7 @@ const MAX_THREADS: usize = 100;
 pub enum ClientOperation {
     ConnectToPeer(Peer),
     SearchResult(FileSearchResult),
+    PeerDisconnected(String),
 }
 struct ClientContext {
     peers: HashMap<String, DefaultPeer>,
@@ -103,6 +104,14 @@ impl Client {
         }
     }
 
+    #[allow(dead_code)]
+    pub fn remove_peer(&self, username: &str) {
+        let mut context = self.context.lock().unwrap();
+        if let Some(peer) = context.peers.remove(username) {
+            drop(peer); // Explicitly drop to trigger cleanup
+        }
+    }
+
     pub fn search(&self, query: &str, timeout: Duration) -> Vec<FileSearchResult> {
         println!("Searching for {}", query);
         if let Some(server) = &self.server {
@@ -139,6 +148,12 @@ impl Client {
                             .search_results
                             .push(file_search);
                     }
+                    ClientOperation::PeerDisconnected(username) => {
+                        let mut context = client_context.lock().unwrap();
+                        if let Some(peer) = context.peers.remove(&username) {
+                            drop(peer); // Explicitly drop to trigger cleanup
+                        }
+                    }
                 }
             }
         });
diff --git a/src/peer/default_peer.rs b/src/peer/default_peer.rs
index 79599c8..a8b08bd 100644
--- a/src/peer/default_peer.rs
+++ b/src/peer/default_peer.rs
@@ -36,7 +36,6 @@ impl DefaultPeer {
             peer,
             peer_channel: None,
             client_channel,
-            // Initialize handles as None
             read_thread: None,
             write_thread: None,
         }
@@ -69,6 +68,7 @@ impl DefaultPeer {
         let mut write_stream = stream; // Use the original stream for writing
 
         let peer = self.peer.clone();
+        let client_channel_for_read = self.client_channel.clone();
 
         // Spawn the reader thread
         self.read_thread = Some(thread::spawn(move || {
@@ -91,6 +91,7 @@ impl DefaultPeer {
                     }
                     Err(e) => {
                         eprintln!("Error reading from peer: {}. Terminating read loop.", e);
+                        let _ = client_channel_for_read.send(ClientOperation::PeerDisconnected(peer.username.clone()));
                         break;
                     }
                 }
@@ -102,6 +103,7 @@ impl DefaultPeer {
                             "Error extracting message in default peer: {}. Terminating read loop.",
                             e
                         );
+                        let _ = client_channel_for_read.send(ClientOperation::PeerDisconnected(peer.username.clone()));
                         break;
                     }
                     Ok(None) => continue,
@@ -110,6 +112,7 @@ impl DefaultPeer {
         }));
 
         let client_channel = self.client_channel.clone();
+        let peer_username = self.peer.username.clone();
         self.write_thread = Some(thread::spawn(move || {
             loop {
                 match peer_reader.recv() {
@@ -118,6 +121,7 @@ impl DefaultPeer {
                             PeerOperation::SendMessage(message) => {
                                 if let Err(e) = write_stream.write_all(&message.get_buffer()) {
                                     eprintln!("Error writing message to stream: {}. Terminating write loop.", e);
+                                    let _ = client_channel.send(ClientOperation::PeerDisconnected(peer_username.clone()));
                                     break;
                                 }
                             }