mirror of
https://github.com/openai/codex.git
synced 2026-11-13 06:45:07 +00:00
Remove WebSocket transport from code-mode hosts (#40692)
## What changed - Accept only `http://` and `https://` gRPC endpoints for app server `--code-mode-host` connections. - Limit the standalone code-mode host to stdio and gRPC listeners. - Remove the WebSocket session provider, dual-WebSocket negotiation, transport lanes, and related dependencies and tests. - Keep coverage for shared and prewarmed gRPC hosts and reject WebSocket host URLs in CLI parsing tests. GitOrigin-RevId: 8792604e79ffd54ce776074a4dc3f7d55e63752a
This commit is contained in:
Ruslan Nigmatullin authored and copyberry committed
1 parent
9be8d6e1c3
commit
725b3a44f8
31 files changed
+153
-2766
No files matched your search
@@ -1,211 +1,22 @@
|
||||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
use std::io::Write as _;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::PoisonError;
|
||||
|
||||
use anyhow::Context;
|
||||
use anyhow::Result;
|
||||
use axum::Router;
|
||||
use axum::body::Body;
|
||||
use axum::extract::ConnectInfo;
|
||||
use axum::extract::Path;
|
||||
use axum::extract::State;
|
||||
use axum::extract::ws::Message;
|
||||
use axum::extract::ws::WebSocket;
|
||||
use axum::extract::ws::WebSocketUpgrade;
|
||||
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 axum::serve::Listener;
|
||||
use axum::serve::ListenerExt;
|
||||
use axum::serve::TapIo;
|
||||
use codex_code_mode_protocol::host::ClientToHost;
|
||||
use codex_code_mode_protocol::host::EncodedFrame;
|
||||
use codex_code_mode_protocol::host::FramedReader;
|
||||
use codex_code_mode_protocol::host::FramedWriter;
|
||||
use codex_code_mode_protocol::host::HostToClient;
|
||||
use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
|
||||
use futures::SinkExt;
|
||||
use futures::StreamExt;
|
||||
use futures::stream::SplitSink;
|
||||
use futures::stream::SplitStream;
|
||||
use tokio::io::AsyncRead;
|
||||
use tokio::io::AsyncWrite;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::oneshot;
|
||||
use tracing::info;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::HostLimits;
|
||||
use crate::MAX_IN_FLIGHT_REQUESTS;
|
||||
use crate::grpc_transport;
|
||||
|
||||
/// The default transport retains the standalone host's original stdio behavior.
|
||||
pub const DEFAULT_LISTEN_URL: &str = "stdio";
|
||||
|
||||
const MAX_WEBSOCKET_FRAME_BYTES: usize = MAX_FRAME_BYTES + std::mem::size_of::<u32>();
|
||||
const MAX_PENDING_BULK_CONNECTIONS: usize = MAX_IN_FLIGHT_REQUESTS;
|
||||
|
||||
type BoxedReader = Box<dyn AsyncRead + Send + Unpin>;
|
||||
type BoxedWriter = Box<dyn AsyncWrite + Send + Unpin>;
|
||||
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
enum ListenTransport {
|
||||
Stdio,
|
||||
WebSocket(SocketAddr),
|
||||
Grpc(SocketAddr),
|
||||
}
|
||||
|
||||
pub(crate) enum ConnectionReader {
|
||||
Framed(FramedReader<BoxedReader>),
|
||||
WebSocket(SplitStream<WebSocket>),
|
||||
}
|
||||
|
||||
pub(crate) enum ConnectionWriter {
|
||||
Framed(FramedWriter<BoxedWriter>),
|
||||
WebSocket(SplitSink<WebSocket, Message>),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct WebSocketListenerState {
|
||||
limits: Arc<HostLimits>,
|
||||
bulk_connections: BulkConnectionRegistry,
|
||||
}
|
||||
|
||||
pub(crate) struct BulkConnection {
|
||||
pub(crate) reader: ConnectionReader,
|
||||
pub(crate) writer: ConnectionWriter,
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub(crate) struct BulkConnectionRegistry {
|
||||
pending: Arc<Mutex<HashMap<Uuid, oneshot::Sender<BulkConnection>>>>,
|
||||
}
|
||||
|
||||
pub(crate) struct BulkConnectionRegistration {
|
||||
registry: BulkConnectionRegistry,
|
||||
token: Uuid,
|
||||
receiver: oneshot::Receiver<BulkConnection>,
|
||||
}
|
||||
|
||||
impl BulkConnectionRegistry {
|
||||
pub(crate) fn reserve(&self) -> Option<BulkConnectionRegistration> {
|
||||
let mut pending = self.pending.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
if pending.len() >= MAX_PENDING_BULK_CONNECTIONS {
|
||||
return None;
|
||||
}
|
||||
|
||||
let token = Uuid::new_v4();
|
||||
let (sender, receiver) = oneshot::channel();
|
||||
pending.insert(token, sender);
|
||||
Some(BulkConnectionRegistration {
|
||||
registry: self.clone(),
|
||||
token,
|
||||
receiver,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn remove(&self, token: Uuid) -> Option<oneshot::Sender<BulkConnection>> {
|
||||
self.pending
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.remove(&token)
|
||||
}
|
||||
}
|
||||
|
||||
impl BulkConnectionRegistration {
|
||||
pub(crate) fn token(&self) -> Uuid {
|
||||
self.token
|
||||
}
|
||||
|
||||
pub(crate) async fn receive(&mut self) -> Result<BulkConnection, oneshot::error::RecvError> {
|
||||
(&mut self.receiver).await
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for BulkConnectionRegistration {
|
||||
fn drop(&mut self) {
|
||||
self.registry.remove(self.token);
|
||||
}
|
||||
}
|
||||
|
||||
impl ConnectionReader {
|
||||
pub(crate) fn from_reader<R>(reader: R) -> Self
|
||||
where
|
||||
R: AsyncRead + Send + Unpin + 'static,
|
||||
{
|
||||
Self::Framed(FramedReader::new(Box::new(reader)))
|
||||
}
|
||||
|
||||
pub(crate) async fn read(&mut self) -> io::Result<Option<ClientToHost>> {
|
||||
match self {
|
||||
Self::Framed(reader) => reader.read().await,
|
||||
Self::WebSocket(reader) => loop {
|
||||
match reader.next().await {
|
||||
Some(Ok(Message::Binary(bytes))) => {
|
||||
return EncodedFrame::decode_framed(&bytes).map(Some);
|
||||
}
|
||||
Some(Ok(Message::Text(_))) => {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
"code-mode websocket messages must be binary framed messages",
|
||||
));
|
||||
}
|
||||
Some(Ok(Message::Ping(_) | Message::Pong(_))) => {}
|
||||
Some(Ok(Message::Close(_))) | None => return Ok(None),
|
||||
Some(Err(err)) => {
|
||||
return Err(io::Error::other(format!(
|
||||
"failed to read code-mode websocket message: {err}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ConnectionWriter {
|
||||
pub(crate) fn from_writer<W>(writer: W) -> Self
|
||||
where
|
||||
W: AsyncWrite + Send + Unpin + 'static,
|
||||
{
|
||||
Self::Framed(FramedWriter::new(Box::new(writer)))
|
||||
}
|
||||
|
||||
pub(crate) async fn write(&mut self, message: &HostToClient) -> io::Result<()> {
|
||||
self.write_frame(EncodedFrame::encode(message)?).await
|
||||
}
|
||||
|
||||
pub(crate) async fn write_frame(&mut self, frame: EncodedFrame) -> io::Result<()> {
|
||||
match self {
|
||||
Self::Framed(writer) => writer.write_frame(&frame).await,
|
||||
Self::WebSocket(writer) => writer
|
||||
.send(Message::Binary(frame.into_framed_bytes().into()))
|
||||
.await
|
||||
.map_err(|err| {
|
||||
io::Error::other(format!(
|
||||
"failed to write code-mode websocket message: {err}"
|
||||
))
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn run_transport(listen_url: &str) -> Result<()> {
|
||||
match parse_listen_url(listen_url)? {
|
||||
ListenTransport::Stdio => crate::run_stdio().await,
|
||||
ListenTransport::WebSocket(bind_address) => run_websocket_listener(bind_address).await,
|
||||
ListenTransport::Grpc(bind_address) => grpc_transport::run_tcp_listener(bind_address).await,
|
||||
}
|
||||
}
|
||||
@@ -215,15 +26,6 @@ fn parse_listen_url(listen_url: &str) -> Result<ListenTransport> {
|
||||
return Ok(ListenTransport::Stdio);
|
||||
}
|
||||
|
||||
if let Some(socket_addr) = listen_url.strip_prefix("ws://") {
|
||||
return socket_addr
|
||||
.parse::<SocketAddr>()
|
||||
.map(ListenTransport::WebSocket)
|
||||
.with_context(|| {
|
||||
format!("invalid websocket --listen URL `{listen_url}`; expected `ws://IP:PORT`")
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(socket_addr) = listen_url.strip_prefix("grpc://") {
|
||||
return socket_addr
|
||||
.parse::<SocketAddr>()
|
||||
@@ -234,129 +36,10 @@ fn parse_listen_url(listen_url: &str) -> Result<ListenTransport> {
|
||||
}
|
||||
|
||||
anyhow::bail!(
|
||||
"unsupported --listen URL `{listen_url}`; expected `ws://IP:PORT`, `grpc://IP:PORT`, `stdio`, or `stdio://`"
|
||||
"unsupported --listen URL `{listen_url}`; expected `grpc://IP:PORT`, `stdio`, or `stdio://`"
|
||||
);
|
||||
}
|
||||
|
||||
async fn bind_websocket_listener(
|
||||
bind_address: SocketAddr,
|
||||
) -> Result<TapIo<TcpListener, impl FnMut(&mut TcpStream) + Send + 'static>> {
|
||||
let listener = TcpListener::bind(bind_address)
|
||||
.await
|
||||
.with_context(|| format!("failed to bind code-mode host websocket to {bind_address}"))?;
|
||||
|
||||
Ok(listener.tap_io(|stream| {
|
||||
if let Err(error) = stream.set_nodelay(/*nodelay*/ true) {
|
||||
warn!(%error, "failed to enable TCP_NODELAY for code-mode host connection");
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
async fn run_websocket_listener(bind_address: SocketAddr) -> Result<()> {
|
||||
let listener = bind_websocket_listener(bind_address).await?;
|
||||
let local_addr = listener
|
||||
.local_addr()
|
||||
.context("failed to read code-mode host websocket listen address")?;
|
||||
let state = WebSocketListenerState {
|
||||
limits: Arc::new(HostLimits::new()),
|
||||
bulk_connections: BulkConnectionRegistry::default(),
|
||||
};
|
||||
info!("codex-code-mode-host listening on ws://{local_addr}");
|
||||
println!("ws://{local_addr}");
|
||||
io::stdout()
|
||||
.flush()
|
||||
.context("failed to publish code-mode host websocket listen address")?;
|
||||
|
||||
let router = Router::new()
|
||||
.route("/", any(websocket_upgrade_handler))
|
||||
.route("/bulk/{token}", any(bulk_websocket_upgrade_handler))
|
||||
.route("/readyz", get(readiness_handler))
|
||||
.layer(middleware::from_fn(reject_requests_with_origin_header))
|
||||
.with_state(state);
|
||||
axum::serve(
|
||||
listener,
|
||||
router.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.await
|
||||
.context("code-mode host websocket listener failed")
|
||||
}
|
||||
|
||||
async fn readiness_handler() -> StatusCode {
|
||||
StatusCode::OK
|
||||
}
|
||||
|
||||
async fn reject_requests_with_origin_header(
|
||||
request: Request<Body>,
|
||||
next: Next,
|
||||
) -> Result<Response, StatusCode> {
|
||||
if request.headers().contains_key(ORIGIN) {
|
||||
warn!(
|
||||
method = %request.method(),
|
||||
uri = %request.uri(),
|
||||
"rejecting code-mode host websocket request with Origin header"
|
||||
);
|
||||
Err(StatusCode::FORBIDDEN)
|
||||
} else {
|
||||
Ok(next.run(request).await)
|
||||
}
|
||||
}
|
||||
|
||||
async fn websocket_upgrade_handler(
|
||||
websocket: WebSocketUpgrade,
|
||||
ConnectInfo(peer_addr): ConnectInfo<SocketAddr>,
|
||||
State(state): State<WebSocketListenerState>,
|
||||
) -> impl IntoResponse {
|
||||
websocket
|
||||
.max_frame_size(MAX_WEBSOCKET_FRAME_BYTES)
|
||||
.max_message_size(MAX_WEBSOCKET_FRAME_BYTES)
|
||||
.on_upgrade(move |stream| async move {
|
||||
info!(%peer_addr, "code-mode host websocket client connected");
|
||||
let (writer, reader) = stream.split();
|
||||
if let Err(err) = crate::run_connection(
|
||||
ConnectionReader::WebSocket(reader),
|
||||
ConnectionWriter::WebSocket(writer),
|
||||
state.limits,
|
||||
Some(state.bulk_connections),
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(%peer_addr, "code-mode host session failed: {err:#}");
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn bulk_websocket_upgrade_handler(
|
||||
websocket: WebSocketUpgrade,
|
||||
Path(token): Path<String>,
|
||||
ConnectInfo(peer_addr): ConnectInfo<SocketAddr>,
|
||||
State(state): State<WebSocketListenerState>,
|
||||
) -> Response {
|
||||
let Ok(token) = Uuid::parse_str(&token) else {
|
||||
return StatusCode::NOT_FOUND.into_response();
|
||||
};
|
||||
let Some(pending) = state.bulk_connections.remove(token) else {
|
||||
return StatusCode::NOT_FOUND.into_response();
|
||||
};
|
||||
|
||||
websocket
|
||||
.max_frame_size(MAX_WEBSOCKET_FRAME_BYTES)
|
||||
.max_message_size(MAX_WEBSOCKET_FRAME_BYTES)
|
||||
.on_upgrade(move |stream| async move {
|
||||
info!(%peer_addr, "code-mode host bulk websocket client connected");
|
||||
let (writer, reader) = stream.split();
|
||||
if pending
|
||||
.send(BulkConnection {
|
||||
reader: ConnectionReader::WebSocket(reader),
|
||||
writer: ConnectionWriter::WebSocket(writer),
|
||||
})
|
||||
.is_err()
|
||||
{
|
||||
warn!(%peer_addr, "code-mode host bulk websocket pairing expired");
|
||||
}
|
||||
})
|
||||
.into_response()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "transport_tests.rs"]
|
||||
mod tests;
|
||||
Reference in new issue
Block a user