Files
codex/codex-rs/arg0/src/openai_api_key_env_var.rs
2025-09-25 13:28:43 -07:00

81 lines
2.2 KiB
Rust

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) {}