diff --git a/flutter/lib/models/native_model.dart b/flutter/lib/models/native_model.dart index 309c30f68..51940bb18 100644 --- a/flutter/lib/models/native_model.dart +++ b/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); diff --git a/src/flutter.rs b/src/flutter.rs index 4008b2420..397d8aa01 100644 --- a/src/flutter.rs +++ b/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}; diff --git a/src/flutter_ffi.rs b/src/flutter_ffi.rs index e54168249..c8de03862 100644 --- a/src/flutter_ffi.rs +++ b/src/flutter_ffi.rs @@ -1448,9 +1448,9 @@ pub fn main_use_texture_render() -> SyncReturn { } } -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. diff --git a/src/platform/windows.rs b/src/platform/windows.rs index 5ca069b80..c708650b8 100644 --- a/src/platform/windows.rs +++ b/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::*; diff --git a/src/privacy_win_mag.rs b/src/privacy_win_mag.rs index fe0ee4f69..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 = Mutex::new(false); static ref CONN_ID: Mutex = Mutex::new(0); static ref CUR_HOOK_THREAD_ID: Mutex = Mutex::new(0); static ref WND_HANDLERS: Mutex = 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 { @@ -85,7 +95,7 @@ pub fn turn_on_privacy(conn_id: i32) -> ResultType { ); } - 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 = 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/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(); + } } } } 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(cm: ConnectionManager) { e ); } - allow_err!(crate::privacy_win_mag::start()); }); #[cfg(target_os = "windows")]