use crate::auth::{AuthProvider, Credential};
use crate::error::{Error, Result};
use crate::fetch::{FetchMeta, FetchOptions, Fetched, Fetcher};
use crate::payload::DocPayload;
use crate::resource::{Resource, ResourceKind};
use bytes::BytesMut;
use compact_str::CompactString;
use futures::StreamExt;
use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, ETAG, IF_NONE_MATCH, USER_AGENT};
use std::time::{Duration, Instant};
pub(crate) const DEFAULT_UA: &str = concat!("link-r/", env!("CARGO_PKG_VERSION"));
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_REDIRECTS: usize = 5;
#[derive(Debug)]
pub struct HttpFetcher<A: AuthProvider> {
client: reqwest::Client,
auth: A,
user_agent: CompactString,
}
pub(crate) fn default_client() -> Result<reqwest::Client> {
reqwest::Client::builder()
.user_agent(DEFAULT_UA)
.connect_timeout(CONNECT_TIMEOUT)
.timeout(REQUEST_TIMEOUT)
.redirect(reqwest::redirect::Policy::limited(MAX_REDIRECTS))
.build()
.map_err(|e| Error::backend("reqwest", e.to_string()))
}
impl<A: AuthProvider> HttpFetcher<A> {
pub fn new(auth: A) -> Result<Self> {
let client = default_client()?;
Ok(Self {
client,
auth,
user_agent: CompactString::from(DEFAULT_UA),
})
}
#[must_use]
pub fn with_client(client: reqwest::Client, auth: A) -> Self {
Self {
client,
auth,
user_agent: CompactString::from(DEFAULT_UA),
}
}
#[must_use]
pub fn user_agent(mut self, ua: impl Into<CompactString>) -> Self {
self.user_agent = ua.into();
self
}
}
fn map_reqwest(e: &reqwest::Error, started: Instant) -> Error {
if e.is_timeout() {
Error::Timeout {
duration_ms: started.elapsed().as_millis() as u64,
}
} else {
Error::backend("reqwest", e.to_string())
}
}
trait StreamErr {
fn into_error(self, started: Instant) -> Error;
}
impl StreamErr for reqwest::Error {
fn into_error(self, started: Instant) -> Error {
map_reqwest(&self, started)
}
}
#[cfg(test)]
impl StreamErr for std::convert::Infallible {
fn into_error(self, _started: Instant) -> Error {
match self {}
}
}
async fn collect_capped<S, E>(
mut stream: S,
hint: Option<u64>,
max: Option<u64>,
status: u16,
started: Instant,
) -> Result<bytes::Bytes>
where
S: futures::Stream<Item = std::result::Result<bytes::Bytes, E>> + Unpin,
E: StreamErr,
{
let cap = match (hint, max) {
(Some(h), Some(m)) => h.min(m),
(Some(h), None) => h,
(None, _) => 0,
};
let cap = usize::try_from(cap).unwrap_or(0).min(1 << 20);
let mut body = BytesMut::with_capacity(cap);
while let Some(chunk) = stream.next().await {
let chunk = chunk.map_err(|e| e.into_error(started))?;
if let Some(max) = max {
if body.len() as u64 + chunk.len() as u64 > max {
return Err(Error::http(status, format!("body exceeds max {max} bytes")));
}
}
body.extend_from_slice(&chunk);
}
Ok(body.freeze())
}
impl<A: AuthProvider> Fetcher for HttpFetcher<A> {
type FetchFuture<'a>
= impl std::future::Future<Output = Result<Fetched<'a>>> + Send + 'a
where
Self: 'a;
fn fetch<'a>(
&'a self,
resource: &'a Resource,
opts: FetchOptions<'a>,
) -> Self::FetchFuture<'a> {
async move {
let started = Instant::now();
if self.auth.needs_refresh() {
self.auth.refresh().await?;
}
let ua = opts.user_agent.unwrap_or(self.user_agent.as_str());
let mut req = self.client.get(resource.url.clone()).header(USER_AGENT, ua);
match self.auth.credential(resource) {
Credential::Bearer(token) => {
req = req.header(AUTHORIZATION, format!("Bearer {token}"));
}
Credential::Header { name, value } => {
req = req.header(name, value.as_ref());
}
Credential::Anonymous | Credential::Presigned => {}
}
if let Some(etag) = opts.if_none_match {
req = req.header(IF_NONE_MATCH, etag);
}
let resp = req.send().await.map_err(|e| map_reqwest(&e, started))?;
let status = resp.status().as_u16();
match status {
304 => return Err(Error::not_modified(resource.url.as_str())),
401 => return Err(Error::unauthenticated(resource.url.as_str())),
403 => return Err(Error::permission_denied(resource.url.as_str())),
429 => return Err(rate_limited(&resp)),
s if s >= 400 => {
return Err(Error::http(s, resource.url.as_str()));
}
_ => {}
}
let final_url = (resp.url() != &resource.url).then(|| resp.url().clone());
let classify_path = resp.url().path().to_owned();
let kind = resp
.headers()
.get(CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.map_or(ResourceKind::Unknown, ResourceKind::from_content_type)
.refine_with_path(&classify_path);
let etag = resp
.headers()
.get(ETAG)
.and_then(|v| v.to_str().ok())
.map(CompactString::from);
let hint = resp.content_length();
if let (Some(max), Some(len)) = (opts.max_bytes, hint) {
if len > max {
return Err(Error::http(status, format!("body {len} exceeds max {max}")));
}
}
let bytes =
collect_capped(resp.bytes_stream(), hint, opts.max_bytes, status, started).await?;
Ok(Fetched {
meta: FetchMeta {
kind,
etag,
status,
final_url,
},
payload: DocPayload::Owned(bytes),
})
}
}
}
const RETRY_AFTER_MIN_MS: u64 = 1_000;
const RETRY_AFTER_MAX_MS: u64 = 300_000;
fn rate_limited(resp: &reqwest::Response) -> Error {
let now_ms = graph_r::fsio::now_ms();
let header = |name: &str| resp.headers().get(name).and_then(|v| v.to_str().ok());
let retry_after_ms = header("retry-after")
.and_then(|v| objstore_r::net::retry_after(v, now_ms))
.map(|d| u64::try_from(d.as_millis()).unwrap_or(u64::MAX))
.or_else(|| {
header("x-ratelimit-reset")
.and_then(|s| s.trim().parse::<u64>().ok())
.map(|epoch_s| epoch_s.saturating_mul(1000).saturating_sub(now_ms))
})
.unwrap_or(RETRY_AFTER_MIN_MS)
.clamp(RETRY_AFTER_MIN_MS, RETRY_AFTER_MAX_MS);
Error::rate_limited(retry_after_ms)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::AnonymousAuth;
#[test]
fn constructs_with_anonymous_auth() {
let fetcher = HttpFetcher::new(AnonymousAuth).unwrap();
assert_eq!(fetcher.user_agent, CompactString::from(DEFAULT_UA));
}
#[test]
fn user_agent_override() {
let fetcher = HttpFetcher::new(AnonymousAuth)
.unwrap()
.user_agent("my-bot/1.0");
assert_eq!(fetcher.user_agent, "my-bot/1.0");
}
fn chunks(
parts: &[&[u8]],
) -> impl futures::Stream<Item = std::result::Result<bytes::Bytes, std::convert::Infallible>>
{
let owned: Vec<_> = parts
.iter()
.map(|p| Ok(bytes::Bytes::copy_from_slice(p)))
.collect();
futures::stream::iter(owned)
}
#[tokio::test]
async fn collect_capped_under_cap_passes() {
let started = Instant::now();
let out = collect_capped(
chunks(&[b"hello ", b"world"]),
None,
Some(1024),
200,
started,
)
.await
.unwrap();
assert_eq!(&out[..], b"hello world");
}
#[tokio::test]
async fn collect_capped_over_cap_errors_without_length_hint() {
let started = Instant::now();
let err = collect_capped(
chunks(&[b"aaaa", b"bbbb", b"cccc"]),
None,
Some(8),
200,
started,
)
.await
.unwrap_err();
match err {
Error::Http { status, .. } => assert_eq!(status, 200),
other => panic!("expected Http cap error, got {other:?}"),
}
}
#[tokio::test]
async fn collect_capped_no_max_reads_all() {
let started = Instant::now();
let out = collect_capped(chunks(&[b"a", b"b", b"c"]), Some(3), None, 200, started)
.await
.unwrap();
assert_eq!(&out[..], b"abc");
}
}