app-server-transport/src/transport/unix_socket_tests.rs
233 lines
use super::AppServerTransport;use super::CHANNEL_CAPACITY;use super::TransportEvent;use super::acquire_app_server_startup_lock;use super::app_server_control_socket_path;use super::start_control_socket_acceptor;use codex_app_server_protocol::JSONRPCMessage;use codex_app_server_protocol::JSONRPCNotification;use codex_core::config::find_codex_home;use codex_uds::UnixStream;use codex_utils_absolute_path::AbsolutePathBuf;use futures::SinkExt;use futures::StreamExt;use pretty_assertions::assert_eq;use std::io::Result as IoResult;use std::path::Path;use tokio::sync::mpsc;use tokio::time::Duration;use tokio::time::timeout;use tokio_tungstenite::client_async;use tokio_tungstenite::tungstenite::Bytes;use tokio_tungstenite::tungstenite::Message as WebSocketMessage;use tokio_util::sync::CancellationToken;#[test]fn listen_unix_socket_parses_as_unix_socket_transport() { assert_eq!( AppServerTransport::from_listen_url("unix://"), Ok(AppServerTransport::UnixSocket { socket_path: default_control_socket_path() }) );}#[test]fn listen_unix_socket_accepts_absolute_custom_path() { assert_eq!( AppServerTransport::from_listen_url("unix:///tmp/codex.sock"), Ok(AppServerTransport::UnixSocket { socket_path: absolute_path("/tmp/codex.sock") }) );}#[test]fn listen_unix_socket_accepts_relative_custom_path() { assert_eq!( AppServerTransport::from_listen_url("unix://codex.sock"), Ok(AppServerTransport::UnixSocket { socket_path: AbsolutePathBuf::relative_to_current_dir("codex.sock") .expect("relative path should resolve") }) );}#[tokio::test]async fn control_socket_acceptor_upgrades_and_forwards_websocket_text_messages_and_pings() { let temp_dir = tempfile::TempDir::new().expect("temp dir"); let socket_path = test_socket_path(temp_dir.path()); let (transport_event_tx, mut transport_event_rx) = mpsc::channel::<TransportEvent>(CHANNEL_CAPACITY); let shutdown_token = CancellationToken::new(); let accept_handle = start_control_socket_acceptor( socket_path.clone(), transport_event_tx, shutdown_token.clone(), ) .await .expect("control socket acceptor should start"); let stream = connect_to_socket(socket_path.as_path()) .await .expect("client should connect"); let (mut websocket, response) = client_async("ws://localhost/rpc", stream) .await .expect("websocket upgrade should complete"); assert_eq!(response.status().as_u16(), 101); let opened = timeout(Duration::from_secs(1), transport_event_rx.recv()) .await .expect("connection opened event should arrive") .expect("connection opened event"); let connection_id = match opened { TransportEvent::ConnectionOpened { connection_id, .. } => connection_id, _ => panic!("expected connection opened event"), }; let notification = JSONRPCMessage::Notification(JSONRPCNotification { method: "initialized".to_string(), params: None, }); websocket .send(WebSocketMessage::Text( serde_json::to_string(¬ification) .expect("notification should serialize") .into(), )) .await .expect("notification should send"); let incoming = timeout(Duration::from_secs(1), transport_event_rx.recv()) .await .expect("incoming message event should arrive") .expect("incoming message event"); assert_eq!( match incoming { TransportEvent::IncomingMessage { connection_id: incoming_connection_id, message, } => (incoming_connection_id, message), _ => panic!("expected incoming message event"), }, (connection_id, notification) ); websocket .send(WebSocketMessage::Ping(Bytes::from_static(b"check"))) .await .expect("ping should send"); let pong = timeout(Duration::from_secs(1), websocket.next()) .await .expect("pong should arrive") .expect("pong frame") .expect("pong should be valid"); assert_eq!(pong, WebSocketMessage::Pong(Bytes::from_static(b"check"))); websocket.close(None).await.expect("close should send"); let closed = timeout(Duration::from_secs(1), transport_event_rx.recv()) .await .expect("connection closed event should arrive") .expect("connection closed event"); assert!(matches!( closed, TransportEvent::ConnectionClosed { connection_id: closed_connection_id, } if closed_connection_id == connection_id )); shutdown_token.cancel(); accept_handle.await.expect("acceptor should join"); assert_socket_path_removed(socket_path.as_path());}#[tokio::test]async fn app_server_startup_lock_serializes_waiters() { let temp_dir = tempfile::TempDir::new().expect("temp dir"); let lock_path = test_startup_lock_path(temp_dir.path()); let first_lock = acquire_app_server_startup_lock(lock_path.clone()) .await .expect("first startup lock should succeed"); let mut second_lock = tokio::spawn(acquire_app_server_startup_lock(lock_path)); assert!( timeout(Duration::from_millis(100), &mut second_lock) .await .is_err() ); drop(first_lock); second_lock .await .expect("second startup lock task should join") .expect("second startup lock should succeed");}#[cfg(unix)]#[tokio::test]async fn control_socket_file_is_private_after_bind() { use std::os::unix::fs::PermissionsExt; let temp_dir = tempfile::TempDir::new().expect("temp dir"); let socket_path = test_socket_path(temp_dir.path()); let (transport_event_tx, _transport_event_rx) = mpsc::channel::<TransportEvent>(CHANNEL_CAPACITY); let shutdown_token = CancellationToken::new(); let accept_handle = start_control_socket_acceptor( socket_path.clone(), transport_event_tx, shutdown_token.clone(), ) .await .expect("control socket acceptor should start"); let metadata = tokio::fs::metadata(socket_path.as_path()) .await .expect("socket metadata should exist"); assert_eq!(metadata.permissions().mode() & 0o777, 0o600); shutdown_token.cancel(); accept_handle.await.expect("acceptor should join");}fn absolute_path(path: &str) -> AbsolutePathBuf { AbsolutePathBuf::from_absolute_path(path).expect("absolute path")}fn default_control_socket_path() -> AbsolutePathBuf { let codex_home = find_codex_home().expect("codex home"); app_server_control_socket_path(&codex_home).expect("default control socket path")}fn test_socket_path(temp_dir: &Path) -> AbsolutePathBuf { AbsolutePathBuf::from_absolute_path( temp_dir .join("app-server-control") .join("app-server-control.sock"), ) .expect("socket path should resolve")}fn test_startup_lock_path(temp_dir: &Path) -> AbsolutePathBuf { AbsolutePathBuf::from_absolute_path( temp_dir .join("app-server-control") .join("app-server-startup.lock"), ) .expect("startup lock path should resolve")}async fn connect_to_socket(socket_path: &Path) -> IoResult<UnixStream> { UnixStream::connect(socket_path).await}#[cfg(unix)]fn assert_socket_path_removed(socket_path: &Path) { assert!(!socket_path.exists());}#[cfg(windows)]fn assert_socket_path_removed(_socket_path: &Path) { // uds_windows uses a regular filesystem path as its rendezvous point, // but there is no Unix socket filesystem node to assert on.}