rustdesk/rustdesk-server
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk-server.git synced 2025-07-01 23:35:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' into fix-clippy-warning

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
libs/hbb_common/src/udp.rs
View File

@ -5,7 +5,7 @@ use futures::{SinkExt, StreamExt};
use protobuf::Message; use protobuf::Message;
use socket2::{Domain, Socket, Type}; use socket2::{Domain, Socket, Type};
use std::net::SocketAddr; 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_socks::{udp::Socks5UdpFramed, IntoTargetAddr, TargetAddr, ToProxyAddrs};
use tokio_util::{codec::BytesCodec, udp::UdpFramed}; use tokio_util::{codec::BytesCodec, udp::UdpFramed};
@ -37,38 +37,28 @@ fn new_socket(addr: SocketAddr, reuse: bool, buf_size: usize) -> Result<Socket,
addr, addr,
socket.recv_buffer_size() 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())?; socket.bind(&addr.into())?;
Ok(socket) Ok(socket)
} }
impl FramedSocket { impl FramedSocket {
pub async fn new<T: ToSocketAddrs>(addr: T) -> ResultType<Self> { pub async fn new<T: ToSocketAddrs>(addr: T) -> ResultType<Self> {
let socket = UdpSocket::bind(addr).await?; Self::new_reuse(addr, false, 0).await
Ok(Self::Direct(UdpFramed::new(socket, BytesCodec::new())))
} }
pub async fn new_reuse<T: std::net::ToSocketAddrs>(addr: T) -> ResultType<Self> { pub async fn new_reuse<T: ToSocketAddrs>(
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>(
addr: T, addr: T,
reuse: bool,
buf_size: usize, buf_size: usize,
) -> ResultType<Self> { ) -> ResultType<Self> {
let addr = addr let addr = lookup_host(&addr).await?
.to_socket_addrs()?
.next() .next()
.context("could not resolve to any address")?; .context("could not resolve to any address")?;
Ok(Self::Direct(UdpFramed::new( 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(), 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 FramedSocket::Direct(x) = self {
if let Ok(v) = x.get_ref().local_addr() { if let Ok(v) = x.get_ref().local_addr() {
return v.is_ipv4(); return Some(v);
} }
} }
true None
} }
} }

17
src/relay_server.rs
View File

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

49
src/rendezvous_server.rs
View File

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