socks5 support

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

libs/hbb_common/Cargo.toml

@@ -28,6 +28,11 @@ confy = { git = "https://github.com/open-trade/confy" }
 dirs-next = "2.0"
 filetime = "0.2"
 sodiumoxide = "0.2"
+tokio-socks = { git = "https://github.com/fufesou/tokio-socks" }
+futures-core = "0.3"
+futures-sink = "0.3"
+async-trait = "0.1"
+pin-project = "1"
 
 [target.'cfg(not(any(target_os = "android", target_os = "ios")))'.dependencies]
 mac_address = "1.1"

libs/hbb_common/src/config.rs

@@ -71,6 +71,16 @@ pub struct Config {
     keys_confirmed: HashMap<String, bool>,
 }
 
+#[derive(Debug, Default, Serialize, Deserialize, Clone)]
+pub struct Socks5Server {
+    #[serde(default)]
+    pub proxy: String,
+    #[serde(default)]
+    pub username: String,
+    #[serde(default)]
+    pub password: String,
+}
+
 // more variable configs
 #[derive(Debug, Default, Serialize, Deserialize, Clone)]
 pub struct Config2 {
@@ -85,6 +95,9 @@ pub struct Config2 {
     #[serde(default)]
     serial: i32,
 
+    #[serde(default)]
+    socks: Option<Socks5Server>,
+
     // the other scalar value must before this
     #[serde(default)]
     pub options: HashMap<String, String>,
@@ -619,6 +632,16 @@ impl Config {
     pub fn get_remote_id() -> String {
         CONFIG2.read().unwrap().remote_id.clone()
     }
+
+    pub fn set_socks(socks: Option<Socks5Server>) {
+        let mut config = CONFIG2.write().unwrap();
+        config.socks = socks;
+        config.store();
+    }
+
+    pub fn get_socks() -> Option<Socks5Server> {
+        CONFIG2.read().unwrap().socks.clone()
+    }
 }
 
 const PEERS: &str = "peers";

libs/hbb_common/src/lib.rs

@@ -17,16 +17,20 @@ pub use tokio;
 pub use tokio_util;
 pub mod tcp;
 pub mod udp;
+pub mod socket_client;
 pub use env_logger;
 pub use log;
 pub mod bytes_codec;
 #[cfg(feature = "quic")]
 pub mod quic;
 pub use anyhow::{self, bail};
+pub use futures_core;
+pub use futures_sink;
 pub use futures_util;
 pub mod config;
 pub mod fs;
 pub use sodiumoxide;
+pub use tokio_socks;
 
 #[cfg(feature = "quic")]
 pub type Stream = quic::Connection;

libs/hbb_common/src/socket_client.rs

@@ -0,0 +1,87 @@
+use crate::{
+    config::{Config, Socks5Server},
+    tcp::FramedStream,
+    udp::{FramedSocket, UdpFramedWrapper},
+    ResultType,
+};
+use anyhow::{bail, Context};
+use std::net::SocketAddr;
+use tokio::net::ToSocketAddrs;
+use tokio_socks::{udp::Socks5UdpFramed, IntoTargetAddr};
+use tokio_util::{codec::BytesCodec, udp::UdpFramed};
+
+pub fn get_socks5_conf() -> Option<Socks5Server> {
+    // Config::set_socks(Some(Socks5Server {
+    //     proxy: "139.186.136.143:1080".to_owned(),
+    //     ..Default::default()
+    // }));
+    Config::get_socks()
+}
+
+pub async fn connect_tcp<'t, T: IntoTargetAddr<'t>>(
+    target: T,
+    local: SocketAddr,
+    ms_timeout: u64,
+) -> ResultType<FramedStream> {
+    let target_addr = target.into_target_addr()?;
+
+    if let Some(conf) = get_socks5_conf() {
+        FramedStream::connect(
+            conf.proxy.as_str(),
+            target_addr,
+            local,
+            conf.username.as_str(),
+            conf.password.as_str(),
+            ms_timeout,
+        )
+        .await
+    } else {
+        let addrs: Vec<SocketAddr> =
+            std::net::ToSocketAddrs::to_socket_addrs(&target_addr)?.collect();
+        if addrs.is_empty() {
+            bail!("Invalid target addr");
+        };
+
+        FramedStream::new(addrs[0], local, ms_timeout)
+            .await
+            .with_context(|| "Failed to connect to rendezvous server")
+    }
+}
+
+// TODO: merge connect_udp and connect_udp_socks5
+pub async fn connect_udp_socket<T1: ToSocketAddrs>(
+    local: T1,
+) -> ResultType<(
+    FramedSocket<UdpFramedWrapper<UdpFramed<BytesCodec>>>,
+    Option<SocketAddr>,
+)> {
+    Ok((FramedSocket::new(local).await?, None))
+}
+
+pub async fn connect_udp_socks5<'t, T1: IntoTargetAddr<'t>, T2: ToSocketAddrs>(
+    target: T1,
+    local: T2,
+    socks5: &Option<Socks5Server>,
+    ms_timeout: u64,
+) -> ResultType<(
+    FramedSocket<UdpFramedWrapper<Socks5UdpFramed>>,
+    Option<SocketAddr>,
+)> {
+    match socks5 {
+        Some(conf) => {
+            let (socket, addr) = FramedSocket::connect(
+                conf.proxy.as_str(),
+                target,
+                local,
+                conf.username.as_str(),
+                conf.password.as_str(),
+                ms_timeout,
+            )
+            .await?;
+            Ok((socket, Some(addr)))
+        }
+        None => {
+            bail!("Nil socks5 server config")
+        }
+    }
+}

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 {}

libs/hbb_common/src/udp.rs

@@ -1,26 +1,62 @@
 use crate::{bail, ResultType};
-use bytes::BytesMut;
+use anyhow::anyhow;
+use bytes::{Bytes, BytesMut};
 use futures::{SinkExt, StreamExt};
+use futures_core::Stream;
+use futures_sink::Sink;
+use pin_project::pin_project;
 use protobuf::Message;
 use socket2::{Domain, Socket, Type};
 use std::{
-    io::Error,
     net::SocketAddr,
     ops::{Deref, DerefMut},
+    pin::Pin,
+    task::{Context, Poll},
+};
+use tokio::net::{ToSocketAddrs, UdpSocket};
+use tokio_socks::{
+    udp::{Socks5UdpFramed, Socks5UdpMessage},
+    IntoTargetAddr, TargetAddr, ToProxyAddrs,
 };
-use tokio::{net::ToSocketAddrs, net::UdpSocket};
 use tokio_util::{codec::BytesCodec, udp::UdpFramed};
 
-pub struct FramedSocket(UdpFramed<BytesCodec>);
+pub struct FramedSocket<F>(F);
 
-impl Deref for FramedSocket {
-    type Target = UdpFramed<BytesCodec>;
+#[pin_project]
+pub struct UdpFramedWrapper<F>(#[pin] F);
+
+pub trait BytesMutGetter<'a> {
+    fn get_bytes_mut(&'a self) -> &'a BytesMut;
+}
+
+impl<F> Deref for FramedSocket<F> {
+    type Target = F;
 
     fn deref(&self) -> &Self::Target {
         &self.0
     }
 }
 
+impl<F> DerefMut for FramedSocket<F> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl<F> Deref for UdpFramedWrapper<F> {
+    type Target = F;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<F> DerefMut for UdpFramedWrapper<F> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
 fn new_socket(addr: SocketAddr, reuse: bool) -> Result<Socket, std::io::Error> {
     let socket = match addr {
         SocketAddr::V4(..) => Socket::new(Domain::ipv4(), Type::dgram(), None),
@@ -38,50 +74,101 @@ fn new_socket(addr: SocketAddr, reuse: bool) -> Result<Socket, std::io::Error> {
     Ok(socket)
 }
 
-impl DerefMut for FramedSocket {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.0
-    }
-}
-
-impl FramedSocket {
+impl FramedSocket<UdpFramedWrapper<UdpFramed<BytesCodec>>> {
     pub async fn new<T: ToSocketAddrs>(addr: T) -> ResultType<Self> {
         let socket = UdpSocket::bind(addr).await?;
-        Ok(Self(UdpFramed::new(socket, BytesCodec::new())))
+        Ok(Self(UdpFramedWrapper(UdpFramed::new(
+            socket,
+            BytesCodec::new(),
+        ))))
     }
 
     #[allow(clippy::never_loop)]
     pub async fn new_reuse<T: std::net::ToSocketAddrs>(addr: T) -> ResultType<Self> {
         for addr in addr.to_socket_addrs()? {
-            return Ok(Self(UdpFramed::new(
-                UdpSocket::from_std(new_socket(addr, true)?.into_udp_socket())?,
+            let socket = new_socket(addr, true)?.into_udp_socket();
+            return Ok(Self(UdpFramedWrapper(UdpFramed::new(
+                UdpSocket::from_std(socket)?,
                 BytesCodec::new(),
-            )));
+            ))));
         }
         bail!("could not resolve to any address");
     }
+}
+
+impl FramedSocket<UdpFramedWrapper<Socks5UdpFramed>> {
+    pub async fn connect<'a, 't, P: ToProxyAddrs, T1: IntoTargetAddr<'t>, T2: ToSocketAddrs>(
+        proxy: P,
+        target: T1,
+        local: T2,
+        username: &'a str,
+        password: &'a str,
+        ms_timeout: u64,
+    ) -> ResultType<(Self, SocketAddr)> {
+        let framed = if username.trim().is_empty() {
+            super::timeout(
+                ms_timeout,
+                Socks5UdpFramed::connect(proxy, target, Some(local)),
+            )
+            .await??
+        } else {
+            super::timeout(
+                ms_timeout,
+                Socks5UdpFramed::connect_with_password(
+                    proxy,
+                    target,
+                    Some(local),
+                    username,
+                    password,
+                ),
+            )
+            .await??
+        };
+        let addr = if let TargetAddr::Ip(c) = framed.target_addr() {
+            c
+        } else {
+            unreachable!()
+        };
+        log::trace!(
+            "Socks5 udp connected, local addr: {}, target addr: {}",
+            framed.local_addr().unwrap(),
+            &addr
+        );
+        Ok((Self(UdpFramedWrapper(framed)), addr))
+    }
+}
+
+// TODO: simplify this constraint
+impl<F> FramedSocket<F>
+where
+    F: Unpin + Stream + Sink<(Bytes, SocketAddr)>,
+    <F as Sink<(Bytes, SocketAddr)>>::Error: Sync + Send + std::error::Error + 'static,
+{
+    pub async fn new_with(self) -> ResultType<Self> {
+        Ok(self)
+    }
 
     #[inline]
     pub async fn send(&mut self, msg: &impl Message, addr: SocketAddr) -> ResultType<()> {
         self.0
-            .send((bytes::Bytes::from(msg.write_to_bytes().unwrap()), addr))
+            .send((Bytes::from(msg.write_to_bytes().unwrap()), addr))
             .await?;
         Ok(())
     }
 
     #[inline]
     pub async fn send_raw(&mut self, msg: &'static [u8], addr: SocketAddr) -> ResultType<()> {
-        self.0.send((bytes::Bytes::from(msg), addr)).await?;
+        self.0.send((Bytes::from(msg), addr)).await?;
         Ok(())
     }
 
     #[inline]
-    pub async fn next(&mut self) -> Option<Result<(BytesMut, SocketAddr), Error>> {
+    pub async fn next(&mut self) -> Option<<F as Stream>::Item> {
         self.0.next().await
     }
 
     #[inline]
-    pub async fn next_timeout(&mut self, ms: u64) -> Option<Result<(BytesMut, SocketAddr), Error>> {
+    pub async fn next_timeout(&mut self, ms: u64) -> Option<<F as Stream>::Item> {
         if let Ok(res) =
             tokio::time::timeout(std::time::Duration::from_millis(ms), self.0.next()).await
         {
@@ -91,3 +178,59 @@ impl FramedSocket {
         }
     }
 }
+
+impl<'a> BytesMutGetter<'a> for BytesMut {
+    fn get_bytes_mut(&'a self) -> &'a BytesMut {
+        self
+    }
+}
+
+impl<'a> BytesMutGetter<'a> for Socks5UdpMessage {
+    fn get_bytes_mut(&'a self) -> &'a BytesMut {
+        &self.data
+    }
+}
+
+impl<F, M, E> Stream for UdpFramedWrapper<F>
+where
+    F: Stream<Item = std::result::Result<(M, SocketAddr), E>>,
+    for<'b> M: BytesMutGetter<'b> + std::fmt::Debug,
+    E: std::error::Error + Into<anyhow::Error>,
+{
+    type Item = ResultType<(BytesMut, SocketAddr)>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match self.project().0.poll_next(cx) {
+            Poll::Ready(Some(Ok((msg, addr)))) => {
+                Poll::Ready(Some(Ok((msg.get_bytes_mut().clone(), addr))))
+            }
+            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(anyhow!(e)))),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+impl<F> Sink<(Bytes, SocketAddr)> for UdpFramedWrapper<F>
+where
+    F: Sink<(Bytes, SocketAddr)>,
+{
+    type Error = <F as Sink<(Bytes, SocketAddr)>>::Error;
+
+    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.project().0.poll_ready(cx)
+    }
+
+    fn start_send(self: Pin<&mut Self>, item: (Bytes, SocketAddr)) -> Result<(), Self::Error> {
+        self.project().0.start_send(item)
+    }
+
+    #[allow(unused_mut)]
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.project().0.poll_flush(cx)
+    }
+
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.project().0.poll_close(cx)
+    }
+}