use crate::net::{PublicEndpoint, PublicHttpClient, PublicHttpError};
use crate::util::throttle::Throttler;
use bytes::{Bytes, BytesMut};
use reqwest::StatusCode;
/// Direct-only public HTTP transport that records 429 rate-limit responses into
/// the shared [`Throttler`] before Jacquard discards the headers.
#[derive(Clone)]
pub struct ThrottledHttpClient {
client: PublicHttpClient,
max_body_bytes: usize,
pub throttler: Throttler,
}
impl ThrottledHttpClient {
#[cfg(test)]
pub(crate) fn new(client: PublicHttpClient, throttler: Throttler) -> Self {
Self::new_with_body_limit(client, throttler, 2 * 1024 * 1024)
}
pub(crate) fn new_with_body_limit(
client: PublicHttpClient,
throttler: Throttler,
max_body_bytes: usize,
) -> Self {
Self {
client,
max_body_bytes,
throttler,
}
}
/// Starts a GET request for an already typed public endpoint. Callers that
/// need streaming (CAR) bodies must finish it with [`Self::send_car`].
pub fn get(&self, endpoint: &PublicEndpoint) -> reqwest::RequestBuilder {
self.client.get(endpoint)
}
/// Executes a typed public request while retaining reqwest's streaming
/// response and the socket-coupled resolver owned by the capability.
pub async fn send_car(
&self,
request: reqwest::Request,
) -> Result<reqwest::Response, PublicHttpError> {
self.client.execute(request).await
}
}
fn uri_to_origin(uri: &http::Uri) -> Option<url::Url> {
let scheme = uri.scheme_str()?;
let host = uri.host()?;
let authority = if host.parse::<std::net::Ipv6Addr>().is_ok() {
format!("[{host}]")
} else {
host.to_owned()
};
let port = uri.port_u16().or_else(|| match scheme {
"http" => Some(80),
"https" => Some(443),
_ => None,
})?;
url::Url::parse(&format!("{scheme}://{authority}:{port}/")).ok()
}
impl ThrottledHttpClient {
pub(crate) async fn send_http(
&self,
request: http::Request<Vec<u8>>,
) -> Result<http::Response<Vec<u8>>, PublicHttpError> {
let uri = request.uri().clone();
let host_url = uri_to_origin(&uri);
let res = self
.client
.send_http_bounded(request, self.max_body_bytes)
.await;
if let Some((resp, host_url)) = res
.as_ref()
.ok()
.filter(|resp| resp.status() == StatusCode::TOO_MANY_REQUESTS)
.zip(host_url.as_ref())
{
let headers = resp.headers();
let retry_after = headers
.get(http::header::RETRY_AFTER)
.and_then(|v| v.to_str().ok())
.and_then(|s| s.parse::<u64>().ok());
let rate_limit_reset = headers
.get("ratelimit-reset")
.and_then(|v| v.to_str().ok())
.and_then(|s| s.parse::<i64>().ok())
.map(|ts| {
let now = chrono::Utc::now().timestamp();
(ts - now).max(1) as u64
});
let delay_secs = retry_after.or(rate_limit_reset);
self.throttler
.get_handle(host_url)
.await
.record_ratelimit(delay_secs);
}
res
}
}
/// streams a response body into a single contiguous [`Bytes`], enforcing a hard byte ceiling.
///
/// returns `Ok(None)` when the body exceeds `max_bytes`: the response is dropped without
/// buffering the remainder, so per-task memory stays bounded regardless of what the PDS sends
/// (a chunked response with no `Content-Length`, or a decompression bomb — `chunk` yields
/// already-decoded bytes, so the ceiling applies to the decompressed size). the buffer is grown
/// in place and never double-buffered: at most `buf + one chunk` is resident at any point.
pub(crate) async fn collect_body_bounded(
mut resp: reqwest::Response,
max_bytes: usize,
) -> Result<Option<Bytes>, reqwest::Error> {
// reject early when a declared (uncompressed) length already exceeds the ceiling. reqwest
// strips `Content-Length` from decoded responses, so a present value is the real body size.
let content_length = resp.content_length();
if content_length.is_some_and(|len| len > max_bytes as u64) {
return Ok(None);
}
let reserve = content_length.unwrap_or(0).min(max_bytes as u64) as usize;
let mut buf = BytesMut::with_capacity(reserve);
while let Some(chunk) = resp.chunk().await? {
if buf.len() + chunk.len() > max_bytes {
return Ok(None);
}
buf.extend_from_slice(&chunk);
}
Ok(Some(buf.freeze()))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::net::PublicHttpError;
use axum::{Router, body::Body, response::IntoResponse, routing::get};
#[tokio::test]
async fn public_xrpc_body_limit_is_enforced_before_caller_materializes() {
let app = Router::new().route(
"/",
get(|| async { Body::from(vec![b'x'; 65]).into_response() }),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
let transport = reqwest::Client::builder()
.resolve("public.example", address)
.build()
.unwrap();
let client = ThrottledHttpClient::new_with_body_limit(
crate::net::PublicHttpClient::test(transport),
Throttler::new(10),
64,
);
let request = http::Request::builder()
.method(http::Method::GET)
.uri("http://public.example/")
.body(Vec::new())
.unwrap();
let result = client.send_http(request).await;
assert!(matches!(
result,
Err(PublicHttpError::BodyTooLarge { limit: 64 })
));
server.abort();
}
#[tokio::test]
async fn public_transport_rejects_private_literal_before_connecting() {
let client = ThrottledHttpClient::new(
crate::net::PublicHttpClient::test(reqwest::Client::new()),
Throttler::new(10),
);
let request = http::Request::builder()
.method(http::Method::GET)
.uri("http://127.0.0.1:9/")
.body(Vec::new())
.unwrap();
let result = client.send_http(request).await;
assert!(matches!(result, Err(PublicHttpError::UnsafeEndpoint(_))));
}
}