rustdesk/src/platform/windows.rs

1748 lines
52 KiB
Rust
Raw Normal View History

2022-05-12 17:35:25 +08:00
use super::{CursorData, ResultType};
use crate::common::PORTABLE_APPNAME_RUNTIME_ENV_KEY;
2022-05-12 17:35:25 +08:00
use crate::ipc;
use crate::license::*;
use hbb_common::{
allow_err, bail,
config::{self, Config},
log, sleep, timeout, tokio,
};
use std::io::prelude::*;
use std::{
ffi::OsString,
fs, io, mem,
path::PathBuf,
2022-05-12 17:35:25 +08:00
sync::{Arc, Mutex},
time::{Duration, Instant},
};
use winapi::{
ctypes::c_void,
2022-05-12 17:35:25 +08:00
shared::{minwindef::*, ntdef::NULL, windef::*},
um::{
errhandlingapi::GetLastError,
handleapi::CloseHandle,
minwinbase::STILL_ACTIVE,
processthreadsapi::{
GetCurrentProcess, GetCurrentProcessId, GetExitCodeProcess, OpenProcess,
OpenProcessToken, PROCESS_INFORMATION, STARTUPINFOW,
},
securitybaseapi::GetTokenInformation,
shellapi::ShellExecuteW,
winbase::*,
wingdi::*,
winnt::{
TokenElevation, HANDLE, PROCESS_QUERY_LIMITED_INFORMATION, TOKEN_ELEVATION, TOKEN_QUERY,
},
winuser::*,
2022-05-12 17:35:25 +08:00
},
};
use windows_service::{
define_windows_service,
service::{
ServiceControl, ServiceControlAccept, ServiceExitCode, ServiceState, ServiceStatus,
ServiceType,
},
service_control_handler::{self, ServiceControlHandlerResult},
};
use winreg::enums::*;
use winreg::RegKey;
pub fn get_cursor_pos() -> Option<(i32, i32)> {
unsafe {
#[allow(invalid_value)]
2022-05-12 17:35:25 +08:00
let mut out = mem::MaybeUninit::uninit().assume_init();
if GetCursorPos(&mut out) == FALSE {
return None;
}
return Some((out.x, out.y));
}
}
pub fn reset_input_cache() {}
pub fn get_cursor() -> ResultType<Option<u64>> {
unsafe {
#[allow(invalid_value)]
2022-05-12 17:35:25 +08:00
let mut ci: CURSORINFO = mem::MaybeUninit::uninit().assume_init();
ci.cbSize = std::mem::size_of::<CURSORINFO>() as _;
if crate::portable_service::client::get_cursor_info(&mut ci) == FALSE {
2022-05-12 17:35:25 +08:00
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<Self> {
unsafe {
#[allow(invalid_value)]
2022-05-12 17:35:25 +08:00
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<CursorData> {
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<u8> = Vec::new();
cbits.resize(cbits_size as _, 0);
let mut mbits: Vec<u8> = 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(),
2022-05-12 17:35:25 +08:00
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<BITMAP> {
unsafe {
let mut bm: BITMAP = mem::zeroed();
if GetObjectA(
handle as _,
std::mem::size_of::<BITMAP>() 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<Self> {
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<Self> {
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<Self> {
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<u8>,
) -> 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: {}", get_error());
}
}
Ok(())
}
fn fix_cursor_mask(
mbits: &mut Vec<u8>,
cbits: &mut Vec<u8>,
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<OsString>) {
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 has_rdp_service() -> BOOL;
fn get_current_session(rdp: BOOL) -> DWORD;
fn LaunchProcessWin(cmd: *const u16, session_id: DWORD, as_user: BOOL) -> HANDLE;
fn GetSessionUserTokenWin(lphUserToken: LPHANDLE, dwSessionId: DWORD, as_user: BOOL) -> BOOL;
2022-05-12 17:35:25 +08:00
fn selectInputDesktop() -> BOOL;
fn inputDesktopSelected() -> BOOL;
fn is_windows_server() -> 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;
2022-12-14 00:51:43 +08:00
fn is_local_system() -> BOOL;
2022-05-12 17:35:25 +08:00
}
extern "system" {
fn BlockInput(v: BOOL) -> BOOL;
}
#[tokio::main(flavor = "current_thread")]
async fn run_service(_arguments: Vec<OsString>) -> 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 => {
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?;
loop {
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();
}
_ => {}
}
}
}
_ => {}
},
Err(_) => {
// timeout
unsafe {
let tmp = get_current_session(share_rdp());
if tmp == 0xFFFFFFFF {
continue;
}
let mut close_sent = false;
if tmp != session_id {
log::info!("session changed from {} to {}", session_id, tmp);
session_id = tmp;
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<HANDLE> {
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("")
);
use std::os::windows::ffi::OsStrExt;
let wstr: Vec<u16> = std::ffi::OsStr::new(&cmd)
.encode_wide()
.chain(Some(0).into_iter())
.collect();
let wstr = wstr.as_ptr();
let h = unsafe { LaunchProcessWin(wstr, session_id, FALSE) };
if h.is_null() {
log::error!("Failed to launch server: {}", get_error());
2022-05-12 17:35:25 +08:00
}
Ok(h)
}
pub fn run_as_user(arg: Vec<&str>) -> ResultType<Option<std::process::Child>> {
2022-05-12 17:35:25 +08:00
let cmd = format!(
"\"{}\" {}",
std::env::current_exe()?.to_str().unwrap_or(""),
arg.join(" "),
2022-05-12 17:35:25 +08:00
);
let session_id = unsafe { get_current_session(share_rdp()) };
use std::os::windows::ffi::OsStrExt;
let wstr: Vec<u16> = std::ffi::OsStr::new(&cmd)
.encode_wide()
.chain(Some(0).into_iter())
.collect();
let wstr = wstr.as_ptr();
let h = unsafe { LaunchProcessWin(wstr, session_id, TRUE) };
if h.is_null() {
bail!(
"Failed to launch {:?} with session id {}: {}",
2022-05-12 17:35:25 +08:00
arg,
session_id,
get_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
fn send_sas() {
#[link(name = "sas")]
extern "system" {
pub fn SendSAS(AsUser: BOOL);
}
unsafe {
log::info!("SAS received");
SendSAS(FALSE);
}
}
lazy_static::lazy_static! {
static ref SUPPRESS: Arc<Mutex<Instant>> = 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: {}", get_error());
*s = Instant::now();
}
} else {
log::info!("Desktop switched");
}
return res;
}
}
return false;
}
fn get_error() -> String {
unsafe {
let buff_size = 256;
let mut buff: Vec<u16> = Vec::with_capacity(buff_size);
buff.resize(buff_size, 0);
let errno = GetLastError();
let chars_copied = FormatMessageW(
FORMAT_MESSAGE_IGNORE_INSERTS
| FORMAT_MESSAGE_FROM_SYSTEM
| FORMAT_MESSAGE_ARGUMENT_ARRAY,
std::ptr::null(),
errno,
0,
buff.as_mut_ptr(),
(buff_size + 1) as u32,
std::ptr::null_mut(),
);
if chars_copied == 0 {
return "".to_owned();
}
let mut curr_char: usize = chars_copied as usize;
while curr_char > 0 {
let ch = buff[curr_char];
if ch >= ' ' as u16 {
break;
}
curr_char -= 1;
}
let sl = std::slice::from_raw_parts(buff.as_ptr(), curr_char);
let err_msg = String::from_utf16(sl);
return err_msg.unwrap_or("".to_owned());
}
}
fn share_rdp() -> BOOL {
if get_reg("share_rdp") != "true" {
FALSE
} else {
TRUE
}
}
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" }
);
2022-06-27 00:36:31 +08:00
run_cmds(cmd, false, "share_rdp").ok();
2022-05-12 17:35:25 +08:00
}
pub fn get_active_username() -> String {
2022-12-14 00:51:43 +08:00
if !is_root() {
return crate::username();
2022-05-12 17:35:25 +08:00
}
2022-12-14 00:51:43 +08:00
2022-05-12 17:35:25 +08:00
extern "C" {
fn get_active_user(path: *mut u16, n: u32, rdp: BOOL) -> u32;
}
let buff_size = 256;
let mut buff: Vec<u16> = 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()
}
pub fn get_active_user_home() -> Option<PathBuf> {
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
}
2022-05-12 17:35:25 +08:00
pub fn is_prelogin() -> bool {
let username = get_active_username();
username.is_empty() || username == "SYSTEM"
}
pub fn is_root() -> bool {
2022-12-14 00:51:43 +08:00
// https://stackoverflow.com/questions/4023586/correct-way-to-find-out-if-a-service-is-running-as-the-system-user
unsafe { is_local_system() == TRUE }
2022-05-12 17:35:25 +08:00
}
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);
}
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;
2022-05-12 17:35:25 +08:00
}
}
#[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 (_, path, _, _) = get_install_info();
let process_exe = crate::ui::win_privacy::INJECTED_PROCESS_EXE;
let origin_process_exe = crate::ui::win_privacy::ORIGIN_PROCESS_EXE;
let exe_file = std::env::current_exe()?;
if exe_file.parent().is_none() {
bail!("Cannot get parent of current exe file");
}
let cur_dir = exe_file.parent().unwrap();
let cur_exe = cur_dir.join(process_exe);
if !std::path::Path::new(&cur_exe).exists() {
std::fs::copy(origin_process_exe, cur_exe)?;
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
);
}
}
}
// Force update broker exe if failed to check modified time.
let cmds = format!(
"
chcp 65001
taskkill /F /IM {broker_exe}
copy /Y \"{origin_process_exe}\" \"{cur_exe}\"
",
broker_exe = process_exe,
origin_process_exe = origin_process_exe,
cur_exe = cur_exe.to_string_lossy().to_string(),
);
2022-06-27 00:36:31 +08:00
run_cmds(cmds, false, "update_broker")?;
Ok(())
}
2022-05-12 17:35:25 +08:00
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_exe_cmd(src_exe: &str, _exe: &str, path: &str) -> String {
#[cfg(feature = "flutter")]
let main_exe = format!(
"XCOPY \"{}\" \"{}\" /Y /E /H /C /I /K /R /Z",
PathBuf::from(src_exe)
.parent()
.unwrap()
.to_string_lossy()
.to_string(),
path
);
#[cfg(not(feature = "flutter"))]
let main_exe = format!(
"copy /Y \"{src_exe}\" \"{exe}\"",
src_exe = src_exe,
exe = _exe
);
return format!(
"
{main_exe}
copy /Y \"{ORIGIN_PROCESS_EXE}\" \"{path}\\{broker_exe}\"
\"{src_exe}\" --extract \"{path}\"
",
main_exe = main_exe,
path = path,
ORIGIN_PROCESS_EXE = crate::ui::win_privacy::ORIGIN_PROCESS_EXE,
broker_exe = crate::ui::win_privacy::INJECTED_PROCESS_EXE,
);
}
2022-05-12 17:35:25 +08:00
pub fn update_me() -> ResultType<()> {
let (_, path, _, exe) = get_install_info();
2022-05-12 17:35:25 +08:00
let src_exe = std::env::current_exe()?.to_str().unwrap_or("").to_owned();
let cmds = format!(
"
chcp 65001
sc stop {app_name}
taskkill /F /IM {broker_exe}
taskkill /F /IM {app_name}.exe /FI \"PID ne {cur_pid}\"
{copy_exe}
2022-05-12 17:35:25 +08:00
sc start {app_name}
{lic}
",
copy_exe = copy_exe_cmd(&src_exe, &exe, &path),
broker_exe = crate::ui::win_privacy::INJECTED_PROCESS_EXE,
2022-05-12 17:35:25 +08:00
app_name = crate::get_app_name(),
lic = register_licence(),
cur_pid = get_current_pid(),
2022-05-12 17:35:25 +08:00
);
std::thread::sleep(std::time::Duration::from_millis(1000));
2022-06-27 00:36:31 +08:00
run_cmds(cmds, false, "update")?;
2022-05-12 17:35:25 +08:00
std::thread::sleep(std::time::Duration::from_millis(2000));
std::process::Command::new(&exe).arg("--tray").spawn().ok();
std::process::Command::new(&exe).spawn().ok();
std::process::Command::new(&exe)
.args(&["--remove", &src_exe])
.spawn()?;
Ok(())
}
fn get_after_install(exe: &str) -> String {
let app_name = crate::get_app_name();
let ext = app_name.to_lowercase();
format!("
chcp 65001
reg add HKEY_CLASSES_ROOT\\.{ext} /f
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\\\"\"
sc create {app_name} binpath= \"\\\"{exe}\\\" --service\" start= auto DisplayName= \"{app_name} Service\"
netsh advfirewall firewall add rule name=\"{app_name} Service\" dir=in action=allow program=\"{exe}\" enable=yes
sc start {app_name}
reg add HKEY_LOCAL_MACHINE\\Software\\Microsoft\\Windows\\CurrentVersion\\Policies\\System /f /v SoftwareSASGeneration /t REG_DWORD /d 1
", ext=ext, exe=exe, app_name=app_name)
}
2022-06-27 00:36:31 +08:00
pub fn install_me(options: &str, path: String, silent: bool, debug: bool) -> ResultType<()> {
2022-05-12 17:35:25 +08:00
let uninstall_str = get_uninstall();
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];
}
2022-06-27 00:36:31 +08:00
let tmp_path = std::env::temp_dir().to_string_lossy().to_string();
2022-05-12 17:35:25 +08:00
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
",
tmp_path = tmp_path,
app_name = crate::get_app_name(),
exe = exe,
),
"vbs",
2022-06-27 00:36:31 +08:00
"mk_shortcut",
2022-05-12 17:35:25 +08:00
)?
.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
",
tmp_path = tmp_path,
app_name = crate::get_app_name(),
exe = exe,
),
"vbs",
2022-06-27 00:36:31 +08:00
"uninstall_shortcut",
2022-05-12 17:35:25 +08:00
)?
.to_str()
.unwrap_or("")
.to_owned();
let tray_shortcut = 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
",
tmp_path = tmp_path,
app_name = crate::get_app_name(),
exe = exe,
),
"vbs",
2022-06-27 00:36:31 +08:00
"tray_shortcut",
2022-05-12 17:35:25 +08:00
)?
.to_str()
.unwrap_or("")
.to_owned();
let mut shortcuts = Default::default();
if options.contains("desktopicon") {
shortcuts = format!(
"copy /Y \"{}\\{}.lnk\" \"%PUBLIC%\\Desktop\\\"",
tmp_path,
crate::get_app_name()
);
}
if options.contains("startmenu") {
shortcuts = format!(
"{}
md \"{start_menu}\"
copy /Y \"{tmp_path}\\{app_name}.lnk\" \"{start_menu}\\\"
copy /Y \"{tmp_path}\\Uninstall {app_name}.lnk\" \"{start_menu}\\\"
",
shortcuts,
start_menu = start_menu,
tmp_path = tmp_path,
app_name = crate::get_app_name(),
);
}
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
2022-07-06 01:33:04 +08:00
// 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\"
",
mk_shortcut = mk_shortcut,
uninstall_shortcut = uninstall_shortcut,
tray_shortcut = tray_shortcut,
tmp_path = tmp_path,
app_name = crate::get_app_name(),
);
let src_exe = std::env::current_exe()?.to_str().unwrap_or("").to_string();
2022-05-12 17:35:25 +08:00
let cmds = format!(
"
{uninstall_str}
chcp 65001
md \"{path}\"
{copy_exe}
2022-05-12 17:35:25 +08:00
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}\"
2022-06-13 19:22:37 +08:00
reg add {subkey} /f /v DisplayVersion /t REG_SZ /d \"{version}\"
2022-05-12 17:35:25 +08:00
reg add {subkey} /f /v Version /t REG_SZ /d \"{version}\"
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 {major}
reg add {subkey} /f /v VersionMinor /t REG_DWORD /d {minor}
reg add {subkey} /f /v VersionBuild /t REG_DWORD /d {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
{lic}
cscript \"{mk_shortcut}\"
cscript \"{uninstall_shortcut}\"
cscript \"{tray_shortcut}\"
2022-06-27 00:36:31 +08:00
copy /Y \"{tmp_path}\\{app_name} Tray.lnk\" \"%PROGRAMDATA%\\Microsoft\\Windows\\Start Menu\\Programs\\Startup\\\"
2022-05-12 17:35:25 +08:00
{shortcuts}
copy /Y \"{tmp_path}\\Uninstall {app_name}.lnk\" \"{path}\\\"
{dels}
2022-05-12 17:35:25 +08:00
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}
{after_install}
{sleep}
2022-05-12 17:35:25 +08:00
",
uninstall_str=uninstall_str,
path=path,
exe=exe,
subkey=subkey,
app_name=crate::get_app_name(),
version=crate::VERSION,
major=version_major,
minor=version_minor,
build=version_build,
size=size,
mk_shortcut=mk_shortcut,
uninstall_shortcut=uninstall_shortcut,
tray_shortcut=tray_shortcut,
tmp_path=tmp_path,
shortcuts=shortcuts,
config_path=Config::file().to_str().unwrap_or(""),
lic=register_licence(),
after_install=get_after_install(&exe),
sleep=if debug {
"timeout 300"
} else {
""
},
dels=if debug {
""
} else {
&dels
},
copy_exe = copy_exe_cmd(&src_exe, &exe, &path),
2022-05-12 17:35:25 +08:00
);
2022-06-27 00:36:31 +08:00
run_cmds(cmds, debug, "install")?;
2022-05-12 17:35:25 +08:00
std::thread::sleep(std::time::Duration::from_millis(2000));
if !silent {
std::process::Command::new(&exe).spawn()?;
std::process::Command::new(&exe).arg("--tray").spawn()?;
std::thread::sleep(std::time::Duration::from_millis(1000));
}
2022-05-12 17:35:25 +08:00
Ok(())
}
pub fn run_after_install() -> ResultType<()> {
let (_, _, _, exe) = get_install_info();
2022-06-27 00:36:31 +08:00
run_cmds(get_after_install(&exe), true, "after_install")
2022-05-12 17:35:25 +08:00
}
pub fn run_before_uninstall() -> ResultType<()> {
2022-06-27 00:36:31 +08:00
run_cmds(get_before_uninstall(), true, "before_install")
2022-05-12 17:35:25 +08:00
}
fn get_before_uninstall() -> String {
let app_name = crate::get_app_name();
let ext = app_name.to_lowercase();
format!(
"
chcp 65001
sc stop {app_name}
sc delete {app_name}
taskkill /F /IM {broker_exe}
taskkill /F /IM {app_name}.exe /FI \"PID ne {cur_pid}\"
2022-05-12 17:35:25 +08:00
reg delete HKEY_CLASSES_ROOT\\.{ext} /f
netsh advfirewall firewall delete rule name=\"{app_name} Service\"
",
app_name = app_name,
broker_exe = crate::ui::win_privacy::INJECTED_PROCESS_EXE,
ext = ext,
cur_pid = get_current_pid(),
2022-05-12 17:35:25 +08:00
)
}
fn get_uninstall() -> String {
let (subkey, path, start_menu, _) = get_install_info();
format!(
"
{before_uninstall}
reg delete {subkey} /f
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\"
2022-05-12 17:35:25 +08:00
",
before_uninstall=get_before_uninstall(),
subkey=subkey,
app_name = crate::get_app_name(),
path = path,
start_menu = start_menu,
)
}
pub fn uninstall_me() -> ResultType<()> {
2022-06-27 00:36:31 +08:00
run_cmds(get_uninstall(), true, "uninstall")
2022-05-12 17:35:25 +08:00
}
2022-06-27 00:36:31 +08:00
fn write_cmds(cmds: String, ext: &str, tip: &str) -> ResultType<std::path::PathBuf> {
2022-05-12 17:35:25 +08:00
let mut tmp = std::env::temp_dir();
2022-06-27 00:36:31 +08:00
tmp.push(format!("{}_{}.{}", crate::get_app_name(), tip, ext));
2022-05-12 17:35:25 +08:00
let mut file = std::fs::File::create(&tmp)?;
// 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
let cmds = cmds.replace("\r\n", "\n").replace("\n", "\r\n");
if ext == "vbs" {
let mut v: Vec<u16> = 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 }
}
2022-06-27 00:36:31 +08:00
fn run_cmds(cmds: String, show: bool, tip: &str) -> ResultType<()> {
let tmp = write_cmds(cmds, "bat", tip)?;
let tmp_fn = tmp.to_str().unwrap_or("");
let res = runas::Command::new("cmd")
.args(&["/C", &tmp_fn])
2022-05-12 17:35:25 +08:00
.show(show)
.force_prompt(true)
.status();
if !show {
allow_err!(std::fs::remove_file(tmp));
}
2022-05-12 17:35:25 +08:00
let _ = res?;
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 {
let v = if v { TRUE } else { FALSE };
unsafe { BlockInput(v) == TRUE }
}
pub fn add_recent_document(path: &str) {
extern "C" {
fn AddRecentDocument(path: *const u16);
}
use std::os::windows::ffi::OsStrExt;
let wstr: Vec<u16> = 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 {
use windows_service::{
service::ServiceAccess,
service_manager::{ServiceManager, ServiceManagerAccess},
};
let (_, _, _, exe) = get_install_info();
if !std::fs::metadata(exe).is_ok() {
return false;
}
let manager_access = ServiceManagerAccess::CONNECT;
if let Ok(service_manager) = ServiceManager::local_computer(None::<&str>, manager_access) {
if let Ok(_) =
service_manager.open_service(crate::get_app_name(), ServiceAccess::QUERY_CONFIG)
{
return true;
}
}
return false;
}
pub fn get_installed_version() -> String {
let (_, _, _, exe) = get_install_info();
if let Ok(output) = std::process::Command::new(exe).arg("--version").output() {
for line in String::from_utf8_lossy(&output.stdout).lines() {
return line.to_owned();
}
}
"".to_owned()
}
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()
}
fn get_license_from_exe_name() -> ResultType<License> {
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;
log::debug!("update portable executable name to {}", exe);
}
2022-05-15 00:25:58 +08:00
get_license_from_string(&exe)
2022-05-12 17:35:25 +08:00
}
#[inline]
pub fn is_win_server() -> bool {
unsafe { is_windows_server() > 0 }
}
pub fn get_license() -> Option<License> {
let mut lic: License = Default::default();
if let Ok(tmp) = get_license_from_exe_name() {
lic = tmp;
} else {
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 fn bootstrap() {
if let Some(lic) = get_license() {
*config::PROD_RENDEZVOUS_SERVER.write().unwrap() = lic.host.clone();
}
}
fn register_licence() -> String {
let (subkey, _, _, _) = get_install_info();
if let Ok(lic) = get_license_from_exe_name() {
format!(
"
reg add {subkey} /f /v Key /t REG_SZ /d \"{key}\"
reg add {subkey} /f /v Host /t REG_SZ /d \"{host}\"
reg add {subkey} /f /v Api /t REG_SZ /d \"{api}\"
",
subkey = subkey,
key = &lic.key,
host = &lic.host,
api = &lic.api,
)
} else {
"".to_owned()
}
}
pub fn is_rdp_service_open() -> bool {
unsafe { has_rdp_service() == TRUE }
}
pub fn create_shortcut(id: &str) -> ResultType<()> {
let exe = std::env::current_exe()?.to_str().unwrap_or("").to_owned();
let shortcut = write_cmds(
format!(
"
Set oWS = WScript.CreateObject(\"WScript.Shell\")
strDesktop = oWS.SpecialFolders(\"Desktop\")
Set objFSO = CreateObject(\"Scripting.FileSystemObject\")
sLinkFile = objFSO.BuildPath(strDesktop, \"{id}.lnk\")
Set oLink = oWS.CreateShortcut(sLinkFile)
oLink.TargetPath = \"{exe}\"
oLink.Arguments = \"--connect {id}\"
oLink.Save
",
exe = exe,
id = id,
),
"vbs",
2022-06-27 00:36:31 +08:00
"connect_shortcut",
2022-05-12 17:35:25 +08:00
)?
.to_str()
.unwrap_or("")
.to_owned();
std::process::Command::new("cscript")
.arg(&shortcut)
.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() {
2022-05-18 16:12:50 +08:00
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 {
if FALSE
== GetSessionUserTokenWin(
&mut token as _,
session_id,
if as_user { TRUE } else { FALSE },
)
{
NULL as _
} else {
token
}
}
}
pub fn run_background(exe: &str, arg: &str) -> ResultType<bool> {
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<bool> {
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<bool> {
run_uac(
std::env::current_exe()?
.to_string_lossy()
.to_string()
.as_str(),
"--version",
)
}
pub fn elevate(arg: &str) -> ResultType<bool> {
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 {
unsafe {
log::error!("Failed to run as system, errno={}", GetLastError());
}
}
}
} else {
if !is_elevate {
if let Ok(true) = elevate(arg_elevate) {
std::process::exit(0);
} else {
unsafe {
log::error!("Failed to elevate, errno={}", GetLastError());
}
}
}
}
}
Err(_) => unsafe {
log::error!("Failed to get elevation status, errno={}", GetLastError());
},
}
}
}
// https://github.com/mgostIH/process_list/blob/master/src/windows/mod.rs
#[repr(transparent)]
pub(self) struct RAIIHandle(pub HANDLE);
impl Drop for RAIIHandle {
fn drop(&mut self) {
// This never gives problem except when running under a debugger.
unsafe { CloseHandle(self.0) };
}
}
pub fn is_elevated(process_id: Option<DWORD>) -> ResultType<bool> {
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, errno {}", GetLastError())
}
let _handle = RAIIHandle(handle);
let mut token: HANDLE = mem::zeroed();
if OpenProcessToken(handle, TOKEN_QUERY, &mut token) == FALSE {
bail!("Failed to open process token, errno {}", GetLastError())
}
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::<TOKEN_ELEVATION>() as _,
&mut size,
) == FALSE
{
bail!("Failed to get token information, errno {}", GetLastError())
}
Ok(token_elevation.TokenIsElevated != 0)
}
}
pub fn is_foreground_window_elevated() -> ResultType<bool> {
unsafe {
let mut process_id: DWORD = 0;
GetWindowThreadProcessId(GetForegroundWindow(), &mut process_id);
if process_id == 0 {
bail!("Failed to get processId, errno {}", GetLastError())
}
is_elevated(Some(process_id))
}
}
fn get_current_pid() -> u32 {
unsafe { GetCurrentProcessId() }
}
pub fn get_double_click_time() -> u32 {
unsafe { GetDoubleClickTime() }
}
fn wide_string(s: &str) -> Vec<u16> {
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 create_process_with_logon(user: &str, pwd: &str, exe: &str, arg: &str) -> ResultType<()> {
unsafe {
let wuser = wide_string(user);
let wpc = wide_string("");
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::<STARTUPINFOW>() 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,
)
{
bail!("CreateProcessWithLogonW failed, errno={}", GetLastError());
}
}
return Ok(());
}