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..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")))] @@ -793,8 +794,22 @@ 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); + #[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; + false + } else { + true + } + } + + async fn check_whitelist(&mut self, addr: &SocketAddr) -> bool { let whitelist: Vec = Config::get_option("whitelist") .split(",") .filter(|x| !x.is_empty()) @@ -819,7 +834,17 @@ impl Connection { true, json!({ "ip":addr.ip() }), ); - 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();