From bb8257b7f29a8a6f3ee7d6f1fd4e7992846a2341 Mon Sep 17 00:00:00 2001 From: rustdesk Date: Thu, 14 Jul 2022 15:21:26 +0800 Subject: [PATCH] https://github.com/rustdesk/rustdesk/issues/899 --- src/client.rs | 22 +++++++++++--- src/common.rs | 62 ++++++++++++++++++++++++++++++-------- src/ipc.rs | 30 +++++++++++++++--- src/mobile.rs | 1 - src/rendezvous_mediator.rs | 7 +++-- 5 files changed, 98 insertions(+), 24 deletions(-) diff --git a/src/client.rs b/src/client.rs index d63ce970c..c6ccd2c26 100644 --- a/src/client.rs +++ b/src/client.rs @@ -148,11 +148,25 @@ impl Client { true, )); } - let rendezvous_server = crate::get_rendezvous_server(1_000).await; - log::info!("rendezvous server: {}", rendezvous_server); - + let (mut rendezvous_server, servers, contained) = crate::get_rendezvous_server(1_000).await; 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; + debug_assert!(!servers.contains(&rendezvous_server)); + if socket.is_err() && !servers.is_empty() { + log::info!("try the other servers: {:?}", servers); + for server in servers { + socket = socket_client::connect_tcp(&*server, any_addr, RENDEZVOUS_TIMEOUT).await; + if socket.is_ok() { + rendezvous_server = server; + break; + } + } + crate::refresh_rendezvous_server(); + } else if !contained { + crate::refresh_rendezvous_server(); + } + log::info!("rendezvous server: {}", rendezvous_server); + let mut socket = socket?; let my_addr = socket.local_addr(); let mut signed_id_pk = Vec::new(); let mut relay_server = "".to_owned(); diff --git a/src/common.rs b/src/common.rs index 667dd3e7e..5ef127f46 100644 --- a/src/common.rs +++ b/src/common.rs @@ -12,7 +12,6 @@ use hbb_common::{ rendezvous_proto::*, sleep, socket_client, tokio, ResultType, }; -#[cfg(any(target_os = "android", target_os = "ios", feature = "cli"))] use hbb_common::{config::RENDEZVOUS_PORT, futures::future::join_all}; use std::sync::{Arc, Mutex}; @@ -247,7 +246,7 @@ async fn test_nat_type_() -> ResultType { return Ok(true); } let start = std::time::Instant::now(); - let rendezvous_server = get_rendezvous_server(1_000).await; + let (rendezvous_server, _, _) = get_rendezvous_server(1_000).await; let server1 = rendezvous_server; let tmp: Vec<&str> = server1.split(":").collect(); if tmp.len() != 2 { @@ -316,31 +315,59 @@ async fn test_nat_type_() -> ResultType { Ok(ok) } -#[cfg(any(target_os = "android", target_os = "ios"))] -pub async fn get_rendezvous_server(_ms_timeout: u64) -> String { - Config::get_rendezvous_server() +pub async fn get_rendezvous_server(ms_timeout: u64) -> (String, Vec, bool) { + let (mut a, mut b) = get_rendezvous_server_(ms_timeout).await; + let mut b: Vec = b + .drain(..) + .map(|x| { + if !x.contains(":") { + format!("{}:{}", x, config::RENDEZVOUS_PORT) + } else { + x + } + }) + .collect(); + let c = if b.contains(&a) { + b = b.drain(..).filter(|x| x != &a).collect(); + true + } else { + a = b.pop().unwrap_or(a); + false + }; + (a, b, c) } +#[inline] +#[cfg(any(target_os = "android", target_os = "ios"))] +fn get_rendezvous_server_(_ms_timeout: u64) -> (String, Vec) { + ( + Config::get_rendezvous_server(), + Config::get_rendezvous_servers(), + ) +} + +#[inline] #[cfg(not(any(target_os = "android", target_os = "ios")))] -pub async fn get_rendezvous_server(ms_timeout: u64) -> String { +async fn get_rendezvous_server_(ms_timeout: u64) -> (String, Vec) { crate::ipc::get_rendezvous_server(ms_timeout).await } +#[inline] #[cfg(any(target_os = "android", target_os = "ios"))] pub async fn get_nat_type(_ms_timeout: u64) -> i32 { Config::get_nat_type() } +#[inline] #[cfg(not(any(target_os = "android", target_os = "ios")))] pub async fn get_nat_type(ms_timeout: u64) -> i32 { crate::ipc::get_nat_type(ms_timeout).await } -#[cfg(any(target_os = "android", target_os = "ios", feature = "cli"))] #[tokio::main(flavor = "current_thread")] async fn test_rendezvous_server_() { let servers = Config::get_rendezvous_servers(); - hbb_common::config::ONLINE.lock().unwrap().clear(); + Config::reset_online(); let mut futs = Vec::new(); for host in servers { futs.push(tokio::spawn(async move { @@ -363,11 +390,21 @@ async fn test_rendezvous_server_() { join_all(futs).await; } -#[cfg(any(target_os = "android", target_os = "ios", feature = "cli"))] pub fn test_rendezvous_server() { std::thread::spawn(test_rendezvous_server_); } +pub fn refresh_rendezvous_server() { + #[cfg(any(target_os = "android", target_os = "ios", feature = "cli"))] + test_rendezvous_server(); + #[cfg(not(any(target_os = "android", target_os = "ios", feature = "cli")))] + std::thread::spawn(|| { + if crate::ipc::test_rendezvous_server().is_err() { + test_rendezvous_server(); + } + }); +} + #[inline] pub fn get_time() -> i64 { std::time::SystemTime::now() @@ -437,14 +474,15 @@ pub fn is_modifier(evt: &KeyEvent) -> bool { } pub fn check_software_update() { - std::thread::spawn(move || allow_err!(_check_software_update())); + std::thread::spawn(move || allow_err!(check_software_update_())); } #[tokio::main(flavor = "current_thread")] -async fn _check_software_update() -> hbb_common::ResultType<()> { +async fn check_software_update_() -> hbb_common::ResultType<()> { sleep(3.).await; - let rendezvous_server = socket_client::get_target_addr(&get_rendezvous_server(1_000).await)?; + let rendezvous_server = + socket_client::get_target_addr(&format!("rs-sg.rustdesk.com:{}", config::RENDEZVOUS_PORT))?; let mut socket = socket_client::new_udp(Config::get_any_listen_addr(), RENDEZVOUS_TIMEOUT).await?; diff --git a/src/ipc.rs b/src/ipc.rs index 24a156eba..7df06cd22 100644 --- a/src/ipc.rs +++ b/src/ipc.rs @@ -127,6 +127,7 @@ pub enum Data { ClipbaordFile(ClipbaordFile), ClipboardFileEnabled(bool), PrivacyModeState((i32, PrivacyModeState)), + TestRendezvousServer, } #[tokio::main(flavor = "current_thread")] @@ -286,7 +287,11 @@ 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()); + value = Some(format!( + "{},{}", + Config::get_rendezvous_server(), + Config::get_rendezvous_servers().join(",") + )); } else if name == "rendezvous_servers" { value = Some(Config::get_rendezvous_servers().join(",")); } else { @@ -336,7 +341,9 @@ async fn handle(data: Data, stream: &mut Connection) { .await ); } - + Data::TestRendezvousServer => { + crate::test_rendezvous_server(); + } _ => {} } } @@ -520,11 +527,17 @@ pub fn get_password() -> String { } } -pub async fn get_rendezvous_server(ms_timeout: u64) -> String { +pub async fn get_rendezvous_server(ms_timeout: u64) -> (String, Vec) { if let Ok(Some(v)) = get_config_async("rendezvous_server", ms_timeout).await { - v + let mut urls = v.split(","); + let a = urls.next().unwrap_or_default().to_owned(); + let b: Vec = urls.map(|x| x.to_owned()).collect(); + (a, b) } else { - Config::get_rendezvous_server() + ( + Config::get_rendezvous_server(), + Config::get_rendezvous_servers(), + ) } } @@ -636,3 +649,10 @@ pub async fn set_socks(value: config::Socks5Server) -> ResultType<()> { .await?; Ok(()) } + +#[tokio::main(flavor = "current_thread")] +pub async fn test_rendezvous_server() -> ResultType<()> { + let mut c = connect(1000, "").await?; + c.send(&Data::TestRendezvousServer).await?; + Ok(()) +} diff --git a/src/mobile.rs b/src/mobile.rs index 5b9651e54..942510b17 100644 --- a/src/mobile.rs +++ b/src/mobile.rs @@ -591,7 +591,6 @@ impl Connection { log::debug!("Exit io_loop of id={}", session.id); } Err(err) => { - crate::common::test_rendezvous_server(); session.msgbox("error", "Connection Error", &err.to_string()); } } diff --git a/src/rendezvous_mediator.rs b/src/rendezvous_mediator.rs index a7f90b977..e2397fa16 100644 --- a/src/rendezvous_mediator.rs +++ b/src/rendezvous_mediator.rs @@ -51,9 +51,12 @@ impl RendezvousMediator { check_zombie(); let server = new_server(); if Config::get_nat_type() == NatType::UNKNOWN_NAT as i32 { - crate::common::test_nat_type(); + crate::test_nat_type(); nat_tested = true; } + if !Config::get_option("stop-service").is_empty() { + crate::test_rendezvous_server(); + } let server_cloned = server.clone(); tokio::spawn(async move { direct_server(server_cloned).await; @@ -68,7 +71,7 @@ impl RendezvousMediator { Config::reset_online(); if Config::get_option("stop-service").is_empty() { if !nat_tested { - crate::common::test_nat_type(); + crate::test_nat_type(); nat_tested = true; } let mut futs = Vec::new();