use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use arc_swap::ArcSwap;
use http::Extensions;
use reqwest::ClientBuilder;
use reqwest_middleware::ClientWithMiddleware;
use reqwest_retry::policies::ExponentialBackoff;
use reqwest_retry::RetryTransientMiddleware;
use reqwest_retry::{Jitter, RetryDecision, RetryPolicy};
use thiserror::Error;
use super::retry_after::{anchor_budget, RetryAfterMiddleware};
const DEFAULT_RETRY_AFTER_CAPACITY: usize = 256;
const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(30);
const DEFAULT_MAX_TOTAL_RETRY_DURATION: Duration = Duration::from_secs(10);
const DEFAULT_RETRY_AFTER_CEILING: Duration = Duration::from_secs(300);
const DEFAULT_STREAM_TOTAL_BYTE_CAP: u64 = 1024 * 1024 * 1024;
const DEFAULT_MAX_RETRIES: u32 = 3;
const RETRY_BACKOFF_MIN_INTERVAL: Duration = Duration::from_millis(100);
const RETRY_BACKOFF_MAX_INTERVAL: Duration = Duration::from_secs(2);
#[derive(Debug, Clone)]
pub struct ClientCertConfig {
pub cert_pem: PathBuf,
pub key_pem: PathBuf,
}
#[derive(Debug, Clone)]
pub struct RetryConfig {
pub max_retries: u32,
pub initial_backoff: Duration,
pub max_retry_interval: Duration,
}
impl Default for RetryConfig {
fn default() -> Self {
Self {
max_retries: DEFAULT_MAX_RETRIES,
initial_backoff: RETRY_BACKOFF_MIN_INTERVAL,
max_retry_interval: RETRY_BACKOFF_MAX_INTERVAL,
}
}
}
#[derive(Debug, Clone)]
pub struct HttpClientConfig {
pub pool_max_idle_per_host: Option<usize>,
pub request_timeout: Option<Duration>,
pub connect_timeout: Option<Duration>,
pub read_timeout: Option<Duration>,
pub stream_total_byte_cap: u64,
pub retry: RetryConfig,
pub max_total_retry_duration: Duration,
pub retry_after_ceiling: Duration,
pub ca_bundle: Option<PathBuf>,
pub client_cert: Option<ClientCertConfig>,
}
impl Default for HttpClientConfig {
fn default() -> Self {
Self {
pool_max_idle_per_host: None,
request_timeout: Some(DEFAULT_REQUEST_TIMEOUT),
connect_timeout: Some(DEFAULT_CONNECT_TIMEOUT),
read_timeout: Some(DEFAULT_READ_TIMEOUT),
stream_total_byte_cap: DEFAULT_STREAM_TOTAL_BYTE_CAP,
retry: RetryConfig::default(),
max_total_retry_duration: DEFAULT_MAX_TOTAL_RETRY_DURATION,
retry_after_ceiling: DEFAULT_RETRY_AFTER_CEILING,
ca_bundle: None,
client_cert: None,
}
}
}
#[derive(Debug, Error)]
pub enum HttpClientBuildError {
#[error("failed to read CA bundle from {path}: {source}")]
CaBundleRead {
path: PathBuf,
#[source]
source: std::io::Error,
},
#[error("failed to parse CA bundle at {path}: {source}")]
CaBundleParse {
path: PathBuf,
#[source]
source: reqwest::Error,
},
#[error("failed to read client cert from {path}: {source}")]
ClientCertRead {
path: PathBuf,
#[source]
source: std::io::Error,
},
#[error("failed to parse client cert at {path}: {source}")]
ClientCertParse {
path: PathBuf,
#[source]
source: reqwest::Error,
},
#[error("failed to build reqwest client: {0}")]
Build(reqwest::Error),
}
pub struct SharedHttpClient {
inner: ArcSwap<SharedHttpInner>,
}
#[derive(Clone)]
struct SharedHttpInner {
client: Arc<ClientWithMiddleware>,
stream_client: Arc<ClientWithMiddleware>,
config: Arc<HttpClientConfig>,
}
impl std::fmt::Debug for SharedHttpClient {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SharedHttpClient")
.field("config", &self.inner.load().config)
.finish_non_exhaustive()
}
}
impl SharedHttpClient {
pub fn new(config: HttpClientConfig) -> Result<Self, HttpClientBuildError> {
let (client, stream_client) = build_clients_sync(&config)?;
Ok(Self {
inner: ArcSwap::from_pointee(SharedHttpInner {
client: Arc::new(client),
stream_client: Arc::new(stream_client),
config: Arc::new(config),
}),
})
}
pub fn client(&self) -> Arc<ClientWithMiddleware> {
Arc::clone(&self.inner.load().client)
}
pub fn stream_client(&self) -> Arc<ClientWithMiddleware> {
Arc::clone(&self.inner.load().stream_client)
}
pub fn config(&self) -> Arc<HttpClientConfig> {
Arc::clone(&self.inner.load().config)
}
pub async fn reload(&self, config: HttpClientConfig) -> Result<(), HttpClientBuildError> {
let (client, stream_client) = build_clients(&config).await?;
self.inner.store(Arc::new(SharedHttpInner {
client: Arc::new(client),
stream_client: Arc::new(stream_client),
config: Arc::new(config),
}));
Ok(())
}
}
fn same_host_redirect_policy() -> reqwest::redirect::Policy {
reqwest::redirect::Policy::custom(|attempt| {
if attempt.previous().len() >= MAX_REDIRECT_HOPS {
return attempt.error("too many redirects");
}
let previous = match attempt.previous().last() {
Some(url) => url.clone(),
None => return attempt.follow(),
};
let next = attempt.url();
let scheme_matches = previous.scheme() == next.scheme();
let host_matches = previous.host_str() == next.host_str();
let port_matches = previous.port_or_known_default() == next.port_or_known_default();
if scheme_matches && host_matches && port_matches {
attempt.follow()
} else {
attempt.stop()
}
})
}
const MAX_REDIRECT_HOPS: usize = 10;
fn is_idempotent(method: &reqwest::Method) -> bool {
matches!(
*method,
reqwest::Method::GET
| reqwest::Method::HEAD
| reqwest::Method::PUT
| reqwest::Method::DELETE
| reqwest::Method::OPTIONS
)
}
struct RetryGateMiddleware {
retry: Arc<RetryTransientMiddleware<TotalRetryBudget<ExponentialBackoff>>>,
max_total_retry_duration: Duration,
}
impl RetryGateMiddleware {
fn new(retry: &RetryConfig, max_total_retry_duration: Duration) -> Self {
let policy = ExponentialBackoff::builder()
.retry_bounds(retry.initial_backoff, retry.max_retry_interval)
.jitter(Jitter::Bounded)
.base(2)
.build_with_max_retries(retry.max_retries);
Self {
retry: Arc::new(RetryTransientMiddleware::new_with_policy(
TotalRetryBudget {
budget: max_total_retry_duration,
inner: policy,
},
)),
max_total_retry_duration,
}
}
}
#[async_trait::async_trait]
impl reqwest_middleware::Middleware for RetryGateMiddleware {
async fn handle(
&self,
req: reqwest::Request,
extensions: &mut Extensions,
next: reqwest_middleware::Next<'_>,
) -> reqwest_middleware::Result<reqwest::Response> {
if is_idempotent(req.method()) {
anchor_budget(extensions, self.max_total_retry_duration);
self.retry.handle(req, extensions, next).await
} else {
next.run(req, extensions).await
}
}
}
struct TotalRetryBudget<P> {
budget: Duration,
inner: P,
}
impl<P: RetryPolicy> RetryPolicy for TotalRetryBudget<P> {
fn should_retry(&self, request_start_time: SystemTime, n_past_retries: u32) -> RetryDecision {
let elapsed = SystemTime::now()
.duration_since(request_start_time)
.unwrap_or_default();
if elapsed >= self.budget {
return RetryDecision::DoNotRetry;
}
match self.inner.should_retry(request_start_time, n_past_retries) {
RetryDecision::DoNotRetry => RetryDecision::DoNotRetry,
RetryDecision::Retry { execute_after } => {
let hard_stop = request_start_time
.checked_add(self.budget)
.unwrap_or(execute_after);
RetryDecision::Retry {
execute_after: execute_after.min(hard_stop),
}
}
}
}
}
async fn build_clients(
config: &HttpClientConfig,
) -> Result<(ClientWithMiddleware, ClientWithMiddleware), HttpClientBuildError> {
let pems = read_pems(config).await?;
let client = build_client_with_pems(config, pems.clone(), false)?;
let stream_client = build_client_with_pems(config, pems, true)?;
Ok((client, stream_client))
}
fn build_clients_sync(
config: &HttpClientConfig,
) -> Result<(ClientWithMiddleware, ClientWithMiddleware), HttpClientBuildError> {
let pems = read_pems_sync(config)?;
let client = build_client_with_pems(config, pems.clone(), false)?;
let stream_client = build_client_with_pems(config, pems, true)?;
Ok((client, stream_client))
}
#[derive(Clone)]
struct ClientPems {
ca: Option<Vec<u8>>,
client: Option<(Vec<u8>, Vec<u8>)>,
}
async fn read_pems(config: &HttpClientConfig) -> Result<ClientPems, HttpClientBuildError> {
let ca = match &config.ca_bundle {
Some(path) => Some(tokio::fs::read(path).await.map_err(|source| {
HttpClientBuildError::CaBundleRead {
path: path.clone(),
source,
}
})?),
None => None,
};
let client = match &config.client_cert {
Some(cfg) => {
let cert_pem = tokio::fs::read(&cfg.cert_pem).await.map_err(|source| {
HttpClientBuildError::ClientCertRead {
path: cfg.cert_pem.clone(),
source,
}
})?;
let key_pem = tokio::fs::read(&cfg.key_pem).await.map_err(|source| {
HttpClientBuildError::ClientCertRead {
path: cfg.key_pem.clone(),
source,
}
})?;
Some((cert_pem, key_pem))
}
None => None,
};
Ok(ClientPems { ca, client })
}
fn read_pems_sync(config: &HttpClientConfig) -> Result<ClientPems, HttpClientBuildError> {
let ca = match &config.ca_bundle {
Some(path) => {
Some(
std::fs::read(path).map_err(|source| HttpClientBuildError::CaBundleRead {
path: path.clone(),
source,
})?,
)
}
None => None,
};
let client = match &config.client_cert {
Some(cfg) => {
let cert_pem = std::fs::read(&cfg.cert_pem).map_err(|source| {
HttpClientBuildError::ClientCertRead {
path: cfg.cert_pem.clone(),
source,
}
})?;
let key_pem = std::fs::read(&cfg.key_pem).map_err(|source| {
HttpClientBuildError::ClientCertRead {
path: cfg.key_pem.clone(),
source,
}
})?;
Some((cert_pem, key_pem))
}
None => None,
};
Ok(ClientPems { ca, client })
}
fn build_client_with_pems(
config: &HttpClientConfig,
pems: ClientPems,
streaming: bool,
) -> Result<ClientWithMiddleware, HttpClientBuildError> {
let ClientPems {
ca: ca_pem,
client: client_pems,
} = pems;
let mut builder = ClientBuilder::new();
builder = builder.redirect(same_host_redirect_policy());
if let Some(pool_max_idle) = config.pool_max_idle_per_host {
builder = builder.pool_max_idle_per_host(pool_max_idle);
}
if !streaming {
if let Some(timeout) = config.request_timeout {
builder = builder.timeout(timeout);
}
}
if let Some(timeout) = config.connect_timeout {
builder = builder.connect_timeout(timeout);
}
if let Some(timeout) = config.read_timeout {
builder = builder.read_timeout(timeout);
}
if let Some(pem) = &ca_pem {
let certs = reqwest::Certificate::from_pem_bundle(pem).map_err(|source| {
HttpClientBuildError::CaBundleParse {
path: config
.ca_bundle
.clone()
.unwrap_or_else(|| PathBuf::from("<ca-bundle>")),
source,
}
})?;
for cert in certs {
builder = builder.add_root_certificate(cert);
}
}
if let Some((cert_pem, key_pem)) = &client_pems {
let identity = reqwest::Identity::from_pem(concat_pem(cert_pem, key_pem).as_slice())
.map_err(|source| HttpClientBuildError::ClientCertParse {
path: config
.client_cert
.as_ref()
.map(|cfg| cfg.cert_pem.clone())
.unwrap_or_else(|| PathBuf::from("<client-cert>")),
source,
})?;
builder = builder.identity(identity);
}
let reqwest_client = builder.build().map_err(HttpClientBuildError::Build)?;
let retry_after = RetryAfterMiddleware::with_capacity_ceiling_and_budget(
DEFAULT_RETRY_AFTER_CAPACITY,
config.retry_after_ceiling,
config.max_total_retry_duration,
);
let client = reqwest_middleware::ClientBuilder::new(reqwest_client)
.with(RetryGateMiddleware::new(
&config.retry,
config.max_total_retry_duration,
))
.with(retry_after)
.build();
Ok(client)
}
fn concat_pem(cert: &[u8], key: &[u8]) -> Vec<u8> {
let mut combined = Vec::with_capacity(cert.len() + key.len() + 1);
combined.extend_from_slice(cert);
if !cert.is_empty() && cert.last() != Some(&b'\n') {
combined.push(b'\n');
}
combined.extend_from_slice(key);
combined
}