Allow setting terminal env vars via PTY options

Closes #7778.

5 files changed, 90 insertions, 7 deletions

diff --git a/alacritty/src/cli.rs b/alacritty/src/cli.rs
index d5e24b4a..91ba2fd6 100644
--- a/alacritty/src/cli.rs
+++ b/alacritty/src/cli.rs
@@ -1,4 +1,5 @@
 use std::cmp::max;
+use std::collections::HashMap;
 use std::ops::{Deref, DerefMut};
 use std::path::PathBuf;
 use std::rc::Rc;
@@ -195,6 +196,7 @@ impl From<TerminalOptions> for PtyOptions {
             working_directory: options.working_directory.take(),
             shell: options.command().map(Into::into),
             hold: options.hold,
+            env: HashMap::new(),
         }
     }
 }
diff --git a/alacritty/src/config/ui_config.rs b/alacritty/src/config/ui_config.rs
index a4b6c2c5..580a3dad 100644
--- a/alacritty/src/config/ui_config.rs
+++ b/alacritty/src/config/ui_config.rs
@@ -167,7 +167,12 @@ impl UiConfig {
     /// Derive [`PtyOptions`] from the config.
     pub fn pty_config(&self) -> PtyOptions {
         let shell = self.shell.clone().map(Into::into);
-        PtyOptions { shell, working_directory: self.working_directory.clone(), hold: false }
+        PtyOptions {
+            shell,
+            working_directory: self.working_directory.clone(),
+            hold: false,
+            env: HashMap::new(),
+        }
     }
 
     /// Generate key bindings for all keyboard hints.
diff --git a/alacritty_terminal/src/tty/mod.rs b/alacritty_terminal/src/tty/mod.rs
index 35d227b8..55d263ca 100644
--- a/alacritty_terminal/src/tty/mod.rs
+++ b/alacritty_terminal/src/tty/mod.rs
@@ -1,5 +1,6 @@
 //! TTY related functionality.
 
+use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::{env, io};
@@ -29,6 +30,9 @@ pub struct Options {
 
     /// Remain open after child process exits.
     pub hold: bool,
+
+    /// Extra environment variables.
+    pub env: HashMap<String, String>,
 }
 
 /// Shell options.
diff --git a/alacritty_terminal/src/tty/unix.rs b/alacritty_terminal/src/tty/unix.rs
index 455abbd2..a4b07b74 100644
--- a/alacritty_terminal/src/tty/unix.rs
+++ b/alacritty_terminal/src/tty/unix.rs
@@ -217,9 +217,11 @@ pub fn new(config: &Options, window_size: WindowSize, window_id: u64) -> Result<
     builder.env("ALACRITTY_WINDOW_ID", &window_id);
     builder.env("USER", user.user);
     builder.env("HOME", user.home);
-
     // Set Window ID for clients relying on X11 hacks.
     builder.env("WINDOWID", window_id);
+    for (key, value) in &config.env {
+        builder.env(key, value);
+    }
 
     unsafe {
         builder.pre_exec(move || {
diff --git a/alacritty_terminal/src/tty/windows/conpty.rs b/alacritty_terminal/src/tty/windows/conpty.rs
index 9731b4f0..244681e7 100644
--- a/alacritty_terminal/src/tty/windows/conpty.rs
+++ b/alacritty_terminal/src/tty/windows/conpty.rs
@@ -1,5 +1,8 @@
-use log::info;
+use log::{info, warn};
+use std::collections::{HashMap, HashSet};
+use std::ffi::OsStr;
 use std::io::Error;
+use std::os::windows::ffi::OsStrExt;
 use std::os::windows::io::IntoRawHandle;
 use std::{mem, ptr};
 
@@ -13,8 +16,8 @@ use windows_sys::{s, w};
 
 use windows_sys::Win32::System::Threading::{
     CreateProcessW, InitializeProcThreadAttributeList, UpdateProcThreadAttribute,
-    EXTENDED_STARTUPINFO_PRESENT, PROCESS_INFORMATION, PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE,
-    STARTF_USESTDHANDLES, STARTUPINFOEXW, STARTUPINFOW,
+    CREATE_UNICODE_ENVIRONMENT, EXTENDED_STARTUPINFO_PRESENT, PROCESS_INFORMATION,
+    PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE, STARTF_USESTDHANDLES, STARTUPINFOEXW, STARTUPINFOW,
 };
 
 use crate::event::{OnResize, WindowSize};
@@ -198,8 +201,18 @@ pub fn new(config: &Options, window_size: WindowSize) -> Option<Pty> {
         }
     }
 
+    // Prepare child process creation arguments.
     let cmdline = win32_string(&cmdline(config));
     let cwd = config.working_directory.as_ref().map(win32_string);
+    let mut creation_flags = EXTENDED_STARTUPINFO_PRESENT;
+    let custom_env_block = convert_custom_env(&config.env);
+    let custom_env_block_pointer = match &custom_env_block {
+        Some(custom_env_block) => {
+            creation_flags |= CREATE_UNICODE_ENVIRONMENT;
+            custom_env_block.as_ptr() as *mut std::ffi::c_void
+        },
+        None => ptr::null_mut(),
+    };
 
     let mut proc_info: PROCESS_INFORMATION = unsafe { mem::zeroed() };
     unsafe {
@@ -209,8 +222,8 @@ pub fn new(config: &Options, window_size: WindowSize) -> Option<Pty> {
             ptr::null_mut(),
             ptr::null_mut(),
             false as i32,
-            EXTENDED_STARTUPINFO_PRESENT,
-            ptr::null_mut(),
+            creation_flags,
+            custom_env_block_pointer,
             cwd.as_ref().map_or_else(ptr::null, |s| s.as_ptr()),
             &mut startup_info_ex.StartupInfo as *mut STARTUPINFOW,
             &mut proc_info as *mut PROCESS_INFORMATION,
@@ -230,6 +243,63 @@ pub fn new(config: &Options, window_size: WindowSize) -> Option<Pty> {
     Some(Pty::new(conpty, conout, conin, child_watcher))
 }
 
+// Windows environment variables are case-insensitive, and the caller is responsible for
+// deduplicating environment variables, so do that here while converting.
+//
+// https://learn.microsoft.com/en-us/previous-versions/troubleshoot/windows/win32/createprocess-cannot-eliminate-duplicate-variables#environment-variables
+fn convert_custom_env(custom_env: &HashMap<String, String>) -> Option<Vec<u16>> {
+    // Windows inherits parent's env when no `lpEnvironment` parameter is specified.
+    if custom_env.is_empty() {
+        return None;
+    }
+
+    let mut converted_block = Vec::new();
+    let mut all_env_keys = HashSet::new();
+    for (custom_key, custom_value) in custom_env {
+        let custom_key_os = OsStr::new(custom_key);
+        if all_env_keys.insert(custom_key_os.to_ascii_uppercase()) {
+            add_windows_env_key_value_to_block(
+                &mut converted_block,
+                custom_key_os,
+                OsStr::new(&custom_value),
+            );
+        } else {
+            warn!(
+                "Omitting environment variable pair with duplicate key: \
+                 '{custom_key}={custom_value}'"
+            );
+        }
+    }
+
+    // Pull the current process environment after, to avoid overwriting the user provided one.
+    for (inherited_key, inherited_value) in std::env::vars_os() {
+        if all_env_keys.insert(inherited_key.to_ascii_uppercase()) {
+            add_windows_env_key_value_to_block(
+                &mut converted_block,
+                &inherited_key,
+                &inherited_value,
+            );
+        }
+    }
+
+    converted_block.push(0);
+    Some(converted_block)
+}
+
+// According to the `lpEnvironment` parameter description:
+// https://learn.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-createprocessa#parameters
+//
+// > An environment block consists of a null-terminated block of null-terminated strings. Each
+// string is in the following form:
+// >
+// > name=value\0
+fn add_windows_env_key_value_to_block(block: &mut Vec<u16>, key: &OsStr, value: &OsStr) {
+    block.extend(key.encode_wide());
+    block.push('=' as u16);
+    block.extend(value.encode_wide());
+    block.push(0);
+}
+
 // Panic with the last os error as message.
 fn panic_shell_spawn() {
     panic!("Unable to spawn shell: {}", Error::last_os_error());