refactor udp framed

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

libs/hbb_common/Cargo.toml

@@ -29,10 +29,6 @@ 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

@@ -55,6 +55,12 @@ pub const RENDEZVOUS_SERVERS: &'static [&'static str] = &[
 pub const RENDEZVOUS_PORT: i32 = 21116;
 pub const RELAY_PORT: i32 = 21117;
 
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub enum NetworkType {
+    Direct,
+    ProxySocks,
+}
+
 #[derive(Debug, Default, Serialize, Deserialize, Clone)]
 pub struct Config {
     #[serde(default)]
@@ -642,6 +648,13 @@ impl Config {
     pub fn get_socks() -> Option<Socks5Server> {
         CONFIG2.read().unwrap().socks.clone()
     }
+
+    pub fn get_network_type() -> NetworkType {
+        match &CONFIG2.read().unwrap().socks {
+            None => NetworkType::Direct,
+            Some(_) => NetworkType::ProxySocks,
+        }
+    }
 }
 
 const PEERS: &str = "peers";

libs/hbb_common/src/lib.rs

@@ -24,8 +24,6 @@ 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;

libs/hbb_common/src/socket_client.rs

@@ -1,22 +1,21 @@
 use crate::{
-    config::{Config, Socks5Server},
+    config::{Config, NetworkType},
     tcp::FramedStream,
-    udp::{FramedSocket, UdpFramedWrapper},
+    udp::FramedSocket,
     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};
+use tokio_socks::IntoTargetAddr;
 
-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()
-}
+// 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,
@@ -25,7 +24,7 @@ pub async fn connect_tcp<'t, T: IntoTargetAddr<'t>>(
 ) -> ResultType<FramedStream> {
     let target_addr = target.into_target_addr()?;
 
-    if let Some(conf) = get_socks5_conf() {
+    if let Some(conf) = Config::get_socks() {
         FramedStream::connect(
             conf.proxy.as_str(),
             target_addr,
@@ -48,26 +47,13 @@ pub async fn connect_tcp<'t, T: IntoTargetAddr<'t>>(
     }
 }
 
-// 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>(
+pub async fn connect_udp<'t, T1: IntoTargetAddr<'t>, T2: ToSocketAddrs>(
     target: T1,
     local: T2,
-    socks5: &Option<Socks5Server>,
     ms_timeout: u64,
-) -> ResultType<(
-    FramedSocket<UdpFramedWrapper<Socks5UdpFramed>>,
-    Option<SocketAddr>,
-)> {
-    match socks5 {
+) -> ResultType<(FramedSocket, Option<SocketAddr>)> {
+    match Config::get_socks() {
+        None => Ok((FramedSocket::new(local).await?, None)),
         Some(conf) => {
             let (socket, addr) = FramedSocket::connect(
                 conf.proxy.as_str(),
@@ -80,8 +66,12 @@ pub async fn connect_udp_socks5<'t, T1: IntoTargetAddr<'t>, T2: ToSocketAddrs>(
             .await?;
             Ok((socket, Some(addr)))
         }
-        None => {
-            bail!("Nil socks5 server config")
-        }
+    }
+}
+
+pub async fn reconnect_udp<T: ToSocketAddrs>(local: T) -> ResultType<Option<FramedSocket>> {
+    match Config::get_network_type() {
+        NetworkType::Direct => Ok(Some(FramedSocket::new(local).await?)),
+        _ => Ok(None),
     }
 }

libs/hbb_common/src/udp.rs

@@ -2,59 +2,16 @@ use crate::{bail, ResultType};
 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::{
-    net::SocketAddr,
-    ops::{Deref, DerefMut},
-    pin::Pin,
-    task::{Context, Poll},
-};
+use std::net::SocketAddr;
 use tokio::net::{ToSocketAddrs, UdpSocket};
-use tokio_socks::{
-    udp::{Socks5UdpFramed, Socks5UdpMessage},
-    IntoTargetAddr, TargetAddr, ToProxyAddrs,
-};
+use tokio_socks::{udp::Socks5UdpFramed, IntoTargetAddr, TargetAddr, ToProxyAddrs};
 use tokio_util::{codec::BytesCodec, udp::UdpFramed};
 
-pub struct FramedSocket<F>(F);
-
-#[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
-    }
+pub enum FramedSocket {
+    Direct(UdpFramed<BytesCodec>),
+    ProxySocks(Socks5UdpFramed),
 }
 
 fn new_socket(addr: SocketAddr, reuse: bool) -> Result<Socket, std::io::Error> {
@@ -74,29 +31,24 @@ fn new_socket(addr: SocketAddr, reuse: bool) -> Result<Socket, std::io::Error> {
     Ok(socket)
 }
 
-impl FramedSocket<UdpFramedWrapper<UdpFramed<BytesCodec>>> {
+impl FramedSocket {
     pub async fn new<T: ToSocketAddrs>(addr: T) -> ResultType<Self> {
         let socket = UdpSocket::bind(addr).await?;
-        Ok(Self(UdpFramedWrapper(UdpFramed::new(
-            socket,
-            BytesCodec::new(),
-        ))))
+        Ok(Self::Direct(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()? {
             let socket = new_socket(addr, true)?.into_udp_socket();
-            return Ok(Self(UdpFramedWrapper(UdpFramed::new(
+            return Ok(Self::Direct(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,
@@ -134,43 +86,48 @@ impl FramedSocket<UdpFramedWrapper<Socks5UdpFramed>> {
             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)
+        Ok((Self::ProxySocks(framed), addr))
     }
 
     #[inline]
     pub async fn send(&mut self, msg: &impl Message, addr: SocketAddr) -> ResultType<()> {
-        self.0
-            .send((Bytes::from(msg.write_to_bytes().unwrap()), addr))
-            .await?;
+        let send_data = (Bytes::from(msg.write_to_bytes().unwrap()), addr);
+        let _ = match self {
+            Self::Direct(f) => f.send(send_data).await?,
+            Self::ProxySocks(f) => f.send(send_data).await?,
+        };
         Ok(())
     }
 
     #[inline]
     pub async fn send_raw(&mut self, msg: &'static [u8], addr: SocketAddr) -> ResultType<()> {
-        self.0.send((Bytes::from(msg), addr)).await?;
+        let _ = match self {
+            Self::Direct(f) => f.send((Bytes::from(msg), addr)).await?,
+            Self::ProxySocks(f) => f.send((Bytes::from(msg), addr)).await?,
+        };
         Ok(())
     }
 
     #[inline]
-    pub async fn next(&mut self) -> Option<<F as Stream>::Item> {
-        self.0.next().await
+    pub async fn next(&mut self) -> Option<ResultType<(BytesMut, SocketAddr)>> {
+        match self {
+            Self::Direct(f) => match f.next().await {
+                Some(Ok((data, addr))) => Some(Ok((data, addr))),
+                Some(Err(e)) => Some(Err(anyhow!(e))),
+                None => None,
+            },
+            Self::ProxySocks(f) => match f.next().await {
+                Some(Ok((data, addr))) => Some(Ok((data.data, addr))),
+                Some(Err(e)) => Some(Err(anyhow!(e))),
+                None => None,
+            },
+        }
     }
 
     #[inline]
-    pub async fn next_timeout(&mut self, ms: u64) -> Option<<F as Stream>::Item> {
+    pub async fn next_timeout(&mut self, ms: u64) -> Option<ResultType<(BytesMut, SocketAddr)>> {
         if let Ok(res) =
-            tokio::time::timeout(std::time::Duration::from_millis(ms), self.0.next()).await
+            tokio::time::timeout(std::time::Duration::from_millis(ms), self.next()).await
         {
             res
         } else {
@@ -178,59 +135,3 @@ where
         }
     }
 }
-
-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)
-    }
-}