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

rendezvous_mediator: ensure once one exits, the others also exit.

Browse Source
This commit is contained in:
open-trade 2022-01-04 17:44:06 +08:00
parent cbcad03bc6
commit fb392b81cb
1 changed files with 11 additions and 30 deletions

41
src/rendezvous_mediator.rs

@ -17,7 +17,10 @@ use hbb_common::{
}; };
use std::{ use std::{
net::SocketAddr, net::SocketAddr,
sync::{Arc, Mutex}, sync::{
atomic::{AtomicBool, Ordering},
Arc, Mutex,
},
time::SystemTime, time::SystemTime,
}; };
use uuid::Uuid; use uuid::Uuid;
@ -27,6 +30,7 @@ type Message = RendezvousMessage;
lazy_static::lazy_static! { lazy_static::lazy_static! {
pub static ref SOLVING_PK_MISMATCH: Arc<Mutex<String>> = Default::default(); pub static ref SOLVING_PK_MISMATCH: Arc<Mutex<String>> = Default::default();
} }
static EXITED: AtomicBool = AtomicBool::new(false);
#[derive(Clone)] #[derive(Clone)]
pub struct RendezvousMediator { pub struct RendezvousMediator {
@ -55,40 +59,14 @@ impl RendezvousMediator {
} }
let mut futs = Vec::new(); let mut futs = Vec::new();
let servers = Config::get_rendezvous_servers(); let servers = Config::get_rendezvous_servers();
EXITED.store(false, Ordering::SeqCst);
for host in servers.clone() { for host in servers.clone() {
let server = server.clone(); let server = server.clone();
let servers = servers.clone(); let servers = servers.clone();
futs.push(tokio::spawn(async move { futs.push(tokio::spawn(async move {
allow_err!(Self::start(server, host, servers).await); allow_err!(Self::start(server, host, servers).await);
// let socks5_conf = socket_client::get_socks5_conf(); // EXITED here is to ensure once one exits, the others also exit.
// if socks5_conf.is_some() { EXITED.store(true, Ordering::SeqCst);
// let target = format!("{}:{}", host, RENDEZVOUS_PORT);
// let conn_fn = |bind_addr: SocketAddr| {
// let target = target.clone();
// let conf_ref = &socks5_conf;
// async move {
// socket_client::connect_udp_socks5(
// target,
// bind_addr,
// conf_ref,
// RENDEZVOUS_TIMEOUT,
// )
// .await
// }
// };
// allow_err!(Self::start(server, host, servers, conn_fn, true).await);
// } else {
// allow_err!(
// Self::start(
// server,
// host,
// servers,
// socket_client::connect_udp_socket,
// false,
// )
// .await
// );
// }
})); }));
} }
join_all(futs).await; join_all(futs).await;
@ -246,6 +224,9 @@ impl RendezvousMediator {
if !Config::get_option("stop-service").is_empty() { if !Config::get_option("stop-service").is_empty() {
break; break;
} }
if EXITED.load(Ordering::SeqCst) {
break;
}
if rz.addr.port() == 0 { if rz.addr.port() == 0 {
// tcp is established to help connecting socks5 // tcp is established to help connecting socks5
allow_err!(rz.dns_check(&mut host_addr)); allow_err!(rz.dns_check(&mut host_addr));