socks5: fix error handling in awaiting phase

Optimize error handling in other states
cyphernet-labs · May 4, 2024 · 873815d · 873815d
1 parent 2db9a90
commit 873815d
Showing 1 changed file with 13 additions and 7 deletions.
diff --git a/socks5-client/src/lib.rs b/socks5-client/src/lib.rs
@@ -83,13 +83,20 @@ impl Socks5 {
     pub fn advance(&mut self, input: &[u8]) -> Result<Vec<u8>, Error> {
         eprintln!("SOCKS5 now {self:?}");
         eprintln!("> Input is {input:02x?}");
+        if !input.is_empty() && input[0] != 0x05 {
+            *self = Socks5::Failed(Error::VersionNotSupported(input[0]));
+            return Err(Error::VersionNotSupported(input[0]));
+        }
         match self {
             Socks5::Initial(addr, false) if !addr.requires_proxy() => {
                 *self = Socks5::Active(addr.clone());
                 Ok(vec![])
             }
             Socks5::Initial(addr, _) => {
-                debug_assert!(input.is_empty());
+                if !input.is_empty() {
+                    *self = Socks5::Failed(Error::InvalidReply);
+                    return Err(Error::InvalidReply);
+                }
                 let out = vec![0x05, 0x01, 0x00];
                 *self = Socks5::Connected(addr.clone());
                 Ok(out)
@@ -99,10 +106,6 @@ impl Socks5 {
                     *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 {
                     *self = Socks5::Failed(Error::AuthRequired);
                     return Err(Error::AuthRequired);
@@ -114,8 +117,11 @@ impl Socks5 {
                 Ok(out)
             }
             Socks5::Awaiting => {
-                debug_assert_eq!(input.len(), 3);
-                if input[0] != 0x00 {
+                if input.len() != 3 {
+                    *self = Socks5::Failed(Error::InvalidReply);
+                    return Err(Error::InvalidReply);
+                }
+                if input[1] != 0x00 {
                     let err = ServerError::from(input[1]);
                     *self = Socks5::Rejected(err);
                     return Err(Error::Closed);