diff --git a/src/server.rs b/src/server.rs index 5d6e8c2..594225d 100644 --- a/src/server.rs +++ b/src/server.rs @@ -13,6 +13,8 @@ use serde_json::{Value, json}; use std::{sync::Arc, time::Duration}; use tokio::sync::Semaphore; +const RESPONSES_BODY_LIMIT: usize = 64 * 1024 * 1024; + pub struct Gateway { pub backend: Backend, pub models: Vec, @@ -42,15 +44,24 @@ pub fn router(gateway: Gateway) -> Router { .route("/codex/images/edits", post(crate::images::edit)) .route("/backend-api/codex/images/generations", post(crate::images::generate)) .route("/backend-api/codex/images/edits", post(crate::images::edit)) - .route("/v1/responses", post(compatible_create).get(crate::websocket::upgrade)) + .route( + "/v1/responses", + post(compatible_create) + .layer(DefaultBodyLimit::max(RESPONSES_BODY_LIMIT)) + .get(crate::websocket::upgrade), + ) .route( "/codex/responses", - post(native_create).get(crate::websocket::upgrade), + post(native_create) + .layer(DefaultBodyLimit::max(RESPONSES_BODY_LIMIT)) + .get(crate::websocket::upgrade), ) .route("/codex/models", get(native_models)) .route( "/backend-api/codex/responses", - post(native_create).get(crate::websocket::upgrade), + post(native_create) + .layer(DefaultBodyLimit::max(RESPONSES_BODY_LIMIT)) + .get(crate::websocket::upgrade), ) .route("/backend-api/codex/models", get(native_models)) .route("/v1/responses/compact", post(compact)) @@ -285,15 +296,15 @@ async fn native_create( .map_err(|_| GatewayError::invalid("invalid zstd request"))?; let mut decoded = Vec::new(); decoder - .take(16 * 1024 * 1024 + 1) + .take(RESPONSES_BODY_LIMIT as u64 + 1) .read_to_end(&mut decoded) .map_err(|_| GatewayError::invalid("invalid zstd request"))?; - if decoded.len() > 16 * 1024 * 1024 { + if decoded.len() > RESPONSES_BODY_LIMIT { return Err(GatewayError { upstream_response: None, status: StatusCode::PAYLOAD_TOO_LARGE, code: "invalid_request_error", - message: "decoded body exceeds 16 MiB".into(), + message: "decoded body exceeds 64 MiB".into(), }); } Ok(decoded) @@ -396,9 +407,8 @@ async fn respond( usage.failure("upstream_timeout"); timeout_error() })? - .map_err(|e| { + .inspect_err(|e| { usage.failure(e.code); - e })?; let upstream_headers = forward_response_headers(&running.headers); let mut accumulator = crate::output::OutputAccumulator::default(); diff --git a/tests/gateway.rs b/tests/gateway.rs index e77aa18..1fe0bd3 100644 --- a/tests/gateway.rs +++ b/tests/gateway.rs @@ -220,7 +220,12 @@ impl Harness { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let upstream_url = format!("http://{}", listener.local_addr().unwrap()); let app = Router::new() - .route("/responses", post(upstream).get(upstream_ws)) + .route( + "/responses", + post(upstream) + .layer(axum::extract::DefaultBodyLimit::max(65 * 1024 * 1024)) + .get(upstream_ws), + ) .route("/realtime/calls", post(realtime_cases::upstream_call)) .route("/realtime", axum::routing::get(upstream_ws)) .route( @@ -886,7 +891,7 @@ async fn native_rejects_unauthorized_and_nonstream_and_decompression_bombs() { .status(), StatusCode::BAD_REQUEST ); - let bomb = zstd::encode_all(vec![b' '; 16 * 1024 * 1024 + 1].as_slice(), 1).unwrap(); + let bomb = zstd::encode_all(vec![b' '; 64 * 1024 * 1024 + 1].as_slice(), 1).unwrap(); assert_eq!( client .post(&url) @@ -899,6 +904,14 @@ async fn native_rejects_unauthorized_and_nonstream_and_decompression_bombs() { .status(), StatusCode::PAYLOAD_TOO_LARGE ); + let oversized = client + .post(&url) + .bearer_auth("client-key") + .body(vec![b' '; 64 * 1024 * 1024 + 1]) + .send() + .await + .unwrap(); + assert_eq!(oversized.status(), StatusCode::PAYLOAD_TOO_LARGE); assert!(h.fake.received.lock().unwrap().is_empty()); } @@ -1583,14 +1596,17 @@ mod response_type_cases; #[tokio::test] async fn public_http_preserves_prepared_codex_requests() { - let h = Harness::new(events(), StatusCode::OK, false, Duration::from_secs(5)).await; + let h = Harness::new(events(), StatusCode::OK, false, Duration::from_secs(30)).await; + let image = format!("data:image/png;base64,{}", "A".repeat(17 * 1024 * 1024)); let body = json!({"model":"future-model", "stream":true, "store":false, - "instructions":"unchanged", "input":[], "tools":null, + "instructions":"unchanged", "tools":null, + "input":[{"role":"user","content":[{"type":"input_image","image_url":image}]}], "client_metadata":{"trace":"keep"}, "future_field":{"keep":true}}); - let raw = serde_json::to_vec(&body).unwrap(); - for compressed in [false, true] { + let mut raw = serde_json::to_vec(&body).unwrap(); + raw.resize(64 * 1024 * 1024, b' '); + for (path, compressed) in [("v1", false), ("codex", true), ("backend-api/codex", false)] { let response = reqwest::Client::new() - .post(format!("{}/v1/responses", h.url)) + .post(format!("{}/{path}/responses", h.url)) .bearer_auth("client-key") .header("content-type", "application/json") .header( @@ -1615,9 +1631,9 @@ async fn public_http_preserves_prepared_codex_requests() { ); } let received = h.fake.received.lock().unwrap(); - assert_eq!(received.len(), 2); + assert_eq!(received.len(), 3); for (_, forwarded) in received.iter() { - assert_eq!(forwarded, &body); + assert!(forwarded == &body); } }