use hashbrown::HashMap;
use std::path::Path;
pub const FILTERED_ENV_VARS: &[&str] = &[
"OPENAI_API_KEY",
"ANTHROPIC_API_KEY",
"GEMINI_API_KEY",
"XAI_API_KEY",
"DEEPSEEK_API_KEY",
"OPENROUTER_API_KEY",
"GROQ_API_KEY",
"MISTRAL_API_KEY",
"COHERE_API_KEY",
"AZURE_OPENAI_API_KEY",
"HUGGINGFACE_API_KEY",
"HF_TOKEN",
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
"AWS_SESSION_TOKEN",
"GOOGLE_APPLICATION_CREDENTIALS",
"GOOGLE_CLOUD_PROJECT",
"AZURE_CLIENT_ID",
"AZURE_CLIENT_SECRET",
"AZURE_TENANT_ID",
"AZURE_SUBSCRIPTION_ID",
"GITHUB_TOKEN",
"GH_TOKEN",
"GITHUB_PAT",
"NPM_TOKEN",
"NPM_AUTH_TOKEN",
"CARGO_REGISTRY_TOKEN",
"PYPI_TOKEN",
"DATABASE_URL",
"DB_PASSWORD",
"PGPASSWORD",
"MYSQL_PWD",
"REDIS_PASSWORD",
"MONGO_PASSWORD",
"SSH_AUTH_SOCK",
"GPG_AGENT_INFO",
"LD_PRELOAD",
"LD_LIBRARY_PATH",
"LD_AUDIT",
"LD_DEBUG",
"LD_PROFILE",
"DYLD_INSERT_LIBRARIES",
"DYLD_LIBRARY_PATH",
"DYLD_FRAMEWORK_PATH",
"DYLD_FALLBACK_LIBRARY_PATH",
"VAULT_TOKEN",
"CONSUL_HTTP_TOKEN",
"DOCKER_AUTH_CONFIG",
"KUBECONFIG",
"KUBE_TOKEN",
"SLACK_TOKEN",
"SLACK_BOT_TOKEN",
"DISCORD_TOKEN",
"TELEGRAM_BOT_TOKEN",
];
pub const PRESERVED_ENV_VARS: &[&str] = &[
"PATH",
"HOME",
"USER",
"SHELL",
"TERM",
"LANG",
"LC_ALL",
"LC_CTYPE",
"TZ",
"XDG_CONFIG_HOME",
"XDG_DATA_HOME",
"XDG_CACHE_HOME",
"XDG_RUNTIME_DIR",
"EDITOR",
"VISUAL",
"PAGER",
"CARGO_HOME",
"RUSTUP_HOME",
"GOPATH",
"GOROOT",
"JAVA_HOME",
"PYTHON",
"PYTHONPATH",
"NODE_PATH",
"COLORTERM",
"FORCE_COLOR",
"NO_COLOR",
"CLICOLOR",
"CLICOLOR_FORCE",
"TMPDIR",
"TEMP",
"TMP",
];
pub const VTCODE_SANDBOX_ACTIVE: &str = "VTCODE_SANDBOX_ACTIVE";
pub const VTCODE_SANDBOX_NETWORK_DISABLED: &str = "VTCODE_SANDBOX_NETWORK_DISABLED";
pub const VTCODE_SANDBOX_TYPE: &str = "VTCODE_SANDBOX_TYPE";
pub const VTCODE_SANDBOX_WRITABLE_ROOTS: &str = "VTCODE_SANDBOX_WRITABLE_ROOTS";
pub fn build_sanitized_env(
current_env: &HashMap<String, String>,
sandbox_active: bool,
network_disabled: bool,
sandbox_type: &str,
writable_roots: &[&Path],
) -> HashMap<String, String> {
let mut sanitized = HashMap::new();
for key in PRESERVED_ENV_VARS {
if let Some(value) = current_env.get(*key) {
sanitized.insert(key.to_string(), value.clone());
}
}
if sandbox_active {
sanitized.insert(VTCODE_SANDBOX_ACTIVE.to_string(), "1".to_string());
sanitized.insert(VTCODE_SANDBOX_TYPE.to_string(), sandbox_type.to_string());
if network_disabled {
sanitized.insert(VTCODE_SANDBOX_NETWORK_DISABLED.to_string(), "1".to_string());
}
if !writable_roots.is_empty() {
let roots: Vec<String> = writable_roots
.iter()
.map(|p| p.display().to_string())
.collect();
sanitized.insert(VTCODE_SANDBOX_WRITABLE_ROOTS.to_string(), roots.join(":"));
}
}
sanitized
}
pub fn should_filter_env_var(key: &str) -> bool {
FILTERED_ENV_VARS.contains(&key)
|| key.starts_with("AWS_")
|| key.starts_with("AZURE_")
|| key.starts_with("GOOGLE_")
|| key.starts_with("GCP_")
|| key.starts_with("LD_")
|| key.starts_with("DYLD_")
|| key.ends_with("_TOKEN")
|| key.ends_with("_KEY")
|| key.ends_with("_SECRET")
|| key.ends_with("_PASS")
|| key.ends_with("_PWD")
|| key.ends_with("_PASSWORD")
|| key.ends_with("_CREDENTIALS")
}
pub fn filter_sensitive_env(env: &HashMap<String, String>) -> HashMap<String, String> {
env.iter()
.filter(|(k, _)| !should_filter_env_var(k))
.map(|(k, v)| (k.clone(), v.clone()))
.collect()
}
#[cfg(target_os = "linux")]
#[allow(unsafe_code)]
pub fn setup_parent_death_signal() -> std::io::Result<()> {
setup_parent_death_signal_with_check(unsafe { libc::getppid() })
}
#[cfg(target_os = "linux")]
#[allow(unsafe_code)]
pub fn setup_parent_death_signal_with_check(
expected_parent_pid: libc::pid_t,
) -> std::io::Result<()> {
use std::io::Error;
let result = unsafe { libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM) };
if result == -1 {
return Err(Error::other(format!(
"prctl(PR_SET_PDEATHSIG) failed: {}",
Error::last_os_error()
)));
}
if unsafe { libc::getppid() } != expected_parent_pid {
unsafe { libc::raise(libc::SIGTERM) };
}
Ok(())
}
#[cfg(not(target_os = "linux"))]
pub fn setup_parent_death_signal() -> std::io::Result<()> {
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
const TEST_API_KEY_VALUE: &str = "test-openai-key";
#[test]
fn test_should_filter_sensitive_vars() {
assert!(should_filter_env_var("OPENAI_API_KEY"));
assert!(should_filter_env_var("AWS_SECRET_ACCESS_KEY"));
assert!(should_filter_env_var("GITHUB_TOKEN"));
assert!(should_filter_env_var("LD_PRELOAD"));
assert!(should_filter_env_var("DYLD_INSERT_LIBRARIES"));
assert!(should_filter_env_var("MY_CUSTOM_TOKEN"));
assert!(should_filter_env_var("MY_CUSTOM_PASS"));
assert!(should_filter_env_var("MYSQL_PWD"));
assert!(should_filter_env_var("DATABASE_PASSWORD"));
assert!(!should_filter_env_var("PATH"));
assert!(!should_filter_env_var("HOME"));
assert!(!should_filter_env_var("TERM"));
}
#[test]
fn test_build_sanitized_env() {
let mut current = HashMap::new();
current.insert("PATH".to_string(), "/usr/bin".to_string());
current.insert("HOME".to_string(), "/home/user".to_string());
current.insert("OPENAI_API_KEY".to_string(), TEST_API_KEY_VALUE.to_string());
current.insert("RANDOM_VAR".to_string(), "value".to_string());
let sanitized = build_sanitized_env(¤t, true, true, "MacosSeatbelt", &[]);
assert_eq!(sanitized.get("PATH"), Some(&"/usr/bin".to_string()));
assert_eq!(sanitized.get("HOME"), Some(&"/home/user".to_string()));
assert!(!sanitized.contains_key("OPENAI_API_KEY"));
assert!(!sanitized.contains_key("RANDOM_VAR"));
assert_eq!(sanitized.get(VTCODE_SANDBOX_ACTIVE), Some(&"1".to_string()));
assert_eq!(
sanitized.get(VTCODE_SANDBOX_NETWORK_DISABLED),
Some(&"1".to_string())
);
assert_eq!(
sanitized.get(VTCODE_SANDBOX_TYPE),
Some(&"MacosSeatbelt".to_string())
);
}
#[test]
fn test_filter_sensitive_env() {
let mut env = HashMap::new();
env.insert("PATH".to_string(), "/usr/bin".to_string());
env.insert("OPENAI_API_KEY".to_string(), TEST_API_KEY_VALUE.to_string());
env.insert("MY_VAR".to_string(), "value".to_string());
env.insert("AWS_ACCESS_KEY_ID".to_string(), "AKIA...".to_string());
env.insert("CUSTOM_PASS".to_string(), "let-me-in".to_string());
env.insert("SERVICE_PWD".to_string(), "super-secret".to_string());
let filtered = filter_sensitive_env(&env);
assert!(filtered.contains_key("PATH"));
assert!(filtered.contains_key("MY_VAR"));
assert!(!filtered.contains_key("OPENAI_API_KEY"));
assert!(!filtered.contains_key("AWS_ACCESS_KEY_ID"));
assert!(!filtered.contains_key("CUSTOM_PASS"));
assert!(!filtered.contains_key("SERVICE_PWD"));
}
}