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

objstore_r/
http.rs
