objstore_r/
http.rs

1// Author: spicyzhug@gmail.com
2// Copyright: lituus-io, all rights reserved
3//
4//! [`Client`]: the reqwest-backed [`Transport`] (feature `http`).
5//!
6//! It applies the same policy as the GCS store: HTTPS only (loopback `http`
7//! for emulators and tests), redirects refused, bounded bodies, bearer tokens
8//! from any [`TokenProvider`] (plus the quota-project header for user
9//! credentials), and jittered exponential backoff that honours `Retry-After`.
10//! One client holds one connection pool; with HTTP/2 many concurrent calls
11//! share a connection.
12//!
13//! ```no_run
14//! # async fn demo() -> objstore_r::Result<()> {
15//! use objstore_r::transport::{Request, Transport};
16//! use objstore_r::{http::Client, token::Adc, Bytes};
17//!
18//! let client = Client::new(Adc::discover()?, Default::default())?;
19//! let url = "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/publishers/google/models/m:generateContent";
20//! let resp = client.send(Request::post_json(url, Bytes::from_static(b"{}"))).await?.error_for_status()?;
21//! # let _ = resp; Ok(()) }
22//! ```
23
24use crate::net::{backoff, is_allowed_endpoint, is_retriable_status, transport_error};
25use crate::token::TokenProvider;
26use crate::transport::{BodyStream, Method, Request, Response, Retry, Transport};
27use crate::{Bytes, Error, Result};
28use std::future::Future;
29use std::time::Duration;
30
31/// Longest server-requested `Retry-After` honoured.
32const MAX_RETRY_AFTER: Duration = Duration::from_secs(30);
33
34/// Settings.
35///
36/// No timeout covers a whole response body: a large or slow body takes as
37/// long as it takes while bytes keep arriving, and fails once none arrive for
38/// [`idle_timeout`](Self::idle_timeout).
39#[derive(Debug, Clone)]
40pub struct HttpConfig {
41    /// The wait for [`Transport::send`]'s response headers, request body
42    /// included (plus `body / min_throughput` for a large request body).
43    pub timeout: Duration,
44    /// The wait for [`Transport::open`]'s response headers: a streaming API
45    /// may think for a while before its first byte.
46    pub stream_timeout: Duration,
47    /// Connection establishment (TCP, TLS and any proxy `CONNECT`).
48    pub connect_timeout: Duration,
49    /// The longest a response body may deliver no bytes, whole or streamed.
50    pub idle_timeout: Duration,
51    /// The slowest request-body upload rate tolerated, in bytes per second:
52    /// it extends the header wait by `body / min_throughput`. Must not be zero.
53    pub min_throughput: u64,
54    /// Retries after the first attempt (per the request's [`Retry`] policy).
55    pub max_retries: u32,
56    /// Largest body accepted, whole or streamed.
57    pub max_response_bytes: u64,
58}
59
60impl Default for HttpConfig {
61    fn default() -> Self {
62        Self {
63            timeout: Duration::from_secs(60),
64            stream_timeout: Duration::from_secs(300),
65            connect_timeout: Duration::from_secs(10),
66            idle_timeout: Duration::from_secs(60),
67            min_throughput: 64 << 10,
68            max_retries: 4,
69            max_response_bytes: 64 << 20,
70        }
71    }
72}
73
74/// An authorised HTTPS client. See the module docs.
75#[derive(Debug)]
76pub struct Client<T: TokenProvider> {
77    http: reqwest::Client,
78    tokens: T,
79    cfg: HttpConfig,
80}
81
82impl<T: TokenProvider> Client<T> {
83    /// A client authenticating with `tokens`.
84    pub fn new(tokens: T, cfg: HttpConfig) -> Result<Self> {
85        if cfg.min_throughput == 0
86            || [cfg.timeout, cfg.stream_timeout, cfg.connect_timeout, cfg.idle_timeout]
87                .contains(&Duration::ZERO)
88        {
89            return Err(Error::protocol(
90                "invalid configuration: timeouts and min_throughput must be positive",
91            ));
92        }
93        // No client-wide timeout and no reqwest `read_timeout` (its timer is
94        // not reset by request-body progress): see `exchange` and `Chunks`.
95        let http = reqwest::Client::builder()
96            .connect_timeout(cfg.connect_timeout)
97            .redirect(reqwest::redirect::Policy::none())
98            .pool_idle_timeout(Duration::from_secs(90))
99            .user_agent(concat!("objstore-r/", env!("CARGO_PKG_VERSION")))
100            .build()
101            .map_err(|e| Error::unavailable(format!("http client: {e}"), false))?;
102        Ok(Self { http, tokens, cfg })
103    }
104
105    /// The token provider.
106    pub fn tokens(&self) -> &T {
107        &self.tokens
108    }
109
110    /// The settings.
111    pub fn config(&self) -> &HttpConfig {
112        &self.cfg
113    }
114
115    fn check(req: &Request<'_>) -> Result<()> {
116        if !is_allowed_endpoint(req.url) {
117            return Err(Error::protocol("the URL must be https:// (http:// only to loopback)"));
118        }
119        if req.headers.iter().any(|(k, _)| k.eq_ignore_ascii_case("authorization")) {
120            return Err(Error::protocol("authorization is set by the transport"));
121        }
122        Ok(())
123    }
124
125    async fn build(&self, req: &Request<'_>) -> Result<reqwest::RequestBuilder> {
126        let method = match req.method {
127            Method::Get => reqwest::Method::GET,
128            Method::Post => reqwest::Method::POST,
129            Method::Put => reqwest::Method::PUT,
130            Method::Delete => reqwest::Method::DELETE,
131            Method::Patch => reqwest::Method::PATCH,
132        };
133        let mut rb = self.http.request(method, req.url);
134        let t = self.tokens.token().await?;
135        if !t.secret().is_empty() {
136            rb = rb.bearer_auth(t.secret());
137        }
138        if let Some(q) = self.tokens.quota_project() {
139            rb = rb.header("x-goog-user-project", q);
140        }
141        if let Some(ct) = req.content_type {
142            rb = rb.header("content-type", ct);
143        }
144        for &(k, v) in req.headers {
145            rb = rb.header(k, v);
146        }
147        if !req.body.is_empty() || matches!(req.method, Method::Post | Method::Put | Method::Patch)
148        {
149            rb = rb.body(req.body.clone());
150        }
151        Ok(rb)
152    }
153
154    /// Send with the request's retry policy; returns once headers arrive,
155    /// each attempt within `wait` plus the request body's time at
156    /// `min_throughput`.
157    async fn exchange(&self, req: &Request<'_>, wait: Duration) -> Result<reqwest::Response> {
158        Self::check(req)?;
159        let ms = (req.body.len() as u128 * 1000) / u128::from(self.cfg.min_throughput);
160        let wait =
161            wait.saturating_add(Duration::from_millis(u64::try_from(ms).unwrap_or(u64::MAX)));
162        let mut attempt = 0;
163        loop {
164            let rb = self.build(req).await?;
165            // The error, and whether the connection itself failed.
166            let outcome = match tokio::time::timeout(wait, rb.send()).await {
167                Ok(Ok(r)) => Ok(r),
168                Ok(Err(e)) => Err((transport_error(&e), e.is_connect())),
169                Err(_) => Err((Error::unavailable("timed out", true), false)),
170            };
171            let retry = match (&outcome, req.retry) {
172                (_, Retry::Never) => false,
173                (Ok(r), Retry::Transient) => is_retriable_status(r.status().as_u16()),
174                (Ok(r), Retry::Throttled) => r.status().as_u16() == 429,
175                (Err((e, _)), Retry::Transient) => e.is_retriable(),
176                (Err((_, connect)), Retry::Throttled) => *connect,
177            };
178            if !retry || attempt >= self.cfg.max_retries {
179                return outcome.map_err(|(e, _)| e);
180            }
181            attempt += 1;
182            let wait = outcome
183                .as_ref()
184                .ok()
185                .and_then(retry_after)
186                .map_or(backoff(attempt), |ra| ra.max(backoff(attempt)));
187            tokio::time::sleep(wait).await;
188        }
189    }
190
191    async fn read_all(&self, resp: reqwest::Response) -> Result<Bytes> {
192        let limit = self.cfg.max_response_bytes;
193        if resp.content_length().is_some_and(|n| n > limit) {
194            return Err(Error::TooLarge { limit });
195        }
196        let mut body = Chunks { resp, left: limit, limit, idle: self.cfg.idle_timeout };
197        // Most responses arrive in one chunk: return it without copying.
198        let Some(first) = body.next().await? else { return Ok(Bytes::new()) };
199        let Some(second) = body.next().await? else { return Ok(first) };
200        let mut buf = bytes::BytesMut::with_capacity(first.len() + second.len());
201        buf.extend_from_slice(&first);
202        buf.extend_from_slice(&second);
203        while let Some(chunk) = body.next().await? {
204            buf.extend_from_slice(&chunk);
205        }
206        Ok(buf.freeze())
207    }
208}
209
210/// The response's `Retry-After` (either form), capped.
211fn retry_after(resp: &reqwest::Response) -> Option<Duration> {
212    let v = resp.headers().get("retry-after")?.to_str().ok()?;
213    Some(crate::net::retry_after(v, graph_r::fsio::now_ms())?.min(MAX_RETRY_AFTER))
214}
215
216/// A bounded chunked body: each chunk must arrive within the idle timeout.
217/// Dropping it aborts the transfer.
218#[derive(Debug)]
219pub struct Chunks {
220    resp: reqwest::Response,
221    left: u64,
222    limit: u64,
223    idle: Duration,
224}
225
226impl Chunks {
227    async fn next(&mut self) -> Result<Option<Bytes>> {
228        let chunk = tokio::time::timeout(self.idle, self.resp.chunk())
229            .await
230            .map_err(|_| Error::unavailable("stalled: no data within the idle timeout", true))?
231            .map_err(|e| transport_error(&e))?;
232        if let Some(c) = &chunk {
233            let n = c.len() as u64;
234            if n > self.left {
235                return Err(Error::TooLarge { limit: self.limit });
236            }
237            self.left -= n;
238        }
239        Ok(chunk)
240    }
241}
242
243impl BodyStream for Chunks {
244    type Next<'a>
245        = impl Future<Output = Result<Option<Bytes>>> + 'a
246    where
247        Self: 'a;
248
249    fn next_chunk(&mut self) -> Self::Next<'_> {
250        self.next()
251    }
252}
253
254impl<T: TokenProvider> Transport for Client<T> {
255    type Send<'a>
256        = impl Future<Output = Result<Response>> + 'a
257    where
258        Self: 'a;
259    type Body<'a>
260        = Chunks
261    where
262        Self: 'a;
263    type Open<'a>
264        = impl Future<Output = Result<(u16, Chunks)>> + 'a
265    where
266        Self: 'a;
267
268    fn send<'a>(&'a self, req: Request<'a>) -> Self::Send<'a> {
269        async move {
270            let resp = self.exchange(&req, self.cfg.timeout).await?;
271            let status = resp.status().as_u16();
272            let body = self.read_all(resp).await?;
273            Ok(Response { status, body })
274        }
275    }
276
277    fn open<'a>(&'a self, req: Request<'a>) -> Self::Open<'a> {
278        async move {
279            let resp = self.exchange(&req, self.cfg.stream_timeout).await?;
280            let limit = self.cfg.max_response_bytes;
281            if resp.content_length().is_some_and(|n| n > limit) {
282                return Err(Error::TooLarge { limit });
283            }
284            let idle = self.cfg.idle_timeout;
285            Ok((resp.status().as_u16(), Chunks { resp, left: limit, limit, idle }))
286        }
287    }
288}
289
290/// Run `fut` to completion on a single-threaded runtime driving this
291/// transport's I/O and timers — for command-line tools with a synchronous
292/// `main`. Libraries should stay runtime-agnostic and leave this to the
293/// binary. Fails only if the runtime cannot start.
294#[cfg(feature = "blocking")]
295pub fn block_on<F: std::future::Future>(fut: F) -> std::io::Result<F::Output> {
296    let rt = tokio::runtime::Builder::new_current_thread().enable_all().build()?;
297    Ok(rt.block_on(fut))
298}
objstore_r/http.rs

objstore_r/
http.rs
