fix(security): make unsafe env::set_var calls safe with explicit invariants (#968)

* fix(security): make unsafe env::set_var calls safe with explicit invariants

`std::env::set_var` is unsafe in Rust 1.82+ because concurrent calls
from multiple threads cause undefined behavior. This commit addresses
the two production-code call sites:

1. `bootstrap.rs:load_ironclaw_env()` -- called before the Tokio
   runtime starts (genuinely single-threaded). Added a `debug_assert!`
   that verifies no Tokio runtime is active, making the safety
   invariant machine-checkable rather than relying on a comment.

2. `llm/session.rs:api_key_login()` -- was calling `set_var` mid-
   execution inside the multi-threaded Tokio runtime (UB risk).
   Replaced with `set_runtime_env()`, a new thread-safe overlay
   backed by `OnceLock<Mutex<HashMap>>`. The overlay integrates with
   the existing `optional_env()` config resolution and a new
   `env_or_override()` reader function.

All call sites that read `NEARAI_API_KEY` via raw `std::env::var()`
(wizard.rs, main.rs, doctor.rs) are updated to use the thread-safe
`env_or_override()` helper instead, so the value set during
interactive login is visible without mutating the process environment.

Test code `set_var`/`remove_var` calls (bootstrap tests, config tests,
shell tests, oauth tests, wizard tests) are left as-is since they run
under `ENV_MUTEX` serialization and are not production paths.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

* fix: address review feedback on thread-safe env overlay PR

- Replace debug_assert! with runtime check in bootstrap.rs so release
  builds skip unsafe set_var when a Tokio runtime is active
- Recover from mutex poison in set_runtime_env instead of silently
  dropping writes (poisoned HashMap is still usable)
- Skip empty override values in env_or_override and optional_env for
  consistency with real env var handling
- Fix doc comment on env_or_override (real env checked first, not
  runtime overrides)
- Update api_key_login doc to describe runtime overlay instead of
  env var mutation

[skip-regression-check]

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

* fix: use LazyLock::lock() for INJECTED_VARS; use set_runtime_env() in bootstrap fallback

- helpers.rs: fix env_or_override() to call INJECTED_VARS.lock() instead
  of .get() — INJECTED_VARS was changed upstream from OnceLock<HashMap>
  to LazyLock<Mutex<HashMap>>; calling .get() caused a compile error
  (E0599: no method named 'get' for LazyLock)

- bootstrap.rs: when load_ironclaw_env() is called with an active Tokio
  runtime, use set_runtime_env("DATABASE_BACKEND", "libsql") instead of
  silently dropping the write. This ensures DATABASE_BACKEND is always
  set regardless of thread context (addresses ilblackdragon review item 1).

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

---------

Co-authored-by: Gabe Hamilton <[email protected]>
Co-authored-by: Claude Sonnet 4.6 <[email protected]>

src/bootstrap.rs

@@ -116,9 +116,18 @@ pub fn load_ironclaw_env() {
             .join(".ironclaw")
             .join("ironclaw.db");
         if default_db.exists() {
-            // SAFETY: `load_ironclaw_env` is called from a synchronous `fn main()`
-            // before the Tokio runtime is started, so no other threads exist yet.
-            unsafe { std::env::set_var("DATABASE_BACKEND", "libsql") };
+            if tokio::runtime::Handle::try_current().is_ok() {
+                // Tokio runtime is active (multi-threaded); std::env::set_var is UB here.
+                // Fall back to the thread-safe runtime overlay so the value is always set.
+                tracing::warn!(
+                    "load_ironclaw_env called with active Tokio runtime; \
+                     using runtime env overlay for DATABASE_BACKEND"
+                );
+                crate::config::set_runtime_env("DATABASE_BACKEND", "libsql");
+            } else {
+                // SAFETY: No Tokio runtime = no other threads = safe to call set_var.
+                unsafe { std::env::set_var("DATABASE_BACKEND", "libsql") };
+            }
         }
     }
 }

src/cli/doctor.rs

@@ -220,7 +220,7 @@ async fn check_nearai_session() -> CheckResult {
     let session_path = crate::config::llm::default_session_path();
     if !session_path.exists() {
         // Check for API key mode
-        if std::env::var("NEARAI_API_KEY").is_ok() {
+        if crate::config::helpers::env_or_override("NEARAI_API_KEY").is_some() {
             return CheckResult::Pass("API key configured".into());
         }
         return CheckResult::Fail(format!(

src/config/helpers.rs

@@ -1,6 +1,9 @@
+use std::collections::HashMap;
+use std::sync::{Mutex, OnceLock};
+
 use crate::error::ConfigError;
 
-use super::INJECTED_VARS;
+use crate::config::INJECTED_VARS;
 
 /// Crate-wide mutex for tests that mutate process environment variables.
 ///
@@ -11,6 +14,73 @@ use super::INJECTED_VARS;
 #[cfg(test)]
 pub(crate) static ENV_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
 
+/// Thread-safe mutable overlay for env vars set at runtime.
+///
+/// Unlike `INJECTED_VARS` (which is set once at startup from the secrets
+/// store), this map supports writes at any point during the process
+/// lifetime. It replaces unsafe `std::env::set_var` calls that would
+/// otherwise be UB in multi-threaded programs (Rust 1.82+).
+///
+/// Priority: real env vars > `RUNTIME_ENV_OVERRIDES` > `INJECTED_VARS`.
+static RUNTIME_ENV_OVERRIDES: OnceLock<Mutex<HashMap<String, String>>> = OnceLock::new();
+
+fn runtime_overrides() -> &'static Mutex<HashMap<String, String>> {
+    RUNTIME_ENV_OVERRIDES.get_or_init(|| Mutex::new(HashMap::new()))
+}
+
+/// Set a runtime environment override (thread-safe alternative to `std::env::set_var`).
+///
+/// Values set here are visible to `optional_env()`, `env_or_override()`, and
+/// all config resolution that goes through those helpers. This avoids the UB
+/// of `std::env::set_var` in multi-threaded programs.
+pub fn set_runtime_env(key: &str, value: &str) {
+    runtime_overrides()
+        .lock()
+        .unwrap_or_else(|e| e.into_inner())
+        .insert(key.to_string(), value.to_string());
+}
+
+/// Read an env var, checking the real environment first, then runtime overrides.
+///
+/// Priority: real env vars > runtime overrides > `INJECTED_VARS`.
+/// Empty values are treated as unset at every layer for consistency with
+/// `optional_env()`.
+///
+/// Use this instead of `std::env::var()` when the value might have been set
+/// via `set_runtime_env()` (e.g., `NEARAI_API_KEY` during interactive login).
+pub fn env_or_override(key: &str) -> Option<String> {
+    // Real env vars always win
+    if let Ok(val) = std::env::var(key)
+        && !val.is_empty()
+    {
+        return Some(val);
+    }
+
+    // Check runtime overrides (skip empty values for consistency with optional_env)
+    if let Some(val) = runtime_overrides()
+        .lock()
+        .unwrap_or_else(|e| e.into_inner())
+        .get(key)
+        .filter(|v| !v.is_empty())
+        .cloned()
+    {
+        return Some(val);
+    }
+
+    // Check INJECTED_VARS (secrets from DB, set once at startup)
+    if let Some(val) = INJECTED_VARS
+        .lock()
+        .unwrap_or_else(|e| e.into_inner())
+        .get(key)
+        .filter(|v| !v.is_empty())
+        .cloned()
+    {
+        return Some(val);
+    }
+
+    None
+}
+
 pub(crate) fn optional_env(key: &str) -> Result<Option<String>, ConfigError> {
     // Check real env vars first (always win over injected secrets)
     match std::env::var(key) {
@@ -24,6 +94,17 @@ pub(crate) fn optional_env(key: &str) -> Result<Option<String>, ConfigError> {
         }
     }
 
+    // Fall back to runtime overrides (set via set_runtime_env)
+    if let Some(val) = runtime_overrides()
+        .lock()
+        .unwrap_or_else(|e| e.into_inner())
+        .get(key)
+        .filter(|v| !v.is_empty())
+        .cloned()
+    {
+        return Ok(Some(val));
+    }
+
     // Fall back to thread-safe overlay (secrets injected from DB)
     if let Some(val) = INJECTED_VARS
         .lock()
@@ -94,3 +175,55 @@ pub(crate) fn parse_string_env(
 ) -> Result<String, ConfigError> {
     Ok(optional_env(key)?.unwrap_or_else(|| default.into()))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn runtime_env_override_is_visible_to_env_or_override() {
+        // Use a unique key that won't collide with real env vars.
+        let key = "IRONCLAW_TEST_RUNTIME_OVERRIDE_42";
+
+        // Not set initially
+        assert!(env_or_override(key).is_none());
+
+        // Set via the thread-safe overlay
+        set_runtime_env(key, "test_value");
+
+        // Now visible
+        assert_eq!(env_or_override(key), Some("test_value".to_string()));
+    }
+
+    #[test]
+    fn runtime_env_override_is_visible_to_optional_env() {
+        let key = "IRONCLAW_TEST_OPTIONAL_ENV_OVERRIDE_42";
+
+        assert_eq!(optional_env(key).unwrap(), None);
+
+        set_runtime_env(key, "hello");
+
+        assert_eq!(optional_env(key).unwrap(), Some("hello".to_string()));
+    }
+
+    #[test]
+    fn real_env_var_takes_priority_over_runtime_override() {
+        let _guard = ENV_MUTEX.lock().unwrap();
+        let key = "IRONCLAW_TEST_ENV_PRIORITY_42";
+
+        // Set runtime override
+        set_runtime_env(key, "override_value");
+
+        // Set real env var (should win)
+        // SAFETY: test runs under ENV_MUTEX
+        unsafe { std::env::set_var(key, "real_value") };
+
+        assert_eq!(env_or_override(key), Some("real_value".to_string()));
+
+        // Clean up
+        unsafe { std::env::remove_var(key) };
+
+        // Now the runtime override is visible again
+        assert_eq!(env_or_override(key), Some("override_value".to_string()));
+    }
+}

src/config/mod.rs

@@ -54,6 +54,10 @@ pub use crate::llm::config::{
 };
 pub use crate::llm::session::SessionConfig;
 
+// Thread-safe env var override helpers (replaces unsafe `std::env::set_var`
+// for mid-process env mutations in multi-threaded contexts).
+pub use self::helpers::{env_or_override, set_runtime_env};
+
 /// Thread-safe overlay for injected env vars (secrets loaded from DB).
 ///
 /// Used by `inject_llm_keys_from_secrets()` to make API keys available to

src/llm/session.rs

@@ -373,9 +373,10 @@ impl SessionManager {
     /// NEAR AI Cloud API key entry flow.
     ///
     /// Prompts the user to enter a NEAR AI Cloud API key from
-    /// cloud.near.ai. The key is set as `NEARAI_API_KEY` env var so
-    /// `LlmConfig::resolve()` auto-selects ChatCompletions mode, and
-    /// saved to `~/.ironclaw/.env` for persistence across restarts.
+    /// cloud.near.ai. The key is stored in the thread-safe runtime
+    /// env overlay (via `set_runtime_env`) so `LlmConfig::resolve()`
+    /// auto-selects ChatCompletions mode, and persisted to
+    /// `~/.ironclaw/.env` for survival across restarts.
     /// No session token is saved and no `/v1/users/me` validation is
     /// performed (different auth model).
     async fn api_key_login(&self) -> Result<(), LlmError> {
@@ -403,15 +404,11 @@ impl SessionManager {
             });
         }
 
-        // Set env var so Config picks it up immediately
-        // (LlmConfig::resolve() auto-selects ChatCompletions mode when
-        // NEARAI_API_KEY is present).
-        //
-        // SAFETY: called during single-threaded interactive login flow.
-        #[allow(unused_unsafe)]
-        unsafe {
-            std::env::set_var("NEARAI_API_KEY", &key);
-        }
+        // Make the key visible to Config resolution and `env_or_override()`
+        // callers for the remainder of this process. Uses a thread-safe
+        // overlay instead of `std::env::set_var`, which is UB in
+        // multi-threaded programs (Rust 1.82+).
+        crate::config::helpers::set_runtime_env("NEARAI_API_KEY", &key);
 
         // Persist to ~/.ironclaw/.env so the key survives restarts
         // (bootstrap layer — available before DB is connected).

src/setup/wizard.rs

@@ -1183,9 +1183,9 @@ impl SetupWizard {
         self.persist_session_to_db().await;
 
         // If the user chose the API key path, NEARAI_API_KEY is now set
-        // in the environment. Persist it to the encrypted secrets store
-        // so inject_llm_keys_from_secrets() can load it on future runs.
-        if let Ok(api_key) = std::env::var("NEARAI_API_KEY")
+        // in the runtime env overlay. Persist it to the encrypted secrets
+        // store so inject_llm_keys_from_secrets() can load it on future runs.
+        if let Some(api_key) = crate::config::helpers::env_or_override("NEARAI_API_KEY")
             && !api_key.is_empty()
             && let Ok(ctx) = self.init_secrets_context().await
         {
@@ -2613,8 +2613,9 @@ impl SetupWizard {
             env_vars.push((base_url_env.clone(), base_url.clone()));
         }
 
-        // Preserve NEARAI_API_KEY if present (set by API key auth flow)
-        if let Ok(api_key) = std::env::var("NEARAI_API_KEY")
+        // Preserve NEARAI_API_KEY if present (set by API key auth flow
+        // via the thread-safe runtime env overlay).
+        if let Some(api_key) = crate::config::helpers::env_or_override("NEARAI_API_KEY")
             && !api_key.is_empty()
         {
             env_vars.push(("NEARAI_API_KEY".to_string(), api_key));