From cd958d0942d6d780952ee0e2b719968fc458da57 Mon Sep 17 00:00:00 2001
From: jbesraa <jbesraa@gmail.com>
Date: Thu, 16 Jan 2025 10:35:28 +0200
Subject: [PATCH] Handle unwrap when shutting down JDC server

..This should prevent the code from panicing if a signal
to shutdown is sent the JDC server
---
 protocols/v2/roles-logic-sv2/src/errors.rs | 2 ++
 roles/jd-client/src/lib/downstream.rs      | 4 +++-
 roles/jd-client/src/lib/mod.rs             | 7 +++++--
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/protocols/v2/roles-logic-sv2/src/errors.rs b/protocols/v2/roles-logic-sv2/src/errors.rs
index 20c4bcd554..6badff21c1 100644
--- a/protocols/v2/roles-logic-sv2/src/errors.rs
+++ b/protocols/v2/roles-logic-sv2/src/errors.rs
@@ -61,6 +61,7 @@ pub enum Error {
     HashrateError(InputError),
     LogicErrorMessage(std::boxed::Box<AllMessages<'static>>),
     JDSMissingTransactions,
+    IoError(std::io::Error),
 }
 
 impl From<BinarySv2Error> for Error {
@@ -153,6 +154,7 @@ impl Display for Error {
             HashrateError(e) => write!(f, "Impossible to get Hashrate: {:?}", e),
             LogicErrorMessage(e) => write!(f, "Message is well formatted but can not be handled: {:?}", e),
             JDSMissingTransactions => write!(f, "JD server cannot propagate the block: missing transactions"),
+            IoError(e) => write!(f, "IO error: {:?}", e),
         }
     }
 }
diff --git a/roles/jd-client/src/lib/downstream.rs b/roles/jd-client/src/lib/downstream.rs
index c5d49d304f..d71473385e 100644
--- a/roles/jd-client/src/lib/downstream.rs
+++ b/roles/jd-client/src/lib/downstream.rs
@@ -682,7 +682,9 @@ pub async fn listen_for_downstream_mining(
     jd: Option<Arc<Mutex<JobDeclarator>>>,
 ) -> Result<Arc<Mutex<DownstreamMiningNode>>, Error> {
     info!("Listening for downstream mining connections on {}", address);
-    let listner = TcpListener::bind(address).await.unwrap();
+    let listner = TcpListener::bind(address)
+        .await
+        .map_err(|e| Error::IoError(e))?;
 
     if let Ok((stream, _)) = listner.accept().await {
         let responder = Responder::from_authority_kp(
diff --git a/roles/jd-client/src/lib/mod.rs b/roles/jd-client/src/lib/mod.rs
index ee8b4c1118..307adeec37 100644
--- a/roles/jd-client/src/lib/mod.rs
+++ b/roles/jd-client/src/lib/mod.rs
@@ -337,7 +337,7 @@ impl JobDeclaratorClient {
         };
 
         // Wait for downstream to connect
-        let downstream = downstream::listen_for_downstream_mining(
+        let downstream = match downstream::listen_for_downstream_mining(
             downstream_addr,
             Some(upstream),
             send_solution,
@@ -351,7 +351,10 @@ impl JobDeclaratorClient {
             Some(jd.clone()),
         )
         .await
-        .unwrap();
+        {
+            Ok(d) => d,
+            Err(_e) => return,
+        };
 
         TemplateRx::connect(
             SocketAddr::new(IpAddr::from_str(ip_tp.as_str()).unwrap(), port_tp),