From a2f1ab80daf5ea407b71f082b45bdbfbbe2ccee6 Mon Sep 17 00:00:00 2001
From: 21pages <pages21@163.com>
Date: Mon, 31 Oct 2022 16:08:51 +0800
Subject: [PATCH] run_uac with wchar, refactor elevate_or_run_as_system

Signed-off-by: 21pages <pages21@163.com>
---
 Cargo.lock              |  2 +-
 src/platform/windows.rs | 84 +++++++++++++++++++++++------------------
 src/tray.rs             |  2 +-
 3 files changed, 49 insertions(+), 39 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0bc360fe7..5735c2042 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2559,7 +2559,7 @@ dependencies = [
 [[package]]
 name = "impersonate_system"
 version = "0.1.0"
-source = "git+https://github.com/21pages/impersonate-system#c48f37a8fd17413b2a4ba655c3873bdc5c8d25aa"
+source = "git+https://github.com/21pages/impersonate-system#84b401893d5b6628c8b33b295328d13fbbe2674b"
 dependencies = [
  "cc",
 ]
diff --git a/src/platform/windows.rs b/src/platform/windows.rs
index b067f514b..cda7b56a8 100644
--- a/src/platform/windows.rs
+++ b/src/platform/windows.rs
@@ -9,7 +9,7 @@ use hbb_common::{
 };
 use std::io::prelude::*;
 use std::{
-    ffi::{CString, OsString},
+    ffi::OsString,
     fs, io, mem,
     path::PathBuf,
     sync::{Arc, Mutex},
@@ -27,7 +27,7 @@ use winapi::{
             OpenProcessToken,
         },
         securitybaseapi::GetTokenInformation,
-        shellapi::ShellExecuteA,
+        shellapi::ShellExecuteW,
         winbase::*,
         wingdi::*,
         winnt::{
@@ -1481,17 +1481,19 @@ pub fn get_user_token(session_id: u32, as_user: bool) -> HANDLE {
 }
 
 pub fn run_uac(exe: &str, arg: &str) -> ResultType<bool> {
+    let wop = wide_string("runas");
+    let wexe = wide_string(exe);
+    let warg;
     unsafe {
-        let cstring;
-        let ret = ShellExecuteA(
+        let ret = ShellExecuteW(
             NULL as _,
-            CString::new("runas")?.as_ptr() as _,
-            CString::new(exe)?.as_ptr() as _,
+            wop.as_ptr() as _,
+            wexe.as_ptr() as _,
             if arg.is_empty() {
                 NULL as _
             } else {
-                cstring = CString::new(arg)?;
-                cstring.as_ptr() as _
+                warg = wide_string(arg);
+                warg.as_ptr() as _
             },
             NULL as _,
             SW_SHOWNORMAL,
@@ -1529,7 +1531,7 @@ pub fn run_as_system(arg: &str) -> ResultType<()> {
 }
 
 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, which hasn't happened yet.
+    // avoid possible run recursively due to failed run.
     let arg_elevate = if is_setup {
         "--noinstall --elevate"
     } else {
@@ -1540,37 +1542,37 @@ pub fn elevate_or_run_as_system(is_setup: bool, is_elevate: bool, is_run_as_syst
     } else {
         "--run-as-system"
     };
-    let rerun_as_system = || {
-        if !is_root() {
-            if run_as_system(arg_run_as_system).is_ok() {
-                std::process::exit(0);
-            } else {
-                log::error!("Failed to run as system");
-            }
-        }
-    };
 
-    if is_elevate {
-        if !is_elevated(None).map_or(true, |b| b) {
-            log::error!("Failed to elevate");
-            return;
-        }
-        rerun_as_system();
-    } else if is_run_as_system {
-        if !is_root() {
-            log::error!("Failed to be system");
-        }
+    if is_root() {
+        log::debug!("portable run as system user");
     } else {
-        if let Ok(true) = is_elevated(None) {
-            // right click
-            rerun_as_system();
-        } else {
-            // left click || run without install
-            if let Ok(true) = elevate(arg_elevate) {
-                std::process::exit(0);
-            } else {
-                // do nothing but prompt
+        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());
+            },
         }
     }
 }
@@ -1636,3 +1638,11 @@ fn get_current_pid() -> u32 {
 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()
+}
diff --git a/src/tray.rs b/src/tray.rs
index a71ade329..1f95c8306 100644
--- a/src/tray.rs
+++ b/src/tray.rs
@@ -27,7 +27,7 @@ pub fn start_tray(options: Arc<Mutex<HashMap<String, String>>>) {
         .build()
         .unwrap();
     let old_state = Arc::new(Mutex::new(0));
-    let _ = crate::ui_interface::SENDER.lock().unwrap();
+    let _sender = crate::ui_interface::SENDER.lock().unwrap();
     event_loop.run(move |event, _, control_flow| {
         if options.lock().unwrap().get("ipc-closed").is_some() {
             *control_flow = ControlFlow::Exit;