diff --git a/codex-rs/code-mode-host/src/grpc_transport.rs b/codex-rs/code-mode-host/src/grpc_transport.rs index 665e8de55f..4cb59985fa 100644 --- a/codex-rs/code-mode-host/src/grpc_transport.rs +++ b/codex-rs/code-mode-host/src/grpc_transport.rs @@ -4,9 +4,17 @@ use std::net::SocketAddr; use anyhow::Context; use anyhow::Result; +use axum::extract::Request; +use axum::http::Method; +use axum::http::StatusCode; +use axum::http::Version; +use axum::middleware; +use axum::middleware::Next; +use axum::routing::get; use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHostServer; use codex_code_mode_protocol::host::MAX_FRAME_BYTES; use tokio::net::TcpListener; +use tonic::service::Routes; use tonic::transport::Server; use tonic::transport::server::TcpIncoming; use tracing::info; @@ -24,12 +32,28 @@ pub(super) async fn run_tcp_listener(bind_address: SocketAddr) -> Result<()> { .flush() .context("failed to publish code-mode gRPC listen address")?; + let routes = Routes::new( + CodeModeHostServer::new(GrpcCodeModeHost::new()) + .max_decoding_message_size(MAX_FRAME_BYTES) + .max_encoding_message_size(MAX_FRAME_BYTES), + ) + .into_axum_router() + .route("/healthz", get(|| async { StatusCode::OK })) + .layer(middleware::from_fn( + |request: Request, next: Next| async move { + if request.version() != Version::HTTP_2 + && (request.method() != Method::GET || request.uri().path() != "/healthz") + { + return Err(StatusCode::HTTP_VERSION_NOT_SUPPORTED); + } + + Ok(next.run(request).await) + }, + )); + Server::builder() - .add_service( - CodeModeHostServer::new(GrpcCodeModeHost::new()) - .max_decoding_message_size(MAX_FRAME_BYTES) - .max_encoding_message_size(MAX_FRAME_BYTES), - ) + .accept_http1(/*accept_http1*/ true) + .add_routes(routes.into()) .serve_with_incoming(listener) .await .context("code-mode gRPC TCP listener failed") diff --git a/codex-rs/code-mode-host/tests/grpc_tcp.rs b/codex-rs/code-mode-host/tests/grpc_tcp.rs index 5532b2b2a9..b15699b944 100644 --- a/codex-rs/code-mode-host/tests/grpc_tcp.rs +++ b/codex-rs/code-mode-host/tests/grpc_tcp.rs @@ -1,18 +1,122 @@ +use std::future::poll_fn; use std::process::Stdio; use std::time::Duration; use anyhow::Context; use anyhow::Result; +use axum::http::Request; +use axum::http::StatusCode; +use axum::http::Version; use codex_code_mode_protocol::grpc; use codex_code_mode_protocol::grpc::code_mode_host_client::CodeModeHostClient; +use pretty_assertions::assert_eq; use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; use tokio::io::BufReader; +use tokio::net::TcpStream; use tokio::process::Command; use tokio::time::timeout; +use tonic::body::Body; +use tonic::codegen::Service; use tonic::transport::Endpoint; +#[path = "support/host.rs"] +mod host; + +use host::HostHarness; + const TEST_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 10); +#[tokio::test] +async fn tcp_listener_serves_http1_healthz() -> Result<()> { + let host = HostHarness::start("grpc://127.0.0.1:0").await?; + let request = "GET /healthz HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n"; + let response = http1_response(&host.endpoint, request.as_bytes()).await?; + + assert_eq!(response.lines().next(), Some("HTTP/1.1 200 OK")); + Ok(()) +} + +#[tokio::test] +async fn tcp_listener_rejects_http1_grpc_requests() -> Result<()> { + let host = HostHarness::start("grpc://127.0.0.1:0").await?; + let mut request = concat!( + "POST /codex.code_mode.v1.CodeModeHost/OpenSession HTTP/1.1\r\n", + "Host: localhost\r\n", + "Content-Type: application/grpc\r\n", + "Content-Length: 5\r\n", + "Connection: close\r\n\r\n" + ) + .as_bytes() + .to_vec(); + request.extend_from_slice(&[0; 5]); + let response = http1_response(&host.endpoint, &request).await?; + + assert_eq!( + response.lines().next(), + Some("HTTP/1.1 505 HTTP Version Not Supported") + ); + Ok(()) +} + +#[tokio::test] +async fn tcp_listener_serves_http2_healthz() -> Result<()> { + let host = HostHarness::start("grpc://127.0.0.1:0").await?; + let mut channel = Endpoint::from_shared(host.endpoint.clone()) + .context("gRPC code-mode host published an invalid endpoint")? + .connect_timeout(TEST_TIMEOUT) + .timeout(TEST_TIMEOUT) + .connect() + .await + .context("failed to connect to gRPC code-mode host health endpoint")?; + let request = Request::builder() + .uri(format!("{}/healthz", host.endpoint)) + .body(Body::empty()) + .context("failed to build gRPC code-mode host health request")?; + + let response = timeout(TEST_TIMEOUT, async { + poll_fn(|context| channel.poll_ready(context)) + .await + .context("gRPC code-mode host health channel is unavailable")?; + channel + .call(request) + .await + .context("failed to request gRPC code-mode host health") + }) + .await + .context("timed out requesting gRPC code-mode host health")??; + + assert_eq!(response.version(), Version::HTTP_2); + assert_eq!(response.status(), StatusCode::OK); + Ok(()) +} + +async fn http1_response(endpoint: &str, request: &[u8]) -> Result { + let address = endpoint + .strip_prefix("http://") + .context("gRPC code-mode host URL should use http://")?; + + timeout(TEST_TIMEOUT, async { + let mut stream = TcpStream::connect(address) + .await + .context("failed to connect to gRPC code-mode host")?; + stream + .write_all(request) + .await + .context("failed to send HTTP/1.1 request to gRPC code-mode host")?; + + let mut response = String::new(); + stream + .read_to_string(&mut response) + .await + .context("failed to read HTTP/1.1 response from gRPC code-mode host")?; + Ok(response) + }) + .await + .context("timed out waiting for HTTP/1.1 response from gRPC code-mode host")? +} + #[tokio::test] async fn tcp_listener_opens_a_grpc_session() -> Result<()> { let mut host = Command::new(codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?)