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;
pub type Addrs = Box<dyn Iterator<Item = SocketAddr> + Send>;
pub type Resolving = Pin<Box<dyn Future<Output = Result<Addrs, BoxError>> + Send>>;
pub trait Resolve: Send + Sync {
fn resolve(&self, name: Name) -> Resolving;
}
#[derive(Debug)]
pub struct Name(pub(super) HyperName);
impl Name {
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.into())
.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 {}
}