From 5e04f48500ceda3369957a8fb338c679e77fdd0e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?=
 <jedrzej.stuczynski@gmail.com>
Date: Thu, 23 Mar 2023 11:09:43 +0000
Subject: [PATCH] bugfix: drop tasks to connections closed by remote (#3190)

* applied patch #3187

* applied the same concept to the verloc listener
---
 common/mixnode-common/src/verloc/listener.rs  | 24 +++++++++++--------
 .../receiver/connection_handler.rs            |  7 +++---
 .../node/listener/connection_handler/mod.rs   |  7 +++---
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/common/mixnode-common/src/verloc/listener.rs b/common/mixnode-common/src/verloc/listener.rs
index a41e4edc8f..b91b7f50bc 100644
--- a/common/mixnode-common/src/verloc/listener.rs
+++ b/common/mixnode-common/src/verloc/listener.rs
@@ -99,17 +99,24 @@ impl ConnectionHandler {
         let mut framed_conn = Framed::new(conn, EchoPacketCodec);
         while !shutdown_listener.is_shutdown() {
             tokio::select! {
-                Some(echo_packet) = framed_conn.next() => {
+                biased;
+               _ = shutdown_listener.recv() => {
+                    trace!("ConnectionHandler: Shutdown received");
+                }
+                maybe_echo_packet = framed_conn.next() => {
                     // handle echo packet
-                    let reply_packet = match echo_packet {
-                        Ok(echo_packet) => self.handle_echo_packet(echo_packet),
-                        Err(err) => {
-                            error!(
-                                "The socket connection got corrupted with error: {}. Closing the socket",
-                                err
+                    let reply_packet = match maybe_echo_packet {
+                        Some(Ok(echo_packet)) => self.handle_echo_packet(echo_packet),
+                        Some(Err(err)) => {
+                             error!(
+                                "The socket connection got corrupted with error: {err}. Closing the socket",
                             );
                             return;
                         }
+                        None => {
+                            error!("The socket connection got terminated by the remote!");
+                            return;
+                        }
                     };
 
                     // write back the reply (note the lack of framing)
@@ -125,9 +132,6 @@ impl ConnectionHandler {
                         return;
                     }
                 },
-                _ = shutdown_listener.recv() => {
-                    trace!("ConnectionHandler: Shutdown received");
-                }
             }
         }
     }
diff --git a/gateway/src/node/mixnet_handling/receiver/connection_handler.rs b/gateway/src/node/mixnet_handling/receiver/connection_handler.rs
index bd566f5c2f..66a7ddfbab 100644
--- a/gateway/src/node/mixnet_handling/receiver/connection_handler.rs
+++ b/gateway/src/node/mixnet_handling/receiver/connection_handler.rs
@@ -189,9 +189,9 @@ impl<St: Storage> ConnectionHandler<St> {
                 _ = shutdown.recv() => {
                     log::trace!("ConnectionHandler: received shutdown");
                 }
-                Some(framed_sphinx_packet) = framed_conn.next() => {
+                framed_sphinx_packet = framed_conn.next() => {
                     match framed_sphinx_packet {
-                        Ok(framed_sphinx_packet) => {
+                        Some(Ok(framed_sphinx_packet)) => {
                             // TODO: benchmark spawning tokio task with full processing vs just processing it
                             // synchronously under higher load in single and multi-threaded situation.
 
@@ -200,12 +200,13 @@ impl<St: Storage> ConnectionHandler<St> {
                             // that change would only slow things down
                             self.handle_received_packet(framed_sphinx_packet).await;
                         }
-                        Err(err) => {
+                        Some(Err(err)) => {
                             error!(
                                 "The socket connection got corrupted with error: {err}. Closing the socket",
                             );
                             return;
                         }
+                        None => break, // stream got closed by remote
                     }
                 }
             }
diff --git a/mixnode/src/node/listener/connection_handler/mod.rs b/mixnode/src/node/listener/connection_handler/mod.rs
index af72f29620..684793b013 100644
--- a/mixnode/src/node/listener/connection_handler/mod.rs
+++ b/mixnode/src/node/listener/connection_handler/mod.rs
@@ -85,9 +85,9 @@ impl ConnectionHandler {
                 _ = shutdown.recv() => {
                     log::trace!("ConnectionHandler: received shutdown");
                 }
-                Some(framed_sphinx_packet) = framed_conn.next() => {
+                framed_sphinx_packet = framed_conn.next() => {
                     match framed_sphinx_packet {
-                        Ok(framed_sphinx_packet) => {
+                        Some(Ok(framed_sphinx_packet)) => {
                             // TODO: benchmark spawning tokio task with full processing vs just processing it
                             // synchronously (without delaying inside of course,
                             // delay is moved to a global DelayQueue)
@@ -98,12 +98,13 @@ impl ConnectionHandler {
                             // that change would only slow things down
                             self.handle_received_packet(framed_sphinx_packet);
                         }
-                        Err(err) => {
+                        Some(Err(err)) => {
                             error!(
                                 "The socket connection got corrupted with error: {err}. Closing the socket",
                             );
                             return;
                         }
+                        None => break, // stream got closed by remote
                     }
                 },
             }