Merge pull request #4921 from dignow/fix/win_privacy_broker

Fix/win privacy broker

flutter/lib/models/native_model.dart

@@ -233,7 +233,7 @@ class PlatformFFI {
             '_appType:$_appType,info1-id:$id,info2-name:$name,dir:$_dir');
       }
       if (desktopType == DesktopType.cm) {
-        await _ffiBind.cmStartListenIpcThread();
+        await _ffiBind.cmInit();
       }
       await _ffiBind.mainDeviceId(id: id);
       await _ffiBind.mainDeviceName(name: name);

src/flutter.rs

@@ -910,7 +910,7 @@ pub mod connection_manager {
 
     #[inline]
     #[cfg(not(any(target_os = "android", target_os = "ios")))]
-    pub fn start_listen_ipc_thread() {
+    fn start_listen_ipc_thread() {
         start_listen_ipc(true);
     }
 
@@ -931,6 +931,12 @@ pub mod connection_manager {
         }
     }
 
+    #[inline]
+    pub fn cm_init() {
+        #[cfg(not(any(target_os = "android", target_os = "ios")))]
+        start_listen_ipc_thread();
+    }
+
     #[cfg(target_os = "android")]
     use hbb_common::tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
 

src/flutter_ffi.rs

@@ -1448,9 +1448,9 @@ pub fn main_use_texture_render() -> SyncReturn<bool> {
     }
 }
 
-pub fn cm_start_listen_ipc_thread() {
+pub fn cm_init() {
     #[cfg(not(any(target_os = "android", target_os = "ios")))]
-    crate::flutter::connection_manager::start_listen_ipc_thread();
+    crate::flutter::connection_manager::cm_init();
 }
 
 /// Start an ipc server for receiving the url scheme.

src/platform/windows.rs

@@ -2161,9 +2161,11 @@ pub fn uninstall_service(show_new_window: bool) -> bool {
     sc stop {app_name}
     sc delete {app_name}
     if exist \"%PROGRAMDATA%\\Microsoft\\Windows\\Start Menu\\Programs\\Startup\\{app_name} Tray.lnk\" del /f /q \"%PROGRAMDATA%\\Microsoft\\Windows\\Start Menu\\Programs\\Startup\\{app_name} Tray.lnk\"
+    taskkill /F /IM {broker_exe}
     taskkill /F /IM {app_name}.exe{filter}
     ",
         app_name = crate::get_app_name(),
+        broker_exe = WIN_MAG_INJECTED_PROCESS_EXE,
     );
     if let Err(err) = run_cmds(cmds, false, "uninstall") {
         Config::set_option("stop-service".into(), "".into());
@@ -2273,6 +2275,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::*;

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});
@@ -59,17 +58,28 @@ struct WindowHandlers {
 
 impl Drop for WindowHandlers {
     fn drop(&mut self) {
+        self.reset();
+    }
+}
+
+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> {
@@ -85,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));
@@ -156,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");
@@ -257,6 +265,11 @@ pub fn start() -> ResultType<()> {
     Ok(())
 }
 
+#[inline]
+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();
 

src/server.rs

@@ -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 {

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();
+            }
         }
     }
 }

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")]