From fdbd791e8585ff172e93998e93ea04a4237de173 Mon Sep 17 00:00:00 2001
From: PinkP4nther <0x0090@protonmail.com>
Date: Fri, 3 Sep 2021 21:24:39 -0700
Subject: [PATCH] Performance & Error handling updates

---
 .gitignore  |  3 ++-
 src/data.rs | 34 ++++++++++++++++++++++++++++------
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/.gitignore b/.gitignore
index 20f520a..b835682 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,3 @@
 /target
-ideas.txt
\ No newline at end of file
+ideas.txt
+Cargo.lock
diff --git a/src/data.rs b/src/data.rs
index 9ea9bfe..225c57f 100644
--- a/src/data.rs
+++ b/src/data.rs
@@ -89,7 +89,7 @@ impl DownStreamInner {
                         println!("[!] Could not receive DataPipe::Finished notifier!");
                     }
 
-                } else if byte_count == 0 {
+                } else if byte_count == 0 || byte_count == -2 {
 
                     let _ = data_out.send(FullDuplexTcpState::DownStreamShutDown);
                     let _ = self.ds_stream.as_ref().unwrap().lock().unwrap().get_ref().shutdown(Shutdown::Both);
@@ -107,11 +107,13 @@ impl DownStreamInner {
 
         let mut data_length: i64 = 0;
 
+        let mut stream_lock = self.ds_stream.as_mut().unwrap().lock().unwrap();
+
         loop {
 
             let mut r_buf = [0; 1024];
 
-            match self.ds_stream.as_mut().unwrap().lock().unwrap().read(&mut r_buf) {
+            match stream_lock.read(&mut r_buf) {
 
                 Ok(bytes_read) => {
 
@@ -120,10 +122,15 @@ impl DownStreamInner {
 
                     } else if bytes_read != 0 && bytes_read <= 1024 {
 
+                        /*
                         let mut tmp_buf = r_buf.to_vec();
                         tmp_buf.truncate(bytes_read);
+                        */
+                        
 
-                        let _bw = self.internal_data_buffer.write(&tmp_buf).unwrap();
+                        //let _bw = self.internal_data_buffer.write(&tmp_buf).unwrap();
+
+                        let _bw = self.internal_data_buffer.write(r_buf.split_at(bytes_read).0).unwrap();
                         data_length += bytes_read as i64;
 
                     } else {
@@ -141,6 +148,10 @@ impl DownStreamInner {
                             break;
 
                         },
+                        io::ErrorKind::ConnectionReset => {
+                            data_length = -2;
+                            break;
+                        },
                         _ => {println!("[!!!] Got error: {}",e);}
                     }
                 },
@@ -214,7 +225,7 @@ impl UpStreamInner {
                         println!("[!] Could not receive DataPipe::Finished notifier!");
                     }
 
-                } else if byte_count == 0 {
+                } else if byte_count == 0 || byte_count == -2 {
 
                     let _ = data_out.send(FullDuplexTcpState::UpStreamShutDown);
                     let _ = self.us_stream.as_ref().unwrap().lock().unwrap().get_ref().shutdown(Shutdown::Both);
@@ -231,11 +242,13 @@ impl UpStreamInner {
 
         let mut data_length: i64 = 0;
 
+        let mut stream_lock = self.us_stream.as_mut().unwrap().lock().unwrap();
+
         loop {
 
             let mut r_buf = [0; 1024];
 
-            match self.us_stream.as_mut().unwrap().lock().unwrap().read(&mut r_buf) {
+            match stream_lock.read(&mut r_buf) {
 
                 Ok(bytes_read) => {
 
@@ -245,10 +258,15 @@ impl UpStreamInner {
 
                     } else if bytes_read != 0 && bytes_read <= 1024 {
 
+                        /*
                         let mut tmp_buf = r_buf.to_vec();
                         tmp_buf.truncate(bytes_read);
+                        */
+                        
 
-                        let _bw = self.internal_data_buffer.write(&tmp_buf).unwrap();
+                        //let _bw = self.internal_data_buffer.write(&tmp_buf).unwrap();
+
+                        let _bw = self.internal_data_buffer.write(r_buf.split_at(bytes_read).0).unwrap();
                         data_length += bytes_read as i64;
 
                     } else {
@@ -264,6 +282,10 @@ impl UpStreamInner {
                             }
                             break;
                         },
+                        io::ErrorKind::ConnectionReset => {
+                            data_length = -2;
+                            break;
+                        },
                         _ => {println!("[!!!] Got error: {}",e);}
                     }
                 },