From e1c2b8de6efe9104576ed7b048dea92e68086c89 Mon Sep 17 00:00:00 2001
From: 21pages <pages21@163.com>
Date: Mon, 26 Sep 2022 16:23:09 +0800
Subject: [PATCH] windows portable: request elevation && run as system

Signed-off-by: 21pages <pages21@163.com>
---
 Cargo.lock              | 19 +++++++++++++++
 Cargo.toml              |  2 ++
 src/client.rs           |  6 ++++-
 src/core_main.rs        |  6 +++++
 src/platform/windows.rs | 53 ++++++++++++++++++++++++++++++++++++++---
 5 files changed, 82 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 60852eba1..48f47a86f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2537,6 +2537,14 @@ dependencies = [
  "tiff",
 ]
 
+[[package]]
+name = "impersonate_system"
+version = "0.1.0"
+source = "git+https://github.com/21pages/impersonate-system#af4a82050580217a434c2024e181a98de24823ec"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "include_dir"
 version = "0.7.2"
@@ -2593,6 +2601,15 @@ version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "879d54834c8c76457ef4293a689b2a8c59b076067ad77b15efafbb05f92a592b"
 
+[[package]]
+name = "is_elevated"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5299060ff5db63e788015dcb9525ad9b84f4fd9717ed2cbdeba5018cbf42f9b5"
+dependencies = [
+ "winapi 0.3.9",
+]
+
 [[package]]
 name = "itertools"
 version = "0.9.0"
@@ -4329,7 +4346,9 @@ dependencies = [
  "flutter_rust_bridge_codegen",
  "hbb_common",
  "hound",
+ "impersonate_system",
  "include_dir",
+ "is_elevated",
  "jni",
  "lazy_static",
  "libc",
diff --git a/Cargo.toml b/Cargo.toml
index 062e32abb..f8b93c2cd 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -91,6 +91,8 @@ winapi = { version = "0.3", features = ["winuser"] }
 winreg = "0.10"
 windows-service = "0.4"
 virtual_display = { path = "libs/virtual_display" }
+is_elevated = "0.1.2"
+impersonate_system = { git = "https://github.com/21pages/impersonate-system" }
 
 [target.'cfg(target_os = "macos")'.dependencies]
 objc = "0.2"
diff --git a/src/client.rs b/src/client.rs
index c70956b63..214af6cf7 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -1354,7 +1354,11 @@ impl LoginConfigHandler {
             username: self.id.clone(),
             password: password.into(),
             my_id,
-            my_name: crate::username(),
+            my_name: if cfg!(windows) {
+                crate::platform::get_active_username()
+            } else {
+                crate::username()
+            },
             option: self.get_option_message(true).into(),
             session_id: self.session_id,
             version: crate::VERSION.to_string(),
diff --git a/src/core_main.rs b/src/core_main.rs
index f514cd790..7c86bdd19 100644
--- a/src/core_main.rs
+++ b/src/core_main.rs
@@ -57,6 +57,12 @@ pub fn core_main() -> Option<Vec<String>> {
                 .ok();
         }
     }
+    #[cfg(windows)]
+    #[cfg(not(debug_assertions))]
+    if !crate::platform::is_installed() && args.is_empty() {
+        let arg = if is_setup { "--noinstall" } else { "" };
+        crate::platform::run_check_elevation(arg);
+    }
     if args.is_empty() {
         std::thread::spawn(move || crate::start_server(false));
     } else {
diff --git a/src/platform/windows.rs b/src/platform/windows.rs
index fa9fb5b10..be549f7e1 100644
--- a/src/platform/windows.rs
+++ b/src/platform/windows.rs
@@ -1420,16 +1420,63 @@ pub fn get_user_token(session_id: u32, as_user: bool) -> HANDLE {
     }
 }
 
-pub fn check_super_user_permission() -> ResultType<bool> {
+pub fn run_uac(exe: &str, arg: &str) -> ResultType<bool> {
     unsafe {
+        let cstring;
         let ret = ShellExecuteA(
             NULL as _,
             CString::new("runas")?.as_ptr() as _,
-            CString::new("cmd")?.as_ptr() as _,
-            CString::new("/c /q")?.as_ptr() as _,
+            CString::new(exe)?.as_ptr() as _,
+            if arg.is_empty() {
+                NULL as _
+            } else {
+                cstring = CString::new(arg)?;
+                cstring.as_ptr() as _
+            },
             NULL as _,
             SW_SHOWNORMAL,
         );
         return Ok(ret as i32 > 32);
     }
 }
+
+pub fn check_super_user_permission() -> ResultType<bool> {
+    run_uac("cmd", "/c /q")
+}
+
+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 run_check_elevation(arg: &str) {
+    if !is_elevated::is_elevated() {
+        if let Ok(true) = elevate(arg) {
+            std::process::exit(0);
+        } else {
+            // do nothing but prompt
+        }
+    } else {
+        if !is_root() {
+            if run_as_system(arg).is_ok() {
+                std::process::exit(0);
+            } else {
+                // to-do: should not happen
+                log::error!("Failed to run as system");
+            }
+        }
+    }
+}