fix: correctly remove device from list when sending

Signed-off-by: Martichou <[email protected]>
Martichou · Jul 8, 2024 · 758ada8 · 758ada8
1 parent 614bcbc
commit 758ada8
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 18 deletions.
diff --git a/app/common/vue_lib/src/index.ts b/app/common/vue_lib/src/index.ts
@@ -126,7 +126,7 @@ export async function sendInfo(vm: TauriVM, eid: string) {
 	if (vm.outboundPayload === undefined) return;
 
 	const ei = vm.endpointsInfo.find((el) => el.id === eid);
-	if (!ei) return;
+	if (!ei || !ei.ip || !ei.port) return;
 
 	const msg: SendInfo = {
 		id: ei.id,

diff --git a/core_lib/src/hdl/mdns_discovery.rs b/core_lib/src/hdl/mdns_discovery.rs
@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use mdns_sd::{ServiceDaemon, ServiceEvent};
 use serde::{Deserialize, Serialize};
 use tokio::net::TcpStream;
@@ -11,6 +13,7 @@ use crate::DeviceType;
 #[derive(Debug, Clone, Default, Deserialize, Serialize, TS)]
 #[ts(export)]
 pub struct EndpointInfo {
+    pub fullname: String,
     pub id: String,
     pub name: Option<String>,
     pub ip: Option<String>,
@@ -37,6 +40,9 @@ impl MDnsDiscovery {
         let service_type = "_FC9F5ED42C8A._tcp.local.";
         let receiver = self.daemon.browse(service_type)?;
 
+        // Map with fullname as key and EndpointInfo as value
+        let mut cache: HashMap<String, EndpointInfo> = HashMap::new();
+
         loop {
             tokio::select! {
                 _ = ctk.cancelled() => {
@@ -77,25 +83,40 @@ impl MDnsDiscovery {
                                         Err(_) => continue
                                     };
 
-                                    if TcpStream::connect(format!("{ip}:{port}")).await.is_ok() {
+                                    let ip_port = format!("{ip}:{port}");
+                                    let fullname = info.get_fullname().to_string();
+                                    if TcpStream::connect(&ip_port).await.is_ok() {
                                         let ei = EndpointInfo {
-                                            id: info.get_fullname().to_string(),
+                                            fullname: fullname.clone(),
+                                            id: ip_port,
                                             name: Some(dn),
                                             ip: Some(ip.to_string()),
                                             port: Some(port.to_string()),
                                             rtype: Some(dt),
                                             present: Some(true),
                                         };
-                                        info!("Resolved a new service: {:?}", ei);
+                                        info!("ServiceResolved: Resolved a new service: {:?}", ei);
+                                        cache.insert(fullname.clone(), ei.clone());
                                         let _ = self.sender.send(ei);
                                     }
                                 }
                                 ServiceEvent::ServiceRemoved(_, fullname) => {
-                                    info!("Remove a previous service: {}", fullname);
-                                    let _ = self.sender.send(EndpointInfo {
-                                        id: fullname,
-                                        ..Default::default()
-                                    });
+                                    trace!("ServiceRemoved: checking if should remove {}", fullname);
+                                    // Only remove if it has not been seen in the last cleanup_threshold
+                                    let should_remove = if let Some(ei) = cache.get(&fullname) {
+                                        Some(ei.id.clone())
+                                    } else {
+                                        None
+                                    };
+
+                                    if let Some(id) = should_remove {
+                                        info!("ServiceRemoved: Remove a previous service: {}", fullname);
+                                        cache.remove(&fullname);
+                                        let _ = self.sender.send(EndpointInfo {
+                                            id,
+                                            ..Default::default()
+                                        });
+                                    }
                                 }
                                 ServiceEvent::SearchStarted(_) | ServiceEvent::SearchStopped(_) => {}
                                 _ => {}

diff --git a/core_lib/src/manager.rs b/core_lib/src/manager.rs
@@ -1,4 +1,3 @@
-use anyhow::anyhow;
 use serde::{Deserialize, Serialize};
 use tokio::net::{TcpListener, TcpStream};
 use tokio::sync::broadcast::Sender;
@@ -58,7 +57,7 @@ impl TcpServer {
                 Some(i) = self.connect_receiver.recv() => {
                     info!("{INNER_NAME}: connect_receiver: got {:?}", i);
                     if let Err(e) = self.connect(cctk, i).await {
-                        error!("{INNER_NAME}: error sending: {:?}", e);
+                        error!("{INNER_NAME}: error sending: {}", e.to_string());
                     }
                 }
                 r = self.tcp_listener.accept() => {
@@ -112,13 +111,7 @@ impl TcpServer {
     /// To be called inside a separate task if we want to handle concurrency
     pub async fn connect(&self, ctk: CancellationToken, si: SendInfo) -> Result<(), anyhow::Error> {
         debug!("{INNER_NAME}: Connecting to: {}", si.addr);
-        let socket = match TcpStream::connect(si.addr.clone()).await {
-            Ok(r) => r,
-            Err(e) => {
-                warn!("Couldn't connect to {}: {}", si.addr, e);
-                return Err(anyhow!("failed to connect to {}", si.addr));
-            }
-        };
+        let socket = TcpStream::connect(si.addr.clone()).await?;
 
         let mut or = OutboundRequest::new(
             self.endpoint_id,