mirror/rustdesk
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk.git synced 2025-06-07 18:02:48 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix possible id mismatch bug because of register_pk udp packet lost

Browse Source
This commit is contained in:
rustdesk 2024-01-07 20:23:00 +08:00
parent 7e93a5d3cf
commit 396d7680d3
1 changed files with 11 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
src/rendezvous_mediator.rs
View File

@ -2,7 +2,7 @@ use std::{
net::SocketAddr, net::SocketAddr,
sync::{ sync::{
atomic::{AtomicBool, Ordering}, atomic::{AtomicBool, Ordering},
Arc, Mutex, Arc,
}, },
time::Instant, time::Instant,
}; };
@ -22,6 +22,7 @@ use hbb_common::{
tcp::FramedStream, tcp::FramedStream,
tokio::{ tokio::{
self, select, self, select,
sync::Mutex,
time::{interval, Duration}, time::{interval, Duration},
}, },
udp::FramedSocket, udp::FramedSocket,
@ -47,7 +48,6 @@ pub struct RendezvousMediator {
addr: TargetAddr<'static>, addr: TargetAddr<'static>,
host: String, host: String,
host_prefix: String, host_prefix: String,
last_id_pk_registry: String,
} }
impl RendezvousMediator { impl RendezvousMediator {
@ -84,6 +84,7 @@ impl RendezvousMediator {
crate::platform::linux_desktop_manager::start_xdesktop(); crate::platform::linux_desktop_manager::start_xdesktop();
loop { loop {
Config::reset_online(); Config::reset_online();
*SOLVING_PK_MISMATCH.lock().await = "".to_owned();
if Config::get_option("stop-service").is_empty() if Config::get_option("stop-service").is_empty()
&& !crate::platform::installing_service() && !crate::platform::installing_service()
{ {
@ -135,7 +136,6 @@ impl RendezvousMediator {
addr: addr.clone(), addr: addr.clone(),
host: host.clone(), host: host.clone(),
host_prefix, host_prefix,
last_id_pk_registry: "".to_owned(),
}; };
let mut timer = interval(TIMER_OUT); let mut timer = interval(TIMER_OUT);
@ -205,7 +205,7 @@ impl RendezvousMediator {
let elapsed_resp = last_register_resp.map(|x| x.elapsed().as_millis() as i64).unwrap_or(REG_INTERVAL); let elapsed_resp = last_register_resp.map(|x| x.elapsed().as_millis() as i64).unwrap_or(REG_INTERVAL);
let timeout = (elapsed_resp - last_register_sent.map(|x| x.elapsed().as_millis() as i64).unwrap_or(REG_INTERVAL)) > REG_TIMEOUT; let timeout = (elapsed_resp - last_register_sent.map(|x| x.elapsed().as_millis() as i64).unwrap_or(REG_INTERVAL)) > REG_TIMEOUT;
if timeout || elapsed_resp >= REG_INTERVAL { if timeout || elapsed_resp >= REG_INTERVAL {
allow_err!(rz.register_peer(Sink::Framed(&mut socket, &addr)).await); rz.register_peer(Sink::Framed(&mut socket, &addr)).await?;
last_register_sent = now; last_register_sent = now;
if timeout { if timeout {
fails += 1; fails += 1;
@ -246,7 +246,7 @@ impl RendezvousMediator {
update_latency(); update_latency();
if rpr.request_pk { if rpr.request_pk {
log::info!("request_pk received from {}", self.host); log::info!("request_pk received from {}", self.host);
allow_err!(self.register_pk(sink).await); self.register_pk(sink).await?;
} }
} }
Some(rendezvous_message::Union::RegisterPkResponse(rpr)) => { Some(rendezvous_message::Union::RegisterPkResponse(rpr)) => {
@ -255,10 +255,10 @@ impl RendezvousMediator {
Ok(register_pk_response::Result::OK) => { Ok(register_pk_response::Result::OK) => {
Config::set_key_confirmed(true); Config::set_key_confirmed(true);
Config::set_host_key_confirmed(&self.host_prefix, true); Config::set_host_key_confirmed(&self.host_prefix, true);
*SOLVING_PK_MISMATCH.lock().unwrap() = "".to_owned(); *SOLVING_PK_MISMATCH.lock().await = "".to_owned();
} }
Ok(register_pk_response::Result::UUID_MISMATCH) => { Ok(register_pk_response::Result::UUID_MISMATCH) => {
allow_err!(self.handle_uuid_mismatch(sink).await); self.handle_uuid_mismatch(sink).await?;
} }
_ => { _ => {
log::error!("unknown RegisterPkResponse"); log::error!("unknown RegisterPkResponse");
@ -310,7 +310,6 @@ impl RendezvousMediator {
addr: conn.local_addr().into_target_addr()?, addr: conn.local_addr().into_target_addr()?,
host: host.clone(), host: host.clone(),
host_prefix: host.clone(), host_prefix: host.clone(),
last_id_pk_registry: "".to_owned(),
}; };
let mut timer = interval(TIMER_OUT); let mut timer = interval(TIMER_OUT);
loop { loop {
@ -483,7 +482,6 @@ impl RendezvousMediator {
let pk = Config::get_key_pair().1; let pk = Config::get_key_pair().1;
let uuid = hbb_common::get_uuid(); let uuid = hbb_common::get_uuid();
let id = Config::get_id(); let id = Config::get_id();
self.last_id_pk_registry = id.clone();
msg_out.set_register_pk(RegisterPk { msg_out.set_register_pk(RegisterPk {
id, id,
uuid: uuid.into(), uuid: uuid.into(),
@ -495,11 +493,8 @@ impl RendezvousMediator {
} }
async fn handle_uuid_mismatch(&mut self, socket: Sink<'_>) -> ResultType<()> { async fn handle_uuid_mismatch(&mut self, socket: Sink<'_>) -> ResultType<()> {
if self.last_id_pk_registry != Config::get_id() {
return Ok(());
}
{ {
let mut solving = SOLVING_PK_MISMATCH.lock().unwrap(); let mut solving = SOLVING_PK_MISMATCH.lock().await;
if solving.is_empty() || *solving == self.host { if solving.is_empty() || *solving == self.host {
log::info!("UUID_MISMATCH received from {}", self.host); log::info!("UUID_MISMATCH received from {}", self.host);
Config::set_key_confirmed(false); Config::set_key_confirmed(false);
@ -513,9 +508,11 @@ impl RendezvousMediator {
} }
async fn register_peer(&mut self, socket: Sink<'_>) -> ResultType<()> { async fn register_peer(&mut self, socket: Sink<'_>) -> ResultType<()> {
if !SOLVING_PK_MISMATCH.lock().unwrap().is_empty() { let solving = SOLVING_PK_MISMATCH.lock().await;
if !(solving.is_empty() || *solving == self.host) {
return Ok(()); return Ok(());
} }
drop(solving);
if !Config::get_key_confirmed() || !Config::get_host_key_confirmed(&self.host_prefix) { if !Config::get_key_confirmed() || !Config::get_host_key_confirmed(&self.host_prefix) {
log::info!( log::info!(
"register_pk of {} due to key not confirmed", "register_pk of {} due to key not confirmed",