From d64bc7248e208de6e95f7f671ca472f55b19e3ee Mon Sep 17 00:00:00 2001
From: Wez Furlong <wez@wezfurlong.org>
Date: Tue, 19 Apr 2022 09:56:23 -0700
Subject: [PATCH] pty: pre-fill base env from registry env settings on Windows

I think I'd like to make a config option for this, but for the moment,
this first pass unconditionally updates the base environment with
data from the registry.

refs: https://github.com/wez/wezterm/issues/1848
---
 Cargo.lock            | 14 ++-----
 pty/Cargo.toml        |  1 +
 pty/src/cmdbuilder.rs | 98 ++++++++++++++++++++++++++++++++++++++++++-
 window/Cargo.toml     |  2 +-
 4 files changed, 101 insertions(+), 14 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index a22a58efa..d9de2b4e6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1041,7 +1041,7 @@ dependencies = [
  "rustc_version 0.4.0",
  "toml",
  "vswhom",
- "winreg 0.10.1",
+ "winreg",
 ]
 
 [[package]]
@@ -2946,6 +2946,7 @@ dependencies = [
  "smol",
  "ssh2",
  "winapi 0.3.9",
+ "winreg",
 ]
 
 [[package]]
@@ -5000,7 +5001,7 @@ dependencies = [
  "wezterm-input-types",
  "winapi 0.3.9",
  "windows",
- "winreg 0.6.2",
+ "winreg",
  "x11",
  "xcb 0.9.0",
  "xcb-imdkit",
@@ -5051,15 +5052,6 @@ version = "0.33.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ff1e4aa646495048ec7f3ffddc411e1d829c026a2ec62b39da15c1055e406eaa"
 
-[[package]]
-name = "winreg"
-version = "0.6.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b2986deb581c4fe11b621998a5e53361efe6b48a151178d0cd9eeffa4dc6acc9"
-dependencies = [
- "winapi 0.3.9",
-]
-
 [[package]]
 name = "winreg"
 version = "0.10.1"
diff --git a/pty/Cargo.toml b/pty/Cargo.toml
index b6dda62cd..2df06b4b9 100644
--- a/pty/Cargo.toml
+++ b/pty/Cargo.toml
@@ -37,6 +37,7 @@ winapi = { version = "0.3", features = [
     "namedpipeapi",
     "synchapi",
 ]}
+winreg = "0.10"
 
 [dev-dependencies]
 smol = "1.2"
diff --git a/pty/src/cmdbuilder.rs b/pty/src/cmdbuilder.rs
index 4e2f07823..eb4a99e3d 100644
--- a/pty/src/cmdbuilder.rs
+++ b/pty/src/cmdbuilder.rs
@@ -37,7 +37,7 @@ impl EnvEntry {
 }
 
 fn get_base_env() -> BTreeMap<OsString, EnvEntry> {
-    std::env::vars_os()
+    let mut env: BTreeMap<OsString, EnvEntry> = std::env::vars_os()
         .map(|(key, value)| {
             (
                 EnvEntry::map_key(key.clone()),
@@ -48,7 +48,101 @@ fn get_base_env() -> BTreeMap<OsString, EnvEntry> {
                 },
             )
         })
-        .collect()
+        .collect();
+
+    #[cfg(windows)]
+    {
+        use std::os::windows::ffi::OsStringExt;
+        use winapi::um::processenv::ExpandEnvironmentStringsW;
+        use winreg::enums::{RegType, HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE};
+        use winreg::types::FromRegValue;
+        use winreg::{RegKey, RegValue};
+
+        fn reg_value_to_string(value: &RegValue) -> anyhow::Result<OsString> {
+            match value.vtype {
+                RegType::REG_EXPAND_SZ => {
+                    let src = unsafe {
+                        std::slice::from_raw_parts(
+                            value.bytes.as_ptr() as *const u16,
+                            value.bytes.len() / 2,
+                        )
+                    };
+                    let size =
+                        unsafe { ExpandEnvironmentStringsW(src.as_ptr(), std::ptr::null_mut(), 0) };
+                    let mut buf = vec![0u16; size as usize + 1];
+                    unsafe {
+                        ExpandEnvironmentStringsW(src.as_ptr(), buf.as_mut_ptr(), buf.len() as u32)
+                    };
+
+                    let mut buf = buf.as_slice();
+                    while let Some(0) = buf.last() {
+                        buf = &buf[0..buf.len() - 1];
+                    }
+                    Ok(OsString::from_wide(buf))
+                }
+                _ => Ok(OsString::from_reg_value(value)?),
+            }
+        }
+
+        if let Ok(sys_env) = RegKey::predef(HKEY_LOCAL_MACHINE)
+            .open_subkey("System\\CurrentControlSet\\Control\\Session Manager\\Environment")
+        {
+            for res in sys_env.enum_values() {
+                if let Ok((name, value)) = res {
+                    if name.to_ascii_lowercase() == "username" {
+                        continue;
+                    }
+                    if let Ok(value) = reg_value_to_string(&value) {
+                        log::trace!("adding SYS env: {:?} {:?}", name, value);
+                        env.insert(
+                            EnvEntry::map_key(name.clone().into()),
+                            EnvEntry {
+                                is_from_base_env: true,
+                                preferred_key: name.into(),
+                                value,
+                            },
+                        );
+                    }
+                }
+            }
+        }
+
+        if let Ok(sys_env) = RegKey::predef(HKEY_CURRENT_USER).open_subkey("Environment") {
+            for res in sys_env.enum_values() {
+                if let Ok((name, value)) = res {
+                    if let Ok(value) = reg_value_to_string(&value) {
+                        // Merge the system and user paths together
+                        let value = if name.to_ascii_lowercase() == "path" {
+                            match env.get(&EnvEntry::map_key(name.clone().into())) {
+                                Some(entry) => {
+                                    let mut result = OsString::new();
+                                    result.push(&entry.value);
+                                    result.push(";");
+                                    result.push(&value);
+                                    result
+                                }
+                                None => value,
+                            }
+                        } else {
+                            value
+                        };
+
+                        log::trace!("adding USER env: {:?} {:?}", name, value);
+                        env.insert(
+                            EnvEntry::map_key(name.clone().into()),
+                            EnvEntry {
+                                is_from_base_env: true,
+                                preferred_key: name.into(),
+                                value,
+                            },
+                        );
+                    }
+                }
+            }
+        }
+    }
+
+    env
 }
 
 /// `CommandBuilder` is used to prepare a command to be spawned into a pty.
diff --git a/window/Cargo.toml b/window/Cargo.toml
index f3cfe02d2..fbe834db0 100644
--- a/window/Cargo.toml
+++ b/window/Cargo.toml
@@ -61,7 +61,7 @@ winapi = { version = "0.3", features = [
 windows = { version="0.33.0", features = [
     "UI_ViewManagement",
 ]}
-winreg = "0.6"
+winreg = "0.10"
 
 [target.'cfg(all(unix, not(target_os = "macos")))'.dependencies]
 dirs-next = "2.0"