From 3d3b9963ff161c53f2cee9083b996da629feacc4 Mon Sep 17 00:00:00 2001
From: Tanner Sommers <tanner@teamhydra.dev>
Date: Wed, 6 Mar 2024 15:03:39 -0500
Subject: [PATCH] fix ssdp hanging if it can't find any packets in time

---
 src-tauri/src/handlers/bambu/mod.rs | 45 ++++++++++++++++++++---------
 src-tauri/src/handlers/ssdp/mod.rs  | 23 +++++++++++----
 2 files changed, 48 insertions(+), 20 deletions(-)

diff --git a/src-tauri/src/handlers/bambu/mod.rs b/src-tauri/src/handlers/bambu/mod.rs
index 080d757..b89f4f0 100644
--- a/src-tauri/src/handlers/bambu/mod.rs
+++ b/src-tauri/src/handlers/bambu/mod.rs
@@ -617,21 +617,35 @@ impl BambuClient {
         for listener in ssdp_listeners {
             println!("[BambuClient::get_device_ips] Running SSDP Discovery ...");
 
-            let messages = listener.listen(Duration::from_secs(5)).await.map_err(|e| {
-                std::io::Error::new(
-                    std::io::ErrorKind::Other,
-                    format!(
-                        "[BambuClient::get_device_ips] Failed to listen for SSDP messages: {}",
+            // Wrap in a timeout to ensure we don't hang forever
+            let messages = tokio::time::timeout(
+                Duration::from_secs(5),
+                listener.listen(Duration::from_secs(5)),
+            );
+
+            match messages.await {
+                Ok(Ok(messages)) => {
+                    println!(
+                        "[BambuClient::get_device_ips] Successfully discovered {} SSDP messages.",
+                        messages.len()
+                    );
+                    ssdp_messages.extend(messages);
+                }
+                Ok(Err(e)) => {
+                    println!(
+                        "[BambuClient::get_device_ips] Failed to discover SSDP messages: {}",
                         e
-                    ),
-                )
-            })?;
+                    );
+                }
+                Err(_) => {
+                    println!(
+                        "[BambuClient::get_device_ips] Timed out while discovering SSDP messages."
+                    );
 
-            ssdp_messages.extend(messages);
-            println!(
-                "[BambuClient::get_device_ips] Found {} SSDP messages so far ...",
-                ssdp_messages.len(),
-            );
+                    // Continue to the next listener
+                    continue;
+                }
+            }
         }
 
         // de-dupe the messages by location
@@ -647,7 +661,10 @@ impl BambuClient {
 
         if unique_messages.len() == 0 {
             println!("[BambuClient::get_device_ips] No unique messages found. Exiting ...");
-            return Ok(vec![]);
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::Other,
+                "No unique messages found during SSDP discovery. Please ensure your devices are connected to the network and try again.",
+            ));
         }
 
         println!(
diff --git a/src-tauri/src/handlers/ssdp/mod.rs b/src-tauri/src/handlers/ssdp/mod.rs
index bff318f..cab7cf8 100644
--- a/src-tauri/src/handlers/ssdp/mod.rs
+++ b/src-tauri/src/handlers/ssdp/mod.rs
@@ -93,6 +93,10 @@ impl SsdpListener {
     ) -> Result<Vec<SsdpMessage>, Box<dyn std::error::Error>> {
         // Create a UDP socket bound to the specified port
         let socket = UdpSocket::bind(format!("0.0.0.0:{}", self.port))?;
+        socket.set_read_timeout(Some(duration)).map_err(|e| {
+            eprintln!("Error setting read timeout: {}", e);
+            e
+        })?;
 
         // Join the SSDP multicast group
         socket.join_multicast_v4(
@@ -113,15 +117,22 @@ impl SsdpListener {
 
         // Receive messages until the specified duration elapses
         while Instant::now() - start_time < duration {
-            match socket.recv_from(&mut buf[..]) {
+            match socket.recv_from(&mut buf) {
                 Ok((size, _)) => {
-                    // Parse and handle the SSDP NOTIFY message
-                    let message = std::str::from_utf8(&buf[..size])?.to_string();
-                    messages.push(SsdpMessage::from_message(&message)?);
+                    println!("Received message of size {}", size);
+                    let message = std::str::from_utf8(&buf[..size])?;
+                    let ssdp_message = SsdpMessage::from_message(message)?;
+                    messages.push(ssdp_message);
                 }
                 Err(e) => {
-                    eprintln!("Error receiving SSDP message: {}", e);
-                    break; // Exit loop on error
+                    if e.kind() == std::io::ErrorKind::WouldBlock
+                        || e.kind() == std::io::ErrorKind::TimedOut
+                    {
+                        println!("No more messages received within the specified duration");
+                        break;
+                    } else {
+                        eprintln!("Error receiving message: {}", e);
+                    }
                 }
             }
         }