Codex Handbook
app-server-transport/src/transport/unix_socket.rs 190 lines
use std::fs::OpenOptions;use std::io::ErrorKind;use std::io::Result as IoResult;use std::path::Path;use super::TransportEvent;use crate::transport::websocket::run_websocket_connection;use codex_uds::UnixListener;use codex_uds::UnixStream;use codex_utils_absolute_path::AbsolutePathBuf;use futures::StreamExt;use tokio::sync::mpsc;use tokio::task::JoinHandle;use tokio::time::Duration;use tokio_tungstenite::accept_async;use tokio_util::sync::CancellationToken;use tracing::error;use tracing::info;use tracing::warn;#[cfg(unix)]const CONTROL_SOCKET_MODE: u32 = 0o600;pub async fn start_control_socket_acceptor(    socket_path: AbsolutePathBuf,    transport_event_tx: mpsc::Sender<TransportEvent>,    shutdown_token: CancellationToken,) -> IoResult<JoinHandle<()>> {    prepare_control_socket_path(socket_path.as_path()).await?;    let listener = UnixListener::bind(socket_path.as_path()).await?;    let socket_guard = ControlSocketFileGuard { socket_path };    set_control_socket_permissions(socket_guard.socket_path.as_path()).await?;    info!(        socket_path = %socket_guard.socket_path.display(),        "app-server control socket listening"    );    Ok(tokio::spawn(run_control_socket_acceptor(        listener,        transport_event_tx,        shutdown_token,        socket_guard,    )))}async fn run_control_socket_acceptor(    mut listener: UnixListener,    transport_event_tx: mpsc::Sender<TransportEvent>,    shutdown_token: CancellationToken,    socket_guard: ControlSocketFileGuard,) {    let _socket_guard = socket_guard;    loop {        let stream = tokio::select! {            _ = shutdown_token.cancelled() => {                break;            }            result = listener.accept() => {                match result {                    Ok(stream) => stream,                    Err(err) => {                        if matches!(                            err.kind(),                            ErrorKind::ConnectionAborted | ErrorKind::ConnectionReset | ErrorKind::Interrupted                        ) {                            warn!("recoverable control socket accept error: {err}");                            continue;                        }                        error!("control socket accept error: {err}");                        tokio::time::sleep(Duration::from_secs(1)).await;                        continue;                    }                }            }        };        let transport_event_tx = transport_event_tx.clone();        tokio::spawn(async move {            let websocket_stream = match accept_async(stream).await {                Ok(websocket_stream) => websocket_stream,                Err(err) => {                    warn!("failed to upgrade control socket websocket connection: {err}");                    return;                }            };            let (websocket_writer, websocket_reader) = websocket_stream.split();            run_websocket_connection(websocket_writer, websocket_reader, transport_event_tx).await;        });    }    info!("control socket acceptor shutting down");}pub async fn prepare_control_socket_path(socket_path: &Path) -> IoResult<()> {    if let Some(parent) = socket_path.parent() {        codex_uds::prepare_private_socket_directory(parent).await?;    }    match UnixStream::connect(socket_path).await {        Ok(_stream) => {            return Err(std::io::Error::new(                ErrorKind::AddrInUse,                format!(                    "app-server control socket is already in use at {}",                    socket_path.display()                ),            ));        }        Err(err) if err.kind() == ErrorKind::NotFound => return Ok(()),        Err(err) if err.kind() == ErrorKind::ConnectionRefused => {}        Err(err) => {            if !socket_path.exists() {                return Ok(());            }            return Err(err);        }    }    if !socket_path.try_exists()? {        return Ok(());    }    if !codex_uds::is_stale_socket_path(socket_path).await? {        return Err(std::io::Error::new(            ErrorKind::AlreadyExists,            format!(                "app-server control socket path exists and is not a socket: {}",                socket_path.display()            ),        ));    }    tokio::fs::remove_file(socket_path).await}pub struct AppServerStartupLock {    _file: std::fs::File,}pub async fn acquire_app_server_startup_lock(    startup_lock_path: AbsolutePathBuf,) -> IoResult<AppServerStartupLock> {    if let Some(parent) = startup_lock_path.as_path().parent() {        codex_uds::prepare_private_socket_directory(parent).await?;    }    tokio::task::spawn_blocking(move || {        let file = OpenOptions::new()            .create(true)            .truncate(false)            .read(true)            .write(true)            .open(startup_lock_path.as_path())?;        file.lock()?;        Ok(AppServerStartupLock { _file: file })    })    .await    .map_err(|err| std::io::Error::other(format!("startup lock task failed: {err}")))?}#[cfg(unix)]async fn set_control_socket_permissions(socket_path: &Path) -> IoResult<()> {    use std::os::unix::fs::PermissionsExt;    tokio::fs::set_permissions(        socket_path,        std::fs::Permissions::from_mode(CONTROL_SOCKET_MODE),    )    .await}#[cfg(not(unix))]async fn set_control_socket_permissions(_socket_path: &Path) -> IoResult<()> {    Ok(())}struct ControlSocketFileGuard {    socket_path: AbsolutePathBuf,}impl Drop for ControlSocketFileGuard {    fn drop(&mut self) {        if let Err(err) = std::fs::remove_file(self.socket_path.as_path())            && err.kind() != ErrorKind::NotFound        {            warn!(                socket_path = %self.socket_path.display(),                %err,                "failed to remove app-server control socket file"            );        }    }}