From 89b4ed48614adf9cead06c716f97ac930c36fbca Mon Sep 17 00:00:00 2001
From: Dr Maxim Orlovsky <orlovsky@cyphernet.io>
Date: Sat, 4 May 2024 19:14:41 +0200
Subject: [PATCH] socks5: fix error handling in awaiting phase

Optimize error handling in other states
---
 socks5-client/src/lib.rs | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/socks5-client/src/lib.rs b/socks5-client/src/lib.rs
index ee99b88..a2a898a 100644
--- a/socks5-client/src/lib.rs
+++ b/socks5-client/src/lib.rs
@@ -114,13 +114,20 @@ impl Socks5 {
                 Ok(out)
             }
             Socks5::Awaiting => {
-                debug_assert_eq!(input.len(), 3);
-                if input[0] != 0x00 {
+                if input.len() != 4 {
+                    *self = Socks5::Failed(Error::InvalidReply);
+                    return Err(Error::InvalidReply);
+                }
+                if input[0] != 0x05 {
+                    *self = Socks5::Failed(Error::VersionNotSupported(input[0]));
+                    return Err(Error::VersionNotSupported(input[0]));
+                }
+                if input[1] != 0x00 {
                     let err = ServerError::from(input[1]);
                     *self = Socks5::Rejected(err);
                     return Err(Error::Closed);
                 }
-                *self = Socks5::Reading(input[1], input[2]);
+                *self = Socks5::Reading(input[2], input[3]);
                 Ok(vec![])
             }
             Socks5::Reading(code1, code2) => {
@@ -141,7 +148,7 @@ impl Socks5 {
         match self {
             Socks5::Initial(_, _) => 0,
             Socks5::Connected(_) => 2,
-            Socks5::Awaiting => 3,
+            Socks5::Awaiting => 4,
             Socks5::Reading(ty, _) if *ty == IPV4 => 4,
             Socks5::Reading(ty, _) if *ty == IPV6 => 16,
             Socks5::Reading(ty, len) if *ty == DOMAIN => *len as usize,