From 97d4b425831987eebe7a80e0edc6192fb45c59e2 Mon Sep 17 00:00:00 2001 From: Ruslan Nigmatullin Date: Mon, 20 Apr 2026 15:59:05 -0700 Subject: [PATCH] uds: add async Unix socket crate (#18254) ## Summary - add a codex-uds crate with async UnixListener and UnixStream wrappers - expose helpers for private socket directory setup and stale socket path checks - migrate codex-stdio-to-uds onto codex-uds and Tokio-based stdio/socket relaying - update the CLI stdio-to-uds command path for the async runner ## Tests - cargo test -p codex-uds -p codex-stdio-to-uds - cargo test -p codex-cli - just fmt - just fix -p codex-uds - just fix -p codex-stdio-to-uds - just fix -p codex-cli - just bazel-lock-check - git diff --check --- codex-rs/Cargo.lock | 16 +- codex-rs/Cargo.toml | 3 + codex-rs/cli/src/main.rs | 3 +- codex-rs/stdio-to-uds/Cargo.toml | 10 +- codex-rs/stdio-to-uds/README.md | 2 +- codex-rs/stdio-to-uds/src/lib.rs | 66 ++-- codex-rs/stdio-to-uds/src/main.rs | 5 +- codex-rs/stdio-to-uds/tests/stdio_to_uds.rs | 182 ++++++----- codex-rs/uds/BUILD.bazel | 6 + codex-rs/uds/Cargo.toml | 29 ++ codex-rs/uds/src/lib.rs | 331 ++++++++++++++++++++ codex-rs/uds/src/lib_tests.rs | 121 +++++++ 12 files changed, 641 insertions(+), 133 deletions(-) create mode 100644 codex-rs/uds/BUILD.bazel create mode 100644 codex-rs/uds/Cargo.toml create mode 100644 codex-rs/uds/src/lib.rs create mode 100644 codex-rs/uds/src/lib_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index e7eb7fc80..7d3720a96 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2891,10 +2891,11 @@ name = "codex-stdio-to-uds" version = "0.0.0" dependencies = [ "anyhow", + "codex-uds", "codex-utils-cargo-bin", "pretty_assertions", "tempfile", - "uds_windows", + "tokio", ] [[package]] @@ -3058,6 +3059,18 @@ dependencies = [ "winsplit", ] +[[package]] +name = "codex-uds" +version = "0.0.0" +dependencies = [ + "async-io", + "pretty_assertions", + "tempfile", + "tokio", + "tokio-util", + "uds_windows", +] + [[package]] name = "codex-utils-absolute-path" version = "0.0.0" @@ -10970,6 +10983,7 @@ checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" dependencies = [ "bytes", "futures-core", + "futures-io", "futures-sink", "futures-util", "pin-project-lite", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index b1a8ae000..fd203e625 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -90,6 +90,7 @@ members = [ "terminal-detection", "test-binary-support", "thread-store", + "uds", "codex-experimental-api-macros", "plugin", "model-provider", @@ -175,6 +176,7 @@ codex-test-binary-support = { path = "test-binary-support" } codex-thread-store = { path = "thread-store" } codex-tools = { path = "tools" } codex-tui = { path = "tui" } +codex-uds = { path = "uds" } codex-utils-absolute-path = { path = "utils/absolute-path" } codex-utils-approval-presets = { path = "utils/approval-presets" } codex-utils-cache = { path = "utils/cache" } @@ -212,6 +214,7 @@ arc-swap = "1.9.0" assert_cmd = "2" assert_matches = "1.5.0" async-channel = "2.3.1" +async-io = "2.6.0" async-stream = "0.3.6" async-trait = "0.1.89" axum = { version = "0.8", default-features = false } diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 849e8b104..eface8ec1 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -1090,8 +1090,7 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { "stdio-to-uds", )?; let socket_path = cmd.socket_path; - tokio::task::spawn_blocking(move || codex_stdio_to_uds::run(socket_path.as_path())) - .await??; + codex_stdio_to_uds::run(socket_path.as_path()).await?; } Some(Subcommand::ExecServer(cmd)) => { reject_remote_mode_for_subcommand( diff --git a/codex-rs/stdio-to-uds/Cargo.toml b/codex-rs/stdio-to-uds/Cargo.toml index bc62ee8d2..76d217692 100644 --- a/codex-rs/stdio-to-uds/Cargo.toml +++ b/codex-rs/stdio-to-uds/Cargo.toml @@ -17,9 +17,13 @@ workspace = true [dependencies] anyhow = { workspace = true } - -[target.'cfg(target_os = "windows")'.dependencies] -uds_windows = { workspace = true } +codex-uds = { workspace = true } +tokio = { workspace = true, features = [ + "io-std", + "io-util", + "macros", + "rt-multi-thread", +] } [dev-dependencies] codex-utils-cargo-bin = { workspace = true } diff --git a/codex-rs/stdio-to-uds/README.md b/codex-rs/stdio-to-uds/README.md index 96ac67960..9fe838e30 100644 --- a/codex-rs/stdio-to-uds/README.md +++ b/codex-rs/stdio-to-uds/README.md @@ -17,4 +17,4 @@ Unfortunately, the Rust standard library does not provide support for UNIX domai https://github.com/rust-lang/rust/issues/56533 -As a workaround, this crate leverages https://crates.io/crates/uds_windows as a dependency on Windows. +As a workaround, this crate uses `codex-uds`, which provides a cross-platform async UDS API backed by https://crates.io/crates/uds_windows on Windows. diff --git a/codex-rs/stdio-to-uds/src/lib.rs b/codex-rs/stdio-to-uds/src/lib.rs index 9d2cb6c0f..a103de331 100644 --- a/codex-rs/stdio-to-uds/src/lib.rs +++ b/codex-rs/stdio-to-uds/src/lib.rs @@ -1,56 +1,46 @@ #![deny(clippy::print_stdout)] use std::io; -use std::io::Write; -use std::net::Shutdown; use std::path::Path; -use std::thread; use anyhow::Context; -use anyhow::anyhow; - -#[cfg(unix)] -use std::os::unix::net::UnixStream; - -#[cfg(windows)] -use uds_windows::UnixStream; +use codex_uds::UnixStream; +use tokio::io::AsyncWriteExt; /// Connects to the Unix Domain Socket at `socket_path` and relays data between /// standard input/output and the socket. -pub fn run(socket_path: &Path) -> anyhow::Result<()> { - let mut stream = UnixStream::connect(socket_path) +pub async fn run(socket_path: &Path) -> anyhow::Result<()> { + let stream = UnixStream::connect(socket_path) + .await .with_context(|| format!("failed to connect to socket at {}", socket_path.display()))?; + let (mut socket_reader, mut socket_writer) = tokio::io::split(stream); - let mut reader = stream - .try_clone() - .context("failed to clone socket for reading")?; - - let stdout_thread = thread::spawn(move || -> io::Result<()> { - let stdout = io::stdout(); - let mut handle = stdout.lock(); - io::copy(&mut reader, &mut handle)?; - handle.flush()?; + let copy_socket_to_stdout = async { + let mut stdout = tokio::io::stdout(); + tokio::io::copy(&mut socket_reader, &mut stdout).await?; + stdout.flush().await?; Ok(()) - }); + }; + let copy_stdin_to_socket = async { + let mut stdin = tokio::io::stdin(); + tokio::io::copy(&mut stdin, &mut socket_writer) + .await + .context("failed to copy data from stdin to socket")?; - let stdin = io::stdin(); - { - let mut handle = stdin.lock(); - io::copy(&mut handle, &mut stream).context("failed to copy data from stdin to socket")?; - } + // The peer can close immediately after sending its response; in that + // race, half-closing our write side can report NotConnected on some + // platforms. + if let Err(err) = socket_writer.shutdown().await + && err.kind() != io::ErrorKind::NotConnected + { + return Err(err).context("failed to shutdown socket writer"); + } - // The peer can close immediately after sending its response; in that race, - // half-closing our write side can report NotConnected on some platforms. - if let Err(err) = stream.shutdown(Shutdown::Write) - && err.kind() != io::ErrorKind::NotConnected - { - return Err(err).context("failed to shutdown socket writer"); - } + anyhow::Ok(()) + }; - let stdout_result = stdout_thread - .join() - .map_err(|_| anyhow!("thread panicked while copying socket data to stdout"))?; - stdout_result.context("failed to copy data from socket to stdout")?; + tokio::try_join!(copy_stdin_to_socket, copy_socket_to_stdout) + .context("failed to relay data between stdio and socket")?; Ok(()) } diff --git a/codex-rs/stdio-to-uds/src/main.rs b/codex-rs/stdio-to-uds/src/main.rs index 69987c31b..2cbf04155 100644 --- a/codex-rs/stdio-to-uds/src/main.rs +++ b/codex-rs/stdio-to-uds/src/main.rs @@ -2,7 +2,8 @@ use std::env; use std::path::PathBuf; use std::process; -fn main() -> anyhow::Result<()> { +#[tokio::main] +async fn main() -> anyhow::Result<()> { let mut args = env::args_os().skip(1); let Some(socket_path) = args.next() else { eprintln!("Usage: codex-stdio-to-uds "); @@ -15,5 +16,5 @@ fn main() -> anyhow::Result<()> { } let socket_path = PathBuf::from(socket_path); - codex_stdio_to_uds::run(&socket_path) + codex_stdio_to_uds::run(&socket_path).await } diff --git a/codex-rs/stdio-to-uds/tests/stdio_to_uds.rs b/codex-rs/stdio-to-uds/tests/stdio_to_uds.rs index af8fd5922..0bc4b0c7f 100644 --- a/codex-rs/stdio-to-uds/tests/stdio_to_uds.rs +++ b/codex-rs/stdio-to-uds/tests/stdio_to_uds.rs @@ -1,7 +1,7 @@ use std::io::ErrorKind; use std::io::Read; -use std::io::Write; use std::process::Command; +use std::process::ExitStatus; use std::process::Stdio; use std::sync::mpsc; use std::thread; @@ -9,17 +9,13 @@ use std::time::Duration; use std::time::Instant; use anyhow::Context; -use anyhow::anyhow; +use codex_uds::UnixListener; use pretty_assertions::assert_eq; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; -#[cfg(unix)] -use std::os::unix::net::UnixListener; - -#[cfg(windows)] -use uds_windows::UnixListener; - -#[test] -fn pipes_stdin_and_stdout_through_socket() -> anyhow::Result<()> { +#[tokio::test] +async fn pipes_stdin_and_stdout_through_socket() -> anyhow::Result<()> { // This test intentionally avoids `read_to_end()` on the server side because // waiting for EOF can race with socket half-close behavior on slower runners. // Reading the exact request length keeps the test deterministic. @@ -32,7 +28,7 @@ fn pipes_stdin_and_stdout_through_socket() -> anyhow::Result<()> { let request = b"request"; let request_path = dir.path().join("request.txt"); std::fs::write(&request_path, request).context("failed to write child stdin fixture")?; - let listener = match UnixListener::bind(&socket_path) { + let listener = match UnixListener::bind(&socket_path).await { Ok(listener) => listener, Err(err) if err.kind() == ErrorKind::PermissionDenied => { eprintln!("skipping test: failed to bind unix socket: {err}"); @@ -43,105 +39,119 @@ fn pipes_stdin_and_stdout_through_socket() -> anyhow::Result<()> { } }; - let (tx, rx) = mpsc::channel(); let (event_tx, event_rx) = mpsc::channel(); - let server_thread = thread::spawn(move || -> anyhow::Result<()> { + let server_task = tokio::spawn(async move { + let mut listener = listener; let _ = event_tx.send("waiting for accept".to_string()); - let (mut connection, _) = listener + let mut connection = listener .accept() + .await .context("failed to accept test connection")?; let _ = event_tx.send("accepted connection".to_string()); let mut received = vec![0; request.len()]; connection .read_exact(&mut received) + .await .context("failed to read data from client")?; let _ = event_tx.send(format!("read {} bytes", received.len())); - tx.send(received) - .map_err(|_| anyhow!("failed to send received bytes to test thread"))?; connection .write_all(b"response") + .await .context("failed to write response to client")?; let _ = event_tx.send("wrote response".to_string()); - Ok(()) + anyhow::Ok(received) }); - let stdin = std::fs::File::open(&request_path).context("failed to open child stdin fixture")?; - let mut child = Command::new(codex_utils_cargo_bin::cargo_bin("codex-stdio-to-uds")?) - .arg(&socket_path) - .stdin(Stdio::from(stdin)) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .context("failed to spawn codex-stdio-to-uds")?; + struct ChildOutput { + status: ExitStatus, + stdout: Vec, + stderr: Vec, + server_events: Vec, + } - let mut child_stdout = child.stdout.take().context("missing child stdout")?; - let mut child_stderr = child.stderr.take().context("missing child stderr")?; - let (stdout_tx, stdout_rx) = mpsc::channel(); - let (stderr_tx, stderr_rx) = mpsc::channel(); - thread::spawn(move || { - let mut stdout = Vec::new(); - let result = child_stdout.read_to_end(&mut stdout).map(|_| stdout); - let _ = stdout_tx.send(result); - }); - thread::spawn(move || { - let mut stderr = Vec::new(); - let result = child_stderr.read_to_end(&mut stderr).map(|_| stderr); - let _ = stderr_tx.send(result); + let child_task = tokio::task::spawn_blocking(move || -> anyhow::Result { + let stdin = + std::fs::File::open(&request_path).context("failed to open child stdin fixture")?; + let mut child = Command::new(codex_utils_cargo_bin::cargo_bin("codex-stdio-to-uds")?) + .arg(&socket_path) + .stdin(Stdio::from(stdin)) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .context("failed to spawn codex-stdio-to-uds")?; + + let mut child_stdout = child.stdout.take().context("missing child stdout")?; + let mut child_stderr = child.stderr.take().context("missing child stderr")?; + let (stdout_tx, stdout_rx) = mpsc::channel(); + let (stderr_tx, stderr_rx) = mpsc::channel(); + thread::spawn(move || { + let mut stdout = Vec::new(); + let result = child_stdout.read_to_end(&mut stdout).map(|_| stdout); + let _ = stdout_tx.send(result); + }); + thread::spawn(move || { + let mut stderr = Vec::new(); + let result = child_stderr.read_to_end(&mut stderr).map(|_| stderr); + let _ = stderr_tx.send(result); + }); + + let mut server_events = Vec::new(); + let deadline = Instant::now() + Duration::from_secs(5); + let status = loop { + while let Ok(event) = event_rx.try_recv() { + server_events.push(event); + } + + if let Some(status) = child.try_wait().context("failed to poll child status")? { + break status; + } + + if Instant::now() >= deadline { + let _ = child.kill(); + let _ = child.wait(); + let stderr = stderr_rx + .recv_timeout(Duration::from_secs(1)) + .context("timed out waiting for child stderr after kill")? + .context("failed to read child stderr")?; + anyhow::bail!( + "codex-stdio-to-uds did not exit in time; server events: {:?}; stderr: {}", + server_events, + String::from_utf8_lossy(&stderr).trim_end() + ); + } + + thread::sleep(Duration::from_millis(25)); + }; + + let stdout = stdout_rx + .recv_timeout(Duration::from_secs(1)) + .context("timed out waiting for child stdout")? + .context("failed to read child stdout")?; + let stderr = stderr_rx + .recv_timeout(Duration::from_secs(1)) + .context("timed out waiting for child stderr")? + .context("failed to read child stderr")?; + + Ok(ChildOutput { + status, + stdout, + stderr, + server_events, + }) }); - let mut server_events = Vec::new(); - let deadline = Instant::now() + Duration::from_secs(5); - let status = loop { - while let Ok(event) = event_rx.try_recv() { - server_events.push(event); - } - - if let Some(status) = child.try_wait().context("failed to poll child status")? { - break status; - } - - if Instant::now() >= deadline { - let _ = child.kill(); - let _ = child.wait(); - let stderr = stderr_rx - .recv_timeout(Duration::from_secs(1)) - .context("timed out waiting for child stderr after kill")? - .context("failed to read child stderr")?; - anyhow::bail!( - "codex-stdio-to-uds did not exit in time; server events: {:?}; stderr: {}", - server_events, - String::from_utf8_lossy(&stderr).trim_end() - ); - } - - thread::sleep(Duration::from_millis(25)); - }; - - let stdout = stdout_rx - .recv_timeout(Duration::from_secs(1)) - .context("timed out waiting for child stdout")? - .context("failed to read child stdout")?; - let stderr = stderr_rx - .recv_timeout(Duration::from_secs(1)) - .context("timed out waiting for child stderr")? - .context("failed to read child stderr")?; + let child_output = child_task.await.context("child task panicked")??; assert!( - status.success(), + child_output.status.success(), "codex-stdio-to-uds exited with {status}; server events: {:?}; stderr: {}", - server_events, - String::from_utf8_lossy(&stderr).trim_end() + child_output.server_events, + String::from_utf8_lossy(&child_output.stderr).trim_end(), + status = child_output.status ); - assert_eq!(stdout, b"response"); + assert_eq!(child_output.stdout, b"response"); - let received = rx - .recv_timeout(Duration::from_secs(1)) - .context("server did not receive data in time")?; + let received = server_task.await.context("server task panicked")??; assert_eq!(received, request); - let server_result = server_thread - .join() - .map_err(|_| anyhow!("server thread panicked"))?; - server_result.context("server failed")?; - Ok(()) } diff --git a/codex-rs/uds/BUILD.bazel b/codex-rs/uds/BUILD.bazel new file mode 100644 index 000000000..a539147dd --- /dev/null +++ b/codex-rs/uds/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "uds", + crate_name = "codex_uds", +) diff --git a/codex-rs/uds/Cargo.toml b/codex-rs/uds/Cargo.toml new file mode 100644 index 000000000..bc3670891 --- /dev/null +++ b/codex-rs/uds/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "codex-uds" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lib] +name = "codex_uds" +path = "src/lib.rs" + +[lints] +workspace = true + +[dependencies] +tokio = { workspace = true, features = ["fs", "net", "rt"] } + +[target.'cfg(windows)'.dependencies] +async-io = { workspace = true } +tokio-util = { workspace = true, features = ["compat"] } +uds_windows = { workspace = true } + +[dev-dependencies] +pretty_assertions = { workspace = true } +tempfile = { workspace = true } +tokio = { workspace = true, features = [ + "io-util", + "macros", + "rt-multi-thread", +] } diff --git a/codex-rs/uds/src/lib.rs b/codex-rs/uds/src/lib.rs new file mode 100644 index 000000000..bb6dd31bf --- /dev/null +++ b/codex-rs/uds/src/lib.rs @@ -0,0 +1,331 @@ +//! Cross-platform async Unix domain socket helpers. + +use std::io::Result as IoResult; +use std::path::Path; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; + +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::io::ReadBuf; + +/// Creates `socket_dir` if needed and restricts it to the current user where +/// the platform exposes Unix permissions. +pub async fn prepare_private_socket_directory(socket_dir: impl AsRef) -> IoResult<()> { + platform::prepare_private_socket_directory(socket_dir.as_ref()).await +} + +/// Returns whether `socket_path` points at a stale Unix socket rendezvous path. +/// +/// On Unix this checks the file type. On Windows, `uds_windows` represents the +/// rendezvous as a regular path, so existence is the only useful stale-path +/// signal available. +pub async fn is_stale_socket_path(socket_path: impl AsRef) -> IoResult { + platform::is_stale_socket_path(socket_path.as_ref()).await +} + +/// Async Unix domain socket listener. +pub struct UnixListener { + inner: platform::Listener, +} + +impl UnixListener { + /// Binds a new listener at `socket_path`. + pub async fn bind(socket_path: impl AsRef) -> IoResult { + platform::bind_listener(socket_path.as_ref()) + .await + .map(|inner| Self { inner }) + } + + /// Accepts the next incoming stream. + pub async fn accept(&mut self) -> IoResult { + self.inner.accept().await.map(|inner| UnixStream { inner }) + } +} + +/// Async Unix domain socket stream. +pub struct UnixStream { + inner: platform::Stream, +} + +impl UnixStream { + /// Connects to `socket_path`. + pub async fn connect(socket_path: impl AsRef) -> IoResult { + platform::connect_stream(socket_path.as_ref()) + .await + .map(|inner| Self { inner }) + } +} + +impl AsyncRead for UnixStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for UnixStream { + fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) + } +} + +#[cfg(unix)] +mod platform { + use std::io; + use std::io::ErrorKind; + use std::io::Result as IoResult; + use std::os::unix::fs::FileTypeExt; + use std::os::unix::fs::PermissionsExt; + use std::path::Path; + + use tokio::fs; + use tokio::net::UnixListener; + use tokio::net::UnixStream; + + /// Owner-only access keeps the control socket directory private while + /// preserving owner traversal and socket path creation. + const SOCKET_DIR_MODE: u32 = 0o700; + const SOCKET_DIR_PERMISSION_BITS: u32 = 0o777; + + pub(super) type Stream = UnixStream; + + pub(super) struct Listener(UnixListener); + + pub(super) async fn prepare_private_socket_directory(socket_dir: &Path) -> IoResult<()> { + let mut dir_builder = fs::DirBuilder::new(); + dir_builder.mode(SOCKET_DIR_MODE); + match dir_builder.create(socket_dir).await { + Ok(()) => return Ok(()), + Err(err) if err.kind() == ErrorKind::AlreadyExists => {} + Err(err) => return Err(err), + } + + let metadata = fs::symlink_metadata(socket_dir).await?; + if !metadata.is_dir() { + return Err(io::Error::new( + ErrorKind::AlreadyExists, + format!( + "socket directory path exists and is not a directory: {}", + socket_dir.display() + ), + )); + } + + let permissions = metadata.permissions(); + // The SSH-over-UDS control socket is reachable by path, so the + // rendezvous directory must be owner-traversable while denying + // group/other access; exact 0700 fixes insecure modes and unusable + // owner-only modes like 0600. + if permissions.mode() & SOCKET_DIR_PERMISSION_BITS != SOCKET_DIR_MODE { + fs::set_permissions(socket_dir, std::fs::Permissions::from_mode(SOCKET_DIR_MODE)) + .await?; + } + + Ok(()) + } + + pub(super) async fn bind_listener(socket_path: &Path) -> IoResult { + UnixListener::bind(socket_path).map(Listener) + } + + impl Listener { + pub(super) async fn accept(&mut self) -> IoResult { + self.0.accept().await.map(|(stream, _addr)| stream) + } + } + + pub(super) async fn connect_stream(socket_path: &Path) -> IoResult { + UnixStream::connect(socket_path).await + } + + pub(super) async fn is_stale_socket_path(socket_path: &Path) -> IoResult { + Ok(fs::symlink_metadata(socket_path) + .await? + .file_type() + .is_socket()) + } +} + +#[cfg(windows)] +mod platform { + use std::io; + use std::io::Result as IoResult; + use std::net::Shutdown; + use std::ops::Deref; + use std::os::windows::io::AsRawSocket; + use std::os::windows::io::AsSocket; + use std::os::windows::io::BorrowedSocket; + use std::path::Path; + use std::pin::Pin; + use std::task::Context; + use std::task::Poll; + use std::task::ready; + + use async_io::Async; + use tokio::io::AsyncRead; + use tokio::io::AsyncWrite; + use tokio::io::ReadBuf; + use tokio::task; + use tokio_util::compat::Compat; + use tokio_util::compat::FuturesAsyncReadCompatExt; + + pub(super) struct Stream(Compat>); + + pub(super) async fn prepare_private_socket_directory(socket_dir: &Path) -> IoResult<()> { + tokio::fs::create_dir_all(socket_dir).await + } + + pub(super) struct Listener(Async); + + pub(super) async fn bind_listener(socket_path: &Path) -> IoResult { + let socket_path = socket_path.to_path_buf(); + let listener = + spawn_blocking_io(move || uds_windows::UnixListener::bind(socket_path)).await?; + Async::new(WindowsUnixListener::from(listener)).map(Listener) + } + + impl Listener { + pub(super) async fn accept(&mut self) -> IoResult { + let (stream, _addr) = self.0.read_with(|listener| listener.accept()).await?; + Async::new(WindowsUnixStream::from(stream)) + .map(FuturesAsyncReadCompatExt::compat) + .map(Stream) + } + } + + pub(super) async fn connect_stream(socket_path: &Path) -> IoResult { + let socket_path = socket_path.to_path_buf(); + let stream = + spawn_blocking_io(move || uds_windows::UnixStream::connect(socket_path)).await?; + Async::new(WindowsUnixStream::from(stream)) + .map(FuturesAsyncReadCompatExt::compat) + .map(Stream) + } + + pub(super) async fn is_stale_socket_path(socket_path: &Path) -> IoResult { + tokio::fs::try_exists(socket_path).await + } + + async fn spawn_blocking_io( + operation: impl FnOnce() -> IoResult + Send + 'static, + ) -> IoResult + where + T: Send + 'static, + { + task::spawn_blocking(operation) + .await + .map_err(|err| io::Error::other(format!("blocking socket task failed: {err}")))? + } + + pub(super) struct WindowsUnixListener(uds_windows::UnixListener); + + impl From for WindowsUnixListener { + fn from(listener: uds_windows::UnixListener) -> Self { + Self(listener) + } + } + + impl Deref for WindowsUnixListener { + type Target = uds_windows::UnixListener; + + fn deref(&self) -> &Self::Target { + &self.0 + } + } + + impl AsSocket for WindowsUnixListener { + fn as_socket(&self) -> BorrowedSocket<'_> { + unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) } + } + } + + pub(super) struct WindowsUnixStream(uds_windows::UnixStream); + + impl From for WindowsUnixStream { + fn from(stream: uds_windows::UnixStream) -> Self { + Self(stream) + } + } + + impl Deref for WindowsUnixStream { + type Target = uds_windows::UnixStream; + + fn deref(&self) -> &Self::Target { + &self.0 + } + } + + impl AsSocket for WindowsUnixStream { + fn as_socket(&self) -> BorrowedSocket<'_> { + unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) } + } + } + + impl io::Read for WindowsUnixStream { + fn read(&mut self, buf: &mut [u8]) -> IoResult { + io::Read::read(&mut self.0, buf) + } + } + + impl io::Write for WindowsUnixStream { + fn write(&mut self, buf: &[u8]) -> IoResult { + io::Write::write(&mut self.0, buf) + } + + fn flush(&mut self) -> IoResult<()> { + io::Write::flush(&mut self.0) + } + } + + impl AsyncRead for Stream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.get_mut().0).poll_read(cx, buf) + } + } + + impl AsyncWrite for Stream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.get_mut().0).poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().0).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let stream = &mut self.get_mut().0; + ready!(Pin::new(&mut *stream).poll_flush(cx))?; + // `Compat>` maps shutdown to `poll_close()`, which only + // flushes for `async_io::Async`; call the socket shutdown directly. + stream.get_ref().get_ref().shutdown(Shutdown::Write)?; + Poll::Ready(Ok(())) + } + } + + unsafe impl async_io::IoSafe for WindowsUnixListener {} + unsafe impl async_io::IoSafe for WindowsUnixStream {} +} + +#[cfg(test)] +mod lib_tests; diff --git a/codex-rs/uds/src/lib_tests.rs b/codex-rs/uds/src/lib_tests.rs new file mode 100644 index 000000000..1861cb500 --- /dev/null +++ b/codex-rs/uds/src/lib_tests.rs @@ -0,0 +1,121 @@ +use std::io::ErrorKind; + +use pretty_assertions::assert_eq; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; + +use super::*; + +#[tokio::test] +async fn prepare_private_socket_directory_creates_directory() { + let temp_dir = tempfile::TempDir::new().expect("temp dir"); + let socket_dir = temp_dir.path().join("app-server-control"); + + prepare_private_socket_directory(&socket_dir) + .await + .expect("socket dir should be created"); + + assert!(socket_dir.is_dir()); +} + +#[cfg(unix)] +#[tokio::test] +async fn prepare_private_socket_directory_sets_existing_permissions_to_owner_only() { + use std::os::unix::fs::PermissionsExt; + + let temp_dir = tempfile::TempDir::new().expect("temp dir"); + for mode in [0o755, 0o600] { + let socket_dir = temp_dir.path().join(format!("app-server-control-{mode:o}")); + std::fs::create_dir(&socket_dir).expect("socket dir should be created"); + std::fs::set_permissions(&socket_dir, std::fs::Permissions::from_mode(mode)) + .expect("socket dir permissions should be changed"); + + prepare_private_socket_directory(&socket_dir) + .await + .expect("socket dir permissions should be set exactly"); + + let mode = std::fs::metadata(&socket_dir) + .expect("socket dir metadata") + .permissions() + .mode(); + assert_eq!(mode & 0o777, 0o700); + } +} + +#[cfg(unix)] +#[tokio::test] +async fn regular_file_path_is_not_stale_socket_path() { + let temp_dir = tempfile::TempDir::new().expect("temp dir"); + let regular_file = temp_dir.path().join("not-a-socket"); + std::fs::write(®ular_file, b"not a socket").expect("regular file should be created"); + + assert!( + !is_stale_socket_path(®ular_file) + .await + .expect("stale socket check should succeed") + ); +} + +#[tokio::test] +async fn bound_listener_path_is_stale_socket_path() { + let temp_dir = tempfile::TempDir::new().expect("temp dir"); + let socket_path = temp_dir.path().join("socket"); + let _listener = match UnixListener::bind(&socket_path).await { + Ok(listener) => listener, + Err(err) if err.kind() == ErrorKind::PermissionDenied => { + eprintln!("skipping test: failed to bind unix socket: {err}"); + return; + } + Err(err) => panic!("failed to bind test socket: {err}"), + }; + + assert!( + is_stale_socket_path(&socket_path) + .await + .expect("stale socket check should succeed") + ); +} + +#[tokio::test] +async fn stream_round_trips_data_between_listener_and_client() { + let temp_dir = tempfile::TempDir::new().expect("temp dir"); + let socket_path = temp_dir.path().join("socket"); + let mut listener = match UnixListener::bind(&socket_path).await { + Ok(listener) => listener, + Err(err) if err.kind() == ErrorKind::PermissionDenied => { + eprintln!("skipping test: failed to bind unix socket: {err}"); + return; + } + Err(err) => panic!("failed to bind test socket: {err}"), + }; + + let server_task = tokio::spawn(async move { + let mut server_stream = listener.accept().await.expect("connection should accept"); + let mut request = [0; 7]; + server_stream + .read_exact(&mut request) + .await + .expect("server should read request"); + assert_eq!(&request, b"request"); + server_stream + .write_all(b"response") + .await + .expect("server should write response"); + }); + + let mut client_stream = UnixStream::connect(&socket_path) + .await + .expect("client should connect"); + client_stream + .write_all(b"request") + .await + .expect("client should write request"); + let mut response = [0; 8]; + client_stream + .read_exact(&mut response) + .await + .expect("client should read response"); + assert_eq!(&response, b"response"); + + server_task.await.expect("server task should join"); +}