Connect command on exising TCP connection

src/tcp.rs

@@ -37,7 +37,7 @@ impl Socks5Stream {
         P: ToProxyAddrs,
         T: IntoTargetAddr<'t>,
     {
-        Self::connect_raw(proxy, target, Authentication::None, Command::Connect)
+        Self::connect_raw(proxy, None, target, Authentication::None, Command::Connect)
     }
 
     /// Connects to a target server through a SOCKS5 proxy using given username and password.
@@ -57,6 +57,45 @@ impl Socks5Stream {
     {
         Self::connect_raw(
             proxy,
+            None,
+            target,
+            Authentication::Password { username, password },
+            Command::Connect,
+        )
+    }
+
+    /// Connects to a target server through a SOCKS5 proxy on an existing TCP connection
+    /// to the proxy
+    pub fn connect_existing<'a, 't, T>(
+        proxy_conn: TcpStream,
+        target: T,
+    ) -> Result<ConnectFuture<'static, 't, stream::Once<SocketAddr, Error>>>
+    where
+        T: IntoTargetAddr<'t>,
+    {
+        Self::connect_raw(
+            proxy_conn.peer_addr()?,
+            Some(proxy_conn),
+            target,
+            Authentication::None,
+            Command::Connect,
+        )
+    }
+
+    /// Connects to a target server through a SOCKS5 proxy on an existing TCP connection
+    /// to the proxy using given username and password
+    pub fn connect_existing_with_password<'a, 't, T>(
+        proxy_conn: TcpStream,
+        target: T,
+        username: &'a str,
+        password: &'a str,
+    ) -> Result<ConnectFuture<'a, 't, stream::Once<SocketAddr, Error>>>
+    where
+        T: IntoTargetAddr<'t>,
+    {
+        Self::connect_raw(
+            proxy_conn.peer_addr()?,
+            Some(proxy_conn),
             target,
             Authentication::Password { username, password },
             Command::Connect,
@@ -65,6 +104,7 @@ impl Socks5Stream {
 
     fn connect_raw<'a, 't, P, T>(
         proxy: P,
+        proxy_conn: Option<TcpStream>,
         target: T,
         auth: Authentication<'a>,
         command: Command,
@@ -91,6 +131,7 @@ impl Socks5Stream {
             auth,
             command,
             proxy.to_proxy_addrs(),
+            proxy_conn,
             target.into_target_addr()?,
         ))
     }
@@ -131,13 +172,24 @@ impl<'a, 't, S> ConnectFuture<'a, 't, S>
 where
     S: Stream<Item = SocketAddr, Error = Error>,
 {
-    fn new(auth: Authentication<'a>, command: Command, proxy: S, target: TargetAddr<'t>) -> Self {
+    fn new(
+        auth: Authentication<'a>,
+        command: Command,
+        proxy: S,
+        proxy_conn: Option<TcpStream>,
+        target: TargetAddr<'t>,
+    ) -> Self {
+        let state = if let Some(c) = proxy_conn {
+            ConnectState::Connected(Some(c))
+        } else {
+            ConnectState::Uninitialized
+        };
         ConnectFuture {
             auth,
             command,
             proxy,
             target,
-            state: ConnectState::Uninitialized,
+            state,
             buf: [0; 513],
             ptr: 0,
             len: 0,
@@ -238,17 +290,17 @@ where
                 ConnectState::Created(ref mut conn_fut) => match conn_fut.poll() {
                     Ok(Async::Ready(tcp)) => {
                         self.state = ConnectState::Connected(Some(tcp));
-                        self.prepare_send_method_selection()
                     }
                     Ok(Async::NotReady) => return Ok(Async::NotReady),
                     Err(_e) => self.state = ConnectState::Uninitialized,
                 },
                 ConnectState::Connected(ref mut opt) => {
-                    let tcp = opt.as_mut().unwrap();
+                    let mut conn = opt.take();
+                    let tcp = conn.as_mut().unwrap();
+                    self.prepare_send_method_selection();
                     self.ptr += try_ready!(tcp.poll_write(&self.buf[self.ptr..self.len]));
-                    ;
                     if self.ptr == self.len {
-                        self.state = ConnectState::MethodSent(opt.take());
+                        self.state = ConnectState::MethodSent(conn);
                         self.prepare_recv_method_selection();
                     }
                 }
@@ -260,22 +312,23 @@ where
                             Err(Error::InvalidResponseVersion)?
                         }
                         match self.buf[1] {
-                            0x00 => self.state = ConnectState::PrepareRequest(opt.take()),
+                            0x00 => self.state = ConnectState::SendRequest(opt.take()),
                             0xff => Err(Error::NoAcceptableAuthMethods)?,
                             0x02 => {
                                 self.state = ConnectState::PasswordAuth(opt.take());
-                                self.prepare_send_password_auth();
                             }
                             m if m != self.auth.id() => Err(Error::UnknownAuthMethod)?,
                             _ => unimplemented!(),
                         }
                     }
                 }
                 ConnectState::PasswordAuth(ref mut opt) => {
-                    let tcp = opt.as_mut().unwrap();
+                    let mut conn = opt.take();
+                    let tcp = conn.as_mut().unwrap();
+                    self.prepare_send_password_auth();
                     self.ptr += try_ready!(tcp.poll_write(&self.buf[self.ptr..self.len]));
                     if self.ptr == self.len {
-                        self.state = ConnectState::PasswordAuthSent(opt.take());
+                        self.state = ConnectState::PasswordAuthSent(conn);
                         self.prepare_recv_password_auth();
                     }
                 }
@@ -289,18 +342,16 @@ where
                         if self.buf[1] != 0x00 {
                             Err(Error::PasswordAuthFailure(self.buf[1]))?
                         }
-                        self.state = ConnectState::PrepareRequest(opt.take());
+                        self.state = ConnectState::SendRequest(opt.take());
                     }
                 }
-                ConnectState::PrepareRequest(ref mut opt) => {
-                    self.state = ConnectState::SendRequest(opt.take());
-                    self.prepare_send_request();
-                }
                 ConnectState::SendRequest(ref mut opt) => {
-                    let tcp = opt.as_mut().unwrap();
+                    let mut conn = opt.take();
+                    let tcp = conn.as_mut().unwrap();
+                    self.prepare_send_request();
                     self.ptr += try_ready!(tcp.poll_write(&self.buf[self.ptr..self.len]));
                     if self.ptr == self.len {
-                        self.state = ConnectState::RequestSent(opt.take());
+                        self.state = ConnectState::RequestSent(conn);
                         self.prepare_recv_reply();
                     }
                 }
@@ -408,7 +459,6 @@ enum ConnectState {
     MethodSent(Option<TcpStream>),
     PasswordAuth(Option<TcpStream>),
     PasswordAuthSent(Option<TcpStream>),
-    PrepareRequest(Option<TcpStream>),
     SendRequest(Option<TcpStream>),
     RequestSent(Option<TcpStream>),
     PrepareReadAddress(Option<TcpStream>),
@@ -442,6 +492,7 @@ impl Socks5Listener {
             Authentication::None,
             Command::Bind,
             proxy.to_proxy_addrs(),
+            None,
             target.into_target_addr()?,
         )))
     }
@@ -469,6 +520,7 @@ impl Socks5Listener {
             Authentication::Password { username, password },
             Command::Bind,
             proxy.to_proxy_addrs(),
+            None,
             target.into_target_addr()?,
         )))
     }

tests/common.rs

@@ -8,7 +8,7 @@ use tokio::{
     runtime::Runtime,
 };
 use tokio_socks::{
-    tcp::{BindFuture, ConnectFuture},
+    tcp::{BindFuture, Socks5Stream},
     Error,
 };
 
@@ -32,9 +32,9 @@ pub fn echo_server(runtime: &mut Runtime) -> Result<()> {
     Ok(())
 }
 
-pub fn test_connect<S>(conn: ConnectFuture<'static, 'static, S>) -> Result<()>
+pub fn test_connect<F>(conn: F) -> Result<()>
 where
-    S: Stream<Item = SocketAddr, Error = Error> + Send + 'static,
+    F: Future<Item = Socks5Stream, Error = Error> + Send + 'static,
 {
     let fut = conn
         .and_then(|tcp| write_all(tcp, MSG).map_err(Into::into))
@@ -46,6 +46,7 @@ where
     Ok(())
 }
 
+#[allow(dead_code)]
 pub fn test_bind<S>(bind: BindFuture<'static, 'static, S>) -> Result<()>
 where
     S: Stream<Item = SocketAddr, Error = Error> + Send + 'static,

tests/integration_tests.sh

@@ -19,6 +19,7 @@ for test in ${list}; do
     test_exit_code=$?
 
     pkill -F /tmp/3proxy-test.pid
+    sleep 1
 
     if test "$test_exit_code" -ne 0; then
         break
@@ -27,4 +28,4 @@ done
 
 
 #pkill -F /tmp/socat-test.pid
-exit ${test_exit_code}
+exit ${test_exit_code}

tests/long_username_password_auth.rs

@@ -1,10 +1,13 @@
 mod common;
 
 use common::{test_bind, test_connect, ECHO_SERVER_ADDR, PROXY_ADDR};
+use futures::future::{err, Future};
+use std::{io, net::SocketAddr};
 use tokio_socks::{
     tcp::{Socks5Listener, Socks5Stream},
     Error,
 };
+use tokio_tcp::TcpStream;
 
 type Result<T> = std::result::Result<T, Error>;
 
@@ -21,3 +24,24 @@ fn bind() -> Result<()> {
         Socks5Listener::bind_with_password(PROXY_ADDR, ECHO_SERVER_ADDR, "mylonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglogin", "longlonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglongpassword")?;
     test_bind(bind)
 }
+
+#[test]
+fn connect_existing() -> Result<()> {
+    let addr = PROXY_ADDR
+        .parse::<SocketAddr>()
+        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+    let conn = TcpStream::connect(&addr).map_err(Error::Io).and_then(
+        |c| -> Box<Future<Item = Socks5Stream, Error = Error> + Send + 'static> {
+            match Socks5Stream::connect_existing_with_password(
+                c,
+                ECHO_SERVER_ADDR,
+                "mylonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglogin",
+                "longlonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglongpassword",
+            ) {
+                Ok(f) => Box::new(f),
+                Err(e) => Box::new(err(e)),
+            }
+        },
+    );
+    test_connect(conn)
+}

tests/no_auth.rs

@@ -1,10 +1,13 @@
 mod common;
 
 use common::{test_bind, test_connect, ECHO_SERVER_ADDR, PROXY_ADDR};
+use futures::future::{err, Future};
+use std::{io, net::SocketAddr};
 use tokio_socks::{
     tcp::{Socks5Listener, Socks5Stream},
     Error,
 };
+use tokio_tcp::TcpStream;
 
 type Result<T> = std::result::Result<T, Error>;
 
@@ -19,3 +22,19 @@ fn bind() -> Result<()> {
     let bind = Socks5Listener::bind(PROXY_ADDR, ECHO_SERVER_ADDR)?;
     test_bind(bind)
 }
+
+#[test]
+fn connect_existing() -> Result<()> {
+    let addr = PROXY_ADDR
+        .parse::<SocketAddr>()
+        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+    let conn = TcpStream::connect(&addr).map_err(Error::Io).and_then(
+        |c| -> Box<Future<Item = Socks5Stream, Error = Error> + Send + 'static> {
+            match Socks5Stream::connect_existing(c, ECHO_SERVER_ADDR) {
+                Ok(f) => Box::new(f),
+                Err(e) => Box::new(err(e)),
+            }
+        },
+    );
+    test_connect(conn)
+}

tests/username_auth.rs

@@ -1,10 +1,13 @@
 mod common;
 
 use common::{test_bind, test_connect, ECHO_SERVER_ADDR, PROXY_ADDR};
+use futures::future::{err, Future};
+use std::{io, net::SocketAddr};
 use tokio_socks::{
     tcp::{Socks5Listener, Socks5Stream},
     Error,
 };
+use tokio_tcp::TcpStream;
 
 type Result<T> = std::result::Result<T, Error>;
 
@@ -21,3 +24,24 @@ fn bind() -> Result<()> {
         Socks5Listener::bind_with_password(PROXY_ADDR, ECHO_SERVER_ADDR, "mylogin", "mypassword")?;
     test_bind(bind)
 }
+
+#[test]
+fn connect_existing() -> Result<()> {
+    let addr = PROXY_ADDR
+        .parse::<SocketAddr>()
+        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+    let conn = TcpStream::connect(&addr).map_err(Error::Io).and_then(
+        |c| -> Box<Future<Item = Socks5Stream, Error = Error> + Send + 'static> {
+            match Socks5Stream::connect_existing_with_password(
+                c,
+                ECHO_SERVER_ADDR,
+                "mylogin",
+                "mypassword",
+            ) {
+                Ok(f) => Box::new(f),
+                Err(e) => Box::new(err(e)),
+            }
+        },
+    );
+    test_connect(conn)
+}