socks5 support

Signed-off-by: fufesou <shuanglongchen@yeah.net>

libs/hbb_common/src/tcp.rs

@@ -4,16 +4,31 @@ use futures::{SinkExt, StreamExt};
 use protobuf::Message;
 use sodiumoxide::crypto::secretbox::{self, Key, Nonce};
 use std::{
-    io::{Error, ErrorKind},
+    io::{self, Error, ErrorKind},
+    net::SocketAddr,
     ops::{Deref, DerefMut},
+    pin::Pin,
+    task::{Context, Poll},
 };
-use tokio::net::{lookup_host, TcpListener, TcpSocket, TcpStream, ToSocketAddrs};
+use tokio::{
+    io::{AsyncRead, AsyncWrite, ReadBuf},
+    net::{lookup_host, TcpListener, TcpSocket, ToSocketAddrs},
+};
+use tokio_socks::{tcp::Socks5Stream, IntoTargetAddr, ToProxyAddrs};
 use tokio_util::codec::Framed;
 
-pub struct FramedStream(Framed<TcpStream, BytesCodec>, Option<(Key, u64, u64)>, u64);
+pub trait TcpStreamTrait: AsyncRead + AsyncWrite + Unpin {}
+pub struct DynTcpStream(Box<dyn TcpStreamTrait + Send>);
+
+pub struct FramedStream(
+    Framed<DynTcpStream, BytesCodec>,
+    SocketAddr,
+    Option<(Key, u64, u64)>,
+    u64,
+);
 
 impl Deref for FramedStream {
-    type Target = Framed<TcpStream, BytesCodec>;
+    type Target = Framed<DynTcpStream, BytesCodec>;
 
     fn deref(&self) -> &Self::Target {
         &self.0
@@ -26,6 +41,20 @@ impl DerefMut for FramedStream {
     }
 }
 
+impl Deref for DynTcpStream {
+    type Target = Box<dyn TcpStreamTrait + Send>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for DynTcpStream {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
 fn new_socket(addr: std::net::SocketAddr, reuse: bool) -> Result<TcpSocket, std::io::Error> {
     let socket = match addr {
         std::net::SocketAddr::V4(..) => TcpSocket::new_v4()?,
@@ -44,8 +73,8 @@ fn new_socket(addr: std::net::SocketAddr, reuse: bool) -> Result<TcpSocket, std:
 }
 
 impl FramedStream {
-    pub async fn new<T: ToSocketAddrs, T2: ToSocketAddrs>(
-        remote_addr: T,
+    pub async fn new<T1: ToSocketAddrs, T2: ToSocketAddrs>(
+        remote_addr: T1,
         local_addr: T2,
         ms_timeout: u64,
     ) -> ResultType<Self> {
@@ -56,27 +85,86 @@ impl FramedStream {
                     new_socket(local_addr, true)?.connect(remote_addr),
                 )
                 .await??;
-                return Ok(Self(Framed::new(stream, BytesCodec::new()), None, 0));
+                let addr = stream.local_addr()?;
+                return Ok(Self(
+                    Framed::new(DynTcpStream(Box::new(stream)), BytesCodec::new()),
+                    addr,
+                    None,
+                    0,
+                ));
             }
         }
         bail!("could not resolve to any address");
     }
 
-    pub fn set_send_timeout(&mut self, ms: u64) {
-        self.2 = ms;
+    pub async fn connect<'a, 't, P, T1, T2>(
+        proxy: P,
+        target: T1,
+        local: T2,
+        username: &'a str,
+        password: &'a str,
+        ms_timeout: u64,
+    ) -> ResultType<Self>
+    where
+        P: ToProxyAddrs,
+        T1: IntoTargetAddr<'t>,
+        T2: ToSocketAddrs,
+    {
+        if let Some(local) = lookup_host(&local).await?.next() {
+            if let Some(proxy) = proxy.to_proxy_addrs().next().await {
+                let stream =
+                    super::timeout(ms_timeout, new_socket(local, true)?.connect(proxy?)).await??;
+                let stream = if username.trim().is_empty() {
+                    super::timeout(
+                        ms_timeout,
+                        Socks5Stream::connect_with_socket(stream, target),
+                    )
+                    .await??
+                } else {
+                    super::timeout(
+                        ms_timeout,
+                        Socks5Stream::connect_with_password_and_socket(
+                            stream, target, username, password,
+                        ),
+                    )
+                    .await??
+                };
+                let addr = stream.local_addr()?;
+                return Ok(Self(
+                    Framed::new(DynTcpStream(Box::new(stream)), BytesCodec::new()),
+                    addr,
+                    None,
+                    0,
+                ));
+            };
+        };
+        bail!("could not resolve to any address");
     }
 
-    pub fn from(stream: TcpStream) -> Self {
-        Self(Framed::new(stream, BytesCodec::new()), None, 0)
+    pub fn local_addr(&self) -> SocketAddr {
+        self.1
+    }
+
+    pub fn set_send_timeout(&mut self, ms: u64) {
+        self.3 = ms;
+    }
+
+    pub fn from(stream: impl TcpStreamTrait + Send + 'static, addr: SocketAddr) -> Self {
+        Self(
+            Framed::new(DynTcpStream(Box::new(stream)), BytesCodec::new()),
+            addr,
+            None,
+            0,
+        )
     }
 
     pub fn set_raw(&mut self) {
         self.0.codec_mut().set_raw();
-        self.1 = None;
+        self.2 = None;
     }
 
     pub fn is_secured(&self) -> bool {
-        self.1.is_some()
+        self.2.is_some()
     }
 
     #[inline]
@@ -87,7 +175,7 @@ impl FramedStream {
     #[inline]
     pub async fn send_raw(&mut self, msg: Vec<u8>) -> ResultType<()> {
         let mut msg = msg;
-        if let Some(key) = self.1.as_mut() {
+        if let Some(key) = self.2.as_mut() {
             key.1 += 1;
             let nonce = Self::get_nonce(key.1);
             msg = secretbox::seal(&msg, &nonce, &key.0);
@@ -98,8 +186,8 @@ impl FramedStream {
 
     #[inline]
     pub async fn send_bytes(&mut self, bytes: Bytes) -> ResultType<()> {
-        if self.2 > 0 {
-            super::timeout(self.2, self.0.send(bytes)).await??;
+        if self.3 > 0 {
+            super::timeout(self.3, self.0.send(bytes)).await??;
         } else {
             self.0.send(bytes).await?;
         }
@@ -109,7 +197,7 @@ impl FramedStream {
     #[inline]
     pub async fn next(&mut self) -> Option<Result<BytesMut, Error>> {
         let mut res = self.0.next().await;
-        if let Some(key) = self.1.as_mut() {
+        if let Some(key) = self.2.as_mut() {
             if let Some(Ok(bytes)) = res.as_mut() {
                 key.2 += 1;
                 let nonce = Self::get_nonce(key.2);
@@ -137,7 +225,7 @@ impl FramedStream {
     }
 
     pub fn set_key(&mut self, key: Key) {
-        self.1 = Some((key, 0, 0));
+        self.2 = Some((key, 0, 0));
     }
 
     fn get_nonce(seqnum: u64) -> Nonce {
@@ -161,3 +249,35 @@ pub async fn new_listener<T: ToSocketAddrs>(addr: T, reuse: bool) -> ResultType<
         bail!("could not resolve to any address");
     }
 }
+
+impl Unpin for DynTcpStream {}
+
+impl AsyncRead for DynTcpStream {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        AsyncRead::poll_read(Pin::new(&mut self.0), cx, buf)
+    }
+}
+
+impl AsyncWrite for DynTcpStream {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        AsyncWrite::poll_write(Pin::new(&mut self.0), cx, buf)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        AsyncWrite::poll_flush(Pin::new(&mut self.0), cx)
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        AsyncWrite::poll_shutdown(Pin::new(&mut self.0), cx)
+    }
+}
+
+impl<R: AsyncRead + AsyncWrite + Unpin> TcpStreamTrait for R {}