| use hyper::client::connect::dns::Name; |
| use hyper::service::Service; |
| |
| use std::collections::HashMap; |
| use std::future::Future; |
| use std::net::SocketAddr; |
| use std::pin::Pin; |
| 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`. |
| fn resolve(&self, name: Name) -> Resolving; |
| } |
| |
| #[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<Name> 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: Name) -> Self::Future { |
| self.resolver.resolve(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), |
| } |
| } |
| } |