From 9335259aff558dd55f86f68ef6b266ee033d32e0 Mon Sep 17 00:00:00 2001
From: Dr Maxim Orlovsky <orlovsky@cyphernet.io>
Date: Sun, 5 May 2024 22:23:16 +0200
Subject: [PATCH] socks5: read remanining data from stream upon failure

---
 socks5-client/src/lib.rs | 22 ++++++++++++----------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/socks5-client/src/lib.rs b/socks5-client/src/lib.rs
index 1663242..8b4434c 100644
--- a/socks5-client/src/lib.rs
+++ b/socks5-client/src/lib.rs
@@ -71,7 +71,7 @@ pub enum Socks5 {
     Awaiting,
     Reading(u8, u8),
     Active(NetAddr<HostName>),
-    Rejected(ServerError),
+    Rejected(ServerError, u8, u8),
     Failed(Error),
 }
 
@@ -116,10 +116,10 @@ impl Socks5 {
                 }
                 if input[1] != 0x00 {
                     let err = ServerError::from(input[1]);
-                    *self = Socks5::Rejected(err);
-                    return Err(Error::Closed);
+                    *self = Socks5::Rejected(err, input[3], input[4]);
+                } else {
+                    *self = Socks5::Reading(input[3], input[4]);
                 }
-                *self = Socks5::Reading(input[3], input[4]);
                 Ok(vec![])
             }
             Socks5::Reading(code1, code2) => {
@@ -132,7 +132,7 @@ impl Socks5 {
                 Ok(vec![])
             }
             Socks5::Active(_) => Err(Error::Completed),
-            Socks5::Rejected(_) | Socks5::Failed(_) => Err(Error::Closed),
+            Socks5::Rejected(_, _, _) | Socks5::Failed(_) => Err(Error::Closed),
         }
     }
 
@@ -141,11 +141,13 @@ impl Socks5 {
             Socks5::Initial(_, _) => 0,
             Socks5::Connected(_) => 2,
             Socks5::Awaiting => 5,
-            Socks5::Reading(ty, _) if *ty == IPV4 => 5,
-            Socks5::Reading(ty, _) if *ty == IPV6 => 17,
-            Socks5::Reading(ty, len) if *ty == DOMAIN => *len as usize + 1,
-            Socks5::Reading(_, _) => 1,
-            Socks5::Active(_) | Socks5::Rejected(_) | Socks5::Failed(_) => 0,
+            Socks5::Reading(ty, _) | Socks5::Rejected(_, ty, _) if *ty == IPV4 => 5,
+            Socks5::Reading(ty, _) | Socks5::Rejected(_, ty, _) if *ty == IPV6 => 17,
+            Socks5::Reading(ty, len) | Socks5::Rejected(_, ty, len) if *ty == DOMAIN => {
+                *len as usize + 1
+            }
+            Socks5::Reading(_, _) | Socks5::Rejected(_, _, _) => 1,
+            Socks5::Active(_) | Socks5::Failed(_) => 0,
         }
     }
 }