feat: remove OPENAI_API_KEY from env and mlock if set

codex-rs/arg0/Cargo.toml

@@ -16,5 +16,6 @@ codex-apply-patch = { workspace = true }
 codex-core = { workspace = true }
 codex-linux-sandbox = { workspace = true }
 dotenvy = { workspace = true }
+libc = { workspace = true }
 tempfile = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread"] }

codex-rs/arg0/src/lib.rs

@@ -7,10 +7,31 @@ use codex_core::CODEX_APPLY_PATCH_ARG1;
 use std::os::unix::fs::symlink;
 use tempfile::TempDir;
 
+mod openai_api_key_env_var;
+use openai_api_key_env_var::extract_locked_openai_api_key;
+
 const LINUX_SANDBOX_ARG0: &str = "codex-linux-sandbox";
 const APPLY_PATCH_ARG0: &str = "apply_patch";
 const MISSPELLED_APPLY_PATCH_ARG0: &str = "applypatch";
 
+/// Arguments supplied to the async entrypoint invoked by [`arg0_dispatch_or_else`].
+///
+/// Currently, these args are not technically computed "pre-main," but they are
+/// computed before spawning any threads / the tokio runtime.
+#[derive(Debug)]
+pub struct PreMainArgs {
+    pub codex_linux_sandbox_exe: Option<PathBuf>,
+
+    /// Value of the `OPENAI_API_KEY` environment variable that was set at
+    /// startup, if any.
+    ///
+    /// If `Some`, the key has already been removed from the environment and an
+    /// attempt was made to `mlock(2)` the string's allocation to keep it
+    /// resident in memory. The reference is leaked for the lifetime of the
+    /// process so it can be safely shared across threads without copying.
+    pub openai_api_key: Option<&'static str>,
+}
+
 /// While we want to deploy the Codex CLI as a single executable for simplicity,
 /// we also want to expose some of its functionality as distinct CLIs, so we use
 /// the "arg0 trick" to determine which CLI to dispatch. This effectively allows
@@ -22,19 +43,20 @@ const MISSPELLED_APPLY_PATCH_ARG0: &str = "applypatch";
 /// [`codex_linux_sandbox::run_main`] (which never returns). Otherwise we:
 ///
 /// 1.  Load `.env` values from `~/.codex/.env` before creating any threads.
-/// 2.  Construct a Tokio multi-thread runtime.
+/// 2.  Depending on `extract_openai_api_key`, extract and lock the
+///     `OPENAI_API_KEY` environment variable, if present.
 /// 3.  Derive the path to the current executable (so children can re-invoke the
 ///     sandbox) when running on Linux.
-/// 4.  Execute the provided async `main_fn` inside that runtime, forwarding any
-///     error. Note that `main_fn` receives `codex_linux_sandbox_exe:
-///     Option<PathBuf>`, as an argument, which is generally needed as part of
-///     constructing [`codex_core::config::Config`].
+/// 4.  Construct a Tokio multi-thread runtime.
+/// 5.  Execute the provided async `main_fn` inside that runtime, forwarding any
+///     error. The closure receives [`PreMainArgs`], which includes the optional
+///     path to the linux sandbox helper as well as the OpenAI API key.
 ///
 /// This function should be used to wrap any `main()` function in binary crates
 /// in this workspace that depends on these helper CLIs.
-pub fn arg0_dispatch_or_else<F, Fut>(main_fn: F) -> anyhow::Result<()>
+pub fn arg0_dispatch_or_else<F, Fut>(extract_openai_api_key: bool, main_fn: F) -> anyhow::Result<()>
 where
-    F: FnOnce(Option<PathBuf>) -> Fut,
+    F: FnOnce(PreMainArgs) -> Fut,
     Fut: Future<Output = anyhow::Result<()>>,
 {
     // Determine if we were invoked via the special alias.
@@ -76,6 +98,13 @@ where
     // before creating any threads/the Tokio runtime.
     load_dotenv();
 
+    // Perform the OPENAI_API_KEY check after loading the .env file.
+    let openai_api_key = if extract_openai_api_key {
+        extract_locked_openai_api_key()
+    } else {
+        None
+    };
+
     // Retain the TempDir so it exists for the lifetime of the invocation of
     // this executable. Admittedly, we could invoke `keep()` on it, but it
     // would be nice to avoid leaving temporary directories behind, if possible.
@@ -99,7 +128,12 @@ where
             None
         };
 
-        main_fn(codex_linux_sandbox_exe).await
+        let args = PreMainArgs {
+            codex_linux_sandbox_exe,
+            openai_api_key,
+        };
+
+        main_fn(args).await
     })
 }
 

codex-rs/arg0/src/openai_api_key_env_var.rs

@@ -0,0 +1,80 @@
+const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY";
+
+pub(crate) fn extract_locked_openai_api_key() -> Option<&'static str> {
+    match std::env::var(OPENAI_API_KEY_ENV_VAR) {
+        Ok(key) => {
+            if key.is_empty() {
+                return None;
+            }
+
+            // Safety: modifying environment variables is only done before new
+            // threads are spawned.
+            clear_api_key_env_var();
+
+            // into_boxed_str() may reallocate, so only lock the memory after
+            // the final allocation is known.
+            let leaked: &'static mut str = Box::leak(key.into_boxed_str());
+            mlock_str(leaked);
+            Some(leaked)
+        }
+        Err(std::env::VarError::NotPresent) => None,
+        Err(std::env::VarError::NotUnicode(_)) => {
+            // Cannot possibly be a valid API key, but we will clear it anyway.
+            clear_api_key_env_var();
+
+            None
+        }
+    }
+}
+
+/// Note this does not guarantee that the memory is wiped, only that the
+/// environment variable is removed from this process's environment.
+fn clear_api_key_env_var() {
+    unsafe {
+        std::env::remove_var(OPENAI_API_KEY_ENV_VAR);
+    }
+}
+
+#[cfg(unix)]
+fn mlock_str(value: &str) {
+    use libc::_SC_PAGESIZE;
+    use libc::c_void;
+    use libc::mlock;
+    use libc::sysconf;
+
+    if value.is_empty() {
+        return;
+    }
+
+    // Safety: we only read the pointer and length for mlock bookkeeping.
+    let page_size = unsafe { sysconf(_SC_PAGESIZE) };
+    if page_size <= 0 {
+        return;
+    }
+    let page_size = page_size as usize;
+    if page_size == 0 {
+        return;
+    }
+
+    let addr = value.as_ptr() as usize;
+    let len = value.len();
+    let start = addr & !(page_size - 1);
+    let addr_end = match addr.checked_add(len) {
+        Some(v) => match v.checked_add(page_size - 1) {
+            Some(total) => total,
+            None => return,
+        },
+        None => return,
+    };
+    let end = addr_end & !(page_size - 1);
+    let size = end.saturating_sub(start);
+    if size == 0 {
+        return;
+    }
+
+    // Best-effort; ignore failures because mlock may require privileges.
+    let _ = unsafe { mlock(start as *const c_void, size) };
+}
+
+#[cfg(not(unix))]
+fn mlock_str(_value: &str) {}

codex-rs/cli/src/main.rs

@@ -2,6 +2,7 @@ use clap::CommandFactory;
 use clap::Parser;
 use clap_complete::Shell;
 use clap_complete::generate;
+use codex_arg0::PreMainArgs;
 use codex_arg0::arg0_dispatch_or_else;
 use codex_chatgpt::apply_command::ApplyCommand;
 use codex_chatgpt::apply_command::run_apply_command;
@@ -224,13 +225,18 @@ fn pre_main_hardening() {
 }
 
 fn main() -> anyhow::Result<()> {
-    arg0_dispatch_or_else(|codex_linux_sandbox_exe| async move {
-        cli_main(codex_linux_sandbox_exe).await?;
+    arg0_dispatch_or_else(true, |pre_main_args| async move {
+        cli_main(pre_main_args).await?;
         Ok(())
     })
 }
 
-async fn cli_main(codex_linux_sandbox_exe: Option<PathBuf>) -> anyhow::Result<()> {
+async fn cli_main(pre_main_args: PreMainArgs) -> anyhow::Result<()> {
+    let PreMainArgs {
+        codex_linux_sandbox_exe,
+        openai_api_key: _,
+    } = pre_main_args;
+
     let MultitoolCli {
         config_overrides: root_config_overrides,
         mut interactive,

codex-rs/exec/src/main.rs

@@ -10,6 +10,7 @@
 //! This allows us to ship a completely separate set of functionality as part
 //! of the `codex-exec` binary.
 use clap::Parser;
+use codex_arg0::PreMainArgs;
 use codex_arg0::arg0_dispatch_or_else;
 use codex_common::CliConfigOverrides;
 use codex_exec::Cli;
@@ -25,7 +26,11 @@ struct TopCli {
 }
 
 fn main() -> anyhow::Result<()> {
-    arg0_dispatch_or_else(|codex_linux_sandbox_exe| async move {
+    arg0_dispatch_or_else(false, |pre_main_args| async move {
+        let PreMainArgs {
+            codex_linux_sandbox_exe,
+            openai_api_key: _,
+        } = pre_main_args;
         let top_cli = TopCli::parse();
         // Merge root-level overrides into inner CLI struct so downstream logic remains unchanged.
         let mut inner = top_cli.inner;

codex-rs/mcp-server/src/main.rs

@@ -1,9 +1,14 @@
+use codex_arg0::PreMainArgs;
 use codex_arg0::arg0_dispatch_or_else;
 use codex_common::CliConfigOverrides;
 use codex_mcp_server::run_main;
 
 fn main() -> anyhow::Result<()> {
-    arg0_dispatch_or_else(|codex_linux_sandbox_exe| async move {
+    arg0_dispatch_or_else(false, |pre_main_args| async move {
+        let PreMainArgs {
+            codex_linux_sandbox_exe,
+            openai_api_key: _,
+        } = pre_main_args;
         run_main(codex_linux_sandbox_exe, CliConfigOverrides::default()).await?;
         Ok(())
     })

codex-rs/tui/src/main.rs

@@ -1,4 +1,5 @@
 use clap::Parser;
+use codex_arg0::PreMainArgs;
 use codex_arg0::arg0_dispatch_or_else;
 use codex_common::CliConfigOverrides;
 use codex_tui::Cli;
@@ -14,7 +15,11 @@ struct TopCli {
 }
 
 fn main() -> anyhow::Result<()> {
-    arg0_dispatch_or_else(|codex_linux_sandbox_exe| async move {
+    arg0_dispatch_or_else(false, |pre_main_args| async move {
+        let PreMainArgs {
+            codex_linux_sandbox_exe,
+            openai_api_key: _,
+        } = pre_main_args;
         let top_cli = TopCli::parse();
         let mut inner = top_cli.inner;
         inner