uu-objstore-r 0.4.1

// Author: spicyzhug@gmail.com
// Copyright: lituus-io, all rights reserved
//
//! [`Client`]: the reqwest-backed [`Transport`] (feature `http`).
//!
//! It applies the same policy as the GCS store: HTTPS only (loopback `http`
//! for emulators and tests), redirects refused, bounded bodies, bearer tokens
//! from any [`TokenProvider`] (plus the quota-project header for user
//! credentials), and jittered exponential backoff that honours `Retry-After`.
//! One client holds one connection pool; with HTTP/2 many concurrent calls
//! share a connection.
//!
//! Every request first passes the client's [`RequestPolicy`] (default
//! [`Permissive`]): [`RequestPolicy::check`] on the method and parsed URL,
//! then, for a `BigQuery` query job, [`RequestPolicy::check_query`] on its
//! SQL text, which the client isolates itself (see [`crate::policy`]). A
//! refusal is [`Error::Policy`], before a token is fetched or a byte sent.
//!
//! ```no_run
//! # async fn demo() -> objstore_r::Result<()> {
//! use objstore_r::transport::{Request, Transport};
//! use objstore_r::{http::Client, token::Adc, Bytes};
//!
//! let client = Client::new(Adc::discover()?, Default::default())?;
//! let url = "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/publishers/google/models/m:generateContent";
//! let resp = client.send(Request::post_json(url, Bytes::from_static(b"{}"))).await?.error_for_status()?;
//! # let _ = resp; Ok(()) }
//! ```

use crate::net::{backoff, is_allowed_endpoint, is_retriable_status, transport_error};
use crate::policy::{query_sql, Permissive, RequestPolicy, Url};
use crate::token::TokenProvider;
use crate::transport::{BodyStream, Method, Request, Response, Retry, Transport};
use crate::{Bytes, Error, Result};
use std::future::Future;
use std::time::Duration;

/// Longest server-requested `Retry-After` honoured.
const MAX_RETRY_AFTER: Duration = Duration::from_secs(30);

/// Settings.
///
/// No timeout covers a whole response body: a large or slow body takes as
/// long as it takes while bytes keep arriving, and fails once none arrive for
/// [`idle_timeout`](Self::idle_timeout).
#[derive(Debug, Clone)]
pub struct HttpConfig {
    /// The wait for [`Transport::send`]'s response headers, request body
    /// included (plus `body / min_throughput` for a large request body).
    pub timeout: Duration,
    /// The wait for [`Transport::open`]'s response headers: a streaming API
    /// may think for a while before its first byte.
    pub stream_timeout: Duration,
    /// Connection establishment (TCP, TLS and any proxy `CONNECT`).
    pub connect_timeout: Duration,
    /// The longest a response body may deliver no bytes, whole or streamed.
    pub idle_timeout: Duration,
    /// The slowest request-body upload rate tolerated, in bytes per second:
    /// it extends the header wait by `body / min_throughput`. Must not be zero.
    pub min_throughput: u64,
    /// Retries after the first attempt (per the request's [`Retry`] policy).
    pub max_retries: u32,
    /// Largest body accepted, whole or streamed.
    pub max_response_bytes: u64,
}

impl Default for HttpConfig {
    fn default() -> Self {
        Self {
            timeout: Duration::from_secs(60),
            stream_timeout: Duration::from_secs(300),
            connect_timeout: Duration::from_secs(10),
            idle_timeout: Duration::from_secs(60),
            min_throughput: 64 << 10,
            max_retries: 4,
            max_response_bytes: 64 << 20,
        }
    }
}

/// An authorised HTTPS client under request policy `P`. See the module docs.
#[derive(Debug)]
pub struct Client<T: TokenProvider, P: RequestPolicy = Permissive> {
    http: reqwest::Client,
    tokens: T,
    policy: P,
    cfg: HttpConfig,
}

impl<T: TokenProvider> Client<T> {
    /// A client authenticating with `tokens`, under [`Permissive`].
    pub fn new(tokens: T, cfg: HttpConfig) -> Result<Self> {
        Client::with_policy(tokens, Permissive, cfg)
    }
}

impl<T: TokenProvider, P: RequestPolicy> Client<T, P> {
    /// A client authenticating with `tokens` that sends only what `policy`
    /// allows.
    pub fn with_policy(tokens: T, policy: P, cfg: HttpConfig) -> Result<Self> {
        if cfg.min_throughput == 0
            || [cfg.timeout, cfg.stream_timeout, cfg.connect_timeout, cfg.idle_timeout]
                .contains(&Duration::ZERO)
        {
            return Err(Error::protocol(
                "invalid configuration: timeouts and min_throughput must be positive",
            ));
        }
        // No client-wide timeout and no reqwest `read_timeout` (its timer is
        // not reset by request-body progress): see `exchange` and `Chunks`.
        let http = reqwest::Client::builder()
            .connect_timeout(cfg.connect_timeout)
            .redirect(reqwest::redirect::Policy::none())
            .pool_idle_timeout(Duration::from_secs(90))
            .user_agent(concat!("objstore-r/", env!("CARGO_PKG_VERSION")))
            .build()
            .map_err(|e| Error::unavailable(format!("http client: {e}"), false))?;
        Ok(Self { http, tokens, policy, cfg })
    }

    /// The request policy.
    pub fn policy(&self) -> &P {
        &self.policy
    }

    /// The token provider.
    pub fn tokens(&self) -> &T {
        &self.tokens
    }

    /// The settings.
    pub fn config(&self) -> &HttpConfig {
        &self.cfg
    }

    /// The transport's own rules, then the policy's; returns the URL as it
    /// will be sent.
    fn check(&self, req: &Request<'_>) -> Result<Url> {
        if !is_allowed_endpoint(req.url) {
            return Err(Error::protocol("the URL must be https:// (http:// only to loopback)"));
        }
        for (k, _) in req.headers {
            if k.eq_ignore_ascii_case("authorization") {
                return Err(Error::protocol("authorization is set by the transport"));
            }
            // A method override would make the checked method a fiction.
            let k = k.to_ascii_lowercase();
            if k.contains("method-override") || k == "x-http-method" {
                return Err(Error::protocol("method overrides are not sent"));
            }
        }
        let url = Url::parse(req.url).map_err(|_| Error::protocol("unparseable URL"))?;
        self.policy.check(req.method, &url)?;
        if let Some(sql) = query_sql(req.method, &url, &req.body)? {
            self.policy.check_query(&sql)?;
        }
        Ok(url)
    }

    async fn build(&self, req: &Request<'_>, url: &Url) -> Result<reqwest::RequestBuilder> {
        let method = match req.method {
            Method::Get => reqwest::Method::GET,
            Method::Post => reqwest::Method::POST,
            Method::Put => reqwest::Method::PUT,
            Method::Delete => reqwest::Method::DELETE,
            Method::Patch => reqwest::Method::PATCH,
        };
        let mut rb = self.http.request(method, url.clone());
        let t = self.tokens.token().await?;
        if !t.secret().is_empty() {
            rb = rb.bearer_auth(t.secret());
        }
        if let Some(q) = self.tokens.quota_project() {
            rb = rb.header("x-goog-user-project", q);
        }
        if let Some(ct) = req.content_type {
            rb = rb.header("content-type", ct);
        }
        for &(k, v) in req.headers {
            rb = rb.header(k, v);
        }
        if !req.body.is_empty() || matches!(req.method, Method::Post | Method::Put | Method::Patch)
        {
            rb = rb.body(req.body.clone());
        }
        Ok(rb)
    }

    /// Send with the request's retry policy; returns once headers arrive,
    /// each attempt within `wait` plus the request body's time at
    /// `min_throughput`.
    async fn exchange(&self, req: &Request<'_>, wait: Duration) -> Result<reqwest::Response> {
        let url = self.check(req)?;
        let ms = (req.body.len() as u128 * 1000) / u128::from(self.cfg.min_throughput);
        let wait =
            wait.saturating_add(Duration::from_millis(u64::try_from(ms).unwrap_or(u64::MAX)));
        let mut attempt = 0;
        loop {
            let rb = self.build(req, &url).await?;
            // The error, and whether the connection itself failed.
            let outcome = match tokio::time::timeout(wait, rb.send()).await {
                Ok(Ok(r)) => Ok(r),
                Ok(Err(e)) => Err((transport_error(&e), e.is_connect())),
                Err(_) => Err((Error::unavailable("timed out", true), false)),
            };
            let retry = match (&outcome, req.retry) {
                (_, Retry::Never) => false,
                (Ok(r), Retry::Transient) => is_retriable_status(r.status().as_u16()),
                (Ok(r), Retry::Throttled) => r.status().as_u16() == 429,
                (Err((e, _)), Retry::Transient) => e.is_retriable(),
                (Err((_, connect)), Retry::Throttled) => *connect,
            };
            if !retry || attempt >= self.cfg.max_retries {
                return outcome.map_err(|(e, _)| e);
            }
            attempt += 1;
            let wait = outcome
                .as_ref()
                .ok()
                .and_then(retry_after)
                .map_or(backoff(attempt), |ra| ra.max(backoff(attempt)));
            tokio::time::sleep(wait).await;
        }
    }

    async fn read_all(&self, resp: reqwest::Response) -> Result<Bytes> {
        let limit = self.cfg.max_response_bytes;
        if resp.content_length().is_some_and(|n| n > limit) {
            return Err(Error::TooLarge { limit });
        }
        let mut body = Chunks { resp, left: limit, limit, idle: self.cfg.idle_timeout };
        // Most responses arrive in one chunk: return it without copying.
        let Some(first) = body.next().await? else { return Ok(Bytes::new()) };
        let Some(second) = body.next().await? else { return Ok(first) };
        let mut buf = bytes::BytesMut::with_capacity(first.len() + second.len());
        buf.extend_from_slice(&first);
        buf.extend_from_slice(&second);
        while let Some(chunk) = body.next().await? {
            buf.extend_from_slice(&chunk);
        }
        Ok(buf.freeze())
    }
}

/// The response's `Retry-After` (either form), capped.
fn retry_after(resp: &reqwest::Response) -> Option<Duration> {
    let v = resp.headers().get("retry-after")?.to_str().ok()?;
    Some(crate::net::retry_after(v, graph_r::fsio::now_ms())?.min(MAX_RETRY_AFTER))
}

/// A bounded chunked body: each chunk must arrive within the idle timeout.
/// Dropping it aborts the transfer.
#[derive(Debug)]
pub struct Chunks {
    resp: reqwest::Response,
    left: u64,
    limit: u64,
    idle: Duration,
}

impl Chunks {
    async fn next(&mut self) -> Result<Option<Bytes>> {
        let chunk = tokio::time::timeout(self.idle, self.resp.chunk())
            .await
            .map_err(|_| Error::unavailable("stalled: no data within the idle timeout", true))?
            .map_err(|e| transport_error(&e))?;
        if let Some(c) = &chunk {
            let n = c.len() as u64;
            if n > self.left {
                return Err(Error::TooLarge { limit: self.limit });
            }
            self.left -= n;
        }
        Ok(chunk)
    }
}

impl BodyStream for Chunks {
    type Next<'a>
        = impl Future<Output = Result<Option<Bytes>>> + 'a
    where
        Self: 'a;

    fn next_chunk(&mut self) -> Self::Next<'_> {
        self.next()
    }
}

impl<T: TokenProvider, P: RequestPolicy> Transport for Client<T, P> {
    type Send<'a>
        = impl Future<Output = Result<Response>> + 'a
    where
        Self: 'a;
    type Body<'a>
        = Chunks
    where
        Self: 'a;
    type Open<'a>
        = impl Future<Output = Result<(u16, Chunks)>> + 'a
    where
        Self: 'a;

    fn send<'a>(&'a self, req: Request<'a>) -> Self::Send<'a> {
        async move {
            let resp = self.exchange(&req, self.cfg.timeout).await?;
            let status = resp.status().as_u16();
            let body = self.read_all(resp).await?;
            Ok(Response { status, body })
        }
    }

    fn open<'a>(&'a self, req: Request<'a>) -> Self::Open<'a> {
        async move {
            let resp = self.exchange(&req, self.cfg.stream_timeout).await?;
            let limit = self.cfg.max_response_bytes;
            if resp.content_length().is_some_and(|n| n > limit) {
                return Err(Error::TooLarge { limit });
            }
            let idle = self.cfg.idle_timeout;
            Ok((resp.status().as_u16(), Chunks { resp, left: limit, limit, idle }))
        }
    }
}

/// Run `fut` to completion on a single-threaded runtime driving this
/// transport's I/O and timers — for command-line tools with a synchronous
/// `main`. Libraries should stay runtime-agnostic and leave this to the
/// binary. Fails only if the runtime cannot start.
#[cfg(feature = "blocking")]
pub fn block_on<F: std::future::Future>(fut: F) -> std::io::Result<F::Output> {
    let rt = tokio::runtime::Builder::new_current_thread().enable_all().build()?;
    Ok(rt.block_on(fut))
}