From e80da4af0d043e7e8ec87cc080156791a3a16f0d Mon Sep 17 00:00:00 2001
From: dignow <linlong1265@gmail.com>
Date: Sun, 9 Jul 2023 02:05:25 +0800
Subject: [PATCH] debug, win, privacy mode, broker

Signed-off-by: dignow <linlong1265@gmail.com>
---
 src/flutter.rs           | 13 -------------
 src/flutter_ffi.rs       |  5 -----
 src/privacy_win_mag.rs   | 21 ++++++++++++---------
 src/server/connection.rs |  4 ++++
 4 files changed, 16 insertions(+), 27 deletions(-)

diff --git a/src/flutter.rs b/src/flutter.rs
index d37e7d56c..397d8aa01 100644
--- a/src/flutter.rs
+++ b/src/flutter.rs
@@ -931,23 +931,10 @@ pub mod connection_manager {
         }
     }
 
-    extern "C" fn cm_clear_hook() {
-        #[cfg(target_os = "windows")]
-        crate::privacy_win_mag::stop();
-    }
-
     #[inline]
     pub fn cm_init() {
         #[cfg(not(any(target_os = "android", target_os = "ios")))]
         start_listen_ipc_thread();
-        #[cfg(target_os = "windows")]
-        shutdown_hooks::add_shutdown_hook(cm_clear_hook);
-    }
-
-    #[inline]
-    pub fn cm_clear() {
-        #[cfg(target_os = "windows")]
-        crate::privacy_win_mag::stop();
     }
 
     #[cfg(target_os = "android")]
diff --git a/src/flutter_ffi.rs b/src/flutter_ffi.rs
index 3730a1b44..c8de03862 100644
--- a/src/flutter_ffi.rs
+++ b/src/flutter_ffi.rs
@@ -1453,11 +1453,6 @@ pub fn cm_init() {
     crate::flutter::connection_manager::cm_init();
 }
 
-pub fn cm_clear() {
-    #[cfg(not(any(target_os = "android", target_os = "ios")))]
-    crate::flutter::connection_manager::cm_clear();
-}
-
 /// Start an ipc server for receiving the url scheme.
 ///
 /// * Should only be called in the main flutter window.
diff --git a/src/privacy_win_mag.rs b/src/privacy_win_mag.rs
index 33fc0faf0..446cba260 100644
--- a/src/privacy_win_mag.rs
+++ b/src/privacy_win_mag.rs
@@ -21,7 +21,7 @@ use winapi::{
         libloaderapi::{GetModuleHandleA, GetModuleHandleExA, GetProcAddress},
         memoryapi::{VirtualAllocEx, WriteProcessMemory},
         processthreadsapi::{
-            CreateProcessAsUserW, GetCurrentThreadId, QueueUserAPC, ResumeThread,
+            CreateProcessAsUserW, GetCurrentThreadId, QueueUserAPC, ResumeThread, TerminateProcess,
             PROCESS_INFORMATION, STARTUPINFOW,
         },
         winbase::{WTSGetActiveConsoleSessionId, CREATE_SUSPENDED, DETACHED_PROCESS},
@@ -46,7 +46,6 @@ pub const GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS: u32 = 4;
 const WM_USER_EXIT_HOOK: u32 = WM_USER + 1;
 
 lazy_static::lazy_static! {
-    static ref DLL_FOUND: Mutex<bool> = Mutex::new(false);
     static ref CONN_ID: Mutex<i32> = Mutex::new(0);
     static ref CUR_HOOK_THREAD_ID: Mutex<DWORD> = Mutex::new(0);
     static ref WND_HANDLERS: Mutex<WindowHandlers> = Mutex::new(WindowHandlers{hthread: 0, hprocess: 0});
@@ -66,16 +65,21 @@ impl Drop for WindowHandlers {
 impl WindowHandlers {
     fn reset(&mut self) {
         unsafe {
+            if self.hprocess != 0 {
+                let _res = TerminateProcess(self.hprocess as _, 0);
+                CloseHandle(self.hprocess as _);
+            }
+            self.hprocess = 0;
             if self.hthread != 0 {
                 CloseHandle(self.hthread as _);
             }
             self.hthread = 0;
-            if self.hprocess != 0 {
-                CloseHandle(self.hprocess as _);
-            }
-            self.hprocess = 0;
         }
     }
+
+    fn is_default(&self) -> bool {
+        self.hthread == 0 && self.hprocess == 0
+    }
 }
 
 pub fn turn_on_privacy(conn_id: i32) -> ResultType<bool> {
@@ -91,7 +95,7 @@ pub fn turn_on_privacy(conn_id: i32) -> ResultType<bool> {
         );
     }
 
-    if !*DLL_FOUND.lock().unwrap() {
+    if WND_HANDLERS.lock().unwrap().is_default() {
         log::info!("turn_on_privacy, dll not found when started, try start");
         start()?;
         std::thread::sleep(std::time::Duration::from_millis(1_000));
@@ -162,8 +166,6 @@ pub fn start() -> ResultType<()> {
         );
     }
 
-    *DLL_FOUND.lock().unwrap() = true;
-
     let hwnd = wait_find_privacy_hwnd(1_000)?;
     if !hwnd.is_null() {
         log::info!("Privacy window is ready");
@@ -263,6 +265,7 @@ pub fn start() -> ResultType<()> {
     Ok(())
 }
 
+#[inline]
 pub fn stop() {
     WND_HANDLERS.lock().unwrap().reset();
 }
diff --git a/src/server/connection.rs b/src/server/connection.rs
index f276be439..4224d670a 100644
--- a/src/server/connection.rs
+++ b/src/server/connection.rs
@@ -2559,6 +2559,10 @@ mod raii {
             if active_conns_lock.is_empty() {
                 video_service::try_plug_out_virtual_display();
             }
+            #[cfg(all(windows))]
+            if active_conns_lock.is_empty() {
+                crate::privacy_win_mag::stop();
+            }
         }
     }
 }