Merge bb7732ada7 into sapling-pr-archive-bolinfest

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<OwnedF
         .with_context(|| format!("failed to duplicate {name} for escalation transfer"))
 }
 
+async fn connect_escalation_stream(
+    handshake_client: AsyncDatagramSocket,
+) -> 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<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 (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::<EscalateResponse>()
         .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::<EscalateResponse>().await?;
+            drop(server_stream_guard);
+            assert_eq!(-1, unsafe { libc::fcntl(guard_fd, libc::F_GETFD) });
+            Ok::<EscalateResponse, anyhow::Error>(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::<EscalateRequest>().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(())
+    }
 }