Merge branch 'master' into fix-clippy-warning

This commit is contained in:
RustDesk 2023-01-07 00:28:18 +08:00 committed by GitHub
commit 0e01cfcd3a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 52 additions and 48 deletions

View File

@ -5,7 +5,7 @@ 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,38 +37,28 @@ 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
}
pub async fn new_reuse<T: std::net::ToSocketAddrs>(addr: T) -> ResultType<Self> {
let addr = addr
.to_socket_addrs()?
.next()
.context("could not resolve to any address")?;
let socket = new_socket(addr, true, 0)?.into_udp_socket();
Ok(Self::Direct(UdpFramed::new(
UdpSocket::from_std(socket)?,
BytesCodec::new(),
)))
}
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> {
let addr = addr
.to_socket_addrs()?
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, false, buf_size)?.into_udp_socket())?,
UdpSocket::from_std(new_socket(addr, reuse, buf_size)?.into_udp_socket())?,
BytesCodec::new(),
)))
}
@ -168,12 +158,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
}
}

View File

@ -8,7 +8,7 @@ use hbb_common::{
protobuf::Message as _,
rendezvous_proto::*,
sleep,
tcp::{new_listener, FramedStream},
tcp::{listen_any, FramedStream},
timeout,
tokio::{
self,
@ -77,19 +77,14 @@ pub async fn start(port: &str, key: &str) -> ResultType<()> {
BLOCKLIST_FILE,
BLOCKLIST.read().await.len()
);
let addr = format!("0.0.0.0:{}", port);
log::info!("Listening on tcp {}", addr);
let addr2 = format!("0.0.0.0:{}", port.parse::<u16>().unwrap() + 2);
log::info!("Listening on websocket {}", addr2);
let port: u16 = port.parse()?;
log::info!("Listening on tcp :{}", port);
let port2 = port + 2;
log::info!("Listening on websocket :{}", port2);
let main_task = async move {
loop {
log::info!("Start");
io_loop(
new_listener(&addr, false).await?,
new_listener(&addr2, false).await?,
&key,
)
.await;
io_loop(listen_any(port).await?, listen_any(port2).await?, &key).await;
}
};
let listen_signal = crate::common::listen_signal();

View File

@ -15,7 +15,7 @@ use hbb_common::{
register_pk_response::Result::{TOO_FREQUENT, UUID_MISMATCH},
*,
},
tcp::{new_listener, FramedStream},
tcp::{listen_any, FramedStream},
timeout,
tokio::{
self,
@ -32,7 +32,7 @@ use ipnetwork::Ipv4Network;
use sodiumoxide::crypto::sign;
use std::{
collections::HashMap,
net::{IpAddr, Ipv4Addr, SocketAddr},
net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
sync::Arc,
time::Instant,
};
@ -92,15 +92,15 @@ impl RendezvousServer {
pub async fn start(port: i32, serial: i32, key: &str, rmem: usize) -> ResultType<()> {
let (key, sk) = Self::get_server_sk(key);
let addr = format!("0.0.0.0:{}", port);
let addr2 = format!("0.0.0.0:{}", port - 1);
let addr3 = format!("0.0.0.0:{}", port + 2);
let nat_port = port - 1;
let ws_port = port + 2;
let pm = PeerMap::new().await?;
log::info!("serial={}", serial);
let rendezvous_servers = get_servers(&get_arg("rendezvous-servers"), "rendezvous-servers");
log::info!("Listening on tcp/udp {}", addr);
log::info!("Listening on tcp {}, extra port for NAT test", addr2);
log::info!("Listening on websocket {}", addr3);
let mut socket = FramedSocket::new_with_buf_size(&addr, rmem).await?;
log::info!("Listening on tcp/udp :{}", port);
log::info!("Listening on tcp :{}, extra port for NAT test", nat_port);
log::info!("Listening on websocket :{}", ws_port);
let mut socket = create_udp_listener(port, rmem).await?;
let (tx, mut rx) = mpsc::unbounded_channel::<Data>();
let software_url = get_arg("software-url");
let version = hbb_common::get_version_from_url(&software_url);
@ -138,9 +138,9 @@ impl RendezvousServer {
log::info!("local-ip: {:?}", rs.inner.local_ip);
std::env::set_var("PORT_FOR_API", port.to_string());
rs.parse_relay_servers(&get_arg("relay-servers"));
let mut listener = new_listener(&addr, false).await?;
let mut listener2 = new_listener(&addr2, false).await?;
let mut listener3 = new_listener(&addr3, false).await?;
let mut listener = create_tcp_listener(port).await?;
let mut listener2 = create_tcp_listener(nat_port).await?;
let mut listener3 = create_tcp_listener(ws_port).await?;
let test_addr = std::env::var("TEST_HBBS").unwrap_or_default();
if std::env::var("ALWAYS_USE_RELAY")
.unwrap_or_default()
@ -186,19 +186,19 @@ impl RendezvousServer {
{
LoopFailure::UdpSocket => {
drop(socket);
socket = FramedSocket::new_with_buf_size(&addr, rmem).await?;
socket = create_udp_listener(port, rmem).await?;
}
LoopFailure::Listener => {
drop(listener);
listener = new_listener(&addr, false).await?;
listener = create_tcp_listener(port).await?;
}
LoopFailure::Listener2 => {
drop(listener2);
listener2 = new_listener(&addr2, false).await?;
listener2 = create_tcp_listener(nat_port).await?;
}
LoopFailure::Listener3 => {
drop(listener3);
listener3 = new_listener(&addr3, false).await?;
listener3 = create_tcp_listener(ws_port).await?;
}
}
}
@ -1267,3 +1267,22 @@ async fn send_rk_res(
});
socket.send(&msg_out, addr).await
}
async fn create_udp_listener(port: i32, rmem: usize) -> ResultType<FramedSocket> {
let addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), port as _);
if let Ok(s) = FramedSocket::new_reuse(&addr, false, rmem).await {
log::debug!("listen on udp {:?}", s.local_addr());
return Ok(s);
}
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), port as _);
let s = FramedSocket::new_reuse(&addr, false, rmem).await?;
log::debug!("listen on udp {:?}", s.local_addr());
return Ok(s);
}
#[inline]
async fn create_tcp_listener(port: i32) -> ResultType<TcpListener> {
let s = listen_any(port as _).await?;
log::debug!("listen on tcp {:?}", s.local_addr());
Ok(s)
}