shell-escalation/src/unix/escalate_client.rs
144 lines
use std::io;use std::os::fd::AsFd;use std::os::fd::AsRawFd;use std::os::fd::OwnedFd;use anyhow::Context as _;use codex_utils_absolute_path::AbsolutePathBuf;use crate::unix::escalate_protocol::ESCALATE_SOCKET_ENV_VAR;use crate::unix::escalate_protocol::EXEC_WRAPPER_ENV_VAR;use crate::unix::escalate_protocol::EscalateAction;use crate::unix::escalate_protocol::EscalateRequest;use crate::unix::escalate_protocol::EscalateResponse;use crate::unix::escalate_protocol::SuperExecMessage;use crate::unix::escalate_protocol::SuperExecResult;use crate::unix::socket::AsyncDatagramSocket;use crate::unix::socket::AsyncSocket;fn get_escalate_client() -> anyhow::Result<AsyncDatagramSocket> { // TODO: we should defensively require only calling this once, since AsyncSocket will take ownership of the fd. let client_fd = std::env::var(ESCALATE_SOCKET_ENV_VAR)?.parse::<i32>()?; if client_fd < 0 { return Err(anyhow::anyhow!( "{ESCALATE_SOCKET_ENV_VAR} is not a valid file descriptor: {client_fd}" )); } Ok(unsafe { AsyncDatagramSocket::from_raw_fd(client_fd) }?)}fn duplicate_fd_for_transfer(fd: impl AsFd, name: &str) -> anyhow::Result<OwnedFd> { fd.as_fd() .try_clone_to_owned() .with_context(|| format!("failed to duplicate {name} for escalation transfer"))}pub async fn run_shell_escalation_execve_wrapper( file: String, argv: Vec<String>,) -> anyhow::Result<i32> { let handshake_client = get_escalate_client()?; let (server, client) = AsyncSocket::pair()?; const HANDSHAKE_MESSAGE: [u8; 1] = [0]; handshake_client .send_with_fds(&HANDSHAKE_MESSAGE, &[server.into_inner().into()]) .await .context("failed to send handshake datagram")?; let env = std::env::vars() .filter(|(k, _)| !matches!(k.as_str(), ESCALATE_SOCKET_ENV_VAR | EXEC_WRAPPER_ENV_VAR)) .collect(); client .send(EscalateRequest { file: file.clone().into(), argv: argv.clone(), workdir: AbsolutePathBuf::current_dir()?, env, }) .await .context("failed to send EscalateRequest")?; let message = client .receive::<EscalateResponse>() .await .context("failed to receive EscalateResponse")?; match message.action { EscalateAction::Escalate => { // Duplicate stdio before transferring ownership to the server. The // wrapper must keep using its own stdin/stdout/stderr until the // escalated child takes over. let destination_fds = [ io::stdin().as_raw_fd(), io::stdout().as_raw_fd(), io::stderr().as_raw_fd(), ]; let fds_to_send = [ duplicate_fd_for_transfer(io::stdin(), "stdin")?, duplicate_fd_for_transfer(io::stdout(), "stdout")?, duplicate_fd_for_transfer(io::stderr(), "stderr")?, ]; // TODO: also forward signals over the super-exec socket client .send_with_fds( SuperExecMessage { fds: destination_fds.into_iter().collect(), }, &fds_to_send, ) .await .context("failed to send SuperExecMessage")?; let SuperExecResult { exit_code } = client.receive::<SuperExecResult>().await?; Ok(exit_code) } EscalateAction::Run => { // We avoid std::process::Command here because we want to be as transparent as // possible. std::os::unix::process::CommandExt has .exec() but it does some funky // stuff with signal masks and dup2() on its standard FDs, which we don't want. use std::ffi::CString; let file = CString::new(file).context("NUL in file")?; let argv_cstrs: Vec<CString> = argv .iter() .map(|s| CString::new(s.as_str()).context("NUL in argv")) .collect::<Result<Vec<_>, _>>()?; let mut argv: Vec<*const libc::c_char> = argv_cstrs.iter().map(|s| s.as_ptr()).collect(); argv.push(std::ptr::null()); let err = unsafe { libc::execv(file.as_ptr(), argv.as_ptr()); std::io::Error::last_os_error() }; Err(err.into()) } EscalateAction::Deny { reason } => { match reason { Some(reason) => eprintln!("Execution denied: {reason}"), None => eprintln!("Execution denied"), } Ok(1) } }}#[cfg(test)]mod tests { use super::*; use std::os::fd::AsRawFd; use std::os::unix::net::UnixStream; #[test] fn duplicate_fd_for_transfer_does_not_close_original() { let (left, _right) = UnixStream::pair().expect("socket pair"); let original_fd = left.as_raw_fd(); let duplicate = duplicate_fd_for_transfer(&left, "test fd").expect("duplicate fd"); assert_ne!(duplicate.as_raw_fd(), original_fd); drop(duplicate); assert_ne!(unsafe { libc::fcntl(original_fd, libc::F_GETFD) }, -1); }}