use std::collections::HashMap;
use std::fmt;
use std::sync::{Arc, Mutex};
use spire_core::identity::{Identity, WeakIdentity};
use spire_core::{Error, Result};
pub(crate) type Builder = dyn Fn() -> reqwest::ClientBuilder + Send + Sync;
#[derive(Clone)]
pub(crate) struct Clients {
builder: Arc<Builder>,
direct: reqwest::Client,
identities: Arc<Mutex<HashMap<Box<str>, Built>>>,
}
struct Built {
identity: WeakIdentity,
client: reqwest::Client,
}
impl Default for Clients {
fn default() -> Self {
Self {
builder: Arc::new(reqwest::Client::builder),
direct: reqwest::Client::default(),
identities: Arc::default(),
}
}
}
impl Clients {
pub(crate) fn new(builder: Arc<Builder>) -> Result<Self> {
let direct = builder().build().map_err(|error| {
Error::backend("failed to build the HTTP client").with_source(error)
})?;
Ok(Self {
builder,
direct,
identities: Arc::default(),
})
}
pub(crate) const fn direct(&self) -> &reqwest::Client {
&self.direct
}
pub(crate) fn of(&self, identity: &Identity) -> Result<reqwest::Client> {
let mut identities = self
.identities
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(built) = identities.get(identity.name())
&& built.identity.refers_to(identity)
{
return Ok(built.client.clone());
}
let client = self.build_for(identity)?;
identities.retain(|_, built| built.identity.is_alive());
let built = Built {
identity: identity.downgrade(),
client: client.clone(),
};
identities.insert(identity.name().into(), built);
Ok(client)
}
fn build_for(&self, identity: &Identity) -> Result<reqwest::Client> {
let mut builder = (self.builder)();
if let Some(proxy) = identity.proxy() {
let socks = matches!(proxy.scheme_str(), Some("socks5" | "socks5h"));
if socks && !cfg!(feature = "socks") {
let message = format!(
"identity `{}` has a SOCKS proxy, which needs the `socks` feature",
identity.name()
);
return Err(Error::invalid(message));
}
let proxy = reqwest::Proxy::all(proxy.to_string()).map_err(|error| {
let message = format!("unsupported proxy for identity `{}`", identity.name());
Error::invalid(message).with_source(error)
})?;
builder = builder.proxy(proxy);
}
#[cfg(feature = "cookies")]
let builder = builder.cookie_store(true);
builder.build().map_err(|error| {
let message = format!(
"failed to build the client of identity `{}`",
identity.name()
);
Error::backend(message).with_source(error)
})
}
#[cfg(test)]
pub(crate) fn len(&self) -> usize {
self.identities
.lock()
.map_or(0, |identities| identities.len())
}
}
impl fmt::Debug for Clients {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let identities = self
.identities
.lock()
.map_or(0, |identities| identities.len());
f.debug_struct("Clients")
.field("identities", &identities)
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn clients() -> Clients {
Clients::default()
}
#[test]
fn clients_are_kept_while_their_identity_is_in_use() {
let clients = clients();
let a = Identity::new("a");
clients.of(&a).unwrap();
clients.of(&a).unwrap();
assert_eq!(clients.len(), 1);
drop(a);
clients.of(&Identity::new("b")).unwrap();
assert_eq!(clients.len(), 1);
}
#[test]
fn a_new_identity_with_an_old_name_gets_a_new_client() {
let clients = clients();
let old = Identity::new("a");
clients.of(&old).unwrap();
let new = Identity::new("a")
.with_proxy("http://127.0.0.1:3128")
.unwrap();
clients.of(&new).unwrap();
assert_eq!(clients.len(), 1);
}
}