Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192use crate::net::{PublicEndpoint, PublicHttpClient, PublicHttpError};use crate::util::throttle::Throttler;use bytes::{Bytes, BytesMut};use jacquard_common::http_client::HttpClient;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 HttpClient for ThrottledHttpClient { type Error = PublicHttpError;
async fn send_http( &self, request: http::Request<Vec<u8>>, ) -> Result<http::Response<Vec<u8>>, Self::Error> { 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(1, 10), 64, ); let request = http::Request::builder() .method(http::Method::GET) .uri("http://public.example/") .body(Vec::new()) .unwrap(); let result = HttpClient::send_http(&client, 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(1, 10), ); let request = http::Request::builder() .method(http::Method::GET) .uri("http://127.0.0.1:9/") .body(Vec::new()) .unwrap(); let result = HttpClient::send_http(&client, request).await; assert!(matches!(result, Err(PublicHttpError::UnsafeEndpoint(_)))); }}