use super::CHANNEL_CAPACITY; use super::TransportEvent; use super::auth::WebsocketAuthPolicy; use super::auth::authorize_upgrade; use super::auth::should_warn_about_unauthenticated_non_loopback_listener; use super::forward_incoming_message; use super::serialize_outgoing_message; use crate::outgoing_message::ConnectionId; use crate::outgoing_message::QueuedOutgoingMessage; use axum::Router; use axum::body::Body; use axum::extract::ConnectInfo; use axum::extract::State; use axum::extract::ws::Message as WebSocketMessage; use axum::extract::ws::WebSocket; use axum::extract::ws::WebSocketUpgrade; use axum::http::HeaderMap; use axum::http::Request; use axum::http::StatusCode; use axum::http::header::ORIGIN; use axum::middleware; use axum::middleware::Next; use axum::response::IntoResponse; use axum::response::Response; use axum::routing::any; use axum::routing::get; use futures::SinkExt; use futures::StreamExt; use owo_colors::OwoColorize; use owo_colors::Stream; use owo_colors::Style; use std::io::Result as IoResult; use std::net::SocketAddr; use std::sync::Arc; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use tokio::net::TcpListener; use tokio::sync::mpsc; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use tracing::error; use tracing::info; use tracing::warn; fn colorize(text: &str, style: Style) -> String { text.if_supports_color(Stream::Stderr, |value| value.style(style)) .to_string() } #[allow(clippy::print_stderr)] fn print_websocket_startup_banner(addr: SocketAddr) { let title = colorize("codex app-server (WebSockets)", Style::new().bold().cyan()); let listening_label = colorize("listening on:", Style::new().dimmed()); let listen_url = colorize(&format!("ws://{addr}"), Style::new().green()); let ready_label = colorize("readyz:", Style::new().dimmed()); let ready_url = colorize(&format!("http://{addr}/readyz"), Style::new().green()); let health_label = colorize("healthz:", Style::new().dimmed()); let health_url = colorize(&format!("http://{addr}/healthz"), Style::new().green()); let note_label = colorize("note:", Style::new().dimmed()); eprintln!("{title}"); eprintln!(" {listening_label} {listen_url}"); eprintln!(" {ready_label} {ready_url}"); eprintln!(" {health_label} {health_url}"); if addr.ip().is_loopback() { eprintln!( " {note_label} binds localhost only (use SSH port-forwarding for remote access)" ); } else { eprintln!( " {note_label} websocket auth is opt-in in this build; configure `--ws-auth ...` before real remote use" ); } } #[derive(Clone)] struct WebSocketListenerState { transport_event_tx: mpsc::Sender, connection_counter: Arc, auth_policy: Arc, } async fn health_check_handler() -> StatusCode { StatusCode::OK } async fn reject_requests_with_origin_header( request: Request, next: Next, ) -> Result { if request.headers().contains_key(ORIGIN) { warn!( method = %request.method(), uri = %request.uri(), "rejecting websocket listener request with Origin header" ); Err(StatusCode::FORBIDDEN) } else { Ok(next.run(request).await) } } async fn websocket_upgrade_handler( websocket: WebSocketUpgrade, ConnectInfo(peer_addr): ConnectInfo, State(state): State, headers: HeaderMap, ) -> impl IntoResponse { if let Err(err) = authorize_upgrade(&headers, state.auth_policy.as_ref()) { warn!( %peer_addr, message = err.message(), "rejecting websocket client during upgrade" ); return (err.status_code(), err.message()).into_response(); } let connection_id = ConnectionId(state.connection_counter.fetch_add(1, Ordering::Relaxed)); info!(%peer_addr, "websocket client connected"); websocket .on_upgrade(move |stream| async move { run_websocket_connection(connection_id, stream, state.transport_event_tx).await; }) .into_response() } pub(crate) async fn start_websocket_acceptor( bind_address: SocketAddr, transport_event_tx: mpsc::Sender, shutdown_token: CancellationToken, auth_policy: WebsocketAuthPolicy, ) -> IoResult> { if should_warn_about_unauthenticated_non_loopback_listener(bind_address, &auth_policy) { warn!( %bind_address, "starting non-loopback websocket listener without auth; websocket auth is opt-in for now and will become the default in a future release" ); } let listener = TcpListener::bind(bind_address).await?; let local_addr = listener.local_addr()?; print_websocket_startup_banner(local_addr); info!("app-server websocket listening on ws://{local_addr}"); let router = Router::new() .route("/readyz", get(health_check_handler)) .route("/healthz", get(health_check_handler)) .fallback(any(websocket_upgrade_handler)) .layer(middleware::from_fn(reject_requests_with_origin_header)) .with_state(WebSocketListenerState { transport_event_tx, connection_counter: Arc::new(AtomicU64::new(1)), auth_policy: Arc::new(auth_policy), }); let server = axum::serve( listener, router.into_make_service_with_connect_info::(), ) .with_graceful_shutdown(async move { shutdown_token.cancelled().await; }); Ok(tokio::spawn(async move { if let Err(err) = server.await { error!("websocket acceptor failed: {err}"); } info!("websocket acceptor shutting down"); })) } async fn run_websocket_connection( connection_id: ConnectionId, websocket_stream: WebSocket, transport_event_tx: mpsc::Sender, ) { let (writer_tx, writer_rx) = mpsc::channel::(CHANNEL_CAPACITY); let writer_tx_for_reader = writer_tx.clone(); let disconnect_token = CancellationToken::new(); if transport_event_tx .send(TransportEvent::ConnectionOpened { connection_id, writer: writer_tx, disconnect_sender: Some(disconnect_token.clone()), }) .await .is_err() { return; } let (websocket_writer, websocket_reader) = websocket_stream.split(); let (writer_control_tx, writer_control_rx) = mpsc::channel::(CHANNEL_CAPACITY); let mut outbound_task = tokio::spawn(run_websocket_outbound_loop( websocket_writer, writer_rx, writer_control_rx, disconnect_token.clone(), )); let mut inbound_task = tokio::spawn(run_websocket_inbound_loop( websocket_reader, transport_event_tx.clone(), writer_tx_for_reader, writer_control_tx, connection_id, disconnect_token.clone(), )); tokio::select! { _ = &mut outbound_task => { disconnect_token.cancel(); inbound_task.abort(); } _ = &mut inbound_task => { disconnect_token.cancel(); outbound_task.abort(); } } let _ = transport_event_tx .send(TransportEvent::ConnectionClosed { connection_id }) .await; } async fn run_websocket_outbound_loop( mut websocket_writer: futures::stream::SplitSink, mut writer_rx: mpsc::Receiver, mut writer_control_rx: mpsc::Receiver, disconnect_token: CancellationToken, ) { loop { tokio::select! { _ = disconnect_token.cancelled() => { break; } message = writer_control_rx.recv() => { let Some(message) = message else { break; }; if websocket_writer.send(message).await.is_err() { break; } } queued_message = writer_rx.recv() => { let Some(queued_message) = queued_message else { break; }; let Some(json) = serialize_outgoing_message(queued_message.message) else { continue; }; if websocket_writer.send(WebSocketMessage::Text(json.into())).await.is_err() { break; } if let Some(write_complete_tx) = queued_message.write_complete_tx { let _ = write_complete_tx.send(()); } } } } } async fn run_websocket_inbound_loop( mut websocket_reader: futures::stream::SplitStream, transport_event_tx: mpsc::Sender, writer_tx_for_reader: mpsc::Sender, writer_control_tx: mpsc::Sender, connection_id: ConnectionId, disconnect_token: CancellationToken, ) { loop { tokio::select! { _ = disconnect_token.cancelled() => { break; } incoming_message = websocket_reader.next() => { match incoming_message { Some(Ok(WebSocketMessage::Text(text))) => { if !forward_incoming_message( &transport_event_tx, &writer_tx_for_reader, connection_id, text.as_ref(), ) .await { break; } } Some(Ok(WebSocketMessage::Ping(payload))) => { match writer_control_tx.try_send(WebSocketMessage::Pong(payload)) { Ok(()) => {} Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => break, Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { warn!("websocket control queue full while replying to ping; closing connection"); break; } } } Some(Ok(WebSocketMessage::Pong(_))) => {} Some(Ok(WebSocketMessage::Close(_))) | None => break, Some(Ok(WebSocketMessage::Binary(_))) => { warn!("dropping unsupported binary websocket message"); } Some(Err(err)) => { warn!("websocket receive error: {err}"); break; } } } } } }