socks5: fix error handling in awaiting phase

Optimize error handling in other states
cyphernet-labs · May 4, 2024 · 89b4ed4 · 89b4ed4
1 parent 2db9a90
commit 89b4ed4
Showing 1 changed file with 11 additions and 4 deletions.
diff --git 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,