mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +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:
1 parent
49025589b0
commit
2494d939cf
2 files changed
+76
No files matched your search
@@ -6,6 +6,8 @@ use std::sync::Arc;
|
|||||||
use std::sync::Mutex;
|
use std::sync::Mutex;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use bytes::Bytes;
|
||||||
|
use futures::TryStream;
|
||||||
use http::HeaderMap;
|
use http::HeaderMap;
|
||||||
use http::HeaderName;
|
use http::HeaderName;
|
||||||
use http::HeaderValue;
|
use http::HeaderValue;
|
||||||
@@ -100,6 +102,24 @@ impl RouteAwareRequestError {
|
|||||||
pub fn is_connect(&self) -> bool {
|
pub fn is_connect(&self) -> bool {
|
||||||
matches!(self, Self::Request(error) if error.is_connect())
|
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"]
|
#[must_use = "requests are not sent unless `send` is awaited"]
|
||||||
@@ -211,6 +231,19 @@ impl RouteAwareRequestBuilder {
|
|||||||
self
|
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> {
|
pub async fn send(self) -> Result<reqwest::Response, RouteAwareRequestError> {
|
||||||
self.pool.send(self.request?).await
|
self.pool.send(self.request?).await
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -10,6 +10,8 @@ use std::sync::atomic::Ordering;
|
|||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
|
|
||||||
|
use bytes::Bytes;
|
||||||
|
use futures::stream;
|
||||||
use pretty_assertions::assert_eq;
|
use pretty_assertions::assert_eq;
|
||||||
use tracing_subscriber::Layer;
|
use tracing_subscriber::Layer;
|
||||||
use tracing_subscriber::layer::SubscriberExt;
|
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]
|
#[tokio::test]
|
||||||
async fn forwards_exact_urls_and_caches_clients_by_resolved_route() {
|
async fn forwards_exact_urls_and_caches_clients_by_resolved_route() {
|
||||||
let pool = RouteAwareClientPool::with_builder(
|
let pool = RouteAwareClientPool::with_builder(
|
||||||
|
|||||||
Reference in new issue
Block a user