Skip to main content

objstore_r/
http.rs

1// Author: support@lituus.io
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#[derive(Debug, Clone)]
36pub struct HttpConfig {
37    /// Whole-request timeout for [`Transport::send`].
38    pub timeout: Duration,
39    /// Whole-request timeout for [`Transport::open`] (long streams).
40    pub stream_timeout: Duration,
41    /// Connection-establishment timeout.
42    pub connect_timeout: Duration,
43    /// Retries after the first attempt (per the request's [`Retry`] policy).
44    pub max_retries: u32,
45    /// Largest body accepted, whole or streamed.
46    pub max_response_bytes: u64,
47}
48
49impl Default for HttpConfig {
50    fn default() -> Self {
51        Self {
52            timeout: Duration::from_secs(60),
53            stream_timeout: Duration::from_secs(300),
54            connect_timeout: Duration::from_secs(10),
55            max_retries: 4,
56            max_response_bytes: 64 << 20,
57        }
58    }
59}
60
61/// An authorised HTTPS client. See the module docs.
62#[derive(Debug)]
63pub struct Client<T: TokenProvider> {
64    http: reqwest::Client,
65    tokens: T,
66    cfg: HttpConfig,
67}
68
69impl<T: TokenProvider> Client<T> {
70    /// A client authenticating with `tokens`.
71    pub fn new(tokens: T, cfg: HttpConfig) -> Result<Self> {
72        let http = reqwest::Client::builder()
73            .connect_timeout(cfg.connect_timeout)
74            .redirect(reqwest::redirect::Policy::none())
75            .pool_idle_timeout(Duration::from_secs(90))
76            .user_agent(concat!("objstore-r/", env!("CARGO_PKG_VERSION")))
77            .build()
78            .map_err(|e| Error::unavailable(format!("http client: {e}"), false))?;
79        Ok(Self { http, tokens, cfg })
80    }
81
82    /// The token provider.
83    pub fn tokens(&self) -> &T {
84        &self.tokens
85    }
86
87    /// The settings.
88    pub fn config(&self) -> &HttpConfig {
89        &self.cfg
90    }
91
92    fn check(req: &Request<'_>) -> Result<()> {
93        if !is_allowed_endpoint(req.url) {
94            return Err(Error::protocol("the URL must be https:// (http:// only to loopback)"));
95        }
96        if req.headers.iter().any(|(k, _)| k.eq_ignore_ascii_case("authorization")) {
97            return Err(Error::protocol("authorization is set by the transport"));
98        }
99        Ok(())
100    }
101
102    async fn build(&self, req: &Request<'_>, timeout: Duration) -> Result<reqwest::RequestBuilder> {
103        let method = match req.method {
104            Method::Get => reqwest::Method::GET,
105            Method::Post => reqwest::Method::POST,
106            Method::Put => reqwest::Method::PUT,
107            Method::Delete => reqwest::Method::DELETE,
108            Method::Patch => reqwest::Method::PATCH,
109        };
110        let mut rb = self.http.request(method, req.url).timeout(timeout);
111        let t = self.tokens.token().await?;
112        if !t.secret().is_empty() {
113            rb = rb.bearer_auth(t.secret());
114        }
115        if let Some(q) = self.tokens.quota_project() {
116            rb = rb.header("x-goog-user-project", q);
117        }
118        if let Some(ct) = req.content_type {
119            rb = rb.header("content-type", ct);
120        }
121        for &(k, v) in req.headers {
122            rb = rb.header(k, v);
123        }
124        if !req.body.is_empty() || matches!(req.method, Method::Post | Method::Put | Method::Patch)
125        {
126            rb = rb.body(req.body.clone());
127        }
128        Ok(rb)
129    }
130
131    /// Send with the request's retry policy; returns once headers arrive.
132    async fn exchange(&self, req: &Request<'_>, timeout: Duration) -> Result<reqwest::Response> {
133        Self::check(req)?;
134        let mut attempt = 0;
135        loop {
136            let outcome = self.build(req, timeout).await?.send().await;
137            let retry = match (&outcome, req.retry) {
138                (_, Retry::Never) => false,
139                (Ok(r), Retry::Transient) => is_retriable_status(r.status().as_u16()),
140                (Ok(r), Retry::Throttled) => r.status().as_u16() == 429,
141                (Err(e), Retry::Transient) => e.is_timeout() || e.is_connect() || e.is_request(),
142                (Err(e), Retry::Throttled) => e.is_connect(),
143            };
144            if !retry || attempt >= self.cfg.max_retries {
145                return outcome.map_err(|e| transport_error(&e));
146            }
147            attempt += 1;
148            let wait = outcome
149                .as_ref()
150                .ok()
151                .and_then(retry_after)
152                .map_or(backoff(attempt), |ra| ra.max(backoff(attempt)));
153            tokio::time::sleep(wait).await;
154        }
155    }
156
157    async fn read_all(&self, resp: reqwest::Response) -> Result<Bytes> {
158        let limit = self.cfg.max_response_bytes;
159        if resp.content_length().is_some_and(|n| n > limit) {
160            return Err(Error::TooLarge { limit });
161        }
162        let mut body = Chunks { resp, left: limit, limit };
163        // Most responses arrive in one chunk: return it without copying.
164        let Some(first) = body.next().await? else { return Ok(Bytes::new()) };
165        let Some(second) = body.next().await? else { return Ok(first) };
166        let mut buf = bytes::BytesMut::with_capacity(first.len() + second.len());
167        buf.extend_from_slice(&first);
168        buf.extend_from_slice(&second);
169        while let Some(chunk) = body.next().await? {
170            buf.extend_from_slice(&chunk);
171        }
172        Ok(buf.freeze())
173    }
174}
175
176/// The response's `Retry-After` (either form), capped.
177fn retry_after(resp: &reqwest::Response) -> Option<Duration> {
178    let v = resp.headers().get("retry-after")?.to_str().ok()?;
179    Some(crate::net::retry_after(v, graph_r::fsio::now_ms())?.min(MAX_RETRY_AFTER))
180}
181
182/// A bounded chunked body. Dropping it aborts the transfer.
183#[derive(Debug)]
184pub struct Chunks {
185    resp: reqwest::Response,
186    left: u64,
187    limit: u64,
188}
189
190impl Chunks {
191    async fn next(&mut self) -> Result<Option<Bytes>> {
192        let chunk = self.resp.chunk().await.map_err(|e| transport_error(&e))?;
193        if let Some(c) = &chunk {
194            let n = c.len() as u64;
195            if n > self.left {
196                return Err(Error::TooLarge { limit: self.limit });
197            }
198            self.left -= n;
199        }
200        Ok(chunk)
201    }
202}
203
204impl BodyStream for Chunks {
205    type Next<'a>
206        = impl Future<Output = Result<Option<Bytes>>> + 'a
207    where
208        Self: 'a;
209
210    fn next_chunk(&mut self) -> Self::Next<'_> {
211        self.next()
212    }
213}
214
215impl<T: TokenProvider> Transport for Client<T> {
216    type Send<'a>
217        = impl Future<Output = Result<Response>> + 'a
218    where
219        Self: 'a;
220    type Body<'a>
221        = Chunks
222    where
223        Self: 'a;
224    type Open<'a>
225        = impl Future<Output = Result<(u16, Chunks)>> + 'a
226    where
227        Self: 'a;
228
229    fn send<'a>(&'a self, req: Request<'a>) -> Self::Send<'a> {
230        async move {
231            let resp = self.exchange(&req, self.cfg.timeout).await?;
232            let status = resp.status().as_u16();
233            let body = self.read_all(resp).await?;
234            Ok(Response { status, body })
235        }
236    }
237
238    fn open<'a>(&'a self, req: Request<'a>) -> Self::Open<'a> {
239        async move {
240            let resp = self.exchange(&req, self.cfg.stream_timeout).await?;
241            let limit = self.cfg.max_response_bytes;
242            if resp.content_length().is_some_and(|n| n > limit) {
243                return Err(Error::TooLarge { limit });
244            }
245            Ok((resp.status().as_u16(), Chunks { resp, left: limit, limit }))
246        }
247    }
248}
249
250/// Run `fut` to completion on a single-threaded runtime driving this
251/// transport's I/O and timers — for command-line tools with a synchronous
252/// `main`. Libraries should stay runtime-agnostic and leave this to the
253/// binary. Fails only if the runtime cannot start.
254#[cfg(feature = "blocking")]
255pub fn block_on<F: std::future::Future>(fut: F) -> std::io::Result<F::Output> {
256    let rt = tokio::runtime::Builder::new_current_thread().enable_all().build()?;
257    Ok(rt.block_on(fut))
258}