use super::{CursorData, ResultType}; use crate::{ common::PORTABLE_APPNAME_RUNTIME_ENV_KEY, custom_server::*, ipc, privacy_mode::win_topmost_window::{self, WIN_TOPMOST_INJECTED_PROCESS_EXE}, }; use hbb_common::{ allow_err, anyhow::anyhow, bail, config::{self, Config}, libc::{c_int, wchar_t}, log, message_proto::{DisplayInfo, Resolution, WindowsSession}, sleep, sysinfo::{Pid, System}, timeout, tokio, }; use std::{ collections::HashMap, ffi::{CString, OsString}, fs, io::{self, prelude::*}, mem, os::{ raw::c_ulong, windows::{ffi::OsStringExt, process::CommandExt}, }, path::*, ptr::null_mut, sync::{atomic::Ordering, Arc, Mutex}, time::{Duration, Instant}, }; use wallpaper; #[cfg(not(debug_assertions))] use winapi::um::libloaderapi::{LoadLibraryExW, LOAD_LIBRARY_SEARCH_USER_DIRS}; use winapi::{ ctypes::c_void, shared::{minwindef::*, ntdef::NULL, windef::*, winerror::*}, um::{ errhandlingapi::GetLastError, handleapi::{CloseHandle, INVALID_HANDLE_VALUE}, libloaderapi::{ GetProcAddress, LoadLibraryA, LoadLibraryExA, LOAD_LIBRARY_SEARCH_SYSTEM32, }, minwinbase::STILL_ACTIVE, processthreadsapi::{ GetCurrentProcess, GetCurrentProcessId, GetExitCodeProcess, OpenProcess, OpenProcessToken, ProcessIdToSessionId, PROCESS_INFORMATION, STARTUPINFOW, }, securitybaseapi::{ AllocateAndInitializeSid, DuplicateToken, EqualSid, FreeSid, GetTokenInformation, }, shellapi::ShellExecuteW, sysinfoapi::{GetNativeSystemInfo, SYSTEM_INFO}, winbase::*, wingdi::*, winnt::{ SecurityImpersonation, TokenElevation, TokenGroups, TokenImpersonation, TokenType, DOMAIN_ALIAS_RID_ADMINS, ES_AWAYMODE_REQUIRED, ES_CONTINUOUS, ES_DISPLAY_REQUIRED, ES_SYSTEM_REQUIRED, HANDLE, PROCESS_ALL_ACCESS, PROCESS_QUERY_LIMITED_INFORMATION, PSID, SECURITY_BUILTIN_DOMAIN_RID, SECURITY_NT_AUTHORITY, SID_IDENTIFIER_AUTHORITY, TOKEN_ELEVATION, TOKEN_GROUPS, TOKEN_QUERY, TOKEN_TYPE, }, winreg::HKEY_CURRENT_USER, winspool::{ EnumPrintersW, GetDefaultPrinterW, PRINTER_ENUM_CONNECTIONS, PRINTER_ENUM_LOCAL, PRINTER_INFO_1W, }, winuser::*, }, }; use windows::Win32::{ Foundation::{CloseHandle as WinCloseHandle, HANDLE as WinHANDLE}, System::Diagnostics::ToolHelp::{ CreateToolhelp32Snapshot, Process32FirstW, Process32NextW, PROCESSENTRY32W, TH32CS_SNAPPROCESS, }, }; use windows_service::{ define_windows_service, service::{ ServiceControl, ServiceControlAccept, ServiceExitCode, ServiceState, ServiceStatus, ServiceType, }, service_control_handler::{self, ServiceControlHandlerResult}, }; use winreg::{enums::*, RegKey}; pub const FLUTTER_RUNNER_WIN32_WINDOW_CLASS: &'static str = "FLUTTER_RUNNER_WIN32_WINDOW"; // main window, install window pub const EXPLORER_EXE: &'static str = "explorer.exe"; pub const SET_FOREGROUND_WINDOW: &'static str = "SET_FOREGROUND_WINDOW"; const REG_NAME_INSTALL_DESKTOPSHORTCUTS: &str = "DESKTOPSHORTCUTS"; const REG_NAME_INSTALL_STARTMENUSHORTCUTS: &str = "STARTMENUSHORTCUTS"; pub const REG_NAME_INSTALL_PRINTER: &str = "PRINTER"; pub fn get_focused_display(displays: Vec) -> Option { unsafe { let hwnd = GetForegroundWindow(); let mut rect: RECT = mem::zeroed(); if GetWindowRect(hwnd, &mut rect as *mut RECT) == 0 { return None; } displays.iter().position(|display| { let center_x = rect.left + (rect.right - rect.left) / 2; let center_y = rect.top + (rect.bottom - rect.top) / 2; center_x >= display.x && center_x <= display.x + display.width && center_y >= display.y && center_y <= display.y + display.height }) } } pub fn get_cursor_pos() -> Option<(i32, i32)> { unsafe { let mut out = mem::MaybeUninit::::uninit(); if GetCursorPos(out.as_mut_ptr()) == FALSE { return None; } let out = out.assume_init(); Some((out.x, out.y)) } } pub fn set_cursor_pos(x: i32, y: i32) -> bool { unsafe { if SetCursorPos(x, y) == FALSE { let err = GetLastError(); log::warn!("SetCursorPos failed: x={}, y={}, error_code={}", x, y, err); return false; } true } } /// Clip cursor to a rectangle. Pass None to unclip. pub fn clip_cursor(rect: Option<(i32, i32, i32, i32)>) -> bool { unsafe { let result = match rect { Some((left, top, right, bottom)) => { let r = RECT { left, top, right, bottom, }; ClipCursor(&r) } None => ClipCursor(std::ptr::null()), }; if result == FALSE { let err = GetLastError(); log::warn!( "ClipCursor failed: rect={:?}, error_code={}", rect, err ); return false; } true } } pub fn reset_input_cache() {} pub fn get_cursor() -> ResultType> { unsafe { #[allow(invalid_value)] let mut ci: CURSORINFO = mem::MaybeUninit::uninit().assume_init(); ci.cbSize = std::mem::size_of::() as _; if crate::portable_service::client::get_cursor_info(&mut ci) == FALSE { return Err(io::Error::last_os_error().into()); } if ci.flags & CURSOR_SHOWING == 0 { Ok(None) } else { Ok(Some(ci.hCursor as _)) } } } struct IconInfo(ICONINFO); impl IconInfo { fn new(icon: HICON) -> ResultType { unsafe { #[allow(invalid_value)] let mut ii = mem::MaybeUninit::uninit().assume_init(); if GetIconInfo(icon, &mut ii) == FALSE { Err(io::Error::last_os_error().into()) } else { let ii = Self(ii); if ii.0.hbmMask.is_null() { bail!("Cursor bitmap handle is NULL"); } return Ok(ii); } } } fn is_color(&self) -> bool { !self.0.hbmColor.is_null() } } impl Drop for IconInfo { fn drop(&mut self) { unsafe { if !self.0.hbmColor.is_null() { DeleteObject(self.0.hbmColor as _); } if !self.0.hbmMask.is_null() { DeleteObject(self.0.hbmMask as _); } } } } // https://github.com/TurboVNC/tightvnc/blob/a235bae328c12fd1c3aed6f3f034a37a6ffbbd22/vnc_winsrc/winvnc/vncEncoder.cpp // https://github.com/TigerVNC/tigervnc/blob/master/win/rfb_win32/DeviceFrameBuffer.cxx pub fn get_cursor_data(hcursor: u64) -> ResultType { unsafe { let mut ii = IconInfo::new(hcursor as _)?; let bm_mask = get_bitmap(ii.0.hbmMask)?; let mut width = bm_mask.bmWidth; let mut height = if ii.is_color() { bm_mask.bmHeight } else { bm_mask.bmHeight / 2 }; let cbits_size = width * height * 4; if cbits_size < 16 { bail!("Invalid icon: too small"); // solve some crash } let mut cbits: Vec = Vec::new(); cbits.resize(cbits_size as _, 0); let mut mbits: Vec = Vec::new(); mbits.resize((bm_mask.bmWidthBytes * bm_mask.bmHeight) as _, 0); let r = GetBitmapBits(ii.0.hbmMask, mbits.len() as _, mbits.as_mut_ptr() as _); if r == 0 { bail!("Failed to copy bitmap data"); } if r != (mbits.len() as i32) { bail!( "Invalid mask cursor buffer size, got {} bytes, expected {}", r, mbits.len() ); } let do_outline; if ii.is_color() { get_rich_cursor_data(ii.0.hbmColor, width, height, &mut cbits)?; do_outline = fix_cursor_mask( &mut mbits, &mut cbits, width as _, height as _, bm_mask.bmWidthBytes as _, ); } else { do_outline = handleMask( cbits.as_mut_ptr(), mbits.as_ptr(), width, height, bm_mask.bmWidthBytes, bm_mask.bmHeight, ) > 0; } if do_outline { let mut outline = Vec::new(); outline.resize(((width + 2) * (height + 2) * 4) as _, 0); drawOutline( outline.as_mut_ptr(), cbits.as_ptr(), width, height, outline.len() as _, ); cbits = outline; width += 2; height += 2; ii.0.xHotspot += 1; ii.0.yHotspot += 1; } Ok(CursorData { id: hcursor, colors: cbits.into(), hotx: ii.0.xHotspot as _, hoty: ii.0.yHotspot as _, width: width as _, height: height as _, ..Default::default() }) } } #[inline] fn get_bitmap(handle: HBITMAP) -> ResultType { unsafe { let mut bm: BITMAP = mem::zeroed(); if GetObjectA( handle as _, std::mem::size_of::() as _, &mut bm as *mut BITMAP as *mut _, ) == FALSE { return Err(io::Error::last_os_error().into()); } if bm.bmPlanes != 1 { bail!("Unsupported multi-plane cursor"); } if bm.bmBitsPixel != 1 { bail!("Unsupported cursor mask format"); } Ok(bm) } } struct DC(HDC); impl DC { fn new() -> ResultType { unsafe { let dc = GetDC(0 as _); if dc.is_null() { bail!("Failed to get a drawing context"); } Ok(Self(dc)) } } } impl Drop for DC { fn drop(&mut self) { unsafe { if !self.0.is_null() { ReleaseDC(0 as _, self.0); } } } } struct CompatibleDC(HDC); impl CompatibleDC { fn new(existing: HDC) -> ResultType { unsafe { let dc = CreateCompatibleDC(existing); if dc.is_null() { bail!("Failed to get a compatible drawing context"); } Ok(Self(dc)) } } } impl Drop for CompatibleDC { fn drop(&mut self) { unsafe { if !self.0.is_null() { DeleteDC(self.0); } } } } struct BitmapDC(CompatibleDC, HBITMAP); impl BitmapDC { fn new(hdc: HDC, hbitmap: HBITMAP) -> ResultType { unsafe { let dc = CompatibleDC::new(hdc)?; let oldbitmap = SelectObject(dc.0, hbitmap as _) as HBITMAP; if oldbitmap.is_null() { bail!("Failed to select CompatibleDC"); } Ok(Self(dc, oldbitmap)) } } fn dc(&self) -> HDC { (self.0).0 } } impl Drop for BitmapDC { fn drop(&mut self) { unsafe { if !self.1.is_null() { SelectObject((self.0).0, self.1 as _); } } } } #[inline] fn get_rich_cursor_data( hbm_color: HBITMAP, width: i32, height: i32, out: &mut Vec, ) -> ResultType<()> { unsafe { let dc = DC::new()?; let bitmap_dc = BitmapDC::new(dc.0, hbm_color)?; if get_di_bits(out.as_mut_ptr(), bitmap_dc.dc(), hbm_color, width, height) > 0 { bail!("Failed to get di bits: {}", io::Error::last_os_error()); } } Ok(()) } fn fix_cursor_mask( mbits: &mut Vec, cbits: &mut Vec, width: usize, height: usize, bm_width_bytes: usize, ) -> bool { let mut pix_idx = 0; for _ in 0..height { for _ in 0..width { if cbits[pix_idx + 3] != 0 { return false; } pix_idx += 4; } } let packed_width_bytes = (width + 7) >> 3; let bm_size = mbits.len(); let c_size = cbits.len(); // Pack and invert bitmap data (mbits) // borrow from tigervnc for y in 0..height { for x in 0..packed_width_bytes { let a = y * packed_width_bytes + x; let b = y * bm_width_bytes + x; if a < bm_size && b < bm_size { mbits[a] = !mbits[b]; } } } // Replace "inverted background" bits with black color to ensure // cross-platform interoperability. Not beautiful but necessary code. // borrow from tigervnc let bytes_row = width << 2; for y in 0..height { let mut bitmask: u8 = 0x80; for x in 0..width { let mask_idx = y * packed_width_bytes + (x >> 3); if mask_idx < bm_size { let pix_idx = y * bytes_row + (x << 2); if (mbits[mask_idx] & bitmask) == 0 { for b1 in 0..4 { let a = pix_idx + b1; if a < c_size { if cbits[a] != 0 { mbits[mask_idx] ^= bitmask; for b2 in b1..4 { let b = pix_idx + b2; if b < c_size { cbits[b] = 0x00; } } break; } } } } } bitmask >>= 1; if bitmask == 0 { bitmask = 0x80; } } } // borrow from noVNC let mut pix_idx = 0; for y in 0..height { for x in 0..width { let mask_idx = y * packed_width_bytes + (x >> 3); let mut alpha = 255; if mask_idx < bm_size { if (mbits[mask_idx] << (x & 0x7)) & 0x80 == 0 { alpha = 0; } } let a = cbits[pix_idx + 2]; let b = cbits[pix_idx + 1]; let c = cbits[pix_idx]; cbits[pix_idx] = a; cbits[pix_idx + 1] = b; cbits[pix_idx + 2] = c; cbits[pix_idx + 3] = alpha; pix_idx += 4; } } return true; } define_windows_service!(ffi_service_main, service_main); fn service_main(arguments: Vec) { if let Err(e) = run_service(arguments) { log::error!("run_service failed: {}", e); } } pub fn start_os_service() { if let Err(e) = windows_service::service_dispatcher::start(crate::get_app_name(), ffi_service_main) { log::error!("start_service failed: {}", e); } } const SERVICE_TYPE: ServiceType = ServiceType::OWN_PROCESS; extern "C" { fn get_current_session(rdp: BOOL) -> DWORD; fn LaunchProcessWin( cmd: *const u16, session_id: DWORD, as_user: BOOL, show: BOOL, token_pid: &mut DWORD, ) -> HANDLE; fn GetSessionUserTokenWin( lphUserToken: LPHANDLE, dwSessionId: DWORD, as_user: BOOL, token_pid: &mut DWORD, ) -> BOOL; fn selectInputDesktop() -> BOOL; fn inputDesktopSelected() -> BOOL; fn is_windows_server() -> BOOL; fn is_windows_10_or_greater() -> BOOL; fn handleMask( out: *mut u8, mask: *const u8, width: i32, height: i32, bmWidthBytes: i32, bmHeight: i32, ) -> i32; fn drawOutline(out: *mut u8, in_: *const u8, width: i32, height: i32, out_size: i32); fn get_di_bits(out: *mut u8, dc: HDC, hbmColor: HBITMAP, width: i32, height: i32) -> i32; fn blank_screen(v: BOOL); fn win32_enable_lowlevel_keyboard(hwnd: HWND) -> i32; fn win32_disable_lowlevel_keyboard(hwnd: HWND); fn win_stop_system_key_propagate(v: BOOL); fn is_win_down() -> BOOL; fn is_local_system() -> BOOL; fn alloc_console_and_redirect(); fn is_service_running_w(svc_name: *const u16) -> bool; } pub fn get_current_session_id(share_rdp: bool) -> DWORD { unsafe { get_current_session(if share_rdp { TRUE } else { FALSE }) } } extern "system" { fn BlockInput(v: BOOL) -> BOOL; } #[tokio::main(flavor = "current_thread")] async fn run_service(_arguments: Vec) -> ResultType<()> { let event_handler = move |control_event| -> ServiceControlHandlerResult { log::info!("Got service control event: {:?}", control_event); match control_event { ServiceControl::Interrogate => ServiceControlHandlerResult::NoError, ServiceControl::Stop | ServiceControl::Preshutdown | ServiceControl::Shutdown => { send_close(crate::POSTFIX_SERVICE).ok(); ServiceControlHandlerResult::NoError } _ => ServiceControlHandlerResult::NotImplemented, } }; // Register system service event handler let status_handle = service_control_handler::register(crate::get_app_name(), event_handler)?; let next_status = ServiceStatus { // Should match the one from system service registry service_type: SERVICE_TYPE, // The new state current_state: ServiceState::Running, // Accept stop events when running controls_accepted: ServiceControlAccept::STOP, // Used to report an error when starting or stopping only, otherwise must be zero exit_code: ServiceExitCode::Win32(0), // Only used for pending states, otherwise must be zero checkpoint: 0, // Only used for pending states, otherwise must be zero wait_hint: Duration::default(), process_id: None, }; // Tell the system that the service is running now status_handle.set_service_status(next_status)?; let mut session_id = unsafe { get_current_session(share_rdp()) }; log::info!("session id {}", session_id); let mut h_process = launch_server(session_id, true).await.unwrap_or(NULL); let mut incoming = ipc::new_listener(crate::POSTFIX_SERVICE).await?; let mut stored_usid = None; loop { let sids: Vec<_> = get_available_sessions(false) .iter() .map(|e| e.sid) .collect(); if !sids.contains(&session_id) || !is_share_rdp() { let current_active_session = unsafe { get_current_session(share_rdp()) }; if session_id != current_active_session { session_id = current_active_session; // https://github.com/rustdesk/rustdesk/discussions/10039 let count = ipc::get_port_forward_session_count(1000).await.unwrap_or(0); if count == 0 { h_process = launch_server(session_id, true).await.unwrap_or(NULL); } } } let res = timeout(super::SERVICE_INTERVAL, incoming.next()).await; match res { Ok(res) => match res { Some(Ok(stream)) => { let mut stream = ipc::Connection::new(stream); if let Ok(Some(data)) = stream.next_timeout(1000).await { match data { ipc::Data::Close => { log::info!("close received"); break; } ipc::Data::SAS => { send_sas(); } ipc::Data::UserSid(usid) => { if let Some(usid) = usid { if session_id != usid { log::info!( "session changed from {} to {}", session_id, usid ); session_id = usid; stored_usid = Some(session_id); h_process = launch_server(session_id, true).await.unwrap_or(NULL); } } } _ => {} } } } _ => {} }, Err(_) => { // timeout unsafe { let tmp = get_current_session(share_rdp()); if tmp == 0xFFFFFFFF { continue; } let mut close_sent = false; if tmp != session_id && stored_usid != Some(session_id) { log::info!("session changed from {} to {}", session_id, tmp); session_id = tmp; let count = ipc::get_port_forward_session_count(1000).await.unwrap_or(0); if count == 0 { send_close_async("").await.ok(); close_sent = true; } } let mut exit_code: DWORD = 0; if h_process.is_null() || (GetExitCodeProcess(h_process, &mut exit_code) == TRUE && exit_code != STILL_ACTIVE && CloseHandle(h_process) == TRUE) { match launch_server(session_id, !close_sent).await { Ok(ptr) => { h_process = ptr; } Err(err) => { log::error!("Failed to launch server: {}", err); } } } } } } } if !h_process.is_null() { send_close_async("").await.ok(); unsafe { CloseHandle(h_process) }; } status_handle.set_service_status(ServiceStatus { service_type: SERVICE_TYPE, current_state: ServiceState::Stopped, controls_accepted: ServiceControlAccept::empty(), exit_code: ServiceExitCode::Win32(0), checkpoint: 0, wait_hint: Duration::default(), process_id: None, })?; Ok(()) } async fn launch_server(session_id: DWORD, close_first: bool) -> ResultType { if close_first { // in case started some elsewhere send_close_async("").await.ok(); } let cmd = format!( "\"{}\" --server", std::env::current_exe()?.to_str().unwrap_or("") ); launch_privileged_process(session_id, &cmd) } pub fn launch_privileged_process(session_id: DWORD, cmd: &str) -> ResultType { use std::os::windows::ffi::OsStrExt; let wstr: Vec = std::ffi::OsStr::new(&cmd) .encode_wide() .chain(Some(0).into_iter()) .collect(); let wstr = wstr.as_ptr(); let mut token_pid = 0; let h = unsafe { LaunchProcessWin(wstr, session_id, FALSE, FALSE, &mut token_pid) }; if h.is_null() { log::error!( "Failed to launch privileged process: {}", io::Error::last_os_error() ); if token_pid == 0 { log::error!("No process winlogon.exe"); } } Ok(h) } pub fn run_as_user(arg: Vec<&str>) -> ResultType> { run_exe_in_cur_session(std::env::current_exe()?.to_str().unwrap_or(""), arg, false) } pub fn run_exe_in_cur_session( exe: &str, arg: Vec<&str>, show: bool, ) -> ResultType> { let Some(session_id) = get_current_process_session_id() else { bail!("Failed to get current process session id"); }; run_exe_in_session(exe, arg, session_id, show) } pub fn run_exe_in_session( exe: &str, arg: Vec<&str>, session_id: DWORD, show: bool, ) -> ResultType> { use std::os::windows::ffi::OsStrExt; let cmd = format!("\"{}\" {}", exe, arg.join(" "),); let wstr: Vec = std::ffi::OsStr::new(&cmd) .encode_wide() .chain(Some(0).into_iter()) .collect(); let wstr = wstr.as_ptr(); let mut token_pid = 0; let h = unsafe { LaunchProcessWin( wstr, session_id, TRUE, if show { TRUE } else { FALSE }, &mut token_pid, ) }; if h.is_null() { if token_pid == 0 { bail!( "Failed to launch {:?} with session id {}: no process {}", arg, session_id, EXPLORER_EXE ); } bail!( "Failed to launch {:?} with session id {}: {}", arg, session_id, io::Error::last_os_error() ); } Ok(None) } #[tokio::main(flavor = "current_thread")] async fn send_close(postfix: &str) -> ResultType<()> { send_close_async(postfix).await } async fn send_close_async(postfix: &str) -> ResultType<()> { ipc::connect(1000, postfix) .await? .send(&ipc::Data::Close) .await?; // sleep a while to wait for closing and exit sleep(0.1).await; Ok(()) } // https://docs.microsoft.com/en-us/windows/win32/api/sas/nf-sas-sendsas // https://www.cnblogs.com/doutu/p/4892726.html pub fn send_sas() { #[link(name = "sas")] extern "system" { pub fn SendSAS(AsUser: BOOL); } unsafe { log::info!("SAS received"); // Check and temporarily set SoftwareSASGeneration if needed let mut original_value: Option = None; let hklm = RegKey::predef(HKEY_LOCAL_MACHINE); if let Ok(policy_key) = hklm.open_subkey_with_flags( "Software\\Microsoft\\Windows\\CurrentVersion\\Policies\\System", KEY_READ | KEY_WRITE, ) { // Read current value match policy_key.get_value::("SoftwareSASGeneration") { Ok(value) => { /* - 0 = None (disabled) - 1 = Services - 2 = Ease of Access applications - 3 = Services and Ease of Access applications (Both) */ if value != 1 && value != 3 { original_value = Some(value); log::info!("SoftwareSASGeneration is {}, setting to 1", value); // Set to 1 for SendSAS to work if let Err(e) = policy_key.set_value("SoftwareSASGeneration", &1u32) { log::error!("Failed to set SoftwareSASGeneration: {}", e); } } } Err(e) => { log::info!( "SoftwareSASGeneration not found or error reading: {}, setting to 1", e ); original_value = Some(0); // Mark that we need to restore (delete) it // Create and set to 1 if let Err(e) = policy_key.set_value("SoftwareSASGeneration", &1u32) { log::error!("Failed to set SoftwareSASGeneration: {}", e); } } } } else { log::error!("Failed to open registry key for SoftwareSASGeneration"); } // Send SAS SendSAS(FALSE); // Restore original value if we changed it if let Some(original) = original_value { if let Ok(policy_key) = hklm.open_subkey_with_flags( "Software\\Microsoft\\Windows\\CurrentVersion\\Policies\\System", KEY_WRITE, ) { if original == 0 { // It didn't exist before, delete it if let Err(e) = policy_key.delete_value("SoftwareSASGeneration") { log::error!("Failed to delete SoftwareSASGeneration: {}", e); } else { log::info!("Deleted SoftwareSASGeneration (restored to original state)"); } } else { // Restore the original value if let Err(e) = policy_key.set_value("SoftwareSASGeneration", &original) { log::error!( "Failed to restore SoftwareSASGeneration to {}: {}", original, e ); } else { log::info!("Restored SoftwareSASGeneration to {}", original); } } } } } } lazy_static::lazy_static! { static ref SUPPRESS: Arc> = Arc::new(Mutex::new(Instant::now())); } pub fn desktop_changed() -> bool { unsafe { inputDesktopSelected() == FALSE } } pub fn try_change_desktop() -> bool { unsafe { if inputDesktopSelected() == FALSE { let res = selectInputDesktop() == TRUE; if !res { let mut s = SUPPRESS.lock().unwrap(); if s.elapsed() > std::time::Duration::from_secs(3) { log::error!("Failed to switch desktop: {}", io::Error::last_os_error()); *s = Instant::now(); } } else { log::info!("Desktop switched"); } return res; } } return false; } fn share_rdp() -> BOOL { if get_reg("share_rdp") != "false" { TRUE } else { FALSE } } pub fn is_share_rdp() -> bool { share_rdp() == TRUE } pub fn set_share_rdp(enable: bool) { let (subkey, _, _, _) = get_install_info(); let cmd = format!( "reg add {} /f /v share_rdp /t REG_SZ /d \"{}\"", subkey, if enable { "true" } else { "false" } ); run_cmds(cmd, false, "share_rdp").ok(); } pub fn get_current_process_session_id() -> Option { get_session_id_of_process(unsafe { GetCurrentProcessId() }) } pub fn get_session_id_of_process(pid: DWORD) -> Option { let mut sid = 0; if unsafe { ProcessIdToSessionId(pid, &mut sid) == TRUE } { Some(sid) } else { None } } pub fn is_physical_console_session() -> Option { if let Some(sid) = get_current_process_session_id() { let physical_console_session_id = unsafe { get_current_session(FALSE) }; if physical_console_session_id == u32::MAX { return None; } return Some(physical_console_session_id == sid); } None } pub fn get_active_username() -> String { // get_active_user will give console username higher priority if let Some(name) = get_current_session_username() { return name; } if !is_root() { return crate::username(); } extern "C" { fn get_active_user(path: *mut u16, n: u32, rdp: BOOL) -> u32; } let buff_size = 256; let mut buff: Vec = Vec::with_capacity(buff_size); buff.resize(buff_size, 0); let n = unsafe { get_active_user(buff.as_mut_ptr(), buff_size as _, share_rdp()) }; if n == 0 { return "".to_owned(); } let sl = unsafe { std::slice::from_raw_parts(buff.as_ptr(), n as _) }; String::from_utf16(sl) .unwrap_or("??".to_owned()) .trim_end_matches('\0') .to_owned() } fn get_current_session_username() -> Option { let Some(sid) = get_current_process_session_id() else { log::error!("get_current_process_session_id failed"); return None; }; Some(get_session_username(sid)) } fn get_session_username(session_id: u32) -> String { extern "C" { fn get_session_user_info(path: *mut u16, n: u32, session_id: u32) -> u32; } let buff_size = 256; let mut buff: Vec = Vec::with_capacity(buff_size); buff.resize(buff_size, 0); let n = unsafe { get_session_user_info(buff.as_mut_ptr(), buff_size as _, session_id) }; if n == 0 { return "".to_owned(); } let sl = unsafe { std::slice::from_raw_parts(buff.as_ptr(), n as _) }; String::from_utf16(sl) .unwrap_or("".to_owned()) .trim_end_matches('\0') .to_owned() } pub fn get_available_sessions(name: bool) -> Vec { extern "C" { fn get_available_session_ids(buf: *mut wchar_t, buf_size: c_int, include_rdp: bool); } const BUF_SIZE: c_int = 1024; let mut buf: Vec = vec![0; BUF_SIZE as usize]; let station_session_id_array = unsafe { get_available_session_ids(buf.as_mut_ptr(), BUF_SIZE, true); let session_ids = String::from_utf16_lossy(&buf); session_ids.trim_matches(char::from(0)).trim().to_string() }; let mut v: Vec = vec![]; // https://learn.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-wtsgetactiveconsolesessionid let physical_console_sid = unsafe { get_current_session(FALSE) }; if physical_console_sid != u32::MAX { let physical_console_name = if name { let physical_console_username = get_session_username(physical_console_sid); if physical_console_username.is_empty() { "Console".to_owned() } else { format!("Console: {physical_console_username}") } } else { "".to_owned() }; v.push(WindowsSession { sid: physical_console_sid, name: physical_console_name, ..Default::default() }); } // https://learn.microsoft.com/en-us/previous-versions//cc722458(v=technet.10)?redirectedfrom=MSDN for type_session_id in station_session_id_array.split(",") { let split: Vec<_> = type_session_id.split(":").collect(); if split.len() == 2 { if let Ok(sid) = split[1].parse::() { if !v.iter().any(|e| (*e).sid == sid) { let name = if name { let name = get_session_username(sid); if name.is_empty() { split[0].to_string() } else { format!("{}: {}", split[0], name) } } else { "".to_owned() }; v.push(WindowsSession { sid, name, ..Default::default() }); } } } } if name { let mut name_count: HashMap = HashMap::new(); for session in &v { *name_count.entry(session.name.clone()).or_insert(0) += 1; } let current_sid = get_current_process_session_id().unwrap_or_default(); for e in v.iter_mut() { let running = e.sid == current_sid && current_sid != 0; if name_count.get(&e.name).map(|v| *v).unwrap_or_default() > 1 { e.name = format!("{} (sid = {})", e.name, e.sid); } if running { e.name = format!("{} (running)", e.name); } } } v } pub fn get_active_user_home() -> Option { let username = get_active_username(); if !username.is_empty() { let drive = std::env::var("SystemDrive").unwrap_or("C:".to_owned()); let home = PathBuf::from(format!("{}\\Users\\{}", drive, username)); if home.exists() { return Some(home); } } None } pub fn is_prelogin() -> bool { let Some(username) = get_current_session_username() else { return false; }; username.is_empty() || username == "SYSTEM" } // `is_logon_ui()` is regardless of multiple sessions now. // It only check if "LogonUI.exe" exists. // // If there're mulitple sessions (logged in users), // some are in the login screen, while the others are not. // Then this function may not work fine if the session we want to handle(connect) is not in the login screen. // But it's a rare case and cannot be simply handled, so it will not be dealt with for the time being. #[inline] pub fn is_logon_ui() -> ResultType { let pids = get_pids("LogonUI.exe")?; Ok(!pids.is_empty()) } pub fn is_root() -> bool { // https://stackoverflow.com/questions/4023586/correct-way-to-find-out-if-a-service-is-running-as-the-system-user unsafe { is_local_system() == TRUE } } pub fn lock_screen() { extern "system" { pub fn LockWorkStation() -> BOOL; } unsafe { LockWorkStation(); } } const IS1: &str = "{54E86BC2-6C85-41F3-A9EB-1A94AC9B1F93}_is1"; fn get_subkey(name: &str, wow: bool) -> String { let tmp = format!( "HKEY_LOCAL_MACHINE\\Software\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{}", name ); if wow { tmp.replace("Microsoft", "Wow6432Node\\Microsoft") } else { tmp } } fn get_valid_subkey() -> String { let subkey = get_subkey(IS1, false); if !get_reg_of(&subkey, "InstallLocation").is_empty() { return subkey; } let subkey = get_subkey(IS1, true); if !get_reg_of(&subkey, "InstallLocation").is_empty() { return subkey; } let app_name = crate::get_app_name(); let subkey = get_subkey(&app_name, true); if !get_reg_of(&subkey, "InstallLocation").is_empty() { return subkey; } return get_subkey(&app_name, false); } // Return install options other than InstallLocation. pub fn get_install_options() -> String { let app_name = crate::get_app_name(); let subkey = format!(".{}", app_name.to_lowercase()); let mut opts = HashMap::new(); let desktop_shortcuts = get_reg_of_hkcr(&subkey, REG_NAME_INSTALL_DESKTOPSHORTCUTS); if let Some(desktop_shortcuts) = desktop_shortcuts { opts.insert(REG_NAME_INSTALL_DESKTOPSHORTCUTS, desktop_shortcuts); } let start_menu_shortcuts = get_reg_of_hkcr(&subkey, REG_NAME_INSTALL_STARTMENUSHORTCUTS); if let Some(start_menu_shortcuts) = start_menu_shortcuts { opts.insert(REG_NAME_INSTALL_STARTMENUSHORTCUTS, start_menu_shortcuts); } let printer = get_reg_of_hkcr(&subkey, REG_NAME_INSTALL_PRINTER); if let Some(printer) = printer { opts.insert(REG_NAME_INSTALL_PRINTER, printer); } serde_json::to_string(&opts).unwrap_or("{}".to_owned()) } // This function return Option, because some registry value may be empty. fn get_reg_of_hkcr(subkey: &str, name: &str) -> Option { let hkcr = RegKey::predef(HKEY_CLASSES_ROOT); if let Ok(tmp) = hkcr.open_subkey(subkey.replace("HKEY_CLASSES_ROOT\\", "")) { return tmp.get_value(name).ok(); } None } pub fn get_install_info() -> (String, String, String, String) { get_install_info_with_subkey(get_valid_subkey()) } fn get_default_install_info() -> (String, String, String, String) { get_install_info_with_subkey(get_subkey(&crate::get_app_name(), false)) } fn get_default_install_path() -> String { let mut pf = "C:\\Program Files".to_owned(); if let Ok(x) = std::env::var("ProgramFiles") { if std::path::Path::new(&x).exists() { pf = x; } } #[cfg(target_pointer_width = "32")] { let tmp = pf.replace("Program Files", "Program Files (x86)"); if std::path::Path::new(&tmp).exists() { pf = tmp; } } format!("{}\\{}", pf, crate::get_app_name()) } pub fn check_update_broker_process() -> ResultType<()> { let process_exe = win_topmost_window::INJECTED_PROCESS_EXE; let origin_process_exe = win_topmost_window::ORIGIN_PROCESS_EXE; let exe_file = std::env::current_exe()?; let Some(cur_dir) = exe_file.parent() else { bail!("Cannot get parent of current exe file"); }; let cur_exe = cur_dir.join(process_exe); // Force update broker exe if failed to check modified time. let cmds = format!( " chcp 65001 taskkill /F /IM {process_exe} copy /Y \"{origin_process_exe}\" \"{cur_exe}\" ", cur_exe = cur_exe.to_string_lossy(), ); if !std::path::Path::new(&cur_exe).exists() { run_cmds(cmds, false, "update_broker")?; return Ok(()); } let ori_modified = fs::metadata(origin_process_exe)?.modified()?; if let Ok(metadata) = fs::metadata(&cur_exe) { if let Ok(cur_modified) = metadata.modified() { if cur_modified == ori_modified { return Ok(()); } else { log::info!( "broker process updated, modify time from {:?} to {:?}", cur_modified, ori_modified ); } } } run_cmds(cmds, false, "update_broker")?; Ok(()) } fn get_install_info_with_subkey(subkey: String) -> (String, String, String, String) { let mut path = get_reg_of(&subkey, "InstallLocation"); if path.is_empty() { path = get_default_install_path(); } path = path.trim_end_matches('\\').to_owned(); let start_menu = format!( "%ProgramData%\\Microsoft\\Windows\\Start Menu\\Programs\\{}", crate::get_app_name() ); let exe = format!("{}\\{}.exe", path, crate::get_app_name()); (subkey, path, start_menu, exe) } pub fn copy_raw_cmd(src_raw: &str, _raw: &str, _path: &str) -> ResultType { let main_raw = format!( "XCOPY \"{}\" \"{}\" /Y /E /H /C /I /K /R /Z", PathBuf::from(src_raw) .parent() .ok_or(anyhow!("Can't get parent directory of {src_raw}"))? .to_string_lossy() .to_string(), _path ); return Ok(main_raw); } pub fn copy_exe_cmd(src_exe: &str, exe: &str, path: &str) -> ResultType { let main_exe = copy_raw_cmd(src_exe, exe, path)?; Ok(format!( " {main_exe} copy /Y \"{ORIGIN_PROCESS_EXE}\" \"{path}\\{broker_exe}\" ", ORIGIN_PROCESS_EXE = win_topmost_window::ORIGIN_PROCESS_EXE, broker_exe = win_topmost_window::INJECTED_PROCESS_EXE, )) } fn get_after_install( exe: &str, reg_value_start_menu_shortcuts: Option, reg_value_desktop_shortcuts: Option, reg_value_printer: Option, ) -> String { let app_name = crate::get_app_name(); let ext = app_name.to_lowercase(); // reg delete HKEY_CURRENT_USER\Software\Classes for // https://github.com/rustdesk/rustdesk/commit/f4bdfb6936ae4804fc8ab1cf560db192622ad01a // and https://github.com/leanflutter/uni_links_desktop/blob/1b72b0226cec9943ca8a84e244c149773f384e46/lib/src/protocol_registrar_impl_windows.dart#L30 let hcu = RegKey::predef(HKEY_CURRENT_USER); hcu.delete_subkey_all(format!("Software\\Classes\\{}", exe)) .ok(); let desktop_shortcuts = reg_value_desktop_shortcuts .map(|v| { format!("reg add HKEY_CLASSES_ROOT\\.{ext} /f /v {REG_NAME_INSTALL_DESKTOPSHORTCUTS} /t REG_SZ /d \"{v}\"") }) .unwrap_or_default(); let start_menu_shortcuts = reg_value_start_menu_shortcuts .map(|v| { format!( "reg add HKEY_CLASSES_ROOT\\.{ext} /f /v {REG_NAME_INSTALL_STARTMENUSHORTCUTS} /t REG_SZ /d \"{v}\"" ) }) .unwrap_or_default(); let reg_printer = reg_value_printer .map(|v| { format!( "reg add HKEY_CLASSES_ROOT\\.{ext} /f /v {REG_NAME_INSTALL_PRINTER} /t REG_SZ /d \"{v}\"" ) }) .unwrap_or_default(); format!(" chcp 65001 reg add HKEY_CLASSES_ROOT\\.{ext} /f {desktop_shortcuts} {start_menu_shortcuts} {reg_printer} reg add HKEY_CLASSES_ROOT\\.{ext}\\DefaultIcon /f reg add HKEY_CLASSES_ROOT\\.{ext}\\DefaultIcon /f /ve /t REG_SZ /d \"\\\"{exe}\\\",0\" reg add HKEY_CLASSES_ROOT\\.{ext}\\shell /f reg add HKEY_CLASSES_ROOT\\.{ext}\\shell\\open /f reg add HKEY_CLASSES_ROOT\\.{ext}\\shell\\open\\command /f reg add HKEY_CLASSES_ROOT\\.{ext}\\shell\\open\\command /f /ve /t REG_SZ /d \"\\\"{exe}\\\" --play \\\"%%1\\\"\" reg add HKEY_CLASSES_ROOT\\{ext} /f reg add HKEY_CLASSES_ROOT\\{ext} /f /v \"URL Protocol\" /t REG_SZ /d \"\" reg add HKEY_CLASSES_ROOT\\{ext}\\shell /f reg add HKEY_CLASSES_ROOT\\{ext}\\shell\\open /f reg add HKEY_CLASSES_ROOT\\{ext}\\shell\\open\\command /f reg add HKEY_CLASSES_ROOT\\{ext}\\shell\\open\\command /f /ve /t REG_SZ /d \"\\\"{exe}\\\" \\\"%%1\\\"\" netsh advfirewall firewall add rule name=\"{app_name} Service\" dir=out action=allow program=\"{exe}\" enable=yes netsh advfirewall firewall add rule name=\"{app_name} Service\" dir=in action=allow program=\"{exe}\" enable=yes {create_service} reg add HKEY_LOCAL_MACHINE\\Software\\Microsoft\\Windows\\CurrentVersion\\Policies\\System /f /v SoftwareSASGeneration /t REG_DWORD /d 1 ", create_service=get_create_service(&exe)) } pub fn install_me(options: &str, path: String, silent: bool, debug: bool) -> ResultType<()> { let uninstall_str = get_uninstall(false, false); let mut path = path.trim_end_matches('\\').to_owned(); let (subkey, _path, start_menu, exe) = get_default_install_info(); let mut exe = exe; if path.is_empty() { path = _path; } else { exe = exe.replace(&_path, &path); } let mut version_major = "0"; let mut version_minor = "0"; let mut version_build = "0"; let versions: Vec<&str> = crate::VERSION.split(".").collect(); if versions.len() > 0 { version_major = versions[0]; } if versions.len() > 1 { version_minor = versions[1]; } if versions.len() > 2 { version_build = versions[2]; } let app_name = crate::get_app_name(); let tmp_path = std::env::temp_dir().to_string_lossy().to_string(); let mk_shortcut = write_cmds( format!( " Set oWS = WScript.CreateObject(\"WScript.Shell\") sLinkFile = \"{tmp_path}\\{app_name}.lnk\" Set oLink = oWS.CreateShortcut(sLinkFile) oLink.TargetPath = \"{exe}\" oLink.Save " ), "vbs", "mk_shortcut", )? .to_str() .unwrap_or("") .to_owned(); // https://superuser.com/questions/392061/how-to-make-a-shortcut-from-cmd let uninstall_shortcut = write_cmds( format!( " Set oWS = WScript.CreateObject(\"WScript.Shell\") sLinkFile = \"{tmp_path}\\Uninstall {app_name}.lnk\" Set oLink = oWS.CreateShortcut(sLinkFile) oLink.TargetPath = \"{exe}\" oLink.Arguments = \"--uninstall\" oLink.IconLocation = \"msiexec.exe\" oLink.Save " ), "vbs", "uninstall_shortcut", )? .to_str() .unwrap_or("") .to_owned(); let tray_shortcut = get_tray_shortcut(&exe, &tmp_path)?; let mut reg_value_desktop_shortcuts = "0".to_owned(); let mut reg_value_start_menu_shortcuts = "0".to_owned(); let mut reg_value_printer = "0".to_owned(); let mut shortcuts = Default::default(); if options.contains("desktopicon") { shortcuts = format!( "copy /Y \"{}\\{}.lnk\" \"%PUBLIC%\\Desktop\\\"", tmp_path, crate::get_app_name() ); reg_value_desktop_shortcuts = "1".to_owned(); } if options.contains("startmenu") { shortcuts = format!( "{shortcuts} md \"{start_menu}\" copy /Y \"{tmp_path}\\{app_name}.lnk\" \"{start_menu}\\\" copy /Y \"{tmp_path}\\Uninstall {app_name}.lnk\" \"{start_menu}\\\" " ); reg_value_start_menu_shortcuts = "1".to_owned(); } let install_printer = options.contains("printer") && is_win_10_or_greater(); if install_printer { reg_value_printer = "1".to_owned(); } let meta = std::fs::symlink_metadata(std::env::current_exe()?)?; let size = meta.len() / 1024; // https://docs.microsoft.com/zh-cn/windows/win32/msi/uninstall-registry-key?redirectedfrom=MSDNa // https://www.windowscentral.com/how-edit-registry-using-command-prompt-windows-10 // https://www.tenforums.com/tutorials/70903-add-remove-allowed-apps-through-windows-firewall-windows-10-a.html // Note: without if exist, the bat may exit in advance on some Windows7 https://github.com/rustdesk/rustdesk/issues/895 let dels = format!( " if exist \"{mk_shortcut}\" del /f /q \"{mk_shortcut}\" if exist \"{uninstall_shortcut}\" del /f /q \"{uninstall_shortcut}\" if exist \"{tray_shortcut}\" del /f /q \"{tray_shortcut}\" if exist \"{tmp_path}\\{app_name}.lnk\" del /f /q \"{tmp_path}\\{app_name}.lnk\" if exist \"{tmp_path}\\Uninstall {app_name}.lnk\" del /f /q \"{tmp_path}\\Uninstall {app_name}.lnk\" if exist \"{tmp_path}\\{app_name} Tray.lnk\" del /f /q \"{tmp_path}\\{app_name} Tray.lnk\" " ); let src_exe = std::env::current_exe()?.to_str().unwrap_or("").to_string(); // potential bug here: if run_cmd cancelled, but config file is changed. if let Some(lic) = get_license() { Config::set_option("key".into(), lic.key); Config::set_option("custom-rendezvous-server".into(), lic.host); Config::set_option("api-server".into(), lic.api); } let tray_shortcuts = if config::is_outgoing_only() { "".to_owned() } else { format!(" cscript \"{tray_shortcut}\" copy /Y \"{tmp_path}\\{app_name} Tray.lnk\" \"%PROGRAMDATA%\\Microsoft\\Windows\\Start Menu\\Programs\\Startup\\\" ") }; let install_remote_printer = if install_printer { // No need to use `|| true` here. // The script will not exit even if `--install-remote-printer` panics. format!("\"{}\" --install-remote-printer", &src_exe) } else if is_win_10_or_greater() { format!("\"{}\" --uninstall-remote-printer", &src_exe) } else { "".to_owned() }; // Remember to check if `update_me` need to be changed if changing the `cmds`. // No need to merge the existing dup code, because the code in these two functions are too critical. // New code should be written in a common function. let cmds = format!( " {uninstall_str} chcp 65001 md \"{path}\" {copy_exe} reg add {subkey} /f reg add {subkey} /f /v DisplayIcon /t REG_SZ /d \"{exe}\" reg add {subkey} /f /v DisplayName /t REG_SZ /d \"{app_name}\" reg add {subkey} /f /v DisplayVersion /t REG_SZ /d \"{version}\" reg add {subkey} /f /v Version /t REG_SZ /d \"{version}\" reg add {subkey} /f /v BuildDate /t REG_SZ /d \"{build_date}\" reg add {subkey} /f /v InstallLocation /t REG_SZ /d \"{path}\" reg add {subkey} /f /v Publisher /t REG_SZ /d \"{app_name}\" reg add {subkey} /f /v VersionMajor /t REG_DWORD /d {version_major} reg add {subkey} /f /v VersionMinor /t REG_DWORD /d {version_minor} reg add {subkey} /f /v VersionBuild /t REG_DWORD /d {version_build} reg add {subkey} /f /v UninstallString /t REG_SZ /d \"\\\"{exe}\\\" --uninstall\" reg add {subkey} /f /v EstimatedSize /t REG_DWORD /d {size} reg add {subkey} /f /v WindowsInstaller /t REG_DWORD /d 0 cscript \"{mk_shortcut}\" cscript \"{uninstall_shortcut}\" {tray_shortcuts} {shortcuts} copy /Y \"{tmp_path}\\Uninstall {app_name}.lnk\" \"{path}\\\" {dels} {import_config} {after_install} {install_remote_printer} {sleep} ", version = crate::VERSION.replace("-", "."), build_date = crate::BUILD_DATE, after_install = get_after_install( &exe, Some(reg_value_start_menu_shortcuts), Some(reg_value_desktop_shortcuts), Some(reg_value_printer) ), sleep = if debug { "timeout 300" } else { "" }, dels = if debug { "" } else { &dels }, copy_exe = copy_exe_cmd(&src_exe, &exe, &path)?, import_config = get_import_config(&exe), ); run_cmds(cmds, debug, "install")?; run_after_run_cmds(silent); Ok(()) } pub fn run_after_install() -> ResultType<()> { let (_, _, _, exe) = get_install_info(); run_cmds( get_after_install(&exe, None, None, None), true, "after_install", ) } pub fn run_before_uninstall() -> ResultType<()> { run_cmds(get_before_uninstall(true), true, "before_install") } fn get_before_uninstall(kill_self: bool) -> String { let app_name = crate::get_app_name(); let ext = app_name.to_lowercase(); let filter = if kill_self { "".to_string() } else { format!(" /FI \"PID ne {}\"", get_current_pid()) }; format!( " chcp 65001 sc stop {app_name} sc delete {app_name} taskkill /F /IM {broker_exe} taskkill /F /IM {app_name}.exe{filter} reg delete HKEY_CLASSES_ROOT\\.{ext} /f reg delete HKEY_CLASSES_ROOT\\{ext} /f netsh advfirewall firewall delete rule name=\"{app_name} Service\" ", broker_exe = WIN_TOPMOST_INJECTED_PROCESS_EXE, ) } /// Constructs the uninstall command string for the application. /// /// # Parameters /// - `kill_self`: The command will kill the process of current app name. If `true`, it will kill /// the current process as well. If `false`, it will exclude the current process from the kill /// command. /// - `uninstall_printer`: If `true`, includes commands to uninstall the remote printer. /// /// # Details /// The `uninstall_printer` parameter determines whether the command to uninstall the remote printer /// is included in the generated uninstall script. If `uninstall_printer` is `false`, the printer /// related command is omitted from the script. fn get_uninstall(kill_self: bool, uninstall_printer: bool) -> String { let reg_uninstall_string = get_reg("UninstallString"); if reg_uninstall_string.to_lowercase().contains("msiexec.exe") { return reg_uninstall_string; } let mut uninstall_cert_cmd = "".to_string(); let mut uninstall_printer_cmd = "".to_string(); if let Ok(exe) = std::env::current_exe() { if let Some(exe_path) = exe.to_str() { uninstall_cert_cmd = format!("\"{}\" --uninstall-cert", exe_path); if uninstall_printer { uninstall_printer_cmd = format!("\"{}\" --uninstall-remote-printer", &exe_path); } } } let (subkey, path, start_menu, _) = get_install_info(); format!( " {before_uninstall} {uninstall_printer_cmd} {uninstall_cert_cmd} reg delete {subkey} /f {uninstall_amyuni_idd} if exist \"{path}\" rd /s /q \"{path}\" if exist \"{start_menu}\" rd /s /q \"{start_menu}\" if exist \"%PUBLIC%\\Desktop\\{app_name}.lnk\" del /f /q \"%PUBLIC%\\Desktop\\{app_name}.lnk\" 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\" ", before_uninstall=get_before_uninstall(kill_self), uninstall_amyuni_idd=get_uninstall_amyuni_idd(), app_name = crate::get_app_name(), ) } pub fn uninstall_me(kill_self: bool) -> ResultType<()> { run_cmds(get_uninstall(kill_self, true), true, "uninstall") } fn write_cmds(cmds: String, ext: &str, tip: &str) -> ResultType { let mut cmds = cmds; let mut tmp = std::env::temp_dir(); // When dir contains these characters, the bat file will not execute in elevated mode. if vec!["&", "@", "^"] .drain(..) .any(|s| tmp.to_string_lossy().to_string().contains(s)) { if let Ok(dir) = user_accessible_folder() { tmp = dir; } } tmp.push(format!("{}_{}.{}", crate::get_app_name(), tip, ext)); let mut file = std::fs::File::create(&tmp)?; if ext == "bat" { let tmp2 = get_undone_file(&tmp)?; std::fs::File::create(&tmp2).ok(); cmds = format!( " {cmds} if exist \"{path}\" del /f /q \"{path}\" ", path = tmp2.to_string_lossy() ); } // in case cmds mixed with \r\n and \n, make sure all ending with \r\n // in some windows, \r\n required for cmd file to run cmds = cmds.replace("\r\n", "\n").replace("\n", "\r\n"); if ext == "vbs" { let mut v: Vec = cmds.encode_utf16().collect(); // utf8 -> utf16le which vbs support it only file.write_all(to_le(&mut v))?; } else { file.write_all(cmds.as_bytes())?; } file.sync_all()?; return Ok(tmp); } fn to_le(v: &mut [u16]) -> &[u8] { for b in v.iter_mut() { *b = b.to_le() } unsafe { v.align_to().1 } } fn get_undone_file(tmp: &Path) -> ResultType { Ok(tmp.with_file_name(format!( "{}.undone", tmp.file_name() .ok_or(anyhow!("Failed to get filename of {:?}", tmp))? .to_string_lossy() ))) } fn run_cmds(cmds: String, show: bool, tip: &str) -> ResultType<()> { let tmp = write_cmds(cmds, "bat", tip)?; let tmp2 = get_undone_file(&tmp)?; let tmp_fn = tmp.to_str().unwrap_or(""); // https://github.com/rustdesk/rustdesk/issues/6786#issuecomment-1879655410 // Specify cmd.exe explicitly to avoid the replacement of cmd commands. let res = runas::Command::new("cmd.exe") .args(&["/C", &tmp_fn]) .show(show) .force_prompt(true) .status(); if !show { allow_err!(std::fs::remove_file(tmp)); } let _ = res?; if tmp2.exists() { allow_err!(std::fs::remove_file(tmp2)); bail!("{} failed", tip); } Ok(()) } pub fn toggle_blank_screen(v: bool) { let v = if v { TRUE } else { FALSE }; unsafe { blank_screen(v); } } pub fn block_input(v: bool) -> (bool, String) { let v = if v { TRUE } else { FALSE }; unsafe { if BlockInput(v) == TRUE { (true, "".to_owned()) } else { (false, format!("Error: {}", io::Error::last_os_error())) } } } pub fn add_recent_document(path: &str) { extern "C" { fn AddRecentDocument(path: *const u16); } use std::os::windows::ffi::OsStrExt; let wstr: Vec = std::ffi::OsStr::new(path) .encode_wide() .chain(Some(0).into_iter()) .collect(); let wstr = wstr.as_ptr(); unsafe { AddRecentDocument(wstr); } } pub fn is_installed() -> bool { let (_, _, _, exe) = get_install_info(); std::fs::metadata(exe).is_ok() } pub fn get_reg(name: &str) -> String { let (subkey, _, _, _) = get_install_info(); get_reg_of(&subkey, name) } fn get_reg_of(subkey: &str, name: &str) -> String { let hklm = RegKey::predef(HKEY_LOCAL_MACHINE); if let Ok(tmp) = hklm.open_subkey(subkey.replace("HKEY_LOCAL_MACHINE\\", "")) { if let Ok(v) = tmp.get_value(name) { return v; } } "".to_owned() } pub fn get_license_from_exe_name() -> ResultType { let mut exe = std::env::current_exe()?.to_str().unwrap_or("").to_owned(); // if defined portable appname entry, replace original executable name with it. if let Ok(portable_exe) = std::env::var(PORTABLE_APPNAME_RUNTIME_ENV_KEY) { exe = portable_exe; } get_custom_server_from_string(&exe) } // We can't directly use `RegKey::set_value` to update the registry value, because it will fail with `ERROR_ACCESS_DENIED` // So we have to use `run_cmds` to update the registry value. pub fn update_install_option(k: &str, v: &str) -> ResultType<()> { // Don't update registry if not installed or not server process. if !is_installed() || !crate::is_server() { return Ok(()); } let app_name = crate::get_app_name(); let ext = app_name.to_lowercase(); let cmds = format!("chcp 65001 && reg add HKEY_CLASSES_ROOT\\.{ext} /f /v {k} /t REG_SZ /d \"{v}\""); run_cmds(cmds, false, "update_install_option")?; Ok(()) } #[inline] pub fn is_win_server() -> bool { unsafe { is_windows_server() > 0 } } #[inline] pub fn is_win_10_or_greater() -> bool { unsafe { is_windows_10_or_greater() > 0 } } pub fn bootstrap() -> bool { if let Ok(lic) = get_license_from_exe_name() { *config::EXE_RENDEZVOUS_SERVER.write().unwrap() = lic.host.clone(); } #[cfg(debug_assertions)] { true } #[cfg(not(debug_assertions))] { // This function will cause `'sciter.dll' was not found neither in PATH nor near the current executable.` when debugging RustDesk. // Only call set_safe_load_dll() on Windows 10 or greater if is_win_10_or_greater() { set_safe_load_dll() } else { true } } } #[cfg(not(debug_assertions))] fn set_safe_load_dll() -> bool { if !unsafe { set_default_dll_directories() } { return false; } // `SetDllDirectoryW` should never fail. // https://docs.microsoft.com/en-us/windows/win32/api/winbase/nf-winbase-setdlldirectoryw if unsafe { SetDllDirectoryW(wide_string("").as_ptr()) == FALSE } { eprintln!("SetDllDirectoryW failed: {}", io::Error::last_os_error()); return false; } true } // https://docs.microsoft.com/en-us/windows/win32/api/libloaderapi/nf-libloaderapi-setdefaultdlldirectories #[cfg(not(debug_assertions))] unsafe fn set_default_dll_directories() -> bool { let module = LoadLibraryExW( wide_string("Kernel32.dll").as_ptr(), 0 as _, LOAD_LIBRARY_SEARCH_SYSTEM32, ); if module.is_null() { return false; } match CString::new("SetDefaultDllDirectories") { Err(e) => { eprintln!("CString::new failed: {}", e); return false; } Ok(func_name) => { let func = GetProcAddress(module, func_name.as_ptr()); if func.is_null() { eprintln!("GetProcAddress failed: {}", io::Error::last_os_error()); return false; } type SetDefaultDllDirectories = unsafe extern "system" fn(DWORD) -> BOOL; let func: SetDefaultDllDirectories = std::mem::transmute(func); if func(LOAD_LIBRARY_SEARCH_SYSTEM32 | LOAD_LIBRARY_SEARCH_USER_DIRS) == FALSE { eprintln!( "SetDefaultDllDirectories failed: {}", io::Error::last_os_error() ); return false; } } } true } pub fn create_shortcut(id: &str) -> ResultType<()> { let exe = std::env::current_exe()?.to_str().unwrap_or("").to_owned(); // https://github.com/rustdesk/rustdesk/issues/13735 // Replace ':' with '_' for filename since ':' is not allowed in Windows filenames // https://github.com/rustdesk/hbb_common/blob/8b0e25867375ba9e6bff548acf44fe6d6ffa7c0e/src/config.rs#L1384 let filename = id.replace(':', "_"); let shortcut = write_cmds( format!( " Set oWS = WScript.CreateObject(\"WScript.Shell\") strDesktop = oWS.SpecialFolders(\"Desktop\") Set objFSO = CreateObject(\"Scripting.FileSystemObject\") sLinkFile = objFSO.BuildPath(strDesktop, \"{filename}.lnk\") Set oLink = oWS.CreateShortcut(sLinkFile) oLink.TargetPath = \"{exe}\" oLink.Arguments = \"--connect {id}\" oLink.Save " ), "vbs", "connect_shortcut", )? .to_str() .unwrap_or("") .to_owned(); std::process::Command::new("cscript") .arg(&shortcut) .creation_flags(CREATE_NO_WINDOW) .output()?; allow_err!(std::fs::remove_file(shortcut)); Ok(()) } pub fn enable_lowlevel_keyboard(hwnd: HWND) { let ret = unsafe { win32_enable_lowlevel_keyboard(hwnd) }; if ret != 0 { log::error!("Failure grabbing keyboard"); return; } } pub fn disable_lowlevel_keyboard(hwnd: HWND) { unsafe { win32_disable_lowlevel_keyboard(hwnd) }; } pub fn stop_system_key_propagate(v: bool) { unsafe { win_stop_system_key_propagate(if v { TRUE } else { FALSE }) }; } pub fn get_win_key_state() -> bool { unsafe { is_win_down() == TRUE } } pub fn quit_gui() { std::process::exit(0); // unsafe { PostQuitMessage(0) }; // some how not work } pub fn get_user_token(session_id: u32, as_user: bool) -> HANDLE { let mut token = NULL as HANDLE; unsafe { let mut _token_pid = 0; if FALSE == GetSessionUserTokenWin( &mut token as _, session_id, if as_user { TRUE } else { FALSE }, &mut _token_pid, ) { NULL as _ } else { token } } } pub fn run_background(exe: &str, arg: &str) -> ResultType { let wexe = wide_string(exe); let warg; unsafe { let ret = ShellExecuteW( NULL as _, NULL as _, wexe.as_ptr() as _, if arg.is_empty() { NULL as _ } else { warg = wide_string(arg); warg.as_ptr() as _ }, NULL as _, SW_HIDE, ); return Ok(ret as i32 > 32); } } pub fn run_uac(exe: &str, arg: &str) -> ResultType { let wop = wide_string("runas"); let wexe = wide_string(exe); let warg; unsafe { let ret = ShellExecuteW( NULL as _, wop.as_ptr() as _, wexe.as_ptr() as _, if arg.is_empty() { NULL as _ } else { warg = wide_string(arg); warg.as_ptr() as _ }, NULL as _, SW_SHOWNORMAL, ); return Ok(ret as i32 > 32); } } pub fn check_super_user_permission() -> ResultType { run_uac( std::env::current_exe()? .to_string_lossy() .to_string() .as_str(), "--version", ) } pub fn elevate(arg: &str) -> ResultType { run_uac( std::env::current_exe()? .to_string_lossy() .to_string() .as_str(), arg, ) } pub fn run_as_system(arg: &str) -> ResultType<()> { let exe = std::env::current_exe()?.to_string_lossy().to_string(); if impersonate_system::run_as_system(&exe, arg).is_err() { bail!(format!("Failed to run {} as system", exe)); } Ok(()) } pub fn elevate_or_run_as_system(is_setup: bool, is_elevate: bool, is_run_as_system: bool) { // avoid possible run recursively due to failed run. log::info!( "elevate: {} -> {:?}, run_as_system: {} -> {}", is_elevate, is_elevated(None), is_run_as_system, crate::username(), ); let arg_elevate = if is_setup { "--noinstall --elevate" } else { "--elevate" }; let arg_run_as_system = if is_setup { "--noinstall --run-as-system" } else { "--run-as-system" }; if is_root() { if is_run_as_system { log::info!("run portable service"); crate::portable_service::server::run_portable_service(); } } else { match is_elevated(None) { Ok(elevated) => { if elevated { if !is_run_as_system { if run_as_system(arg_run_as_system).is_ok() { std::process::exit(0); } else { log::error!( "Failed to run as system, error {}", io::Error::last_os_error() ); } } } else { if !is_elevate { if let Ok(true) = elevate(arg_elevate) { std::process::exit(0); } else { log::error!("Failed to elevate, error {}", io::Error::last_os_error()); } } } } Err(_) => log::error!( "Failed to get elevation status, error {}", io::Error::last_os_error() ), } } } pub fn is_elevated(process_id: Option) -> ResultType { use hbb_common::platform::windows::RAIIHandle; unsafe { let handle: HANDLE = match process_id { Some(process_id) => OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, process_id), None => GetCurrentProcess(), }; if handle == NULL { bail!( "Failed to open process, error {}", io::Error::last_os_error() ) } let _handle = RAIIHandle(handle); let mut token: HANDLE = mem::zeroed(); if OpenProcessToken(handle, TOKEN_QUERY, &mut token) == FALSE { bail!( "Failed to open process token, error {}", io::Error::last_os_error() ) } let _token = RAIIHandle(token); let mut token_elevation: TOKEN_ELEVATION = mem::zeroed(); let mut size: DWORD = 0; if GetTokenInformation( token, TokenElevation, (&mut token_elevation) as *mut _ as *mut c_void, mem::size_of::() as _, &mut size, ) == FALSE { bail!( "Failed to get token information, error {}", io::Error::last_os_error() ) } Ok(token_elevation.TokenIsElevated != 0) } } pub fn is_foreground_window_elevated() -> ResultType { unsafe { let mut process_id: DWORD = 0; GetWindowThreadProcessId(GetForegroundWindow(), &mut process_id); if process_id == 0 { bail!( "Failed to get processId, error {}", io::Error::last_os_error() ) } is_elevated(Some(process_id)) } } fn get_current_pid() -> u32 { unsafe { GetCurrentProcessId() } } pub fn get_double_click_time() -> u32 { unsafe { GetDoubleClickTime() } } pub fn wide_string(s: &str) -> Vec { use std::os::windows::prelude::OsStrExt; std::ffi::OsStr::new(s) .encode_wide() .chain(Some(0).into_iter()) .collect() } /// send message to currently shown window pub fn send_message_to_hnwd( class_name: &str, window_name: &str, dw_data: usize, data: &str, show_window: bool, ) -> bool { unsafe { let class_name_utf16 = wide_string(class_name); let window_name_utf16 = wide_string(window_name); let window = FindWindowW(class_name_utf16.as_ptr(), window_name_utf16.as_ptr()); if window.is_null() { log::warn!("no such window {}:{}", class_name, window_name); return false; } let mut data_struct = COPYDATASTRUCT::default(); data_struct.dwData = dw_data; let mut data_zero: String = data.chars().chain(Some('\0').into_iter()).collect(); println!("send {:?}", data_zero); data_struct.cbData = data_zero.len() as _; data_struct.lpData = data_zero.as_mut_ptr() as _; SendMessageW( window, WM_COPYDATA, 0, &data_struct as *const COPYDATASTRUCT as _, ); if show_window { ShowWindow(window, SW_NORMAL); SetForegroundWindow(window); } } return true; } pub fn get_logon_user_token(user: &str, pwd: &str) -> ResultType { let user_split = user.split("\\").collect::>(); let wuser = wide_string(user_split.get(1).unwrap_or(&user)); let wpc = wide_string(user_split.get(0).unwrap_or(&"")); let wpwd = wide_string(pwd); let mut ph_token: HANDLE = std::ptr::null_mut(); let res = unsafe { LogonUserW( wuser.as_ptr(), wpc.as_ptr(), wpwd.as_ptr(), LOGON32_LOGON_INTERACTIVE, LOGON32_PROVIDER_DEFAULT, &mut ph_token as _, ) }; if res == FALSE { bail!( "Failed to log on user {}: {}", user, std::io::Error::last_os_error() ); } else { if ph_token.is_null() { bail!( "Failed to log on user {}: {}", user, std::io::Error::last_os_error() ); } Ok(ph_token) } } // Ensure the token returned is a primary token. // If the provided token is an impersonation token, it duplicates it to a primary token. // If the provided token is already a primary token, it returns it as is. // The caller is responsible for closing the returned token handle. pub fn ensure_primary_token(user_token: HANDLE) -> ResultType { if user_token.is_null() || user_token == INVALID_HANDLE_VALUE { bail!("Invalid user token provided"); } unsafe { let mut token_type: TOKEN_TYPE = 0; let mut return_length: DWORD = 0; if GetTokenInformation( user_token, TokenType, &mut token_type as *mut _ as *mut _, std::mem::size_of::() as DWORD, &mut return_length, ) == FALSE { bail!( "Failed to get token type, error {}", io::Error::last_os_error() ); } if token_type == TokenImpersonation { let mut duplicate_token: HANDLE = std::ptr::null_mut(); let dup_res = DuplicateToken(user_token, SecurityImpersonation, &mut duplicate_token); CloseHandle(user_token); if dup_res == FALSE { bail!( "Failed to duplicate token, error {}", io::Error::last_os_error() ); } Ok(duplicate_token) } else { Ok(user_token) } } } pub fn is_user_token_admin(user_token: HANDLE) -> ResultType { if user_token.is_null() || user_token == INVALID_HANDLE_VALUE { bail!("Invalid user token provided"); } unsafe { let mut dw_size: DWORD = 0; GetTokenInformation( user_token, TokenGroups, std::ptr::null_mut(), 0, &mut dw_size, ); let last_error = GetLastError(); if last_error != ERROR_INSUFFICIENT_BUFFER { bail!( "Failed to get token groups buffer size, error: {}", last_error ); } if dw_size == 0 { bail!("Token groups buffer size is zero"); } let mut buffer = vec![0u8; dw_size as usize]; if GetTokenInformation( user_token, TokenGroups, buffer.as_mut_ptr() as *mut _, dw_size, &mut dw_size, ) == FALSE { bail!( "Failed to get token groups information, error: {}", io::Error::last_os_error() ); } let p_token_groups = buffer.as_ptr() as *const TOKEN_GROUPS; let group_count = (*p_token_groups).GroupCount; if group_count == 0 { return Ok(false); } let mut nt_authority: SID_IDENTIFIER_AUTHORITY = SID_IDENTIFIER_AUTHORITY { Value: SECURITY_NT_AUTHORITY, }; let mut administrators_group: PSID = std::ptr::null_mut(); if AllocateAndInitializeSid( &mut nt_authority, 2, SECURITY_BUILTIN_DOMAIN_RID, DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0, &mut administrators_group, ) == FALSE { bail!( "Failed to allocate administrators group SID, error: {}", io::Error::last_os_error() ); } if administrators_group.is_null() { bail!("Failed to create administrators group SID"); } let mut is_admin = false; let groups = std::slice::from_raw_parts((*p_token_groups).Groups.as_ptr(), group_count as usize); for group in groups { if EqualSid(administrators_group, group.Sid) == TRUE { is_admin = true; break; } } if !administrators_group.is_null() { FreeSid(administrators_group); } Ok(is_admin) } } pub fn create_process_with_logon(user: &str, pwd: &str, exe: &str, arg: &str) -> ResultType<()> { let last_error_table = HashMap::from([ ( ERROR_LOGON_FAILURE, "The user name or password is incorrect.", ), (ERROR_ACCESS_DENIED, "Access is denied."), ]); unsafe { let user_split = user.split("\\").collect::>(); let wuser = wide_string(user_split.get(1).unwrap_or(&user)); let wpc = wide_string(user_split.get(0).unwrap_or(&"")); let wpwd = wide_string(pwd); let cmd = if arg.is_empty() { format!("\"{}\"", exe) } else { format!("\"{}\" {}", exe, arg) }; let mut wcmd = wide_string(&cmd); let mut si: STARTUPINFOW = mem::zeroed(); si.wShowWindow = SW_HIDE as _; si.lpDesktop = NULL as _; si.cb = std::mem::size_of::() as _; si.dwFlags = STARTF_USESHOWWINDOW; let mut pi: PROCESS_INFORMATION = mem::zeroed(); let wexe = wide_string(exe); if FALSE == CreateProcessWithLogonW( wuser.as_ptr(), wpc.as_ptr(), wpwd.as_ptr(), LOGON_WITH_PROFILE, wexe.as_ptr(), wcmd.as_mut_ptr(), CREATE_UNICODE_ENVIRONMENT, NULL, NULL as _, &mut si as *mut STARTUPINFOW, &mut pi as *mut PROCESS_INFORMATION, ) { let last_error = GetLastError(); bail!( "CreateProcessWithLogonW failed : \"{}\", error {}", last_error_table .get(&last_error) .unwrap_or(&"Unknown error"), io::Error::from_raw_os_error(last_error as _) ); } } return Ok(()); } pub fn set_path_permission(dir: &Path, permission: &str) -> ResultType<()> { std::process::Command::new("icacls") .arg(dir.as_os_str()) .arg("/grant") .arg(format!("*S-1-1-0:(OI)(CI){}", permission)) .arg("/T") .spawn()?; Ok(()) } #[inline] fn str_to_device_name(name: &str) -> [u16; 32] { let mut device_name: Vec = wide_string(name); if device_name.len() < 32 { device_name.resize(32, 0); } let mut result = [0; 32]; result.copy_from_slice(&device_name[..32]); result } pub fn resolutions(name: &str) -> Vec { unsafe { let mut dm: DEVMODEW = std::mem::zeroed(); let mut v = vec![]; let mut num = 0; let device_name = str_to_device_name(name); loop { if EnumDisplaySettingsW(device_name.as_ptr(), num, &mut dm) == 0 { break; } let r = Resolution { width: dm.dmPelsWidth as _, height: dm.dmPelsHeight as _, ..Default::default() }; if !v.contains(&r) { v.push(r); } num += 1; } v } } pub fn current_resolution(name: &str) -> ResultType { let device_name = str_to_device_name(name); unsafe { let mut dm: DEVMODEW = std::mem::zeroed(); dm.dmSize = std::mem::size_of::() as _; if EnumDisplaySettingsW(device_name.as_ptr(), ENUM_CURRENT_SETTINGS, &mut dm) == 0 { bail!( "failed to get current resolution, error {}", io::Error::last_os_error() ); } let r = Resolution { width: dm.dmPelsWidth as _, height: dm.dmPelsHeight as _, ..Default::default() }; Ok(r) } } pub(super) fn change_resolution_directly( name: &str, width: usize, height: usize, ) -> ResultType<()> { let device_name = str_to_device_name(name); unsafe { let mut dm: DEVMODEW = std::mem::zeroed(); dm.dmSize = std::mem::size_of::() as _; dm.dmPelsWidth = width as _; dm.dmPelsHeight = height as _; dm.dmFields = DM_PELSHEIGHT | DM_PELSWIDTH; let res = ChangeDisplaySettingsExW( device_name.as_ptr(), &mut dm, NULL as _, CDS_UPDATEREGISTRY | CDS_GLOBAL | CDS_RESET, NULL, ); if res != DISP_CHANGE_SUCCESSFUL { bail!( "ChangeDisplaySettingsExW failed, res={}, error {}", res, io::Error::last_os_error() ); } Ok(()) } } pub fn user_accessible_folder() -> ResultType { let disk = std::env::var("SystemDrive").unwrap_or("C:".to_string()); let dir1 = PathBuf::from(format!("{}\\ProgramData", disk)); // NOTICE: "C:\Windows\Temp" requires permanent authorization. let dir2 = PathBuf::from(format!("{}\\Windows\\Temp", disk)); let dir; if dir1.exists() { dir = dir1; } else if dir2.exists() { dir = dir2; } else { bail!("no valid user accessible folder"); } Ok(dir) } #[inline] pub fn uninstall_cert() -> ResultType<()> { cert::uninstall_cert() } mod cert { use hbb_common::ResultType; extern "C" { fn DeleteRustDeskTestCertsW(); } pub fn uninstall_cert() -> ResultType<()> { unsafe { DeleteRustDeskTestCertsW(); } Ok(()) } } #[inline] pub fn get_char_from_vk(vk: u32) -> Option { get_char_from_unicode(get_unicode_from_vk(vk)?) } pub fn get_char_from_unicode(unicode: u16) -> Option { let buff = [unicode]; if let Some(chr) = String::from_utf16(&buff[..1]).ok()?.chars().next() { if chr.is_control() { return None; } else { Some(chr) } } else { None } } pub fn get_unicode_from_vk(vk: u32) -> Option { const BUF_LEN: i32 = 32; let mut buff = [0_u16; BUF_LEN as usize]; let buff_ptr = buff.as_mut_ptr(); let len = unsafe { let current_window_thread_id = GetWindowThreadProcessId(GetForegroundWindow(), null_mut()); let layout = GetKeyboardLayout(current_window_thread_id); // refs: https://github.com/rustdesk-org/rdev/blob/25a99ce71ab42843ad253dd51e6a35e83e87a8a4/src/windows/keyboard.rs#L115 let press_state = 129; let mut state: [BYTE; 256] = [0; 256]; let shift_left = rdev::get_modifier(rdev::Key::ShiftLeft); let shift_right = rdev::get_modifier(rdev::Key::ShiftRight); if shift_left { state[VK_LSHIFT as usize] = press_state; } if shift_right { state[VK_RSHIFT as usize] = press_state; } if shift_left || shift_right { state[VK_SHIFT as usize] = press_state; } ToUnicodeEx(vk, 0x00, &state as _, buff_ptr, BUF_LEN, 0, layout) }; if len == 1 { Some(buff[0]) } else { None } } pub fn is_process_consent_running() -> ResultType { let output = std::process::Command::new("cmd") .args(&["/C", "tasklist | findstr consent.exe"]) .creation_flags(CREATE_NO_WINDOW) .output()?; Ok(output.status.success() && !output.stdout.is_empty()) } pub struct WakeLock(u32); // Failed to compile keepawake-rs on i686 impl WakeLock { pub fn new(display: bool, idle: bool, sleep: bool) -> Self { let mut flag = ES_CONTINUOUS; if display { flag |= ES_DISPLAY_REQUIRED; } if idle { flag |= ES_SYSTEM_REQUIRED; } if sleep { flag |= ES_AWAYMODE_REQUIRED; } unsafe { SetThreadExecutionState(flag) }; WakeLock(flag) } pub fn set_display(&mut self, display: bool) -> ResultType<()> { let flag = if display { self.0 | ES_DISPLAY_REQUIRED } else { self.0 & !ES_DISPLAY_REQUIRED }; if flag != self.0 { unsafe { SetThreadExecutionState(flag) }; self.0 = flag; } Ok(()) } } impl Drop for WakeLock { fn drop(&mut self) { unsafe { SetThreadExecutionState(ES_CONTINUOUS) }; } } pub fn uninstall_service(show_new_window: bool, _: bool) -> bool { log::info!("Uninstalling service..."); let filter = format!(" /FI \"PID ne {}\"", get_current_pid()); Config::set_option("stop-service".into(), "Y".into()); let cmds = format!( " chcp 65001 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_TOPMOST_INJECTED_PROCESS_EXE, ); if let Err(err) = run_cmds(cmds, false, "uninstall") { Config::set_option("stop-service".into(), "".into()); log::debug!("{err}"); return true; } run_after_run_cmds(!show_new_window); std::process::exit(0); } pub fn install_service() -> bool { log::info!("Installing service..."); let _installing = crate::platform::InstallingService::new(); let (_, _, _, exe) = get_install_info(); let tmp_path = std::env::temp_dir().to_string_lossy().to_string(); let tray_shortcut = get_tray_shortcut(&exe, &tmp_path).unwrap_or_default(); let filter = format!(" /FI \"PID ne {}\"", get_current_pid()); Config::set_option("stop-service".into(), "".into()); crate::ipc::EXIT_RECV_CLOSE.store(false, Ordering::Relaxed); let cmds = format!( " chcp 65001 taskkill /F /IM {app_name}.exe{filter} cscript \"{tray_shortcut}\" copy /Y \"{tmp_path}\\{app_name} Tray.lnk\" \"%PROGRAMDATA%\\Microsoft\\Windows\\Start Menu\\Programs\\Startup\\\" {import_config} {create_service} if exist \"{tray_shortcut}\" del /f /q \"{tray_shortcut}\" ", app_name = crate::get_app_name(), import_config = get_import_config(&exe), create_service = get_create_service(&exe), ); if let Err(err) = run_cmds(cmds, false, "install") { Config::set_option("stop-service".into(), "Y".into()); crate::ipc::EXIT_RECV_CLOSE.store(true, Ordering::Relaxed); log::debug!("{err}"); return true; } run_after_run_cmds(false); std::process::exit(0); } pub fn update_me(debug: bool) -> ResultType<()> { let app_name = crate::get_app_name(); let src_exe = std::env::current_exe()?.to_string_lossy().to_string(); let (subkey, path, _, exe) = get_install_info(); let is_installed = std::fs::metadata(&exe).is_ok(); if !is_installed { bail!("{} is not installed.", &app_name); } let app_exe_name = &format!("{}.exe", &app_name); let main_window_pids = crate::platform::get_pids_of_process_with_args::<_, &str>(&app_exe_name, &[]); let main_window_sessions = main_window_pids .iter() .map(|pid| get_session_id_of_process(pid.as_u32())) .flatten() .collect::>(); kill_process_by_pids(&app_exe_name, main_window_pids)?; let tray_pids = crate::platform::get_pids_of_process_with_args(&app_exe_name, &["--tray"]); let tray_sessions = tray_pids .iter() .map(|pid| get_session_id_of_process(pid.as_u32())) .flatten() .collect::>(); kill_process_by_pids(&app_exe_name, tray_pids)?; let is_service_running = is_self_service_running(); let mut version_major = "0"; let mut version_minor = "0"; let mut version_build = "0"; let versions: Vec<&str> = crate::VERSION.split(".").collect(); if versions.len() > 0 { version_major = versions[0]; } if versions.len() > 1 { version_minor = versions[1]; } if versions.len() > 2 { version_build = versions[2]; } let meta = std::fs::symlink_metadata(std::env::current_exe()?)?; let size = meta.len() / 1024; let reg_cmd = format!( " reg add {subkey} /f /v DisplayIcon /t REG_SZ /d \"{exe}\" reg add {subkey} /f /v DisplayVersion /t REG_SZ /d \"{version}\" reg add {subkey} /f /v Version /t REG_SZ /d \"{version}\" reg add {subkey} /f /v BuildDate /t REG_SZ /d \"{build_date}\" reg add {subkey} /f /v VersionMajor /t REG_DWORD /d {version_major} reg add {subkey} /f /v VersionMinor /t REG_DWORD /d {version_minor} reg add {subkey} /f /v VersionBuild /t REG_DWORD /d {version_build} reg add {subkey} /f /v EstimatedSize /t REG_DWORD /d {size} ", version = crate::VERSION.replace("-", "."), build_date = crate::BUILD_DATE, ); let filter = format!(" /FI \"PID ne {}\"", get_current_pid()); let restore_service_cmd = if is_service_running { format!("sc start {}", &app_name) } else { "".to_owned() }; // No need to check the install option here, `is_rd_printer_installed` rarely fails. let is_printer_installed = remote_printer::is_rd_printer_installed(&app_name).unwrap_or(false); // Do nothing if the printer is not installed or failed to query if the printer is installed. let (uninstall_printer_cmd, install_printer_cmd) = if is_printer_installed { ( format!("\"{}\" --uninstall-remote-printer", &src_exe), format!("\"{}\" --install-remote-printer", &src_exe), ) } else { ("".to_owned(), "".to_owned()) }; // We do not try to remove all files in the old version. // Because I don't know whether additional files will be installed here after installation, such as drivers. // Just copy files to the installation directory works fine. //if exist \"{path}\" rd /s /q \"{path}\" // md \"{path}\" // // We need `taskkill` because: // 1. There may be some other processes like `rustdesk --connect` are running. // 2. Sometimes, the main window and the tray icon are showing // while I cannot find them by `tasklist` or the methods above. // There's should be 4 processes running: service, server, tray and main window. // But only 2 processes are shown in the tasklist. let cmds = format!( " chcp 65001 sc stop {app_name} taskkill /F /IM {app_name}.exe{filter} {reg_cmd} {copy_exe} {restore_service_cmd} {uninstall_printer_cmd} {install_printer_cmd} {sleep} ", app_name = app_name, copy_exe = copy_exe_cmd(&src_exe, &exe, &path)?, sleep = if debug { "timeout 300" } else { "" }, ); run_cmds(cmds, debug, "update")?; std::thread::sleep(std::time::Duration::from_millis(2000)); if tray_sessions.is_empty() { log::info!("No tray process found."); } else { log::info!("Try to restore the tray process..."); log::info!( "Try to restore the tray process..., sessions: {:?}", &tray_sessions ); for s in tray_sessions { if s != 0 { allow_err!(run_exe_in_session(&exe, vec!["--tray"], s, true)); } } } if main_window_sessions.is_empty() { log::info!("No main window process found."); } else { log::info!("Try to restore the main window process..."); std::thread::sleep(std::time::Duration::from_millis(2000)); for s in main_window_sessions { if s != 0 { allow_err!(run_exe_in_session(&exe, vec![], s, true)); } } } std::thread::sleep(std::time::Duration::from_millis(300)); log::info!("Update completed."); Ok(()) } // Double confirm the process name fn kill_process_by_pids(name: &str, pids: Vec) -> ResultType<()> { let name = name.to_lowercase(); let s = System::new_all(); // No need to check all names of `pids` first, and kill them then. // It's rare case that they're not matched. for pid in pids { if let Some(process) = s.process(pid) { if process.name().to_lowercase() != name { bail!("Failed to kill the process, the names are mismatched."); } if !process.kill() { bail!("Failed to kill the process"); } } else { bail!("Failed to kill the process, the pid is not found"); } } Ok(()) } // Don't launch tray app when running with `\qn`. // 1. Because `/qn` requires administrator permission and the tray app should be launched with user permission. // Or launching the main window from the tray app will cause the main window to be launched with administrator permission. // 2. We are not able to launch the tray app if the UI is in the login screen. // `fn update_me()` can handle the above cases, but for msi update, we need to do more work to handle the above cases. // 1. Record the tray app session ids. // 2. Do the update. // 3. Restore the tray app sessions. // `1` and `3` must be done in custom actions. // We need also to handle the command line parsing to find the tray processes. pub fn update_me_msi(msi: &str, quiet: bool) -> ResultType<()> { let cmds = format!( "chcp 65001 && msiexec /i {msi} {}", if quiet { "/qn LAUNCH_TRAY_APP=N" } else { "" } ); run_cmds(cmds, false, "update-msi")?; Ok(()) } pub fn get_tray_shortcut(exe: &str, tmp_path: &str) -> ResultType { Ok(write_cmds( format!( " Set oWS = WScript.CreateObject(\"WScript.Shell\") sLinkFile = \"{tmp_path}\\{app_name} Tray.lnk\" Set oLink = oWS.CreateShortcut(sLinkFile) oLink.TargetPath = \"{exe}\" oLink.Arguments = \"--tray\" oLink.Save ", app_name = crate::get_app_name(), ), "vbs", "tray_shortcut", )? .to_str() .unwrap_or("") .to_owned()) } fn get_import_config(exe: &str) -> String { if config::is_outgoing_only() { return "".to_string(); } format!(" sc stop {app_name} sc delete {app_name} sc create {app_name} binpath= \"\\\"{exe}\\\" --import-config \\\"{config_path}\\\"\" start= auto DisplayName= \"{app_name} Service\" sc start {app_name} sc stop {app_name} sc delete {app_name} ", app_name = crate::get_app_name(), config_path=Config::file().to_str().unwrap_or(""), ) } fn get_create_service(exe: &str) -> String { if config::is_outgoing_only() { return "".to_string(); } let stop = Config::get_option("stop-service") == "Y"; if stop { format!(" 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\" ", app_name = crate::get_app_name()) } else { format!(" sc create {app_name} binpath= \"\\\"{exe}\\\" --service\" start= auto DisplayName= \"{app_name} Service\" sc start {app_name} ", app_name = crate::get_app_name()) } } fn run_after_run_cmds(silent: bool) { let (_, _, _, exe) = get_install_info(); if !silent { log::debug!("Spawn new window"); allow_err!(std::process::Command::new("cmd") .args(&["/c", "timeout", "/t", "2", "&", &format!("{exe}")]) .creation_flags(winapi::um::winbase::CREATE_NO_WINDOW) .spawn()); } if Config::get_option("stop-service") != "Y" { allow_err!(std::process::Command::new(&exe).arg("--tray").spawn()); } std::thread::sleep(std::time::Duration::from_millis(300)); } #[inline] pub fn try_kill_broker() { allow_err!(std::process::Command::new("cmd") .arg("/c") .arg(&format!( "taskkill /F /IM {}", WIN_TOPMOST_INJECTED_PROCESS_EXE )) .creation_flags(winapi::um::winbase::CREATE_NO_WINDOW) .spawn()); } pub fn message_box(text: &str) { let mut text = text.to_owned(); let nodialog = std::env::var("NO_DIALOG").unwrap_or_default() == "Y"; if !text.ends_with("!") || nodialog { use arboard::Clipboard as ClipboardContext; match ClipboardContext::new() { Ok(mut ctx) => { ctx.set_text(&text).ok(); if !nodialog { text = format!("{}\n\nAbove text has been copied to clipboard", &text); } } _ => {} } } if nodialog { if std::env::var("PRINT_OUT").unwrap_or_default() == "Y" { println!("{text}"); } if let Ok(x) = std::env::var("WRITE_TO_FILE") { if !x.is_empty() { allow_err!(std::fs::write(x, text)); } } return; } let text = text .encode_utf16() .chain(std::iter::once(0)) .collect::>(); let caption = "RustDesk Output" .encode_utf16() .chain(std::iter::once(0)) .collect::>(); unsafe { MessageBoxW(std::ptr::null_mut(), text.as_ptr(), caption.as_ptr(), MB_OK) }; } pub fn alloc_console() { unsafe { alloc_console_and_redirect(); } } fn get_license() -> Option { let mut lic: CustomServer = Default::default(); if let Ok(tmp) = get_license_from_exe_name() { lic = tmp; } else { // for back compatibility from migrating from <= 1.2.1 to 1.2.2 lic.key = get_reg("Key"); lic.host = get_reg("Host"); lic.api = get_reg("Api"); } if lic.key.is_empty() || lic.host.is_empty() { return None; } Some(lic) } pub struct WallPaperRemover { old_path: String, } impl WallPaperRemover { pub fn new() -> ResultType { let start = std::time::Instant::now(); if !Self::need_remove() { bail!("already solid color"); } let old_path = match Self::get_recent_wallpaper() { Ok(old_path) => old_path, Err(e) => { log::info!("Failed to get recent wallpaper: {:?}, use fallback", e); wallpaper::get().map_err(|e| anyhow!(e.to_string()))? } }; Self::set_wallpaper(None)?; log::info!( "created wallpaper remover, old_path: {:?}, elapsed: {:?}", old_path, start.elapsed(), ); Ok(Self { old_path }) } pub fn support() -> bool { wallpaper::get().is_ok() || !Self::get_recent_wallpaper().unwrap_or_default().is_empty() } fn get_recent_wallpaper() -> ResultType { // SystemParametersInfoW may return %appdata%\Microsoft\Windows\Themes\TranscodedWallpaper, not real path and may not real cache // https://www.makeuseof.com/find-desktop-wallpapers-file-location-windows-11/ // https://superuser.com/questions/1218413/write-to-current-users-registry-through-a-different-admin-account let (hkcu, sid) = if is_root() { let sid = get_current_process_session_id().ok_or(anyhow!("failed to get sid"))?; (RegKey::predef(HKEY_USERS), format!("{}\\", sid)) } else { (RegKey::predef(HKEY_CURRENT_USER), "".to_string()) }; let explorer_key = hkcu.open_subkey_with_flags( &format!( "{}Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\Wallpapers", sid ), KEY_READ, )?; Ok(explorer_key.get_value("BackgroundHistoryPath0")?) } fn need_remove() -> bool { if let Ok(wallpaper) = wallpaper::get() { return !wallpaper.is_empty(); } false } fn set_wallpaper(path: Option) -> ResultType<()> { wallpaper::set_from_path(&path.unwrap_or_default()).map_err(|e| anyhow!(e.to_string())) } } impl Drop for WallPaperRemover { fn drop(&mut self) { // If the old background is a slideshow, it will be converted into an image. AnyDesk does the same. allow_err!(Self::set_wallpaper(Some(self.old_path.clone()))); } } fn get_uninstall_amyuni_idd() -> String { match std::env::current_exe() { Ok(path) => format!("\"{}\" --uninstall-amyuni-idd", path.to_str().unwrap_or("")), Err(e) => { log::warn!("Failed to get current exe path, cannot get command of uninstalling idd, Zzerror: {:?}", e); "".to_string() } } } #[inline] pub fn is_self_service_running() -> bool { is_service_running(&crate::get_app_name()) } pub fn is_service_running(service_name: &str) -> bool { unsafe { let service_name = wide_string(service_name); is_service_running_w(service_name.as_ptr() as _) } } pub fn is_x64() -> bool { const PROCESSOR_ARCHITECTURE_AMD64: u16 = 9; let mut sys_info = SYSTEM_INFO::default(); unsafe { GetNativeSystemInfo(&mut sys_info as _); } unsafe { sys_info.u.s().wProcessorArchitecture == PROCESSOR_ARCHITECTURE_AMD64 } } pub fn try_kill_rustdesk_main_window_process() -> ResultType<()> { // Kill rustdesk.exe without extra arg, should only be called by --server // We can find the exact process which occupies the ipc, see more from https://github.com/winsiderss/systeminformer log::info!("try kill rustdesk main window process"); use hbb_common::sysinfo::System; let mut sys = System::new(); sys.refresh_processes(); let my_uid = sys .process((std::process::id() as usize).into()) .map(|x| x.user_id()) .unwrap_or_default(); let my_pid = std::process::id(); let app_name = crate::get_app_name().to_lowercase(); if app_name.is_empty() { bail!("app name is empty"); } for (_, p) in sys.processes().iter() { let p_name = p.name().to_lowercase(); // name equal if !(p_name == app_name || p_name == app_name.clone() + ".exe") { continue; } // arg more than 1 if p.cmd().len() < 1 { continue; } // first arg contain app name if !p.cmd()[0].to_lowercase().contains(&p_name) { continue; } // only one arg or the second arg is empty uni link let is_empty_uni = p.cmd().len() == 2 && crate::common::is_empty_uni_link(&p.cmd()[1]); if !(p.cmd().len() == 1 || is_empty_uni) { continue; } // skip self if p.pid().as_u32() == my_pid { continue; } // because we call it with --server, so we can check user_id, remove this if call it with user process if p.user_id() == my_uid { log::info!("user id equal, continue"); continue; } log::info!("try kill process: {:?}, pid = {:?}", p.cmd(), p.pid()); nt_terminate_process(p.pid().as_u32())?; log::info!("kill process success: {:?}, pid = {:?}", p.cmd(), p.pid()); return Ok(()); } bail!("failed to find rustdesk main window process"); } fn nt_terminate_process(process_id: DWORD) -> ResultType<()> { type NtTerminateProcess = unsafe extern "system" fn(HANDLE, DWORD) -> DWORD; unsafe { let h_module = if is_win_10_or_greater() { LoadLibraryExA( CString::new("ntdll.dll")?.as_ptr(), std::ptr::null_mut(), LOAD_LIBRARY_SEARCH_SYSTEM32, ) } else { LoadLibraryA(CString::new("ntdll.dll")?.as_ptr()) }; if !h_module.is_null() { let f_nt_terminate_process: NtTerminateProcess = std::mem::transmute(GetProcAddress( h_module, CString::new("NtTerminateProcess")?.as_ptr(), )); let h_token = OpenProcess(PROCESS_ALL_ACCESS, 0, process_id); if !h_token.is_null() { if f_nt_terminate_process(h_token, 1) == 0 { log::info!("terminate process {} success", process_id); CloseHandle(h_token); return Ok(()); } else { CloseHandle(h_token); bail!("NtTerminateProcess {} failed", process_id); } } else { bail!("OpenProcess {} failed", process_id); } } else { bail!("Failed to load ntdll.dll"); } } } pub fn try_set_window_foreground(window: HWND) { let env_key = SET_FOREGROUND_WINDOW; if let Ok(value) = std::env::var(env_key) { if value == "1" { unsafe { SetForegroundWindow(window); } std::env::remove_var(env_key); } } } pub mod reg_display_settings { use hbb_common::ResultType; use serde_derive::{Deserialize, Serialize}; use std::collections::HashMap; use winreg::{enums::*, RegValue}; const REG_GRAPHICS_DRIVERS_PATH: &str = "SYSTEM\\CurrentControlSet\\Control\\GraphicsDrivers"; const REG_CONNECTIVITY_PATH: &str = "Connectivity"; #[derive(Serialize, Deserialize, Debug)] pub struct RegRecovery { path: String, key: String, old: (Vec, isize), new: (Vec, isize), } pub fn read_reg_connectivity() -> ResultType>> { let hklm = winreg::RegKey::predef(HKEY_LOCAL_MACHINE); let reg_connectivity = hklm.open_subkey_with_flags( format!("{}\\{}", REG_GRAPHICS_DRIVERS_PATH, REG_CONNECTIVITY_PATH), KEY_READ, )?; let mut map_connectivity = HashMap::new(); for key in reg_connectivity.enum_keys() { let key = key?; let mut map_item = HashMap::new(); let reg_item = reg_connectivity.open_subkey_with_flags(&key, KEY_READ)?; for value in reg_item.enum_values() { let (name, value) = value?; map_item.insert(name, value); } map_connectivity.insert(key, map_item); } Ok(map_connectivity) } pub fn diff_recent_connectivity( map1: HashMap>, map2: HashMap>, ) -> Option { for (subkey, map_item2) in map2 { if let Some(map_item1) = map1.get(&subkey) { let key = "Recent"; if let Some(value1) = map_item1.get(key) { if let Some(value2) = map_item2.get(key) { if value1 != value2 { return Some(RegRecovery { path: format!( "{}\\{}\\{}", REG_GRAPHICS_DRIVERS_PATH, REG_CONNECTIVITY_PATH, subkey ), key: key.to_owned(), old: (value1.bytes.clone(), value1.vtype.clone() as isize), new: (value2.bytes.clone(), value2.vtype.clone() as isize), }); } } } } } None } pub fn restore_reg_connectivity(reg_recovery: RegRecovery, force: bool) -> ResultType<()> { let hklm = winreg::RegKey::predef(HKEY_LOCAL_MACHINE); let reg_item = hklm.open_subkey_with_flags(®_recovery.path, KEY_READ | KEY_WRITE)?; if !force { let cur_reg_value = reg_item.get_raw_value(®_recovery.key)?; let new_reg_value = RegValue { bytes: reg_recovery.new.0, vtype: isize_to_reg_type(reg_recovery.new.1), }; // Compare if the current value is the same as the new value. // If they are not the same, the registry value has been changed by other processes. // So we do not restore the registry value. if cur_reg_value != new_reg_value { return Ok(()); } } let reg_value = RegValue { bytes: reg_recovery.old.0, vtype: isize_to_reg_type(reg_recovery.old.1), }; reg_item.set_raw_value(®_recovery.key, ®_value)?; Ok(()) } #[inline] fn isize_to_reg_type(i: isize) -> RegType { match i { 0 => RegType::REG_NONE, 1 => RegType::REG_SZ, 2 => RegType::REG_EXPAND_SZ, 3 => RegType::REG_BINARY, 4 => RegType::REG_DWORD, 5 => RegType::REG_DWORD_BIG_ENDIAN, 6 => RegType::REG_LINK, 7 => RegType::REG_MULTI_SZ, 8 => RegType::REG_RESOURCE_LIST, 9 => RegType::REG_FULL_RESOURCE_DESCRIPTOR, 10 => RegType::REG_RESOURCE_REQUIREMENTS_LIST, 11 => RegType::REG_QWORD, _ => RegType::REG_NONE, } } } pub fn get_printer_names() -> ResultType> { let mut needed_bytes = 0; let mut returned_count = 0; unsafe { // First call to get required buffer size EnumPrintersW( PRINTER_ENUM_LOCAL | PRINTER_ENUM_CONNECTIONS, std::ptr::null_mut(), 1, std::ptr::null_mut(), 0, &mut needed_bytes, &mut returned_count, ); let mut buffer = vec![0u8; needed_bytes as usize]; if EnumPrintersW( PRINTER_ENUM_LOCAL | PRINTER_ENUM_CONNECTIONS, std::ptr::null_mut(), 1, buffer.as_mut_ptr() as *mut _, needed_bytes, &mut needed_bytes, &mut returned_count, ) == 0 { return Err(anyhow!("Failed to enumerate printers")); } let ptr = buffer.as_ptr() as *const PRINTER_INFO_1W; let printers = std::slice::from_raw_parts(ptr, returned_count as usize); Ok(printers .iter() .filter_map(|p| { let name = p.pName; if !name.is_null() { let mut len = 0; while len < 500 { if name.add(len).is_null() || *name.add(len) == 0 { break; } len += 1; } if len > 0 && len < 500 { Some(String::from_utf16_lossy(std::slice::from_raw_parts( name, len, ))) } else { None } } else { None } }) .collect()) } } extern "C" { fn PrintXPSRawData(printer_name: *const u16, raw_data: *const u8, data_size: c_ulong) -> DWORD; } pub fn send_raw_data_to_printer(printer_name: Option, data: Vec) -> ResultType<()> { let mut printer_name = printer_name.unwrap_or_default(); if printer_name.is_empty() { // use GetDefaultPrinter to get the default printer name let mut needed_bytes = 0; unsafe { GetDefaultPrinterW(std::ptr::null_mut(), &mut needed_bytes); } if needed_bytes > 0 { let mut default_printer_name = vec![0u16; needed_bytes as usize]; unsafe { GetDefaultPrinterW( default_printer_name.as_mut_ptr() as *mut _, &mut needed_bytes, ); } printer_name = String::from_utf16_lossy(&default_printer_name[..needed_bytes as usize]); } } else { if let Ok(names) = crate::platform::windows::get_printer_names() { if !names.contains(&printer_name) { // Don't set the first printer as current printer. // It may not be the desired printer. bail!("Printer name \"{}\" not found", &printer_name); } } } if printer_name.is_empty() { return Err(anyhow!("Failed to get printer name")); } log::info!("Sending data to printer: {}", &printer_name); let printer_name = wide_string(&printer_name); unsafe { let res = PrintXPSRawData( printer_name.as_ptr(), data.as_ptr() as *const u8, data.len() as c_ulong, ); if res != 0 { bail!("Failed to send data to the printer, see logs in C:\\Windows\\temp\\test_rustdesk.log for more details."); } else { log::info!("Successfully sent data to the printer"); } } Ok(()) } fn get_pids>(name: S) -> ResultType> { let name = name.as_ref().to_lowercase(); let mut pids = Vec::new(); unsafe { let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)?; if snapshot == WinHANDLE::default() { return Ok(pids); } let mut entry: PROCESSENTRY32W = std::mem::zeroed(); entry.dwSize = std::mem::size_of::() as u32; if Process32FirstW(snapshot, &mut entry).is_ok() { loop { let proc_name = OsString::from_wide(&entry.szExeFile) .to_string_lossy() .to_lowercase(); if proc_name.contains(&name) { pids.push(entry.th32ProcessID); } if !Process32NextW(snapshot, &mut entry).is_ok() { break; } } } let _ = WinCloseHandle(snapshot); } Ok(pids) } pub fn is_msi_installed() -> std::io::Result { let hklm = RegKey::predef(HKEY_LOCAL_MACHINE); let uninstall_key = hklm.open_subkey(format!( "SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{}", crate::get_app_name() ))?; Ok(1 == uninstall_key.get_value::("WindowsInstaller")?) } pub fn is_cur_exe_the_installed() -> bool { let (_, _, _, exe) = get_install_info(); // Check if is installed, because `exe` is the default path if is not installed. if !std::fs::metadata(&exe).is_ok() { return false; } let mut path = std::env::current_exe().unwrap_or_default(); if let Ok(linked) = path.read_link() { path = linked; } let path = path.to_string_lossy().to_lowercase(); path == exe.to_lowercase() } #[cfg(not(target_pointer_width = "64"))] pub fn get_pids_with_first_arg_check_session, S2: AsRef>( name: S1, arg: S2, same_session_id: bool, ) -> ResultType> { // Though `wmic` can return the sessionId, for simplicity we only return processid. let pids = get_pids_with_first_arg_by_wmic(name, arg); if !same_session_id { return Ok(pids); } let Some(cur_sid) = get_current_process_session_id() else { bail!("Can't get current process session id"); }; let mut same_session_pids = vec![]; for pid in pids.into_iter() { let mut sid = 0; if unsafe { ProcessIdToSessionId(pid.as_u32(), &mut sid) == TRUE } { if sid == cur_sid { same_session_pids.push(pid); } } else { // Only log here, because this call almost never fails. log::warn!( "Failed to get session id of the process id, error: {:?}", std::io::Error::last_os_error() ); } } Ok(same_session_pids) } #[cfg(not(target_pointer_width = "64"))] fn get_pids_with_args_from_wmic_output>( output: std::borrow::Cow<'_, str>, name: &str, args: &[S2], ) -> Vec { // CommandLine= // ProcessId=33796 // // CommandLine= // ProcessId=34668 // // CommandLine="C:\Program Files\RustDesk\RustDesk.exe" --tray // ProcessId=13728 // // CommandLine="C:\Program Files\RustDesk\RustDesk.exe" // ProcessId=10136 let mut pids = Vec::new(); let mut proc_found = false; for line in output.lines() { if line.starts_with("ProcessId=") { if proc_found { if let Ok(pid) = line["ProcessId=".len()..].trim().parse::() { pids.push(hbb_common::sysinfo::Pid::from_u32(pid)); } proc_found = false; } } else if line.starts_with("CommandLine=") { proc_found = false; let cmd = line["CommandLine=".len()..].trim().to_lowercase(); if args.is_empty() { if cmd.ends_with(&name) || cmd.ends_with(&format!("{}\"", &name)) { proc_found = true; } } else { proc_found = args.iter().all(|arg| cmd.contains(arg.as_ref())); } } } pids } // Note the args are not compared strictly, only check if the args are contained in the command line. // If we want to check the args strictly, we need to parse the command line and compare each arg. // Maybe we have to introduce some external crate like `shell_words` to do this. #[cfg(not(target_pointer_width = "64"))] pub(super) fn get_pids_with_args_by_wmic, S2: AsRef>( name: S1, args: &[S2], ) -> Vec { let name = name.as_ref().to_lowercase(); std::process::Command::new("wmic.exe") .args([ "process", "where", &format!("name='{}'", name), "get", "commandline,processid", "/value", ]) .creation_flags(CREATE_NO_WINDOW) .output() .map(|output| { get_pids_with_args_from_wmic_output::( String::from_utf8_lossy(&output.stdout), &name, args, ) }) .unwrap_or_default() } #[cfg(not(target_pointer_width = "64"))] fn get_pids_with_first_arg_from_wmic_output( output: std::borrow::Cow<'_, str>, name: &str, arg: &str, ) -> Vec { let mut pids = Vec::new(); let mut proc_found = false; for line in output.lines() { if line.starts_with("ProcessId=") { if proc_found { if let Ok(pid) = line["ProcessId=".len()..].trim().parse::() { pids.push(hbb_common::sysinfo::Pid::from_u32(pid)); } proc_found = false; } } else if line.starts_with("CommandLine=") { proc_found = false; let cmd = line["CommandLine=".len()..].trim().to_lowercase(); if cmd.is_empty() { continue; } if !arg.is_empty() && cmd.starts_with(arg) { proc_found = true; } else { for x in [&format!("{}\"", name), &format!("{}", name)] { if cmd.contains(x) { let cmd = cmd.split(x).collect::>()[1..].join(""); if arg.is_empty() { if cmd.trim().is_empty() { proc_found = true; } } else if cmd.trim().starts_with(arg) { proc_found = true; } break; } } } } } pids } // Note the args are not compared strictly, only check if the args are contained in the command line. // If we want to check the args strictly, we need to parse the command line and compare each arg. // Maybe we have to introduce some external crate like `shell_words` to do this. #[cfg(not(target_pointer_width = "64"))] pub(super) fn get_pids_with_first_arg_by_wmic, S2: AsRef>( name: S1, arg: S2, ) -> Vec { let name = name.as_ref().to_lowercase(); let arg = arg.as_ref().to_lowercase(); std::process::Command::new("wmic.exe") .args([ "process", "where", &format!("name='{}'", name), "get", "commandline,processid", "/value", ]) .creation_flags(CREATE_NO_WINDOW) .output() .map(|output| { get_pids_with_first_arg_from_wmic_output( String::from_utf8_lossy(&output.stdout), &name, &arg, ) }) .unwrap_or_default() } #[cfg(test)] mod tests { use super::*; #[test] fn test_uninstall_cert() { println!("uninstall driver certs: {:?}", cert::uninstall_cert()); } #[test] fn test_get_unicode_char_by_vk() { let chr = get_char_from_vk(0x41); // VK_A assert_eq!(chr, Some('a')); let chr = get_char_from_vk(VK_ESCAPE as u32); // VK_ESC assert_eq!(chr, None) } #[cfg(not(target_pointer_width = "64"))] #[test] fn test_get_pids_with_args_from_wmic_output() { let output = r#" CommandLine= ProcessId=33796 CommandLine= ProcessId=34668 CommandLine="C:\Program Files\testapp\TestApp.exe" --tray ProcessId=13728 CommandLine="C:\Program Files\testapp\TestApp.exe" ProcessId=10136 "#; let name = "testapp.exe"; let args = vec!["--tray"]; let pids = super::get_pids_with_args_from_wmic_output( String::from_utf8_lossy(output.as_bytes()), name, &args, ); assert_eq!(pids.len(), 1); assert_eq!(pids[0].as_u32(), 13728); let args: Vec<&str> = vec![]; let pids = super::get_pids_with_args_from_wmic_output( String::from_utf8_lossy(output.as_bytes()), name, &args, ); assert_eq!(pids.len(), 1); assert_eq!(pids[0].as_u32(), 10136); let args = vec!["--other"]; let pids = super::get_pids_with_args_from_wmic_output( String::from_utf8_lossy(output.as_bytes()), name, &args, ); assert_eq!(pids.len(), 0); } #[cfg(not(target_pointer_width = "64"))] #[test] fn test_get_pids_with_first_arg_from_wmic_output() { let output = r#" CommandLine= ProcessId=33796 CommandLine= ProcessId=34668 CommandLine="C:\Program Files\testapp\TestApp.exe" --tray ProcessId=13728 CommandLine="C:\Program Files\testapp\TestApp.exe" ProcessId=10136 "#; let name = "testapp.exe"; let arg = "--tray"; let pids = super::get_pids_with_first_arg_from_wmic_output( String::from_utf8_lossy(output.as_bytes()), name, arg, ); assert_eq!(pids.len(), 1); assert_eq!(pids[0].as_u32(), 13728); let arg = ""; let pids = super::get_pids_with_first_arg_from_wmic_output( String::from_utf8_lossy(output.as_bytes()), name, arg, ); assert_eq!(pids.len(), 1); assert_eq!(pids[0].as_u32(), 10136); let arg = "--other"; let pids = super::get_pids_with_first_arg_from_wmic_output( String::from_utf8_lossy(output.as_bytes()), name, arg, ); assert_eq!(pids.len(), 0); } }