hongshaorou/rustdesk
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor target address

Browse Source
This commit is contained in:
rustdesk 2022-01-05 13:21:14 +08:00
parent 3e590b8212
commit 1d8cdb5e93
10 changed files with 157 additions and 174 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
Cargo.lock generated
View File

@ -389,9 +389,9 @@ dependencies = [
[[package]]
name = "clipboard-win"
version = "4.2.2"
version = "4.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3db8340083d28acb43451166543b98c838299b7e0863621be53a338adceea0ed"
checksum = "1951fb8aa063a2ee18b4b4d217e4aa2ec9cc4f2430482983f607fa10cd36d7aa"
dependencies = [
"error-code",
"str-buf",
@ -586,9 +586,9 @@ dependencies = [
[[package]]
name = "cstr_core"
version = "0.2.4"
version = "0.2.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "917ba9efe9e1e736671d5a03f006afc4e7e3f32503e2077e0bcaf519c0c8c1d3"
checksum = "644828c273c063ab0d39486ba42a5d1f3a499d35529c759e763a9c6cb8a0fb08"
dependencies = [
"cty",
"memchr",
@ -1238,9 +1238,9 @@ dependencies = [
[[package]]
name = "generic-array"
version = "0.14.4"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "501466ecc8a30d1d3b7fc9229b122b2ce8ed6e9d9223f1138d4babb253e51817"
checksum = "fd48d33ec7f05fbfa152300fdad764757cbded343c1aa1cff2fbaf4134851803"
dependencies = [
"typenum",
"version_check",
@ -3667,7 +3667,7 @@ dependencies = [
[[package]]
name = "tokio-socks"
version = "0.5.1"
source = "git+https://github.com/fufesou/tokio-socks#121a780c7e6a31c3aac70e7234f5c62eecaf0629"
source = "git+https://github.com/fufesou/tokio-socks#83e9e9202b85b887999375e4d86e437e33d0535c"
dependencies = [
"bytes",
"either",

13
libs/hbb_common/src/config.rs
View File

@ -6,14 +6,13 @@ use sodiumoxide::crypto::sign;
use std::{
collections::HashMap,
fs,
net::SocketAddr,
net::{IpAddr, Ipv4Addr, SocketAddr},
path::{Path, PathBuf},
sync::{Arc, Mutex, RwLock},
time::SystemTime,
};
pub const APP_NAME: &str = "RustDesk";
pub const BIND_INTERFACE: &str = "0.0.0.0";
pub const RENDEZVOUS_TIMEOUT: u64 = 12_000;
pub const CONNECT_TIMEOUT: u64 = 18_000;
pub const COMPRESS_LEVEL: i32 = 3;
@ -346,10 +345,10 @@ impl Config {
#[inline]
pub fn get_any_listen_addr() -> SocketAddr {
format!("{}:0", BIND_INTERFACE).parse().unwrap()
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
}
pub async fn get_rendezvous_server() -> SocketAddr {
pub fn get_rendezvous_server() -> String {
let mut rendezvous_server = Self::get_option("custom-rendezvous-server");
if rendezvous_server.is_empty() {
rendezvous_server = CONFIG2.write().unwrap().rendezvous_server.clone();
@ -363,11 +362,7 @@ impl Config {
if !rendezvous_server.contains(":") {
rendezvous_server = format!("{}:{}", rendezvous_server, RENDEZVOUS_PORT);
}
if let Ok(addr) = crate::to_socket_addr(&rendezvous_server).await {
addr
} else {
Self::get_any_listen_addr()
}
rendezvous_server
}
pub fn get_rendezvous_servers() -> Vec<String> {

3
libs/hbb_common/src/lib.rs
View File

@ -27,9 +27,10 @@ pub use anyhow::{self, bail};
pub use futures_util;
pub mod config;
pub mod fs;
pub use socket_client::to_socket_addr;
pub use sodiumoxide;
pub use tokio_socks;
pub use tokio_socks::IntoTargetAddr;
pub use tokio_socks::TargetAddr;
#[cfg(feature = "quic")]
pub type Stream = quic::Connection;

81
libs/hbb_common/src/socket_client.rs
View File

@ -1,5 +1,5 @@
use crate::{
config::{Config, NetworkType, RENDEZVOUS_TIMEOUT},
config::{Config, NetworkType},
tcp::FramedStream,
udp::FramedSocket,
ResultType,
@ -7,15 +7,43 @@ use crate::{
use anyhow::{bail, Context};
use std::net::SocketAddr;
use tokio::net::ToSocketAddrs;
use tokio_socks::IntoTargetAddr;
use tokio_socks::{IntoTargetAddr, TargetAddr};
// 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 to_socket_addr(host: &str) -> ResultType<SocketAddr> {
use std::net::ToSocketAddrs;
let addrs: Vec<SocketAddr> = host.to_socket_addrs()?.collect();
if addrs.is_empty() {
bail!("Failed to solve {}", host);
}
Ok(addrs[0])
}
pub fn get_target_addr(host: &str) -> ResultType<TargetAddr<'static>> {
let addr = match Config::get_network_type() {
NetworkType::Direct => to_socket_addr(&host)?.into_target_addr()?,
NetworkType::ProxySocks => host.into_target_addr()?,
}
.to_owned();
Ok(addr)
}
pub fn test_if_valid_server(host: &str) -> String {
let mut host = host.to_owned();
if !host.contains(":") {
host = format!("{}:{}", host, 0);
}
match Config::get_network_type() {
NetworkType::Direct => match to_socket_addr(&host) {
Err(err) => err.to_string(),
Ok(_) => "".to_owned(),
},
NetworkType::ProxySocks => match &host.into_target_addr() {
Err(err) => err.to_string(),
Ok(_) => "".to_owned(),
},
}
}
pub async fn connect_tcp<'t, T: IntoTargetAddr<'t>>(
target: T,
@ -47,49 +75,24 @@ pub async fn connect_tcp<'t, T: IntoTargetAddr<'t>>(
}
}
fn native_to_socket_addr(host: &str) -> ResultType<SocketAddr> {
use std::net::ToSocketAddrs;
let addrs: Vec<SocketAddr> = host.to_socket_addrs()?.collect();
if addrs.is_empty() {
bail!("Failed to solve {}", host);
}
Ok(addrs[0])
}
pub async fn to_socket_addr(host: &str) -> ResultType<SocketAddr> {
Ok(
new_udp(host, Config::get_any_listen_addr(), RENDEZVOUS_TIMEOUT)
.await?
.1,
)
}
pub async fn new_udp<'t, T1: IntoTargetAddr<'t> + std::fmt::Display, T2: ToSocketAddrs>(
target: T1,
local: T2,
ms_timeout: u64,
) -> ResultType<(FramedSocket, SocketAddr)> {
pub async fn new_udp<T: ToSocketAddrs>(local: T, ms_timeout: u64) -> ResultType<FramedSocket> {
match Config::get_socks() {
None => Ok((
FramedSocket::new(local).await?,
native_to_socket_addr(&target.to_string())?,
)),
None => Ok(FramedSocket::new(local).await?),
Some(conf) => {
let (socket, addr) = FramedSocket::new_proxy(
let socket = FramedSocket::new_proxy(
conf.proxy.as_str(),
target,
local,
conf.username.as_str(),
conf.password.as_str(),
ms_timeout,
)
.await?;
Ok((socket, addr))
Ok(socket)
}
}
}
pub async fn rebind<T: ToSocketAddrs>(local: T) -> ResultType<Option<FramedSocket>> {
pub async fn rebind_udp<T: ToSocketAddrs>(local: T) -> ResultType<Option<FramedSocket>> {
match Config::get_network_type() {
NetworkType::Direct => Ok(Some(FramedSocket::new(local).await?)),
_ => Ok(None),

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

@ -49,75 +49,79 @@ impl FramedSocket {
bail!("could not resolve to any address");
}
pub async fn new_proxy<'a, 't, P: ToProxyAddrs, T1: IntoTargetAddr<'t>, T2: ToSocketAddrs>(
pub async fn new_proxy<'a, 't, P: ToProxyAddrs, T: ToSocketAddrs>(
proxy: P,
target: T1,
local: T2,
local: T,
username: &'a str,
password: &'a str,
ms_timeout: u64,
) -> ResultType<(Self, SocketAddr)> {
) -> ResultType<Self> {
let framed = if username.trim().is_empty() {
super::timeout(
ms_timeout,
Socks5UdpFramed::connect(proxy, target, Some(local)),
)
.await??
super::timeout(ms_timeout, Socks5UdpFramed::connect(proxy, Some(local))).await??
} else {
super::timeout(
ms_timeout,
Socks5UdpFramed::connect_with_password(
proxy,
target,
Some(local),
username,
password,
),
Socks5UdpFramed::connect_with_password(proxy, 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
framed.socks_addr()
);
Ok((Self::ProxySocks(framed), addr))
Ok(Self::ProxySocks(framed))
}
#[inline]
pub async fn send(&mut self, msg: &impl Message, addr: SocketAddr) -> ResultType<()> {
let send_data = (Bytes::from(msg.write_to_bytes().unwrap()), addr);
pub async fn send(
&mut self,
msg: &impl Message,
addr: impl IntoTargetAddr<'_>,
) -> ResultType<()> {
let addr = addr.into_target_addr()?.to_owned();
let send_data = Bytes::from(msg.write_to_bytes().unwrap());
let _ = match self {
Self::Direct(f) => f.send(send_data).await?,
Self::ProxySocks(f) => f.send(send_data).await?,
Self::Direct(f) => match addr {
TargetAddr::Ip(addr) => f.send((send_data, addr)).await?,
_ => unreachable!(),
},
Self::ProxySocks(f) => f.send((send_data, addr)).await?,
};
Ok(())
}
// https://stackoverflow.com/a/68733302/1926020
#[inline]
pub async fn send_raw(&mut self, msg: &'static [u8], addr: SocketAddr) -> ResultType<()> {
pub async fn send_raw(
&mut self,
msg: &'static [u8],
addr: impl IntoTargetAddr<'static>,
) -> ResultType<()> {
let addr = addr.into_target_addr()?.to_owned();
let _ = match self {
Self::Direct(f) => f.send((Bytes::from(msg), addr)).await?,
Self::Direct(f) => match addr {
TargetAddr::Ip(addr) => f.send((Bytes::from(msg), addr)).await?,
_ => unreachable!(),
},
Self::ProxySocks(f) => f.send((Bytes::from(msg), addr)).await?,
};
Ok(())
}
#[inline]
pub async fn next(&mut self) -> Option<ResultType<(BytesMut, SocketAddr)>> {
pub async fn next(&mut self) -> Option<ResultType<(BytesMut, TargetAddr<'static>)>> {
match self {
Self::Direct(f) => match f.next().await {
Some(Ok((data, addr))) => Some(Ok((data, addr))),
Some(Ok((data, addr))) => {
Some(Ok((data, addr.into_target_addr().unwrap().to_owned())))
}
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(Ok((data, _))) => Some(Ok((data.data, data.dst_addr))),
Some(Err(e)) => Some(Err(anyhow!(e))),
None => None,
},
@ -125,7 +129,10 @@ impl FramedSocket {
}
#[inline]
pub async fn next_timeout(&mut self, ms: u64) -> Option<ResultType<(BytesMut, SocketAddr)>> {
pub async fn next_timeout(
&mut self,
ms: u64,
) -> Option<ResultType<(BytesMut, TargetAddr<'static>)>> {
if let Ok(res) =
tokio::time::timeout(std::time::Duration::from_millis(ms), self.next()).await
{

9
src/client.rs
View File

@ -109,7 +109,8 @@ impl Client {
log::info!("rendezvous server: {}", rendezvous_server);
let mut socket =
socket_client::connect_tcp(rendezvous_server, any_addr, RENDEZVOUS_TIMEOUT).await?;
socket_client::connect_tcp(&*rendezvous_server, any_addr, RENDEZVOUS_TIMEOUT)
.await?;
let my_addr = socket.local_addr();
let mut pk = Vec::new();
let mut relay_server = "".to_owned();
@ -205,7 +206,7 @@ impl Client {
peer,
pk,
&relay_server,
rendezvous_server,
&rendezvous_server,
time_used,
peer_nat_type,
my_nat_type,
@ -221,7 +222,7 @@ impl Client {
peer_id: &str,
pk: Vec<u8>,
relay_server: &str,
rendezvous_server: SocketAddr,
rendezvous_server: &str,
punch_time_used: u64,
peer_nat_type: NatType,
my_nat_type: i32,
@ -385,7 +386,7 @@ impl Client {
async fn request_relay(
peer: &str,
relay_server: String,
rendezvous_server: SocketAddr,
rendezvous_server: &str,
secure: bool,
conn_type: ConnType,
) -> ResultType<Stream> {

50
src/common.rs
View File

@ -13,10 +13,7 @@ use hbb_common::{
};
#[cfg(any(target_os = "android", target_os = "ios", feature = "cli"))]
use hbb_common::{config::RENDEZVOUS_PORT, futures::future::join_all};
use std::{
net::SocketAddr,
sync::{Arc, Mutex},
};
use std::sync::{Arc, Mutex};
pub const CLIPBOARD_NAME: &'static str = "clipboard";
pub const CLIPBOARD_INTERVAL: u64 = 333;
@ -244,13 +241,15 @@ async fn test_nat_type_() -> ResultType<bool> {
let start = std::time::Instant::now();
let rendezvous_server = get_rendezvous_server(100).await;
let server1 = rendezvous_server;
let mut server2 = server1;
if server1.port() == 0 {
// offline
// avoid overflow crash
bail!("Offline");
let tmp: Vec<&str> = server1.split(":").collect();
if tmp.len() != 2 {
bail!("Invalid server address: {}", server1);
}
server2.set_port(server1.port() - 1);
let port: u16 = tmp[1].parse()?;
if port == 0 {
bail!("Invalid server address: {}", server1);
}
let server2 = format!("{}:{}", tmp[0], port);
let mut msg_out = RendezvousMessage::new();
let serial = Config::get_serial();
msg_out.set_test_nat_request(TestNatRequest {
@ -262,7 +261,11 @@ async fn test_nat_type_() -> ResultType<bool> {
let mut addr = Config::get_any_listen_addr();
for i in 0..2 {
let mut socket = socket_client::connect_tcp(
if i == 0 { &server1 } else { &server2 },
if i == 0 {
server1.clone()
} else {
server2.clone()
},
addr,
RENDEZVOUS_TIMEOUT,
)
@ -306,12 +309,12 @@ async fn test_nat_type_() -> ResultType<bool> {
}
#[cfg(any(target_os = "android", target_os = "ios"))]
pub async fn get_rendezvous_server(_ms_timeout: u64) -> SocketAddr {
pub async fn get_rendezvous_server(_ms_timeout: u64) -> String {
Config::get_rendezvous_server()
}
#[cfg(not(any(target_os = "android", target_os = "ios")))]
pub async fn get_rendezvous_server(ms_timeout: u64) -> SocketAddr {
pub async fn get_rendezvous_server(ms_timeout: u64) -> String {
crate::ipc::get_rendezvous_server(ms_timeout).await
}
@ -415,18 +418,6 @@ pub fn is_modifier(evt: &KeyEvent) -> bool {
}
}
#[tokio::main(flavor = "current_thread")]
pub async fn test_if_valid_server(host: String) -> String {
let mut host = host;
if !host.contains(":") {
host = format!("{}:{}", host, 0);
}
match hbb_common::to_socket_addr(&host).await {
Err(err) => err.to_string(),
Ok(_) => "".to_owned(),
}
}
pub fn get_version_number(v: &str) -> i64 {
let mut n = 0;
for x in v.split(".") {
@ -444,12 +435,9 @@ async fn _check_software_update() -> hbb_common::ResultType<()> {
sleep(3.).await;
let rendezvous_server = get_rendezvous_server(1_000).await;
let (mut socket, _) = socket_client::new_udp(
rendezvous_server,
Config::get_any_listen_addr(),
RENDEZVOUS_TIMEOUT,
)
.await?;
let mut socket =
socket_client::new_udp(Config::get_any_listen_addr(), RENDEZVOUS_TIMEOUT).await?;
let mut msg_out = RendezvousMessage::new();
msg_out.set_software_update(SoftwareUpdate {
url: crate::VERSION.to_owned(),

13
src/ipc.rs
View File

@ -13,7 +13,7 @@ use parity_tokio_ipc::{
Connection as Conn, ConnectionClient as ConnClient, Endpoint, Incoming, SecurityAttributes,
};
use serde_derive::{Deserialize, Serialize};
use std::{collections::HashMap, net::SocketAddr};
use std::collections::HashMap;
#[cfg(not(windows))]
use std::{fs::File, io::prelude::*};
@ -202,7 +202,7 @@ async fn handle(data: Data, stream: &mut Connection) {
} else if name == "salt" {
value = Some(Config::get_salt());
} else if name == "rendezvous_server" {
value = Some(Config::get_rendezvous_server().await.to_string());
value = Some(Config::get_rendezvous_server());
} else {
value = None;
}
@ -397,13 +397,12 @@ pub fn get_password() -> String {
}
}
pub async fn get_rendezvous_server(ms_timeout: u64) -> SocketAddr {
pub async fn get_rendezvous_server(ms_timeout: u64) -> String {
if let Ok(Some(v)) = get_config_async("rendezvous_server", ms_timeout).await {
if let Ok(v) = v.parse() {
return v;
}
v
} else {
Config::get_rendezvous_server()
}
return Config::get_rendezvous_server().await;
}
async fn get_options_(ms_timeout: u64) -> ResultType<HashMap<String, String>> {

73
src/rendezvous_mediator.rs
View File

@ -13,7 +13,7 @@ use hbb_common::{
time::{interval, Duration},
},
udp::FramedSocket,
AddrMangle, ResultType,
AddrMangle, IntoTargetAddr, ResultType, TargetAddr,
};
use std::{
net::SocketAddr,
@ -32,15 +32,26 @@ lazy_static::lazy_static! {
}
static SHOULD_EXIT: AtomicBool = AtomicBool::new(false);
#[derive(Clone)]
pub struct RendezvousMediator {
addr: SocketAddr,
addr: TargetAddr<'static>,
host: String,
host_prefix: String,
rendezvous_servers: Vec<String>,
last_id_pk_registry: String,
}
impl Clone for RendezvousMediator {
fn clone(&self) -> Self {
Self {
addr: self.addr.to_owned(),
host: self.host.clone(),
host_prefix: self.host_prefix.clone(),
rendezvous_servers: self.rendezvous_servers.clone(),
last_id_pk_registry: self.last_id_pk_registry.clone(),
}
}
}
impl RendezvousMediator {
pub async fn start_all() {
let mut nat_tested = false;
@ -93,19 +104,17 @@ impl RendezvousMediator {
})
.unwrap_or(host.to_owned());
let mut rz = Self {
addr: Config::get_any_listen_addr(),
addr: Config::get_any_listen_addr().into_target_addr()?,
host: host.clone(),
host_prefix,
rendezvous_servers,
last_id_pk_registry: "".to_owned(),
};
let (mut socket, target_addr) = socket_client::new_udp(
crate::check_port(&host, RENDEZVOUS_PORT),
Config::get_any_listen_addr(),
RENDEZVOUS_TIMEOUT,
)
.await?;
rz.addr = target_addr;
rz.addr = socket_client::get_target_addr(&crate::check_port(&host, RENDEZVOUS_PORT))?;
let any_addr = Config::get_any_listen_addr();
let mut socket = socket_client::new_udp(any_addr, RENDEZVOUS_TIMEOUT).await?;
const TIMER_OUT: Duration = Duration::from_secs(1);
let mut timer = interval(TIMER_OUT);
let mut last_timer = SystemTime::UNIX_EPOCH;
@ -222,18 +231,6 @@ impl RendezvousMediator {
if SHOULD_EXIT.load(Ordering::SeqCst) {
break;
}
if rz.addr.port() == 0 {
allow_err!(rz.dns_check().await);
if rz.addr.port() == 0 {
continue;
} else {
// have to do this for osx, to avoid "Can't assign requested address"
// when socket created before OS network ready
if let Some(s) = socket_client::rebind(Config::get_any_listen_addr()).await? {
socket = s;
};
}
}
let now = SystemTime::now();
if now.duration_since(last_timer).map(|d| d < TIMER_OUT).unwrap_or(false) {
// a workaround of tokio timer bug
@ -251,13 +248,12 @@ impl RendezvousMediator {
Config::update_latency(&host, -1);
old_latency = 0;
if now.duration_since(last_dns_check).map(|d| d.as_millis() as i64).unwrap_or(0) > DNS_INTERVAL {
if let Ok(_) = rz.dns_check().await {
// in some case of network reconnect (dial IP network),
// old UDP socket not work any more after network recover
if let Some(s) = socket_client::rebind(Config::get_any_listen_addr()).await? {
socket = s;
};
}
rz.addr = socket_client::get_target_addr(&crate::check_port(&host, RENDEZVOUS_PORT))?;
// in some case of network reconnect (dial IP network),
// old UDP socket not work any more after network recover
if let Some(s) = socket_client::rebind_udp(any_addr).await? {
socket = s;
};
last_dns_check = now;
}
} else if fails > MAX_FAILS1 {
@ -272,13 +268,6 @@ impl RendezvousMediator {
Ok(())
}
async fn dns_check(&mut self) -> ResultType<()> {
self.addr =
hbb_common::to_socket_addr(&crate::check_port(&self.host, RENDEZVOUS_PORT)).await?;
log::debug!("Lookup dns of {}", self.host);
Ok(())
}
async fn handle_request_relay(&self, rr: RequestRelay, server: ServerPtr) -> ResultType<()> {
self.create_relay(
rr.socket_addr,
@ -310,7 +299,7 @@ impl RendezvousMediator {
);
let mut socket = socket_client::connect_tcp(
self.addr,
self.addr.to_owned(),
Config::get_any_listen_addr(),
RENDEZVOUS_TIMEOUT,
)
@ -338,7 +327,7 @@ impl RendezvousMediator {
let peer_addr = AddrMangle::decode(&fla.socket_addr);
log::debug!("Handle intranet from {:?}", peer_addr);
let mut socket = socket_client::connect_tcp(
self.addr,
self.addr.to_owned(),
Config::get_any_listen_addr(),
RENDEZVOUS_TIMEOUT,
)
@ -382,7 +371,7 @@ impl RendezvousMediator {
log::debug!("Punch hole to {:?}", peer_addr);
let mut socket = {
let socket = socket_client::connect_tcp(
self.addr,
self.addr.to_owned(),
Config::get_any_listen_addr(),
RENDEZVOUS_TIMEOUT,
)
@ -425,7 +414,7 @@ impl RendezvousMediator {
pk,
..Default::default()
});
socket.send(&msg_out, self.addr).await?;
socket.send(&msg_out, self.addr.to_owned()).await?;
Ok(())
}
@ -471,7 +460,7 @@ impl RendezvousMediator {
serial,
..Default::default()
});
socket.send(&msg_out, self.addr).await?;
socket.send(&msg_out, self.addr.to_owned()).await?;
Ok(())
}
}

2
src/ui.rs
View File

@ -315,7 +315,7 @@ impl UI {
}
fn test_if_valid_server(&self, host: String) -> String {
crate::common::test_if_valid_server(host)
hbb_common::socket_client::test_if_valid_server(&host)
}
fn get_sound_inputs(&self) -> Value {