From 4cd025fcf305a9a84b534dcb1e3949d4ae983380 Mon Sep 17 00:00:00 2001 From: fufesou Date: Sat, 13 May 2023 18:19:44 +0800 Subject: [PATCH 1/4] try fix check privacy mode on conn Signed-off-by: fufesou --- src/server.rs | 11 ----------- src/server/connection.rs | 25 +++++++++++++++++++++++-- 2 files changed, 23 insertions(+), 13 deletions(-) diff --git a/src/server.rs b/src/server.rs index 89e57c79e..39cb258af 100644 --- a/src/server.rs +++ b/src/server.rs @@ -118,15 +118,6 @@ async fn accept_connection_(server: ServerPtr, socket: Stream, secure: bool) -> Ok(()) } -async fn check_privacy_mode_on(stream: &mut Stream) -> ResultType<()> { - if video_service::get_privacy_mode_conn_id() > 0 { - let msg_out = - crate::common::make_privacy_mode_msg(back_notification::PrivacyModeState::PrvOnByOther); - timeout(CONNECT_TIMEOUT, stream.send(&msg_out)).await??; - } - Ok(()) -} - pub async fn create_tcp_connection( server: ServerPtr, stream: Stream, @@ -134,8 +125,6 @@ pub async fn create_tcp_connection( secure: bool, ) -> ResultType<()> { let mut stream = stream; - check_privacy_mode_on(&mut stream).await?; - let id = { let mut w = server.write().unwrap(); w.id_count += 1; diff --git a/src/server/connection.rs b/src/server/connection.rs index 6c5634cf9..4da191953 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -793,8 +793,18 @@ impl Connection { self.send(msg_out).await; } - async fn on_open(&mut self, addr: SocketAddr) -> bool { - log::debug!("#{} Connection opened from {}.", self.inner.id, addr); + async fn check_privacy_mode_on(&mut self) -> bool { + if video_service::get_privacy_mode_conn_id() > 0 { + self.send_login_error("Someone turns on privacy mode, exit") + .await; + false + } + { + true + } + } + + async fn check_whitelist(&mut self, addr: &SocketAddr) -> bool { let whitelist: Vec = Config::get_option("whitelist") .split(",") .filter(|x| !x.is_empty()) @@ -822,6 +832,17 @@ impl Connection { sleep(1.).await; return false; } + true + } + + async fn on_open(&mut self, addr: SocketAddr) -> bool { + log::debug!("#{} Connection opened from {}.", self.inner.id, addr); + if !self.check_privacy_mode_on().await { + return false; + } + if !self.check_whitelist(&addr).await { + return false; + } self.ip = addr.ip().to_string(); let mut msg_out = Message::new(); msg_out.set_hash(self.hash.clone()); From 75cdab0ab2e8d8954dd76242c68028acd03ff1fe Mon Sep 17 00:00:00 2001 From: fufesou Date: Sat, 13 May 2023 18:36:48 +0800 Subject: [PATCH 2/4] fix build Signed-off-by: fufesou --- src/server/connection.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/server/connection.rs b/src/server/connection.rs index 4da191953..e60f9b82f 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -798,8 +798,7 @@ impl Connection { self.send_login_error("Someone turns on privacy mode, exit") .await; false - } - { + } else { true } } From ff810ac5dcc838e3cc2d8c3c494b7fba0457ebd9 Mon Sep 17 00:00:00 2001 From: fufesou Date: Sat, 13 May 2023 18:51:03 +0800 Subject: [PATCH 3/4] sleep to ensure the message is received Signed-off-by: fufesou --- src/server/connection.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/server/connection.rs b/src/server/connection.rs index e60f9b82f..650d85d28 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -793,10 +793,16 @@ impl Connection { self.send(msg_out).await; } + #[inline] + async fn sleep_to_ensure_msg_recved() { + sleep(1.).await; + } + async fn check_privacy_mode_on(&mut self) -> bool { if video_service::get_privacy_mode_conn_id() > 0 { self.send_login_error("Someone turns on privacy mode, exit") .await; + Self::sleep_to_ensure_msg_recved().await; false } else { true @@ -828,7 +834,7 @@ impl Connection { true, json!({ "ip":addr.ip() }), ); - sleep(1.).await; + Self::sleep_to_ensure_msg_recved().await; return false; } true From daf7bb6193ca14b321e3b499a22fe9310271a143 Mon Sep 17 00:00:00 2001 From: RustDesk <71636191+rustdesk@users.noreply.github.com> Date: Sat, 13 May 2023 18:56:24 +0800 Subject: [PATCH 4/4] Update connection.rs --- src/server/connection.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/server/connection.rs b/src/server/connection.rs index 650d85d28..b49a9059c 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -322,6 +322,7 @@ impl Connection { tx_desktop_ready: _tx_desktop_ready, }; if !conn.on_open(addr).await { + conn.sleep_to_ensure_msg_recved().await; return; } #[cfg(not(any(target_os = "android", target_os = "ios")))] @@ -802,7 +803,6 @@ impl Connection { if video_service::get_privacy_mode_conn_id() > 0 { self.send_login_error("Someone turns on privacy mode, exit") .await; - Self::sleep_to_ensure_msg_recved().await; false } else { true @@ -834,7 +834,6 @@ impl Connection { true, json!({ "ip":addr.ip() }), ); - Self::sleep_to_ensure_msg_recved().await; return false; } true