feat: windows, custom client, update (#13687)

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

src/platform/windows.rs

@@ -153,11 +153,7 @@ pub fn clip_cursor(rect: Option<(i32, i32, i32, i32)>) -> bool {
         };
         if result == FALSE {
             let err = GetLastError();
-            log::warn!(
-                "ClipCursor failed: rect={:?}, error_code={}",
-                rect,
-                err
-            );
+            log::warn!("ClipCursor failed: rect={:?}, error_code={}", rect, err);
             return false;
         }
         true
@@ -757,15 +753,37 @@ pub fn run_as_user(arg: Vec<&str>) -> ResultType<Option<std::process::Child>> {
     run_exe_in_cur_session(std::env::current_exe()?.to_str().unwrap_or(""), arg, false)
 }
 
+pub fn run_exe_direct(
+    exe: &str,
+    arg: Vec<&str>,
+    show: bool,
+) -> ResultType<Option<std::process::Child>> {
+    let mut cmd = std::process::Command::new(exe);
+    for a in arg {
+        cmd.arg(a);
+    }
+    if !show {
+        cmd.creation_flags(CREATE_NO_WINDOW);
+    }
+    match cmd.spawn() {
+        Ok(child) => Ok(Some(child)),
+        Err(e) => bail!("Failed to start process: {}", e),
+    }
+}
+
 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)
+    if is_root() {
+        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)
+    } else {
+        run_exe_direct(exe, arg, show)
+    }
 }
 
 pub fn run_exe_in_session(
@@ -1331,6 +1349,38 @@ pub fn copy_exe_cmd(src_exe: &str, exe: &str, path: &str) -> ResultType<String>
     ))
 }
 
+#[inline]
+pub fn rename_exe_cmd(src_exe: &str, path: &str) -> ResultType<String> {
+    let src_exe_filename = PathBuf::from(src_exe)
+        .file_name()
+        .ok_or(anyhow!("Can't get file name of {src_exe}"))?
+        .to_string_lossy()
+        .to_string();
+    let app_name = crate::get_app_name().to_lowercase();
+    if src_exe_filename.to_lowercase() == format!("{app_name}.exe") {
+        Ok("".to_owned())
+    } else {
+        Ok(format!(
+            "
+        move /Y \"{path}\\{src_exe_filename}\" \"{path}\\{app_name}.exe\"
+        ",
+        ))
+    }
+}
+
+#[inline]
+pub fn remove_meta_toml_cmd(is_msi: bool, path: &str) -> String {
+    if is_msi && crate::is_custom_client() {
+        format!(
+            "
+        del /F /Q \"{path}\\meta.toml\"
+        ",
+        )
+    } else {
+        "".to_owned()
+    }
+}
+
 fn get_after_install(
     exe: &str,
     reg_value_start_menu_shortcuts: Option<String>,
@@ -1417,7 +1467,11 @@ pub fn install_me(options: &str, path: String, silent: bool, debug: bool) -> Res
     }
     let app_name = crate::get_app_name();
 
+    let current_exe = std::env::current_exe()?;
+
     let tmp_path = std::env::temp_dir().to_string_lossy().to_string();
+    let cur_exe = current_exe.to_str().unwrap_or("").to_owned();
+    let shortcut_icon_location = get_shortcut_icon_location(&cur_exe);
     let mk_shortcut = write_cmds(
         format!(
             "
@@ -1426,6 +1480,7 @@ sLinkFile = \"{tmp_path}\\{app_name}.lnk\"
 
 Set oLink = oWS.CreateShortcut(sLinkFile)
     oLink.TargetPath = \"{exe}\"
+    {shortcut_icon_location}
 oLink.Save
         "
         ),
@@ -1482,8 +1537,13 @@ copy /Y \"{tmp_path}\\Uninstall {app_name}.lnk\" \"{start_menu}\\\"
         reg_value_printer = "1".to_owned();
     }
 
-    let meta = std::fs::symlink_metadata(std::env::current_exe()?)?;
-    let size = meta.len() / 1024;
+    let meta = std::fs::symlink_metadata(&current_exe)?;
+    let mut size = meta.len() / 1024;
+    if let Some(parent_dir) = current_exe.parent() {
+        if let Some(d) = parent_dir.to_str() {
+            size = get_directory_size_kb(d);
+        }
+    }
     // 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
@@ -1536,7 +1596,7 @@ 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 DisplayIcon /t REG_SZ /d \"{display_icon}\"
 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}\"
@@ -1560,6 +1620,7 @@ copy /Y \"{tmp_path}\\Uninstall {app_name}.lnk\" \"{path}\\\"
 {install_remote_printer}
 {sleep}
     ",
+        display_icon = get_custom_icon(&cur_exe).unwrap_or(exe.to_string()),
         version = crate::VERSION.replace("-", "."),
         build_date = crate::BUILD_DATE,
         after_install = get_after_install(
@@ -1795,6 +1856,163 @@ fn get_reg_of(subkey: &str, name: &str) -> String {
     "".to_owned()
 }
 
+fn get_public_base_dir() -> PathBuf {
+    if let Ok(allusersprofile) = std::env::var("ALLUSERSPROFILE") {
+        let path = PathBuf::from(&allusersprofile);
+        if path.exists() {
+            return path;
+        }
+    }
+    if let Ok(public) = std::env::var("PUBLIC") {
+        let path = PathBuf::from(public).join("Documents");
+        if path.exists() {
+            return path;
+        }
+    }
+    let program_data_dir = PathBuf::from("C:\\ProgramData");
+    if program_data_dir.exists() {
+        return program_data_dir;
+    }
+    std::env::temp_dir()
+}
+
+#[inline]
+pub fn get_custom_client_staging_dir() -> PathBuf {
+    get_public_base_dir()
+        .join("RustDesk")
+        .join("RustDeskCustomClientStaging")
+}
+
+/// Removes the custom client staging directory.
+///
+/// Current behavior: intentionally a no-op (does not delete).
+///
+/// Rationale
+/// - The staging directory only contains a small `custom.txt`, leaving it is harmless.
+/// - Deleting directories under a public location (e.g., C:\\ProgramData\\RustDesk) is
+///   susceptible to TOCTOU attacks if an unprivileged user can replace the path with a
+///   symlink/junction between checks and deletion.
+///
+/// Future work:
+/// - Use the files (if needed) in the installation directory instead of a public location.
+///   This directory only contains a small `custom.txt` file.
+/// - Pass the custom client name directly via command line
+///   or environment variable during update installation. Then no staging directory is needed.
+#[inline]
+pub fn remove_custom_client_staging_dir(staging_dir: &Path) -> ResultType<bool> {
+    if !staging_dir.exists() {
+        return Ok(false);
+    }
+
+    // First explicitly removes `custom.txt` to ensure stale config is never replayed,
+    // even if the subsequent directory removal fails.
+    //
+    // `std::fs::remove_file` on a symlink removes the symlink itself, not the target,
+    // so this is safe even in a TOCTOU race.
+    let custom_txt_path = staging_dir.join("custom.txt");
+    if custom_txt_path.exists() {
+        allow_err!(std::fs::remove_file(&custom_txt_path));
+    }
+
+    // Intentionally not deleting. See the function docs for rationale.
+    log::debug!(
+        "Skip deleting staging directory {:?} (intentional to avoid TOCTOU)",
+        staging_dir
+    );
+    Ok(false)
+}
+
+// Prepare custom client update by copying staged custom.txt to current directory and loading it.
+// Returns:
+// 1. Ok(true) if preparation was successful or no staging directory exists.
+// 2. Ok(false) if custom.txt file exists but has invalid contents or fails security checks
+//    (e.g., is a symlink or has invalid contents).
+// 3. Err if any unexpected error occurs during file operations.
+pub fn prepare_custom_client_update() -> ResultType<bool> {
+    let custom_client_staging_dir = get_custom_client_staging_dir();
+    let current_exe = std::env::current_exe()?;
+    let current_exe_dir = current_exe
+        .parent()
+        .ok_or(anyhow!("Cannot get parent directory of current exe"))?;
+
+    let staging_dir = custom_client_staging_dir.clone();
+    let clear_staging_on_exit = crate::SimpleCallOnReturn {
+        b: true,
+        f: Box::new(
+            move || match remove_custom_client_staging_dir(&staging_dir) {
+                Ok(existed) => {
+                    if existed {
+                        log::info!("Custom client staging directory removed successfully.");
+                    }
+                }
+                Err(e) => {
+                    log::error!(
+                        "Failed to remove custom client staging directory {:?}: {}",
+                        staging_dir,
+                        e
+                    );
+                }
+            },
+        ),
+    };
+
+    if custom_client_staging_dir.exists() {
+        let custom_txt_path = custom_client_staging_dir.join("custom.txt");
+        if !custom_txt_path.exists() {
+            return Ok(true);
+        }
+
+        let metadata = std::fs::symlink_metadata(&custom_txt_path)?;
+        if metadata.is_symlink() {
+            log::error!(
+                "custom.txt is a symlink. Refusing to load custom client for security reasons."
+            );
+            drop(clear_staging_on_exit);
+            return Ok(false);
+        }
+        if metadata.is_file() {
+            // Copy custom.txt to current directory
+            let local_custom_file_path = current_exe_dir.join("custom.txt");
+            log::debug!(
+                "Copying staged custom file from {:?} to {:?}",
+                custom_txt_path,
+                local_custom_file_path
+            );
+
+            // No need to check symlink before copying.
+            // `load_custom_client()` will fail if the file is not valid.
+            fs::copy(&custom_txt_path, &local_custom_file_path)?;
+            log::info!("Staged custom client file copied to current directory.");
+
+            // Load custom client
+            let is_custom_file_exists =
+                local_custom_file_path.exists() && local_custom_file_path.is_file();
+            crate::load_custom_client();
+
+            // Remove the copied custom.txt file
+            allow_err!(fs::remove_file(&local_custom_file_path));
+
+            // Check if loaded successfully
+            if is_custom_file_exists && !crate::common::is_custom_client() {
+                // The custom.txt file existed, but its contents are invalid.
+                log::error!("Failed to load custom client from custom.txt.");
+                drop(clear_staging_on_exit);
+                // ERROR_INVALID_DATA
+                return Ok(false);
+            }
+        } else {
+            log::info!("No custom client files found in staging directory.");
+        }
+    } else {
+        log::info!(
+            "Custom client staging directory {:?} does not exist.",
+            custom_client_staging_dir
+        );
+    }
+
+    Ok(true)
+}
+
 pub fn get_license_from_exe_name() -> ResultType<CustomServer> {
     let mut exe = std::env::current_exe()?.to_str().unwrap_or("").to_owned();
     // if defined portable appname entry, replace original executable name with it.
@@ -1903,12 +2121,48 @@ unsafe fn set_default_dll_directories() -> bool {
     true
 }
 
+fn get_custom_icon(exe: &str) -> Option<String> {
+    if crate::is_custom_client() {
+        if let Some(p) = PathBuf::from(exe).parent() {
+            let alter_icon_path = p.join("data\\flutter_assets\\assets\\icon.ico");
+            if alter_icon_path.exists() {
+                // Verify that the icon is not a symlink for security
+                if let Ok(metadata) = std::fs::symlink_metadata(&alter_icon_path) {
+                    if metadata.is_symlink() {
+                        log::warn!(
+                            "Custom icon at {:?} is a symlink, refusing to use it.",
+                            alter_icon_path
+                        );
+                        return None;
+                    }
+                    if metadata.is_file() {
+                        return Some(alter_icon_path.to_string_lossy().to_string());
+                    }
+                }
+            }
+        }
+    }
+    None
+}
+
+#[inline]
+fn get_shortcut_icon_location(exe: &str) -> String {
+    if exe.is_empty() {
+        return "".to_owned();
+    }
+
+    get_custom_icon(exe)
+        .map(|p| format!("oLink.IconLocation = \"{}\"", p))
+        .unwrap_or_default()
+}
+
 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_icon_location = get_shortcut_icon_location(&exe);
     let shortcut = write_cmds(
         format!(
             "
@@ -1919,6 +2173,7 @@ sLinkFile = objFSO.BuildPath(strDesktop, \"{filename}.lnk\")
 Set oLink = oWS.CreateShortcut(sLinkFile)
     oLink.TargetPath = \"{exe}\"
     oLink.Arguments = \"--connect {id}\"
+    {shortcut_icon_location}
 oLink.Save
         "
         ),
@@ -2724,6 +2979,44 @@ if exist \"{tray_shortcut}\" del /f /q \"{tray_shortcut}\"
     std::process::exit(0);
 }
 
+/// Calculate the total size of a directory in KB
+/// Does not follow symlinks to prevent directory traversal attacks.
+fn get_directory_size_kb(path: &str) -> u64 {
+    let mut total_size = 0u64;
+    let mut stack = vec![PathBuf::from(path)];
+
+    while let Some(current_path) = stack.pop() {
+        let entries = match std::fs::read_dir(&current_path) {
+            Ok(entries) => entries,
+            Err(_) => continue,
+        };
+
+        for entry in entries {
+            let entry = match entry {
+                Ok(entry) => entry,
+                Err(_) => continue,
+            };
+
+            let metadata = match std::fs::symlink_metadata(entry.path()) {
+                Ok(metadata) => metadata,
+                Err(_) => continue,
+            };
+
+            if metadata.is_symlink() {
+                continue;
+            }
+
+            if metadata.is_dir() {
+                stack.push(entry.path());
+            } else {
+                total_size = total_size.saturating_add(metadata.len());
+            }
+        }
+    }
+
+    total_size / 1024
+}
+
 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();
@@ -2764,12 +3057,35 @@ pub fn update_me(debug: bool) -> ResultType<()> {
     if versions.len() > 2 {
         version_build = versions[2];
     }
-    let meta = std::fs::symlink_metadata(std::env::current_exe()?)?;
-    let size = meta.len() / 1024;
+    let version = crate::VERSION.replace("-", ".");
+    let size = get_directory_size_kb(&path);
+    let build_date = crate::BUILD_DATE;
+    let display_icon = get_custom_icon(&exe).unwrap_or(exe.to_string());
 
-    let reg_cmd = format!(
-        "
-reg add {subkey} /f /v DisplayIcon /t REG_SZ /d \"{exe}\"
+    let is_msi = is_msi_installed().ok();
+
+    fn get_reg_cmd(
+        subkey: &str,
+        is_msi: Option<bool>,
+        display_icon: &str,
+        version: &str,
+        build_date: &str,
+        version_major: &str,
+        version_minor: &str,
+        version_build: &str,
+        size: u64,
+    ) -> String {
+        let reg_display_icon = if is_msi.unwrap_or(false) {
+            "".to_string()
+        } else {
+            format!(
+                "reg add {} /f /v DisplayIcon /t REG_SZ /d \"{}\"",
+                subkey, display_icon
+            )
+        };
+        format!(
+            "
+{reg_display_icon}
 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}\"
@@ -2777,10 +3093,39 @@ 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 reg_cmd = {
+        let reg_cmd_main = get_reg_cmd(
+            &subkey,
+            is_msi,
+            &display_icon,
+            &version,
+            &build_date,
+            &version_major,
+            &version_minor,
+            &version_build,
+            size,
+        );
+        let reg_cmd_msi = if let Some(reg_msi_key) = get_reg_msi_key(&subkey, is_msi) {
+            get_reg_cmd(
+                &reg_msi_key,
+                is_msi,
+                &display_icon,
+                &version,
+                &build_date,
+                &version_major,
+                &version_minor,
+                &version_build,
+                size,
+            )
+        } else {
+            "".to_owned()
+        };
+        format!("{}{}", reg_cmd_main, reg_cmd_msi)
+    };
 
     let filter = format!(" /FI \"PID ne {}\"", get_current_pid());
     let restore_service_cmd = if is_service_running {
@@ -2820,6 +3165,8 @@ sc stop {app_name}
 taskkill /F /IM {app_name}.exe{filter}
 {reg_cmd}
 {copy_exe}
+{rename_exe}
+{remove_meta_toml}
 {restore_service_cmd}
 {uninstall_printer_cmd}
 {install_printer_cmd}
@@ -2827,43 +3174,106 @@ taskkill /F /IM {app_name}.exe{filter}
     ",
         app_name = app_name,
         copy_exe = copy_exe_cmd(&src_exe, &exe, &path)?,
+        rename_exe = rename_exe_cmd(&src_exe, &path)?,
+        remove_meta_toml = remove_meta_toml_cmd(is_msi.unwrap_or(true), &path),
         sleep = if debug { "timeout 300" } else { "" },
     );
 
+    let _restore_session_guard = crate::common::SimpleCallOnReturn {
+        b: true,
+        f: Box::new(move || {
+            let is_root = is_root();
+            if tray_sessions.is_empty() {
+                log::info!("No tray process found.");
+            } else {
+                log::info!(
+                    "Try to restore the tray process..., sessions: {:?}",
+                    &tray_sessions
+                );
+                // When not running as root, only spawn once since run_exe_direct
+                // doesn't target specific sessions.
+                let mut spawned_non_root_tray = false;
+                for s in tray_sessions.clone().into_iter() {
+                    if s != 0 {
+                        // We need to check if is_root here because if `update_me()` is called from
+                        // the main window running with administrator permission,
+                        // `run_exe_in_session()` will fail with error 1314 ("A required privilege is
+                        // not held by the client").
+                        //
+                        // This issue primarily affects the MSI-installed version running in Administrator
+                        // session during testing, but we check permissions here to be safe.
+                        if is_root {
+                            allow_err!(run_exe_in_session(&exe, vec!["--tray"], s, true));
+                        } else if !spawned_non_root_tray {
+                            // Only spawn once for non-root since run_exe_direct doesn't take session parameter
+                            allow_err!(run_exe_direct(&exe, vec!["--tray"], false));
+                            spawned_non_root_tray = 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));
+                // When not running as root, only spawn once since run_exe_direct
+                // doesn't target specific sessions.
+                let mut spawned_non_root_main = false;
+                for s in main_window_sessions.clone().into_iter() {
+                    if s != 0 {
+                        if is_root {
+                            allow_err!(run_exe_in_session(&exe, vec![], s, true));
+                        } else if !spawned_non_root_main {
+                            // Only spawn once for non-root since run_exe_direct doesn't take session parameter
+                            allow_err!(run_exe_direct(&exe, vec![], false));
+                            spawned_non_root_main = true;
+                        }
+                    }
+                }
+            }
+            std::thread::sleep(std::time::Duration::from_millis(300));
+        }),
+    };
+
     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(())
 }
 
+fn get_reg_msi_key(subkey: &str, is_msi: Option<bool>) -> Option<String> {
+    // Only proceed if it's a custom client and MSI is installed.
+    // `is_msi.unwrap_or(true)` is intentional: subsequent code validates the registry,
+    // hence no early return is required upon MSI detection failure.
+    if !(crate::common::is_custom_client() && is_msi.unwrap_or(true)) {
+        return None;
+    }
+
+    // Get the uninstall string from registry
+    let uninstall_string = get_reg_of(subkey, "UninstallString");
+    if uninstall_string.is_empty() {
+        return None;
+    }
+
+    // Find the product code (GUID) in the uninstall string
+    // Handle both quoted and unquoted GUIDs: /X {GUID} or /X "{GUID}"
+    let start = uninstall_string.rfind('{')?;
+    let end = uninstall_string.rfind('}')?;
+    if start >= end {
+        return None;
+    }
+    let product_code = &uninstall_string[start..=end];
+
+    // Build the MSI registry key path
+    let pos = subkey.rfind('\\')?;
+    let reg_msi_key = format!("{}{}", &subkey[..=pos], product_code);
+
+    Some(reg_msi_key)
+}
+
 // Double confirm the process name
 fn kill_process_by_pids(name: &str, pids: Vec<Pid>) -> ResultType<()> {
     let name = name.to_lowercase();
@@ -2885,6 +3295,109 @@ fn kill_process_by_pids(name: &str, pids: Vec<Pid>) -> ResultType<()> {
     Ok(())
 }
 
+pub fn handle_custom_client_staging_dir_before_update(
+    custom_client_staging_dir: &PathBuf,
+) -> ResultType<()> {
+    let Some(current_exe_dir) = std::env::current_exe()
+        .ok()
+        .and_then(|p| p.parent().map(|p| p.to_path_buf()))
+    else {
+        bail!("Failed to get current exe directory");
+    };
+
+    // Clean up existing staging directory
+    if custom_client_staging_dir.exists() {
+        log::debug!(
+            "Removing existing custom client staging directory: {:?}",
+            custom_client_staging_dir
+        );
+        if let Err(e) = remove_custom_client_staging_dir(custom_client_staging_dir) {
+            bail!(
+                "Failed to remove existing custom client staging directory {:?}: {}",
+                custom_client_staging_dir,
+                e
+            );
+        }
+    }
+
+    let src_path = current_exe_dir.join("custom.txt");
+    if src_path.exists() {
+        // Verify that custom.txt is not a symlink before copying
+        let metadata = match std::fs::symlink_metadata(&src_path) {
+            Ok(m) => m,
+            Err(e) => {
+                bail!(
+                    "Failed to read metadata for custom.txt at {:?}: {}",
+                    src_path,
+                    e
+                );
+            }
+        };
+
+        if metadata.is_symlink() {
+            allow_err!(remove_custom_client_staging_dir(&custom_client_staging_dir));
+            bail!(
+                "custom.txt at {:?} is a symlink, refusing to stage for security reasons.",
+                src_path
+            );
+        }
+
+        if metadata.is_file() {
+            if !custom_client_staging_dir.exists() {
+                if let Err(e) = std::fs::create_dir_all(custom_client_staging_dir) {
+                    bail!("Failed to create parent directory {:?} when staging custom client files: {}", custom_client_staging_dir, e);
+                }
+            }
+            let dst_path = custom_client_staging_dir.join("custom.txt");
+            if let Err(e) = std::fs::copy(&src_path, &dst_path) {
+                allow_err!(remove_custom_client_staging_dir(&custom_client_staging_dir));
+                bail!(
+                    "Failed to copy custom txt from {:?} to {:?}: {}",
+                    src_path,
+                    dst_path,
+                    e
+                );
+            }
+        } else {
+            log::warn!(
+                "custom.txt at {:?} is not a regular file, skipping.",
+                src_path
+            );
+        }
+    } else {
+        log::info!("No custom txt found to stage for update.");
+    }
+
+    Ok(())
+}
+
+// Used for auto update and manual update in the main window.
+pub fn update_to(file: &str) -> ResultType<()> {
+    if file.ends_with(".exe") {
+        let custom_client_staging_dir = get_custom_client_staging_dir();
+        if crate::is_custom_client() {
+            handle_custom_client_staging_dir_before_update(&custom_client_staging_dir)?;
+        } else {
+            // Clean up any residual staging directory from previous custom client
+            allow_err!(remove_custom_client_staging_dir(&custom_client_staging_dir));
+        }
+        if !run_uac(file, "--update")? {
+            bail!(
+                "Failed to run the update exe with UAC, error: {:?}",
+                std::io::Error::last_os_error()
+            );
+        }
+    } else if file.ends_with(".msi") {
+        if let Err(e) = update_me_msi(file, false) {
+            bail!("Failed to run the update msi: {}", e);
+        }
+    } else {
+        // unreachable!()
+        bail!("Unsupported update file format: {}", file);
+    }
+    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.
@@ -2905,6 +3418,7 @@ pub fn update_me_msi(msi: &str, quiet: bool) -> ResultType<()> {
 }
 
 pub fn get_tray_shortcut(exe: &str, tmp_path: &str) -> ResultType<String> {
+    let shortcut_icon_location = get_shortcut_icon_location(exe);
     Ok(write_cmds(
         format!(
             "
@@ -2914,6 +3428,7 @@ sLinkFile = \"{tmp_path}\\{app_name} Tray.lnk\"
 Set oLink = oWS.CreateShortcut(sLinkFile)
     oLink.TargetPath = \"{exe}\"
     oLink.Arguments = \"--tray\"
+    {shortcut_icon_location}
 oLink.Save
         ",
             app_name = crate::get_app_name(),
@@ -2976,6 +3491,44 @@ fn run_after_run_cmds(silent: bool) {
     std::thread::sleep(std::time::Duration::from_millis(300));
 }
 
+#[inline]
+pub fn try_remove_temp_update_files() {
+    let temp_dir = std::env::temp_dir();
+    let Ok(entries) = std::fs::read_dir(&temp_dir) else {
+        log::debug!("Failed to read temp directory: {:?}", temp_dir);
+        return;
+    };
+
+    let one_hour = std::time::Duration::from_secs(60 * 60);
+    for entry in entries {
+        if let Ok(entry) = entry {
+            let path = entry.path();
+            if let Some(file_name) = path.file_name().and_then(|n| n.to_str()) {
+                // Match files like rustdesk-*.msi or rustdesk-*.exe
+                if file_name.starts_with("rustdesk-")
+                    && (file_name.ends_with(".msi") || file_name.ends_with(".exe"))
+                {
+                    // Skip files modified within the last hour to avoid deleting files being downloaded
+                    if let Ok(metadata) = std::fs::metadata(&path) {
+                        if let Ok(modified) = metadata.modified() {
+                            if let Ok(elapsed) = modified.elapsed() {
+                                if elapsed < one_hour {
+                                    continue;
+                                }
+                            }
+                        }
+                    }
+                    if let Err(e) = std::fs::remove_file(&path) {
+                        log::debug!("Failed to remove temp update file {:?}: {}", path, e);
+                    } else {
+                        log::info!("Removed temp update file: {:?}", path);
+                    }
+                }
+            }
+        }
+    }
+}
+
 #[inline]
 pub fn try_kill_broker() {
     allow_err!(std::process::Command::new("cmd")
@@ -3151,7 +3704,8 @@ pub fn is_x64() -> bool {
 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");
+    let app_name = crate::get_app_name().to_lowercase();
+    log::info!("try kill main window process");
     use hbb_common::sysinfo::System;
     let mut sys = System::new();
     sys.refresh_processes();
@@ -3160,7 +3714,6 @@ pub fn try_kill_rustdesk_main_window_process() -> ResultType<()> {
         .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");
     }