From 0c296199c0408c20d8aa02ce2d4393148557391a Mon Sep 17 00:00:00 2001
From: Steven Malis <stevenmalis@microsoft.com>
Date: Tue, 17 Dec 2024 15:40:00 -0500
Subject: [PATCH] uh_core: Don't modify the environment for start requests.

---
 openhcl/underhill_core/src/lib.rs     | 11 ++-----
 openhcl/underhill_core/src/options.rs | 42 +++++++++++++++++++--------
 2 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/openhcl/underhill_core/src/lib.rs b/openhcl/underhill_core/src/lib.rs
index 764b42d139..f9356451d7 100644
--- a/openhcl/underhill_core/src/lib.rs
+++ b/openhcl/underhill_core/src/lib.rs
@@ -192,7 +192,7 @@ fn install_task_name_panic_hook() {
 }
 
 async fn do_main(driver: DefaultDriver, mut tracing: TracingBackend) -> anyhow::Result<()> {
-    let opt = Options::parse(Vec::new())?;
+    let opt = Options::parse(Vec::new(), Vec::new())?;
 
     let crate_name = build_info::get().crate_name();
     let crate_revision = build_info::get().scm_revision();
@@ -496,14 +496,7 @@ async fn run_control(
                             if workers.is_some() {
                                 Err(anyhow::anyhow!("workers have already been started"))?;
                             }
-                            for (key, value) in params.env {
-                                if let Some(value) = value {
-                                    std::env::set_var(key, value);
-                                } else {
-                                    std::env::remove_var(key);
-                                }
-                            }
-                            let new_opt = Options::parse(params.args)
+                            let new_opt = Options::parse(params.args, params.env)
                                 .context("failed to parse new options")?;
 
                             workers = Some(
diff --git a/openhcl/underhill_core/src/options.rs b/openhcl/underhill_core/src/options.rs
index d4ea47dad6..f67df04194 100644
--- a/openhcl/underhill_core/src/options.rs
+++ b/openhcl/underhill_core/src/options.rs
@@ -7,6 +7,9 @@
 
 use anyhow::bail;
 use anyhow::Context;
+use std::collections::BTreeMap;
+use std::ffi::OsStr;
+use std::ffi::OsString;
 use std::path::PathBuf;
 
 // We've made our own parser here instead of using something like clap in order
@@ -119,26 +122,41 @@ pub struct Options {
 }
 
 impl Options {
-    pub(crate) fn parse(extra_args: Vec<String>) -> anyhow::Result<Self> {
-        /// Reads an environment variable, falling back to a legacy variable (replacing
-        /// "OPENHCL_" with "UNDERHILL_") if the original is not set.
-        fn legacy_openhcl_env(name: &str) -> Option<std::ffi::OsString> {
-            std::env::var_os(name).or_else(|| {
-                std::env::var_os(format!(
-                    "UNDERHILL_{}",
-                    name.strip_prefix("OPENHCL_").unwrap_or(name)
-                ))
-            })
+    pub(crate) fn parse(
+        extra_args: Vec<String>,
+        extra_env: Vec<(String, Option<String>)>,
+    ) -> anyhow::Result<Self> {
+        // Pull the entire environment into a BTreeMap for manipulation through extra_env.
+        let mut env: BTreeMap<OsString, OsString> = std::env::vars_os().collect();
+        for (key, value) in extra_env {
+            match value {
+                Some(value) => env.insert(key.into(), value.into()),
+                None => env.remove::<OsStr>(key.as_ref()),
+            };
         }
 
-        fn parse_bool(value: Option<std::ffi::OsString>) -> bool {
+        // Reads an environment variable, falling back to a legacy variable (replacing
+        // "OPENHCL_" with "UNDERHILL_") if the original is not set.
+        let legacy_openhcl_env = |name: &str| -> Option<&OsString> {
+            env.get::<OsStr>(name.as_ref()).or_else(|| {
+                env.get::<OsStr>(
+                    format!(
+                        "UNDERHILL_{}",
+                        name.strip_prefix("OPENHCL_").unwrap_or(name)
+                    )
+                    .as_ref(),
+                )
+            })
+        };
+
+        fn parse_bool(value: Option<&OsString>) -> bool {
             value
                 .map(|v| v.to_ascii_lowercase() == "true" || v == "1")
                 .unwrap_or_default()
         }
 
         let parse_legacy_env_bool = |name| parse_bool(legacy_openhcl_env(name));
-        let parse_env_bool = |name| parse_bool(std::env::var_os(name));
+        let parse_env_bool = |name: &str| parse_bool(env.get::<OsStr>(name.as_ref()));
 
         let parse_legacy_env_number = |name| {
             legacy_openhcl_env(name)