From 1648895156f9c36ddd0f0426dc07041802bafc57 Mon Sep 17 00:00:00 2001
From: dignow <linlong1265@gmail.com>
Date: Sat, 8 Jul 2023 15:26:24 +0800
Subject: [PATCH] fix, win privacy mode, broker process

Signed-off-by: dignow <linlong1265@gmail.com>
---
 flutter/lib/desktop/pages/server_page.dart |  1 +
 src/flutter.rs                             |  5 +++++
 src/flutter_ffi.rs                         |  5 +++++
 src/platform/windows.rs                    |  9 +++++++++
 src/privacy_win_mag.rs                     | 10 ++++++++++
 src/server.rs                              |  4 +++-
 src/ui_cm_interface.rs                     |  1 -
 7 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/flutter/lib/desktop/pages/server_page.dart b/flutter/lib/desktop/pages/server_page.dart
index 25be23142..fe6922938 100644
--- a/flutter/lib/desktop/pages/server_page.dart
+++ b/flutter/lib/desktop/pages/server_page.dart
@@ -40,6 +40,7 @@ class _DesktopServerPageState extends State<DesktopServerPage>
 
   @override
   void dispose() {
+    bind.cmClear();
     windowManager.removeListener(this);
     super.dispose();
   }
diff --git a/src/flutter.rs b/src/flutter.rs
index 4008b2420..6159a6de7 100644
--- a/src/flutter.rs
+++ b/src/flutter.rs
@@ -931,6 +931,11 @@ pub mod connection_manager {
         }
     }
 
+    #[cfg(target_os = "windows")]
+    pub fn cm_clear() {
+        crate::privacy_win_mag::stop();
+    }
+
     #[cfg(target_os = "android")]
     use hbb_common::tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
 
diff --git a/src/flutter_ffi.rs b/src/flutter_ffi.rs
index e54168249..3703a0834 100644
--- a/src/flutter_ffi.rs
+++ b/src/flutter_ffi.rs
@@ -1453,6 +1453,11 @@ pub fn cm_start_listen_ipc_thread() {
     crate::flutter::connection_manager::start_listen_ipc_thread();
 }
 
+pub fn cm_clear() {
+    #[cfg(target_os = "windows")]
+    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/platform/windows.rs b/src/platform/windows.rs
index 5ca069b80..67a2af676 100644
--- a/src/platform/windows.rs
+++ b/src/platform/windows.rs
@@ -2273,6 +2273,15 @@ fn run_after_run_cmds(silent: bool) {
     std::thread::sleep(std::time::Duration::from_millis(300));
 }
 
+#[inline]
+pub fn try_kill_broker() {
+    allow_err!(run_cmds(
+        format!("taskkill /F /IM {}", WIN_MAG_INJECTED_PROCESS_EXE),
+        false,
+        "kill_broker"
+    ));
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/privacy_win_mag.rs b/src/privacy_win_mag.rs
index fe0ee4f69..33fc0faf0 100644
--- a/src/privacy_win_mag.rs
+++ b/src/privacy_win_mag.rs
@@ -59,6 +59,12 @@ struct WindowHandlers {
 
 impl Drop for WindowHandlers {
     fn drop(&mut self) {
+        self.reset();
+    }
+}
+
+impl WindowHandlers {
+    fn reset(&mut self) {
         unsafe {
             if self.hthread != 0 {
                 CloseHandle(self.hthread as _);
@@ -257,6 +263,10 @@ pub fn start() -> ResultType<()> {
     Ok(())
 }
 
+pub fn stop() {
+    WND_HANDLERS.lock().unwrap().reset();
+}
+
 unsafe fn inject_dll<'a>(hproc: HANDLE, hthread: HANDLE, dll_file: &'a str) -> ResultType<()> {
     let dll_file_utf16: Vec<u16> = dll_file.encode_utf16().chain(Some(0).into_iter()).collect();
 
diff --git a/src/server.rs b/src/server.rs
index 8a2cf6178..b2d0b593f 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -303,7 +303,7 @@ impl Server {
     // get a new unique id
     pub fn get_new_id(&mut self) -> i32 {
         self.id_count += 1;
-        self.id_count 
+        self.id_count
     }
 }
 
@@ -389,6 +389,8 @@ pub async fn start_server(is_server: bool) {
         }
         #[cfg(any(target_os = "macos", target_os = "linux"))]
         tokio::spawn(async { sync_and_watch_config_dir().await });
+        #[cfg(target_os = "windows")]
+        crate::platform::try_kill_broker();
         crate::RendezvousMediator::start_all().await;
     } else {
         match crate::ipc::connect(1000, "").await {
diff --git a/src/ui_cm_interface.rs b/src/ui_cm_interface.rs
index e5a5c8e3b..9a7556d25 100644
--- a/src/ui_cm_interface.rs
+++ b/src/ui_cm_interface.rs
@@ -533,7 +533,6 @@ pub async fn start_ipc<T: InvokeUiCM>(cm: ConnectionManager<T>) {
                 e
             );
         }
-        allow_err!(crate::privacy_win_mag::start());
     });
 
     #[cfg(target_os = "windows")]