From 4cd025fcf305a9a84b534dcb1e3949d4ae983380 Mon Sep 17 00:00:00 2001
From: fufesou <shuanglongchen@yeah.net>
Date: Sat, 13 May 2023 18:19:44 +0800
Subject: [PATCH] try fix check privacy mode on conn

Signed-off-by: fufesou <shuanglongchen@yeah.net>
---
 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<String> = 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());