From 5e6cfe8955ba2e5949e9e5e1784608d323cc216c Mon Sep 17 00:00:00 2001
From: Aram Peres <aram.peres@wavy.fm>
Date: Sat, 16 Oct 2021 21:12:08 -0400
Subject: [PATCH] Don't break TCP connection until all data is flushed

Potential fix for #14
---
 src/main.rs | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 4337461..8384d3e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -179,10 +179,6 @@ async fn handle_tcp_proxy_connection(
     trace!("[{}] Virtual client is ready to send data", virtual_port);
 
     loop {
-        if abort.load(Ordering::Relaxed) {
-            break;
-        }
-
         tokio::select! {
             readable_result = socket.readable() => {
                 match readable_result {
@@ -234,7 +230,11 @@ async fn handle_tcp_proxy_connection(
                             );
                         }
                         Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
-                            continue;
+                            if abort.load(Ordering::Relaxed) {
+                                break;
+                            } else {
+                                continue;
+                            }
                         }
                         Err(e) => {
                             error!(
@@ -243,7 +243,13 @@ async fn handle_tcp_proxy_connection(
                             );
                         }
                     },
-                    None => continue,
+                    None => {
+                        if abort.load(Ordering::Relaxed) {
+                            break;
+                        } else {
+                            continue;
+                        }
+                    },
                 }
             }
         }