Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 137 additions & 4 deletions codex-rs/arg0/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use std::fs::File;
use std::future::Future;
use std::path::Path;
use std::path::PathBuf;
Expand All @@ -10,8 +11,24 @@ use tempfile::TempDir;
const LINUX_SANDBOX_ARG0: &str = "codex-linux-sandbox";
const APPLY_PATCH_ARG0: &str = "apply_patch";
const MISSPELLED_APPLY_PATCH_ARG0: &str = "applypatch";
const LOCK_FILENAME: &str = ".lock";

pub fn arg0_dispatch() -> Option<TempDir> {
/// Keeps the per-session PATH entry alive and locked for the process lifetime.
pub struct Arg0PathEntryGuard {
_temp_dir: TempDir,
_lock_file: File,
}

impl Arg0PathEntryGuard {
fn new(temp_dir: TempDir, lock_file: File) -> Self {
Self {
_temp_dir: temp_dir,
_lock_file: lock_file,
}
}
}

pub fn arg0_dispatch() -> Option<Arg0PathEntryGuard> {
// Determine if we were invoked via the special alias.
let mut args = std::env::args_os();
let argv0 = args.next().unwrap_or_default();
Expand Down Expand Up @@ -149,7 +166,7 @@ where
///
/// IMPORTANT: This function modifies the PATH environment variable, so it MUST
/// be called before multiple threads are spawned.
pub fn prepend_path_entry_for_codex_aliases() -> std::io::Result<TempDir> {
pub fn prepend_path_entry_for_codex_aliases() -> std::io::Result<Arg0PathEntryGuard> {
let codex_home = codex_core::config::find_codex_home()?;
#[cfg(not(debug_assertions))]
{
Expand All @@ -167,7 +184,7 @@ pub fn prepend_path_entry_for_codex_aliases() -> std::io::Result<TempDir> {

std::fs::create_dir_all(&codex_home)?;
// Use a CODEX_HOME-scoped temp root to avoid cluttering the top-level directory.
let temp_root = codex_home.join("tmp").join("path");
let temp_root = codex_home.join("tmp").join("arg0");
std::fs::create_dir_all(&temp_root)?;
#[cfg(unix)]
{
Expand All @@ -177,11 +194,25 @@ pub fn prepend_path_entry_for_codex_aliases() -> std::io::Result<TempDir> {
std::fs::set_permissions(&temp_root, std::fs::Permissions::from_mode(0o700))?;
}

// Best-effort cleanup of stale per-session dirs. Ignore failures so startup proceeds.
if let Err(err) = janitor_cleanup(&temp_root) {
eprintln!("WARNING: failed to clean up stale arg0 temp dirs: {err}");
}

let temp_dir = tempfile::Builder::new()
.prefix("codex-arg0")
.tempdir_in(&temp_root)?;
let path = temp_dir.path();

let lock_path = path.join(LOCK_FILENAME);
let lock_file = File::options()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&lock_path)?;
lock_file.try_lock()?;

for filename in &[
APPLY_PATCH_ARG0,
MISSPELLED_APPLY_PATCH_ARG0,
Expand Down Expand Up @@ -231,5 +262,107 @@ pub fn prepend_path_entry_for_codex_aliases() -> std::io::Result<TempDir> {
std::env::set_var("PATH", updated_path_env_var);
}

Ok(temp_dir)
Ok(Arg0PathEntryGuard::new(temp_dir, lock_file))
}

fn janitor_cleanup(temp_root: &Path) -> std::io::Result<()> {
let entries = match std::fs::read_dir(temp_root) {
Ok(entries) => entries,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(err) => return Err(err),
};

for entry in entries.flatten() {
let path = entry.path();
if !path.is_dir() {
continue;
}

// Skip the directory if locking fails or the lock is currently held.
let Some(_lock_file) = try_lock_dir(&path)? else {
continue;
};

match std::fs::remove_dir_all(&path) {
Ok(()) => {}
// Expected TOCTOU race: directory can disappear after read_dir/lock checks.
Err(err) if err.kind() == std::io::ErrorKind::NotFound => continue,
Err(err) => return Err(err),
}
}

Ok(())
}

fn try_lock_dir(dir: &Path) -> std::io::Result<Option<File>> {
let lock_path = dir.join(LOCK_FILENAME);
let lock_file = match File::options().read(true).write(true).open(&lock_path) {
Ok(file) => file,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(err) => return Err(err),
};

match lock_file.try_lock() {
Ok(()) => Ok(Some(lock_file)),
Err(std::fs::TryLockError::WouldBlock) => Ok(None),
Err(err) => Err(err.into()),
}
}

#[cfg(test)]
mod tests {
use super::LOCK_FILENAME;
use super::janitor_cleanup;
use std::fs;
use std::fs::File;
use std::path::Path;

fn create_lock(dir: &Path) -> std::io::Result<File> {
let lock_path = dir.join(LOCK_FILENAME);
File::options()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(lock_path)
}

#[test]
fn janitor_skips_dirs_without_lock_file() -> std::io::Result<()> {
let root = tempfile::tempdir()?;
let dir = root.path().join("no-lock");
fs::create_dir(&dir)?;

janitor_cleanup(root.path())?;

assert!(dir.exists());
Ok(())
}

#[test]
fn janitor_skips_dirs_with_held_lock() -> std::io::Result<()> {
let root = tempfile::tempdir()?;
let dir = root.path().join("locked");
fs::create_dir(&dir)?;
let lock_file = create_lock(&dir)?;
lock_file.try_lock()?;

janitor_cleanup(root.path())?;

assert!(dir.exists());
Ok(())
}

#[test]
fn janitor_removes_dirs_with_unlocked_lock() -> std::io::Result<()> {
let root = tempfile::tempdir()?;
let dir = root.path().join("stale");
fs::create_dir(&dir)?;
create_lock(&dir)?;

janitor_cleanup(root.path())?;

assert!(!dir.exists());
Ok(())
}
}
45 changes: 43 additions & 2 deletions codex-rs/core/tests/suite/mod.rs
Original file line number Diff line number Diff line change
@@ -1,16 +1,57 @@
// Aggregates all former standalone integration tests as modules.
use std::ffi::OsString;

use codex_arg0::Arg0PathEntryGuard;
use codex_arg0::arg0_dispatch;
use ctor::ctor;
use tempfile::TempDir;

struct TestCodexAliasesGuard {
_codex_home: TempDir,
_arg0: Arg0PathEntryGuard,
_previous_codex_home: Option<OsString>,
}

const CODEX_HOME_ENV_VAR: &str = "CODEX_HOME";

// This code runs before any other tests are run.
// It allows the test binary to behave like codex and dispatch to apply_patch and codex-linux-sandbox
// based on the arg0.
// NOTE: this doesn't work on ARM
#[ctor]
pub static CODEX_ALIASES_TEMP_DIR: TempDir = unsafe {
pub static CODEX_ALIASES_TEMP_DIR: TestCodexAliasesGuard = unsafe {
#[allow(clippy::unwrap_used)]
arg0_dispatch().unwrap()
let codex_home = tempfile::Builder::new()
.prefix("codex-core-tests")
.tempdir()
.unwrap();
let previous_codex_home = std::env::var_os(CODEX_HOME_ENV_VAR);
// arg0_dispatch() creates helper links under CODEX_HOME/tmp. Point it at a
// test-owned temp dir so startup never mutates the developer's real ~/.codex.
//
// Safety: #[ctor] runs before tests start, so no test threads exist yet.
unsafe {
std::env::set_var(CODEX_HOME_ENV_VAR, codex_home.path());
}

#[allow(clippy::unwrap_used)]
let arg0 = arg0_dispatch().unwrap();
// Restore the process environment immediately so later tests observe the
// same CODEX_HOME state they started with.
match previous_codex_home.as_ref() {
Some(value) => unsafe {
std::env::set_var(CODEX_HOME_ENV_VAR, value);
},
None => unsafe {
std::env::remove_var(CODEX_HOME_ENV_VAR);
},
}

TestCodexAliasesGuard {
_codex_home: codex_home,
_arg0: arg0,
_previous_codex_home: previous_codex_home,
}
};

#[cfg(not(target_os = "windows"))]
Expand Down
Loading