mirror of
https://github.com/openai/codex.git
synced 2026-05-02 04:11:39 +03:00
81 lines
2.2 KiB
Rust
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) {}
|