From ab4e39e1b3a60814b2e846f85d72c1b557f4204e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= <jon.haggblad@gmail.com>
Date: Thu, 1 Dec 2022 14:25:59 +0100
Subject: [PATCH] Fix a few errors in socks5 client and network-requester
 (backport) (#1824)

* Fix two unwraps in socks5 and network-requester

* Make sure client task never sends shutdown signal
---
 clients/socks5/src/socks/client.rs            | 20 +++++++++++++------
 .../network-requester/src/core.rs             | 10 +++++++---
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/clients/socks5/src/socks/client.rs b/clients/socks5/src/socks/client.rs
index 3e3c24347c..739bd84987 100644
--- a/clients/socks5/src/socks/client.rs
+++ b/clients/socks5/src/socks/client.rs
@@ -168,9 +168,13 @@ impl SocksClient {
         controller_sender: ControllerSender,
         self_address: Recipient,
         lane_queue_lengths: LaneQueueLengths,
-        shutdown_listener: ShutdownListener,
+        mut shutdown_listener: ShutdownListener,
     ) -> Self {
+        // If this task fails and exits, we don't want to send shutdown signal
+        shutdown_listener.mark_as_success();
+
         let connection_id = Self::generate_random();
+
         SocksClient {
             controller_sender,
             connection_id,
@@ -252,10 +256,15 @@ impl SocksClient {
             .await;
 
         let stream = self.stream.run_proxy();
-        let local_stream_remote = stream
-            .peer_addr()
-            .expect("failed to extract peer address")
-            .to_string();
+        let peer_addr = match stream.peer_addr() {
+            Ok(peer_addr) => peer_addr,
+            Err(err) => {
+                log::error!("Unable to extract the remote peer address: {err}");
+                return;
+            }
+        };
+        let local_stream_remote = peer_addr.to_string();
+
         let connection_id = self.connection_id;
         let input_sender = self.input_sender.clone();
 
@@ -319,7 +328,6 @@ impl SocksClient {
             SocksCommand::UdpAssociate => unimplemented!(), // not handled
         };
 
-        self.shutdown_listener.mark_as_success();
         Ok(())
     }
 
diff --git a/service-providers/network-requester/src/core.rs b/service-providers/network-requester/src/core.rs
index 870376d896..e58bffaef0 100644
--- a/service-providers/network-requester/src/core.rs
+++ b/service-providers/network-requester/src/core.rs
@@ -157,9 +157,13 @@ impl ServiceProvider {
         lane_queue_lengths: LaneQueueLengths,
     ) -> Option<ReconstructedMessage> {
         while let Some(msg) = websocket_reader.next().await {
-            let data = msg
-                .expect("we failed to read from the websocket!")
-                .into_data();
+            let data = match msg {
+                Ok(msg) => msg.into_data(),
+                Err(err) => {
+                    log::error!("Failed to read from the websocket: {err}");
+                    continue;
+                }
+            };
 
             // try to recover the actual message from the mix network...
             let deserialized_message = match ServerResponse::deserialize(&data) {