diff --git a/codex-rs/shell-escalation/src/unix/escalate_client.rs b/codex-rs/shell-escalation/src/unix/escalate_client.rs index 0f98883865..a8860a0a1d 100644 --- a/codex-rs/shell-escalation/src/unix/escalate_client.rs +++ b/codex-rs/shell-escalation/src/unix/escalate_client.rs @@ -33,17 +33,31 @@ fn duplicate_fd_for_transfer(fd: impl AsFd, name: &str) -> anyhow::Result anyhow::Result<(AsyncSocket, OwnedFd)> { + let (server, client) = AsyncSocket::pair()?; + let server_stream_guard: OwnedFd = server.into_inner().into(); + let transferred_server_stream = + duplicate_fd_for_transfer(&server_stream_guard, "handshake stream")?; + const HANDSHAKE_MESSAGE: [u8; 1] = [0]; + // Keep one local reference to the transferred stream alive until the server + // answers the first request. On macOS, dropping the sender's last local copy + // immediately after the datagram handshake can make the peer observe EOF + // before the received fd is fully servicing the stream. + handshake_client + .send_with_fds(&HANDSHAKE_MESSAGE, &[transferred_server_stream]) + .await + .context("failed to send handshake datagram")?; + Ok((client, server_stream_guard)) +} + pub async fn run_shell_escalation_execve_wrapper( file: String, argv: Vec, ) -> anyhow::Result { 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 (client, server_stream_guard) = connect_escalation_stream(handshake_client).await?; let env = std::env::vars() .filter(|(k, _)| !matches!(k.as_str(), ESCALATE_SOCKET_ENV_VAR | EXEC_WRAPPER_ENV_VAR)) .collect(); @@ -60,6 +74,7 @@ pub async fn run_shell_escalation_execve_wrapper( .receive::() .await .context("failed to receive EscalateResponse")?; + drop(server_stream_guard); match message.action { EscalateAction::Escalate => { // Duplicate stdio before transferring ownership to the server. The @@ -128,6 +143,11 @@ mod tests { use super::*; use std::os::fd::AsRawFd; use std::os::unix::net::UnixStream; + use std::path::PathBuf; + use std::time::Duration; + + use pretty_assertions::assert_eq; + use tokio::time::sleep; #[test] fn duplicate_fd_for_transfer_does_not_close_original() { @@ -141,4 +161,46 @@ mod tests { assert_ne!(unsafe { libc::fcntl(original_fd, libc::F_GETFD) }, -1); } + + #[tokio::test] + async fn connect_escalation_stream_keeps_sender_alive_until_first_response() + -> anyhow::Result<()> { + let (server_datagram, client_datagram) = AsyncDatagramSocket::pair()?; + let client_task = tokio::spawn(async move { + let (client_stream, server_stream_guard) = + connect_escalation_stream(client_datagram).await?; + let guard_fd = server_stream_guard.as_raw_fd(); + assert_ne!(unsafe { libc::fcntl(guard_fd, libc::F_GETFD) }, -1); + client_stream + .send(EscalateRequest { + file: PathBuf::from("/bin/echo"), + argv: vec!["echo".to_string(), "hello".to_string()], + workdir: AbsolutePathBuf::current_dir()?, + env: Default::default(), + }) + .await?; + let response = client_stream.receive::().await?; + drop(server_stream_guard); + assert_eq!(-1, unsafe { libc::fcntl(guard_fd, libc::F_GETFD) }); + Ok::(response) + }); + + let (_, mut fds) = server_datagram.receive_with_fds().await?; + assert_eq!(fds.len(), 1); + sleep(Duration::from_millis(20)).await; + let server_stream = AsyncSocket::from_fd(fds.remove(0))?; + let request = server_stream.receive::().await?; + assert_eq!(request.file, PathBuf::from("/bin/echo")); + assert_eq!(request.argv, vec!["echo".to_string(), "hello".to_string()]); + + let expected = EscalateResponse { + action: EscalateAction::Deny { + reason: Some("not now".to_string()), + }, + }; + server_stream.send(expected.clone()).await?; + let response = client_task.await??; + assert_eq!(response, expected); + Ok(()) + } }