use hyper_util::client::legacy::connect::dns::Name as HyperName;
use tower_service::Service;
use std::collections::HashMap;
use std::future::Future;
use std::net::SocketAddr;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::Arc;
use std::task::{Context, Poll};
use crate::error::BoxError;
/// Alias for an `Iterator` trait object over `SocketAddr`.
pub type Addrs = Box<dyn Iterator<Item = SocketAddr> + Send>;
/// Alias for the `Future` type returned by a DNS resolver.
pub type Resolving = Pin<Box<dyn Future<Output = Result<Addrs, BoxError>> + Send>>;
/// Trait for customizing DNS resolution in reqwest.
pub trait Resolve: Send + Sync {
/// Performs DNS resolution on a `Name`.
/// The return type is a future containing an iterator of `SocketAddr`.
///
/// It differs from `tower_service::Service<Name>` in several ways:
/// * It is assumed that `resolve` will always be ready to poll.
/// * It does not need a mutable reference to `self`.
/// * Since trait objects cannot make use of associated types, it requires
/// wrapping the returned `Future` and its contained `Iterator` with `Box`.
///
/// Explicitly specified port in the URL will override any port in the resolved `SocketAddr`s.
/// Otherwise, port `0` will be replaced by the conventional port for the given scheme (e.g. 80 for http).
fn resolve(&self, name: Name) -> Resolving;
}
/// A name that must be resolved to addresses.
#[derive(Debug)]
pub struct Name(pub(super) HyperName);
impl Name {
/// View the name as a string.
pub fn as_str(&self) -> &str {
self.0.as_str()
}
}
impl FromStr for Name {
type Err = sealed::InvalidNameError;
fn from_str(host: &str) -> Result<Self, Self::Err> {
HyperName::from_str(host)
.map(Name)
.map_err(|_| sealed::InvalidNameError { _ext: () })
}
}
#[derive(Clone)]
pub(crate) struct DynResolver {
resolver: Arc<dyn Resolve>,
}
impl DynResolver {
pub(crate) fn new(resolver: Arc<dyn Resolve>) -> Self {
Self { resolver }
}
}
impl Service<HyperName> for DynResolver {
type Response = Addrs;
type Error = BoxError;
type Future = Resolving;
fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, name: HyperName) -> Self::Future {
self.resolver.resolve(Name(name))
}
}
pub(crate) struct DnsResolverWithOverrides {
dns_resolver: Arc<dyn Resolve>,
overrides: Arc<HashMap<String, Vec<SocketAddr>>>,
}
impl DnsResolverWithOverrides {
pub(crate) fn new(
dns_resolver: Arc<dyn Resolve>,
overrides: HashMap<String, Vec<SocketAddr>>,
) -> Self {
DnsResolverWithOverrides {
dns_resolver,
overrides: Arc::new(overrides),
}
}
}
impl Resolve for DnsResolverWithOverrides {
fn resolve(&self, name: Name) -> Resolving {
match self.overrides.get(name.as_str()) {
Some(dest) => {
let addrs: Addrs = Box::new(dest.clone().into_iter());
Box::pin(futures_util::future::ready(Ok(addrs)))
}
None => self.dns_resolver.resolve(name),
}
}
}
mod sealed {
use std::fmt;
#[derive(Debug)]
pub struct InvalidNameError {
pub(super) _ext: (),
}
impl fmt::Display for InvalidNameError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("invalid DNS name")
}
}
impl std::error::Error for InvalidNameError {}
}