From fbb0f2485fba02f85ade5169769888304d32d94f Mon Sep 17 00:00:00 2001
From: Jedrzej Stuczynski <jedrzej.stuczynski@gmail.com>
Date: Mon, 20 Jan 2020 16:44:57 +0000
Subject: [PATCH] Removing some print statements in favour of log

---
 nym-client/src/sockets/ws.rs | 35 ++++++++++++++++++++++++++---------
 1 file changed, 26 insertions(+), 9 deletions(-)

diff --git a/nym-client/src/sockets/ws.rs b/nym-client/src/sockets/ws.rs
index 578b661e12..222a9baf35 100644
--- a/nym-client/src/sockets/ws.rs
+++ b/nym-client/src/sockets/ws.rs
@@ -6,7 +6,7 @@ use futures::channel::{mpsc, oneshot};
 use futures::future::FutureExt;
 use futures::io::Error;
 use futures::{SinkExt, StreamExt};
-use log::*;
+use log::{debug, error, info, trace, warn};
 use serde::{Deserialize, Serialize};
 use sphinx::route::{Destination, DestinationAddressBytes};
 use std::io;
@@ -168,6 +168,8 @@ enum ServerResponse {
 
 impl Into<Message> for ServerResponse {
     fn into(self) -> Message {
+        // it should be safe to call `unwrap` here as the message is generated by the server
+        // so if it fails (and consequently panics) it's a bug that should be resolved
         let str_res = serde_json::to_string(&self).unwrap();
         Message::Text(str_res)
     }
@@ -210,14 +212,12 @@ async fn accept_connection(
     let address = stream
         .peer_addr()
         .expect("connected streams should have a peer address");
-    println!("Peer address: {}", address);
+    debug!("Peer address: {}", address);
 
     let mut ws_stream = tokio_tungstenite::accept_async(stream)
         .await
         .expect("Error during the websocket handshake occurred");
 
-    println!("New WebSocket connection: {}", address);
-
     // Create a channel for our stream, which other sockets will use to
     // send us messages. Then register our address with the stream to send
     // data to us.
@@ -235,12 +235,29 @@ async fn accept_connection(
     tokio::spawn(handle_connection(conn));
 
     while let Some(message) = ws_stream.next().await {
-        let message = message.expect("Failed to get request");
-        msg_tx
-            .unbounded_send(message)
-            .expect("Failed to forward request");
+        let message = match message {
+            Ok(msg) => msg,
+            Err(err) => {
+                error!("failed to obtain message from websocket stream! stopping connection handler: {}", err);
+                return;
+            }
+        };
+        if let Err(err) = msg_tx.unbounded_send(message) {
+            error!(
+                "Failed to forward request. Closing the socket connection: {}",
+                err
+            );
+            return;
+        }
         if let Some(resp) = response_rx.next().await {
-            ws_stream.send(resp).await.expect("Failed to send response");
+            if let Err(err) = ws_stream.send(resp).await {
+                warn!(
+                    "Failed to send message over websocket: {}. Assuming the connection is dead.",
+                    err
+                );
+                return;
+            }
+        }
         }
     }
 }