socks5: fix error handling in awaiting phase

Optimize error handling in other states
cyphernet-labs · May 4, 2024 · cb1d671 · cb1d671
1 parent 2db9a90
commit cb1d671
Showing 1 changed file with 12 additions and 11 deletions.
diff --git a/socks5-client/src/lib.rs b/socks5-client/src/lib.rs
@@ -95,10 +95,7 @@ impl Socks5 {
                 Ok(out)
             }
             Socks5::Connected(addr) => {
-                if input.len() != 2 {
-                    *self = Socks5::Failed(Error::InvalidReply);
-                    return Err(Error::InvalidReply);
-                }
+                debug_assert_eq!(input.len(), 2);
                 if input[0] != 0x05 {
                     *self = Socks5::Failed(Error::VersionNotSupported(input[0]));
                     return Err(Error::VersionNotSupported(input[0]));
@@ -114,13 +111,17 @@ impl Socks5 {
                 Ok(out)
             }
             Socks5::Awaiting => {
-                debug_assert_eq!(input.len(), 3);
-                if input[0] != 0x00 {
+                debug_assert_eq!(input.len(), 5);
+                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[3], input[4]);
                 Ok(vec![])
             }
             Socks5::Reading(code1, code2) => {
@@ -141,10 +142,10 @@ impl Socks5 {
         match self {
             Socks5::Initial(_, _) => 0,
             Socks5::Connected(_) => 2,
-            Socks5::Awaiting => 3,
-            Socks5::Reading(ty, _) if *ty == IPV4 => 4,
-            Socks5::Reading(ty, _) if *ty == IPV6 => 16,
-            Socks5::Reading(ty, len) if *ty == DOMAIN => *len as usize,
+            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,
         }