sync with rustdesk-server hbb_common

libs/hbb_common/src/udp.rs

@@ -1,11 +1,11 @@
-use crate::{bail, ResultType};
-use anyhow::anyhow;
+use crate::ResultType;
+use anyhow::{anyhow, Context};
 use bytes::{Bytes, BytesMut};
 use futures::{SinkExt, StreamExt};
 use protobuf::Message;
 use socket2::{Domain, Socket, Type};
 use std::net::SocketAddr;
-use tokio::net::{ToSocketAddrs, UdpSocket};
+use tokio::net::{lookup_host, ToSocketAddrs, UdpSocket};
 use tokio_socks::{udp::Socks5UdpFramed, IntoTargetAddr, TargetAddr, ToProxyAddrs};
 use tokio_util::{codec::BytesCodec, udp::UdpFramed};
 
@@ -37,39 +37,31 @@ fn new_socket(addr: SocketAddr, reuse: bool, buf_size: usize) -> Result<Socket,
         addr,
         socket.recv_buffer_size()
     );
+    if addr.is_ipv6() && addr.ip().is_unspecified() && addr.port() > 0 {
+        socket.set_only_v6(false).ok();
+    }
     socket.bind(&addr.into())?;
     Ok(socket)
 }
 
 impl FramedSocket {
     pub async fn new<T: ToSocketAddrs>(addr: T) -> ResultType<Self> {
-        let socket = UdpSocket::bind(addr).await?;
-        Ok(Self::Direct(UdpFramed::new(socket, BytesCodec::new())))
+        Self::new_reuse(addr, false, 0).await
     }
 
-    #[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, 0)?.into_udp_socket();
-            return Ok(Self::Direct(UdpFramed::new(
-                UdpSocket::from_std(socket)?,
-                BytesCodec::new(),
-            )));
-        }
-        bail!("could not resolve to any address");
-    }
-
-    pub async fn new_with_buf_size<T: std::net::ToSocketAddrs>(
+    pub async fn new_reuse<T: ToSocketAddrs>(
         addr: T,
+        reuse: bool,
         buf_size: usize,
     ) -> ResultType<Self> {
-        for addr in addr.to_socket_addrs()? {
-            return Ok(Self::Direct(UdpFramed::new(
-                UdpSocket::from_std(new_socket(addr, false, buf_size)?.into_udp_socket())?,
-                BytesCodec::new(),
-            )));
-        }
-        bail!("could not resolve to any address");
+        let addr = lookup_host(&addr)
+            .await?
+            .next()
+            .context("could not resolve to any address")?;
+        Ok(Self::Direct(UdpFramed::new(
+            UdpSocket::from_std(new_socket(addr, reuse, buf_size)?.into_udp_socket())?,
+            BytesCodec::new(),
+        )))
     }
 
     pub async fn new_proxy<'a, 't, P: ToProxyAddrs, T: ToSocketAddrs>(
@@ -104,11 +96,12 @@ impl FramedSocket {
     ) -> ResultType<()> {
         let addr = addr.into_target_addr()?.to_owned();
         let send_data = Bytes::from(msg.write_to_bytes()?);
-        let _ = match self {
-            Self::Direct(f) => match addr {
-                TargetAddr::Ip(addr) => f.send((send_data, addr)).await?,
-                _ => {}
-            },
+        match self {
+            Self::Direct(f) => {
+                if let TargetAddr::Ip(addr) = addr {
+                    f.send((send_data, addr)).await?
+                }
+            }
             Self::ProxySocks(f) => f.send((send_data, addr)).await?,
         };
         Ok(())
@@ -123,11 +116,12 @@ impl FramedSocket {
     ) -> ResultType<()> {
         let addr = addr.into_target_addr()?.to_owned();
 
-        let _ = match self {
-            Self::Direct(f) => match addr {
-                TargetAddr::Ip(addr) => f.send((Bytes::from(msg), addr)).await?,
-                _ => {}
-            },
+        match self {
+            Self::Direct(f) => {
+                if let TargetAddr::Ip(addr) = addr {
+                    f.send((Bytes::from(msg), addr)).await?
+                }
+            }
             Self::ProxySocks(f) => f.send((Bytes::from(msg), addr)).await?,
         };
         Ok(())
@@ -165,12 +159,12 @@ impl FramedSocket {
         }
     }
 
-    pub fn is_ipv4(&self) -> bool {
+    pub fn local_addr(&self) -> Option<SocketAddr> {
         if let FramedSocket::Direct(x) = self {
             if let Ok(v) = x.get_ref().local_addr() {
-                return v.is_ipv4();
+                return Some(v);
             }
         }
-        true
+        None
     }
 }