pub mod origins;
pub mod reqwest_transport;
pub mod resolver;
use std::fmt;
use std::net::IpAddr;
use std::pin::Pin;
use async_trait::async_trait;
use bytes::{Bytes, BytesMut};
use futures_util::{Stream, StreamExt};
use url::Url;
pub use origins::{OriginKind, OriginSet, UrlRejection};
pub use reqwest_transport::ReqwestTransport;
#[async_trait]
pub trait Transport: Send + Sync + 'static {
async fn fetch_metadata(&self, req: MetadataRequest)
-> Result<MetadataResponse, UpstreamError>;
async fn open_artifact(&self, req: ArtifactRequest) -> Result<ArtifactBody, UpstreamError>;
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct UpstreamValidators {
pub etag: Option<String>,
pub last_modified: Option<String>,
}
impl UpstreamValidators {
pub fn is_empty(&self) -> bool {
self.etag.is_none() && self.last_modified.is_none()
}
}
pub struct MetadataRequest {
pub url: Url,
pub accept: &'static str,
pub validators: Option<UpstreamValidators>,
pub max_bytes: u64,
}
pub enum MetadataResponse {
Fresh {
body: Bytes,
validators: UpstreamValidators,
},
NotModified {
validators: UpstreamValidators,
},
Missing,
}
pub struct ArtifactRequest {
pub url: Url,
pub max_bytes: u64,
}
pub type ByteStream = Pin<Box<dyn Stream<Item = Result<Bytes, UpstreamError>> + Send>>;
pub struct ArtifactBody {
pub declared_length: Option<u64>,
pub stream: ByteStream,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum UpstreamError {
Timeout,
Transport(String),
Status(u16),
TooLarge { limit: u64 },
RejectedUrl(UrlRejection),
RejectedRedirect { to: Url },
RejectedAddress { addr: IpAddr },
}
impl fmt::Display for UpstreamError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
UpstreamError::Timeout => write!(f, "the upstream request timed out"),
UpstreamError::Transport(reason) => write!(f, "upstream transport failure: {reason}"),
UpstreamError::Status(code) => write!(f, "upstream answered {code}"),
UpstreamError::TooLarge { limit } => {
write!(f, "the upstream body exceeds the {limit}-byte cap")
}
UpstreamError::RejectedUrl(rejection) => {
write!(f, "the upstream URL was refused: {rejection}")
}
UpstreamError::RejectedRedirect { to } => {
write!(f, "refused to follow a redirect to {to}")
}
UpstreamError::RejectedAddress { addr } => {
write!(f, "refused to connect to {addr}")
}
}
}
}
impl std::error::Error for UpstreamError {}
pub fn capped(stream: ByteStream, max_bytes: u64) -> ByteStream {
Box::pin(futures_util::stream::unfold(
(stream, 0u64, false),
move |(mut stream, total, finished)| async move {
if finished {
return None;
}
match stream.next().await {
None => None,
Some(Err(err)) => Some((Err(err), (stream, total, true))),
Some(Ok(chunk)) => {
let total = total.saturating_add(chunk.len() as u64);
if total > max_bytes {
let err = UpstreamError::TooLarge { limit: max_bytes };
Some((Err(err), (stream, total, true)))
} else {
Some((Ok(chunk), (stream, total, false)))
}
}
}
},
))
}
pub async fn collect_capped(stream: ByteStream, max_bytes: u64) -> Result<Bytes, UpstreamError> {
let mut stream = capped(stream, max_bytes);
let mut buffer = BytesMut::new();
while let Some(chunk) = stream.next().await {
buffer.extend_from_slice(&chunk?);
}
Ok(buffer.freeze())
}