From 8d835dc96db3b2771d29ace6ec2fa1afa8774e81 Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Wed, 26 Jul 2023 07:01:48 +0800
Subject: [PATCH] Unexpected comsuming (#48)

---
 Cargo.toml   |  8 ++++++--
 src/socks.rs | 22 +++++++++++++---------
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 11adae6..25a01d5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -23,7 +23,7 @@ mio = { version = "0.8", features = ["os-poll", "net", "os-ext"] }
 nix = { version = "0.26", features = ["process", "signal"] }
 prctl = "1.0"
 smoltcp = { version = "0.10.0", features = ["std", "phy-tuntap_interface"] }
-socks5-impl = { version = "0.4", default-features = false }
+socks5-impl = { version = "0.5", default-features = false }
 thiserror = "1.0"
 unicase = "2.6.0"
 url = "2.4"
@@ -34,6 +34,10 @@ jni = { version = "0.21", default-features = false }
 
 [dev-dependencies]
 ctor = "0.2"
-reqwest = { version = "0.11", default-features = false, features = ["blocking", "json", "rustls-tls"] }
+reqwest = { version = "0.11", default-features = false, features = [
+    "blocking",
+    "json",
+    "rustls-tls",
+] }
 serial_test = "2.0"
 test-log = "0.2"
diff --git a/src/socks.rs b/src/socks.rs
index 3c84b20..6c57f26 100644
--- a/src/socks.rs
+++ b/src/socks.rs
@@ -6,7 +6,9 @@ use crate::{
     },
 };
 use smoltcp::wire::IpProtocol;
-use socks5_impl::protocol::{self, handshake, password_method, Address, AuthMethod, UserKey};
+use socks5_impl::protocol::{
+    self, handshake, password_method, Address, AuthMethod, StreamOperation, UserKey,
+};
 use std::{collections::VecDeque, net::SocketAddr, rc::Rc};
 
 #[derive(Eq, PartialEq, Debug)]
@@ -132,16 +134,17 @@ impl SocksConnection {
     }
 
     fn receive_server_hello_socks5(&mut self) -> Result<(), Error> {
-        let response = handshake::Response::rebuild_from_stream(&mut self.server_inbuf);
+        let response = handshake::Response::retrieve_from_stream(&mut self.server_inbuf.clone());
         if let Err(e) = &response {
             if e.kind() == std::io::ErrorKind::UnexpectedEof {
-                log::trace!("receive_server_hello_socks5 need more data \"{}\"...", e);
+                log::trace!("receive_server_hello_socks5 needs more data \"{}\"...", e);
                 return Ok(());
             } else {
                 return Err(e.to_string().into());
             }
         }
         let respones = response?;
+        self.server_inbuf.drain(0..respones.len());
         let auth_method = respones.method;
 
         if auth_method != AuthMethod::NoAuth && self.credentials.is_none()
@@ -176,17 +179,18 @@ impl SocksConnection {
     }
 
     fn receive_auth_data(&mut self) -> Result<(), Error> {
-        let response = password_method::Response::rebuild_from_stream(&mut self.server_inbuf);
+        use password_method::Response;
+        let response = Response::retrieve_from_stream(&mut self.server_inbuf.clone());
         if let Err(e) = &response {
             if e.kind() == std::io::ErrorKind::UnexpectedEof {
-                log::trace!("receive_auth_data need more data \"{}\"...", e);
+                log::trace!("receive_auth_data needs more data \"{}\"...", e);
                 return Ok(());
             } else {
                 return Err(e.to_string().into());
             }
         }
-        assert!(self.server_inbuf.is_empty());
         let response = response?;
+        self.server_inbuf.drain(0..response.len());
         if response.status != password_method::Status::Succeeded {
             return Err(format!("SOCKS authentication failed: {:?}", response.status).into());
         }
@@ -195,20 +199,20 @@ impl SocksConnection {
     }
 
     fn receive_connection_status(&mut self) -> Result<(), Error> {
-        let response = protocol::Response::rebuild_from_stream(&mut self.server_inbuf);
+        let response = protocol::Response::retrieve_from_stream(&mut self.server_inbuf.clone());
         if let Err(e) = &response {
             if e.kind() == std::io::ErrorKind::UnexpectedEof {
-                log::trace!("receive_connection_status need more data \"{}\"...", e);
+                log::trace!("receive_connection_status needs more data \"{}\"...", e);
                 return Ok(());
             } else {
                 return Err(e.to_string().into());
             }
         }
         let response = response?;
+        self.server_inbuf.drain(0..response.len());
         if response.reply != protocol::Reply::Succeeded {
             return Err(format!("SOCKS connection failed: {}", response.reply).into());
         }
-        assert!(self.server_inbuf.is_empty());
         self.server_outbuf.append(&mut self.data_buf);
         self.data_buf.clear();