fix: Allow program run with administrator to start at startup (zzzgyd…

…i#12)
Hypobenthos · Feb 20, 2024 · 47bc9c6 · 47bc9c6
1 parent fedaeba
commit 47bc9c6
Showing 1 changed file with 60 additions and 24 deletions.
diff --git a/src/windows.rs b/src/windows.rs
@@ -1,8 +1,9 @@
 use crate::{AutoLaunch, Result};
 use winreg::enums::RegType::REG_BINARY;
-use winreg::enums::{HKEY_CURRENT_USER, KEY_READ, KEY_SET_VALUE};
+use winreg::enums::{HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE, KEY_READ, KEY_SET_VALUE};
 use winreg::{RegKey, RegValue};
 
+static ADMIN_REGKEY: &str = "SOFTWARE\\WOW6432Node\\Microsoft\\Windows\\CurrentVersion\\Run";
 static AL_REGKEY: &str = "SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run";
 static TASK_MANAGER_OVERRIDE_REGKEY: &str =
     "SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Explorer\\StartupApproved\\Run";
@@ -36,21 +37,42 @@ impl AutoLaunch {
     /// - failed to set value
     pub fn enable(&self) -> Result<()> {
         let hkcu = RegKey::predef(HKEY_CURRENT_USER);
-        hkcu.open_subkey_with_flags(AL_REGKEY, KEY_SET_VALUE)?
-            .set_value::<_, _>(
+        let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
+        if let Ok(reg) = hklm.open_subkey_with_flags(ADMIN_REGKEY, KEY_SET_VALUE) {
+            reg.set_value::<_, _>(
                 &self.app_name,
                 &format!("{} {}", &self.app_path, &self.args.join(" ")),
             )?;
-
-        // this key maybe not found
-        if let Ok(reg) = hkcu.open_subkey_with_flags(TASK_MANAGER_OVERRIDE_REGKEY, KEY_SET_VALUE) {
-            reg.set_raw_value(
-                &self.app_name,
-                &RegValue {
-                    vtype: REG_BINARY,
-                    bytes: TASK_MANAGER_OVERRIDE_ENABLED_VALUE.to_vec(),
-                },
-            )?;
+            // this key maybe not found
+            if let Ok(reg) =
+                hklm.open_subkey_with_flags(TASK_MANAGER_OVERRIDE_REGKEY, KEY_SET_VALUE)
+            {
+                reg.set_raw_value(
+                    &self.app_name,
+                    &RegValue {
+                        vtype: REG_BINARY,
+                        bytes: TASK_MANAGER_OVERRIDE_ENABLED_VALUE.to_vec(),
+                    },
+                )?;
+            }
+        } else {
+            hkcu.open_subkey_with_flags(AL_REGKEY, KEY_SET_VALUE)?
+                .set_value::<_, _>(
+                    &self.app_name,
+                    &format!("{} {}", &self.app_path, &self.args.join(" ")),
+                )?;
+            // this key maybe not found
+            if let Ok(reg) =
+                hkcu.open_subkey_with_flags(TASK_MANAGER_OVERRIDE_REGKEY, KEY_SET_VALUE)
+            {
+                reg.set_raw_value(
+                    &self.app_name,
+                    &RegValue {
+                        vtype: REG_BINARY,
+                        bytes: TASK_MANAGER_OVERRIDE_ENABLED_VALUE.to_vec(),
+                    },
+                )?;
+            }
         }
 
         Ok(())
@@ -64,26 +86,40 @@ impl AutoLaunch {
     /// - failed to delete value
     pub fn disable(&self) -> Result<()> {
         let hkcu = RegKey::predef(HKEY_CURRENT_USER);
-        hkcu.open_subkey_with_flags(AL_REGKEY, KEY_SET_VALUE)?
-            .delete_value(&self.app_name)?;
+        let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
+
+        if let Ok(reg) = hklm.open_subkey_with_flags(ADMIN_REGKEY, KEY_SET_VALUE) {
+            reg.delete_value(&self.app_name)?;
+        } else {
+            hkcu.open_subkey_with_flags(AL_REGKEY, KEY_SET_VALUE)?
+                .delete_value(&self.app_name)?;
+        }
         Ok(())
     }
 
     /// Check whether the AutoLaunch setting is enabled
     pub fn is_enabled(&self) -> Result<bool> {
         let hkcu = RegKey::predef(HKEY_CURRENT_USER);
+        let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
+        // check if the app is enabled in the admin registry
+        // use `KEY_SET_VALUE` to ensure have admin permission
+        if let Ok(reg) = hklm.open_subkey_with_flags(ADMIN_REGKEY, KEY_SET_VALUE) {
+            let adm_enabled = reg.get_value::<String, _>(&self.app_name).is_ok();
+            let task_manager_enabled = self.task_manager_enabled(hklm);
+            Ok(adm_enabled && task_manager_enabled.unwrap_or(true))
+        } else {
+            let al_enabled = hkcu
+                .open_subkey_with_flags(AL_REGKEY, KEY_READ)?
+                .get_value::<String, _>(&self.app_name)
+                .is_ok();
+            let task_manager_enabled = self.task_manager_enabled(hkcu);
 
-        let al_enabled = hkcu
-            .open_subkey_with_flags(AL_REGKEY, KEY_READ)?
-            .get_value::<String, _>(&self.app_name)
-            .is_ok();
-        let task_manager_enabled = self.task_manager_enabled(hkcu);
-
-        Ok(al_enabled && task_manager_enabled.unwrap_or(true))
+            Ok(al_enabled && task_manager_enabled.unwrap_or(true))
+        }
     }
 
-    fn task_manager_enabled(&self, hkcu: RegKey) -> Option<bool> {
-        let task_manager_override_raw_value = hkcu
+    fn task_manager_enabled(&self, hk: RegKey) -> Option<bool> {
+        let task_manager_override_raw_value = hk
             .open_subkey_with_flags(TASK_MANAGER_OVERRIDE_REGKEY, KEY_READ)
             .ok()?
             .get_raw_value(&self.app_name)