feat, update, win, macos (#11618)

Signed-off-by: fufesou <linlong1266@gmail.com>

src/platform/macos.rs

@@ -27,12 +27,18 @@ use include_dir::{include_dir, Dir};
 use objc::rc::autoreleasepool;
 use objc::{class, msg_send, sel, sel_impl};
 use scrap::{libc::c_void, quartz::ffi::*};
-use std::path::{Path, PathBuf};
+use std::{
+    os::unix::process::CommandExt,
+    path::{Path, PathBuf},
+    process::{Command, Stdio},
+};
 
 static PRIVILEGES_SCRIPTS_DIR: Dir =
     include_dir!("$CARGO_MANIFEST_DIR/src/platform/privileges_scripts");
 static mut LATEST_SEED: i32 = 0;
 
+const UPDATE_TEMP_DIR: &str = "/tmp/.rustdeskupdate";
+
 extern "C" {
     fn CGSCurrentCursorSeed() -> i32;
     fn CGEventCreate(r: *const c_void) -> *const c_void;
@@ -155,6 +161,9 @@ pub fn install_service() -> bool {
     is_installed_daemon(false)
 }
 
+// Remember to check if `update_daemon_agent()` need to be changed if changing `is_installed_daemon()`.
+// 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.
 pub fn is_installed_daemon(prompt: bool) -> bool {
     let daemon = format!("{}_service.plist", crate::get_full_name());
     let agent = format!("{}_server.plist", crate::get_full_name());
@@ -218,6 +227,70 @@ pub fn is_installed_daemon(prompt: bool) -> bool {
     false
 }
 
+fn update_daemon_agent(agent_plist_file: String, update_source_dir: String, sync: bool) {
+    let update_script_file = "update.scpt";
+    let Some(update_script) = PRIVILEGES_SCRIPTS_DIR.get_file(update_script_file) else {
+        return;
+    };
+    let Some(update_script_body) = update_script.contents_utf8().map(correct_app_name) else {
+        return;
+    };
+
+    let Some(daemon_plist) = PRIVILEGES_SCRIPTS_DIR.get_file("daemon.plist") else {
+        return;
+    };
+    let Some(daemon_plist_body) = daemon_plist.contents_utf8().map(correct_app_name) else {
+        return;
+    };
+    let Some(agent_plist) = PRIVILEGES_SCRIPTS_DIR.get_file("agent.plist") else {
+        return;
+    };
+    let Some(agent_plist_body) = agent_plist.contents_utf8().map(correct_app_name) else {
+        return;
+    };
+
+    let func = move || {
+        let mut binding = std::process::Command::new("osascript");
+        let mut cmd = binding
+            .arg("-e")
+            .arg(update_script_body)
+            .arg(daemon_plist_body)
+            .arg(agent_plist_body)
+            .arg(&get_active_username())
+            .arg(std::process::id().to_string())
+            .arg(update_source_dir);
+        match cmd.status() {
+            Err(e) => {
+                log::error!("run osascript failed: {}", e);
+            }
+            _ => {
+                let installed = std::path::Path::new(&agent_plist_file).exists();
+                log::info!("Agent file {} installed: {}", &agent_plist_file, installed);
+                if installed {
+                    // Unload first, or load may not work if already loaded.
+                    // We hope that the load operation can immediately trigger a start.
+                    std::process::Command::new("launchctl")
+                        .args(&["unload", "-w", &agent_plist_file])
+                        .stdin(Stdio::null())
+                        .stdout(Stdio::null())
+                        .stderr(Stdio::null())
+                        .status()
+                        .ok();
+                    let status = std::process::Command::new("launchctl")
+                        .args(&["load", "-w", &agent_plist_file])
+                        .status();
+                    log::info!("launch server, status: {:?}", &status);
+                }
+            }
+        }
+    };
+    if sync {
+        func();
+    } else {
+        std::thread::spawn(func);
+    }
+}
+
 fn correct_app_name(s: &str) -> String {
     let s = s.replace("rustdesk", &crate::get_app_name().to_lowercase());
     let s = s.replace("RustDesk", &crate::get_app_name());
@@ -634,6 +707,140 @@ pub fn quit_gui() {
     };
 }
 
+pub fn update_me() -> ResultType<()> {
+    let is_installed_daemon = is_installed_daemon(false);
+    let option_stop_service = "stop-service";
+    let is_service_stopped = hbb_common::config::option2bool(
+        option_stop_service,
+        &crate::ui_interface::get_option(option_stop_service),
+    );
+
+    let cmd = std::env::current_exe()?;
+    // RustDesk.app/Contents/MacOS/RustDesk
+    let app_dir = cmd
+        .parent()
+        .and_then(|p| p.parent())
+        .and_then(|p| p.parent())
+        .map(|d| d.to_string_lossy().to_string());
+    let Some(app_dir) = app_dir else {
+        bail!("Unknown app directory of current exe file: {:?}", cmd);
+    };
+
+    if is_installed_daemon && !is_service_stopped {
+        let agent = format!("{}_server.plist", crate::get_full_name());
+        let agent_plist_file = format!("/Library/LaunchAgents/{}", agent);
+        std::process::Command::new("launchctl")
+            .args(&["unload", "-w", &agent_plist_file])
+            .stdin(Stdio::null())
+            .stdout(Stdio::null())
+            .stderr(Stdio::null())
+            .status()
+            .ok();
+        update_daemon_agent(agent_plist_file, app_dir, true);
+    } else {
+        // `kill -9` may not work without "administrator privileges"
+        let update_body = format!(
+            r#"
+do shell script "
+pgrep -x 'RustDesk' | grep -v {} | xargs kill -9 && rm -rf /Applications/RustDesk.app && cp -R '{}' /Applications/ && chown -R {}:staff /Applications/RustDesk.app
+" with prompt "RustDesk wants to update itself" with administrator privileges
+    "#,
+            std::process::id(),
+            app_dir,
+            get_active_username()
+        );
+        match Command::new("osascript")
+            .arg("-e")
+            .arg(update_body)
+            .status()
+        {
+            Ok(status) if !status.success() => {
+                log::error!("osascript execution failed with status: {}", status);
+            }
+            Err(e) => {
+                log::error!("run osascript failed: {}", e);
+            }
+            _ => {}
+        }
+    }
+    std::process::Command::new("open")
+        .arg("-n")
+        .arg(&format!("/Applications/{}.app", crate::get_app_name()))
+        .spawn()
+        .ok();
+    // leave open a little time
+    std::thread::sleep(std::time::Duration::from_millis(300));
+    Ok(())
+}
+
+pub fn update_to(file: &str) -> ResultType<()> {
+    extract_dmg(file, UPDATE_TEMP_DIR)?;
+    update_extracted(UPDATE_TEMP_DIR)?;
+    Ok(())
+}
+
+fn extract_dmg(dmg_path: &str, target_dir: &str) -> ResultType<()> {
+    let mount_point = "/Volumes/RustDeskUpdate";
+    let target_path = Path::new(target_dir);
+
+    if target_path.exists() {
+        std::fs::remove_dir_all(target_path)?;
+    }
+    std::fs::create_dir_all(target_path)?;
+
+    Command::new("hdiutil")
+        .args(&["attach", "-nobrowse", "-mountpoint", mount_point, dmg_path])
+        .status()?;
+
+    struct DmgGuard(&'static str);
+    impl Drop for DmgGuard {
+        fn drop(&mut self) {
+            let _ = Command::new("hdiutil")
+                .args(&["detach", self.0, "-force"])
+                .status();
+        }
+    }
+    let _guard = DmgGuard(mount_point);
+
+    let app_name = "RustDesk.app";
+    let src_path = format!("{}/{}", mount_point, app_name);
+    let dest_path = format!("{}/{}", target_dir, app_name);
+
+    let copy_status = Command::new("cp")
+        .args(&["-R", &src_path, &dest_path])
+        .status()?;
+
+    if !copy_status.success() {
+        bail!("Failed to copy application {:?}", copy_status);
+    }
+
+    if !Path::new(&dest_path).exists() {
+        bail!(
+            "Copy operation failed - destination not found at {}",
+            dest_path
+        );
+    }
+
+    Ok(())
+}
+
+fn update_extracted(target_dir: &str) -> ResultType<()> {
+    let exe_path = format!("{}/RustDesk.app/Contents/MacOS/RustDesk", target_dir);
+    let _child = unsafe {
+        Command::new(&exe_path)
+            .arg("--update")
+            .stdin(Stdio::null())
+            .stdout(Stdio::null())
+            .stderr(Stdio::null())
+            .pre_exec(|| {
+                hbb_common::libc::setsid();
+                Ok(())
+            })
+            .spawn()?
+    };
+    Ok(())
+}
+
 pub fn get_double_click_time() -> u32 {
     // to-do: https://github.com/servo/core-foundation-rs/blob/786895643140fa0ee4f913d7b4aeb0c4626b2085/cocoa/src/appkit.rs#L2823
     500 as _

src/platform/mod.rs

@@ -27,7 +27,11 @@ pub mod linux_desktop_manager;
 pub mod gtk_sudo;
 
 #[cfg(not(any(target_os = "android", target_os = "ios")))]
-use hbb_common::{message_proto::CursorData, ResultType};
+use hbb_common::{
+    message_proto::CursorData,
+    sysinfo::{Pid, System},
+    ResultType,
+};
 use std::sync::{Arc, Mutex};
 #[cfg(not(any(target_os = "macos", target_os = "android", target_os = "ios")))]
 pub const SERVICE_INTERVAL: u64 = 300;
@@ -137,6 +141,71 @@ pub fn is_prelogin() -> bool {
     false
 }
 
+// Note: This method is inefficient on Windows. It will get all the processes.
+// It should only be called when performance is not critical.
+// If we wanted to get the command line ourselves, there would be a lot of new code.
+#[cfg(not(any(target_os = "android", target_os = "ios")))]
+fn get_pids_of_process_with_args<S1: AsRef<str>, S2: AsRef<str>>(
+    name: S1,
+    args: &[S2],
+) -> Vec<Pid> {
+    // This function does not work when the process is 32-bit and the OS is 64-bit Windows,
+    // `process.cmd()` always returns [] in this case.
+    // So we use `windows::get_pids_with_args_by_wmic()` instead.
+    #[cfg(all(target_os = "windows", not(target_pointer_width = "64")))]
+    {
+        return windows::get_pids_with_args_by_wmic(name, args);
+    }
+    #[cfg(not(all(target_os = "windows", not(target_pointer_width = "64"))))]
+    {
+        let name = name.as_ref().to_lowercase();
+        let system = System::new_all();
+        system
+            .processes()
+            .iter()
+            .filter(|(_, process)| {
+                process.name().to_lowercase() == name
+                    && process.cmd().len() == args.len() + 1
+                    && args.iter().enumerate().all(|(i, arg)| {
+                        process.cmd()[i + 1].to_lowercase() == arg.as_ref().to_lowercase()
+                    })
+            })
+            .map(|(&pid, _)| pid)
+            .collect()
+    }
+}
+
+// Note: This method is inefficient on Windows. It will get all the processes.
+// It should only be called when performance is not critical.
+#[cfg(not(any(target_os = "android", target_os = "ios")))]
+pub fn get_pids_of_process_with_first_arg<S1: AsRef<str>, S2: AsRef<str>>(
+    name: S1,
+    arg: S2,
+) -> Vec<Pid> {
+    // This function does not work when the process is 32-bit and the OS is 64-bit Windows,
+    // `process.cmd()` always returns [] in this case.
+    // So we use `windows::get_pids_with_first_arg_by_wmic()` instead.
+    #[cfg(all(target_os = "windows", not(target_pointer_width = "64")))]
+    {
+        return windows::get_pids_with_first_arg_by_wmic(name, arg);
+    }
+    #[cfg(not(all(target_os = "windows", not(target_pointer_width = "64"))))]
+    {
+        let name = name.as_ref().to_lowercase();
+        let system = System::new_all();
+        system
+            .processes()
+            .iter()
+            .filter(|(_, process)| {
+                process.name().to_lowercase() == name
+                    && process.cmd().len() >= 2
+                    && process.cmd()[1].to_lowercase() == arg.as_ref().to_lowercase()
+            })
+            .map(|(&pid, _)| pid)
+            .collect()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/platform/privileges_scripts/install.scpt

@@ -12,5 +12,5 @@ on run {daemon_file, agent_file, user}
 
   set sh to sh1 & sh2 & sh3 & sh4 & sh5
 
-  do shell script sh with prompt "RustDesk want to install daemon and agent" with administrator privileges
+  do shell script sh with prompt "RustDesk wants to install daemon and agent" with administrator privileges
 end run

src/platform/privileges_scripts/uninstall.scpt

@@ -3,4 +3,4 @@ set sh2 to "/bin/rm /Library/LaunchDaemons/com.carriez.RustDesk_service.plist;"
 set sh3 to "/bin/rm /Library/LaunchAgents/com.carriez.RustDesk_server.plist;"
 
 set sh to sh1 & sh2 & sh3
-do shell script sh with prompt "RustDesk want to unload daemon" with administrator privileges
+do shell script sh with prompt "RustDesk wants to unload daemon" with administrator privileges

src/platform/privileges_scripts/update.scpt

@@ -0,0 +1,18 @@
+on run {daemon_file, agent_file, user, cur_pid, source_dir}
+
+  set unload_service to "launchctl unload -w /Library/LaunchDaemons/com.carriez.RustDesk_service.plist || true;"
+
+  set kill_others to "pgrep -x 'RustDesk' | grep -v " & cur_pid & " | xargs kill -9;"
+
+  set copy_files to "rm -rf /Applications/RustDesk.app && cp -r " & source_dir & " /Applications && chown -R " & quoted form of user & ":staff /Applications/RustDesk.app;"
+
+  set sh1 to "echo " & quoted form of daemon_file & " > /Library/LaunchDaemons/com.carriez.RustDesk_service.plist && chown root:wheel /Library/LaunchDaemons/com.carriez.RustDesk_service.plist;"
+
+  set sh2 to "echo " & quoted form of agent_file & " > /Library/LaunchAgents/com.carriez.RustDesk_server.plist && chown root:wheel /Library/LaunchAgents/com.carriez.RustDesk_server.plist;"
+
+  set sh3 to "launchctl load -w /Library/LaunchDaemons/com.carriez.RustDesk_service.plist;"
+
+  set sh to unload_service & kill_others & copy_files & sh1 & sh2 & sh3
+
+  do shell script sh with prompt "RustDesk wants to update itself" with administrator privileges
+end run

src/platform/windows.cc

@@ -230,7 +230,7 @@ extern "C"
         return IsWindows10OrGreater();
     }
 
-    HANDLE LaunchProcessWin(LPCWSTR cmd, DWORD dwSessionId, BOOL as_user, DWORD *pDwTokenPid)
+    HANDLE LaunchProcessWin(LPCWSTR cmd, DWORD dwSessionId, BOOL as_user, BOOL show, DWORD *pDwTokenPid)
     {
         HANDLE hProcess = NULL;
         HANDLE hToken = NULL;
@@ -240,6 +240,11 @@ extern "C"
             ZeroMemory(&si, sizeof si);
             si.cb = sizeof si;
             si.dwFlags = STARTF_USESHOWWINDOW;
+            if (show)
+            {
+                si.lpDesktop = (LPWSTR)L"winsta0\\default";
+                si.wShowWindow = SW_SHOW;
+            }
             wchar_t buf[MAX_PATH];
             wcscpy_s(buf, sizeof(buf), cmd);
             PROCESS_INFORMATION pi;

src/platform/windows.rs

@@ -13,7 +13,9 @@ use hbb_common::{
     libc::{c_int, wchar_t},
     log,
     message_proto::{DisplayInfo, Resolution, WindowsSession},
-    sleep, timeout, tokio,
+    sleep,
+    sysinfo::{Pid, System},
+    timeout, tokio,
 };
 use std::{
     collections::HashMap,
@@ -484,6 +486,7 @@ extern "C" {
         cmd: *const u16,
         session_id: DWORD,
         as_user: BOOL,
+        show: BOOL,
         token_pid: &mut DWORD,
     ) -> HANDLE;
     fn GetSessionUserTokenWin(
@@ -669,6 +672,10 @@ async fn launch_server(session_id: DWORD, close_first: bool) -> ResultType<HANDL
         "\"{}\" --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<HANDLE> {
     use std::os::windows::ffi::OsStrExt;
     let wstr: Vec<u16> = std::ffi::OsStr::new(&cmd)
         .encode_wide()
@@ -676,9 +683,12 @@ async fn launch_server(session_id: DWORD, close_first: bool) -> ResultType<HANDL
         .collect();
     let wstr = wstr.as_ptr();
     let mut token_pid = 0;
-    let h = unsafe { LaunchProcessWin(wstr, session_id, FALSE, &mut token_pid) };
+    let h = unsafe { LaunchProcessWin(wstr, session_id, FALSE, FALSE, &mut token_pid) };
     if h.is_null() {
-        log::error!("Failed to launch server: {}", io::Error::last_os_error());
+        log::error!(
+            "Failed to launch privileged process: {}",
+            io::Error::last_os_error()
+        );
         if token_pid == 0 {
             log::error!("No process winlogon.exe");
         }
@@ -687,22 +697,43 @@ async fn launch_server(session_id: DWORD, close_first: bool) -> ResultType<HANDL
 }
 
 pub fn run_as_user(arg: Vec<&str>) -> ResultType<Option<std::process::Child>> {
-    let cmd = format!(
-        "\"{}\" {}",
-        std::env::current_exe()?.to_str().unwrap_or(""),
-        arg.join(" "),
-    );
+    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<Option<std::process::Child>> {
     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<Option<std::process::Child>> {
     use std::os::windows::ffi::OsStrExt;
+    let cmd = format!("\"{}\" {}", exe, arg.join(" "),);
     let wstr: Vec<u16> = 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, &mut token_pid) };
+    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!(
@@ -800,8 +831,12 @@ pub fn set_share_rdp(enable: bool) {
 }
 
 pub fn get_current_process_session_id() -> Option<u32> {
+    get_session_id_of_process(unsafe { GetCurrentProcessId() })
+}
+
+pub fn get_session_id_of_process(pid: DWORD) -> Option<u32> {
     let mut sid = 0;
-    if unsafe { ProcessIdToSessionId(GetCurrentProcessId(), &mut sid) == TRUE } {
+    if unsafe { ProcessIdToSessionId(pid, &mut sid) == TRUE } {
         Some(sid)
     } else {
         None
@@ -1348,6 +1383,9 @@ copy /Y \"{tmp_path}\\{app_name} Tray.lnk\" \"%PROGRAMDATA%\\Microsoft\\Windows\
 ")
     };
 
+    // 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}
@@ -2366,6 +2404,171 @@ if exist \"{tray_shortcut}\" del /f /q \"{tray_shortcut}\"
     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::<Vec<_>>();
+    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::<Vec<_>>();
+    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()
+    };
+    // 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}
+{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<Pid>) -> 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<String> {
     Ok(write_cmds(
         format!(
@@ -2450,23 +2653,6 @@ pub fn try_kill_broker() {
         .spawn());
 }
 
-#[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)
-    }
-}
-
 pub fn message_box(text: &str) {
     let mut text = text.to_owned();
     let nodialog = std::env::var("NO_DIALOG").unwrap_or_default() == "Y";
@@ -2974,3 +3160,264 @@ fn get_pids<S: AsRef<str>>(name: S) -> ResultType<Vec<u32>> {
 
     Ok(pids)
 }
+
+pub fn is_msi_installed() -> std::io::Result<bool> {
+    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::<u32, _>("WindowsInstaller")?)
+}
+
+#[cfg(not(target_pointer_width = "64"))]
+fn get_pids_with_args_from_wmic_output<S2: AsRef<str>>(
+    output: std::borrow::Cow<'_, str>,
+    name: &str,
+    args: &[S2],
+) -> Vec<hbb_common::sysinfo::Pid> {
+    // 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::<u32>() {
+                    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<S1: AsRef<str>, S2: AsRef<str>>(
+    name: S1,
+    args: &[S2],
+) -> Vec<hbb_common::sysinfo::Pid> {
+    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::<S2>(
+                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<hbb_common::sysinfo::Pid> {
+    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::<u32>() {
+                    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::<Vec<_>>()[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<S1: AsRef<str>, S2: AsRef<str>>(
+    name: S1,
+    arg: S2,
+) -> Vec<hbb_common::sysinfo::Pid> {
+    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);
+    }
+}