mirror of
https://github.com/openai/codex.git
synced 2026-08-23 13:09:46 +00:00
Support streaming bodies in route-aware HTTP requests (#35715)
## What changed - Add `RouteAwareRequestBuilder::body_stream` for sending fallible byte streams without exposing the underlying HTTP client body type. - Expose request- and body-error classification on `RouteAwareRequestError`. - Add `RouteAwareRequestError::without_url` so callers can remove credential-bearing URLs, such as signed upload URLs, from transport errors. ## Testing - Verify streamed request bytes reach the server. - Verify URL secrets are absent after stripping a transport error's URL. GitOrigin-RevId: 6a100f56a9aa7d79d9b6dd0f103135f5f2060617
This commit is contained in:
@@ -6,6 +6,8 @@ use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures::TryStream;
|
||||
use http::HeaderMap;
|
||||
use http::HeaderName;
|
||||
use http::HeaderValue;
|
||||
@@ -100,6 +102,24 @@ impl RouteAwareRequestError {
|
||||
pub fn is_connect(&self) -> bool {
|
||||
matches!(self, Self::Request(error) if error.is_connect())
|
||||
}
|
||||
|
||||
pub fn is_body(&self) -> bool {
|
||||
matches!(self, Self::Request(error) if error.is_body())
|
||||
}
|
||||
|
||||
pub fn is_request(&self) -> bool {
|
||||
matches!(self, Self::Request(error) if error.is_request())
|
||||
}
|
||||
|
||||
/// Removes a request URL from the underlying transport error before it is logged or returned.
|
||||
///
|
||||
/// Use this for requests whose URL can contain credentials, such as signed blob uploads.
|
||||
pub fn without_url(self) -> Self {
|
||||
match self {
|
||||
Self::Request(error) => Self::Request(error.without_url()),
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use = "requests are not sent unless `send` is awaited"]
|
||||
@@ -211,6 +231,19 @@ impl RouteAwareRequestBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Sets a streaming request body without exposing the underlying HTTP implementation.
|
||||
pub fn body_stream<S>(mut self, stream: S) -> Self
|
||||
where
|
||||
S: TryStream + Send + 'static,
|
||||
S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
|
||||
Bytes: From<S::Ok>,
|
||||
{
|
||||
if let Ok(request) = &mut self.request {
|
||||
*request.body_mut() = Some(reqwest::Body::wrap_stream(stream));
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
pub async fn send(self) -> Result<reqwest::Response, RouteAwareRequestError> {
|
||||
self.pool.send(self.request?).await
|
||||
}
|
||||
|
||||
@@ -10,6 +10,8 @@ use std::sync::atomic::Ordering;
|
||||
use std::time::Duration;
|
||||
use std::time::Instant;
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures::stream;
|
||||
use pretty_assertions::assert_eq;
|
||||
use tracing_subscriber::Layer;
|
||||
use tracing_subscriber::layer::SubscriberExt;
|
||||
@@ -38,6 +40,47 @@ fn request_builder_debug_redacts_url_secrets() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streams_request_bodies_without_exposing_reqwest_body() {
|
||||
let (address, server) = spawn_response_server(vec![
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_string(),
|
||||
]);
|
||||
let pool = RouteAwareClientPool::new(
|
||||
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
|
||||
ClientRouteClass::Api,
|
||||
);
|
||||
let response = pool
|
||||
.put(format!("http://{address}/upload"))
|
||||
.header(http::header::CONTENT_LENGTH, /*value*/ 5)
|
||||
.body_stream(stream::iter(vec![Ok::<_, io::Error>(Bytes::from_static(
|
||||
b"hello",
|
||||
))]))
|
||||
.send()
|
||||
.await
|
||||
.expect("streaming request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let requests = server.join().expect("response server should finish");
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert!(requests[0].ends_with("\r\n\r\nhello"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn without_url_redacts_transport_error_urls() {
|
||||
let listener = TcpListener::bind(("127.0.0.1", 0)).expect("listener should bind");
|
||||
let address = listener.local_addr().expect("listener should have address");
|
||||
drop(listener);
|
||||
let secret = "signed-secret";
|
||||
let error = reqwest::Client::new()
|
||||
.get(format!("http://{address}/upload?sig={secret}"))
|
||||
.send()
|
||||
.await
|
||||
.expect_err("closed listener should reject request");
|
||||
let error = RouteAwareRequestError::from(error).without_url();
|
||||
|
||||
assert!(!error.to_string().contains(secret));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forwards_exact_urls_and_caches_clients_by_resolved_route() {
|
||||
let pool = RouteAwareClientPool::with_builder(
|
||||
|
||||
Reference in New Issue
Block a user