1 | use hyper::client::connect::dns::Name; |
2 | use hyper::service::Service; |
3 | |
4 | use std::collections::HashMap; |
5 | use std::future::Future; |
6 | use std::net::SocketAddr; |
7 | use std::pin::Pin; |
8 | use std::sync::Arc; |
9 | use std::task::{Context, Poll}; |
10 | |
11 | use crate::error::BoxError; |
12 | |
13 | /// Alias for an `Iterator` trait object over `SocketAddr`. |
14 | pub type Addrs = Box<dyn Iterator<Item = SocketAddr> + Send>; |
15 | |
16 | /// Alias for the `Future` type returned by a DNS resolver. |
17 | pub type Resolving = Pin<Box<dyn Future<Output = Result<Addrs, BoxError>> + Send>>; |
18 | |
19 | /// Trait for customizing DNS resolution in reqwest. |
20 | pub trait Resolve: Send + Sync { |
21 | /// Performs DNS resolution on a `Name`. |
22 | /// The return type is a future containing an iterator of `SocketAddr`. |
23 | /// |
24 | /// It differs from `tower_service::Service<Name>` in several ways: |
25 | /// * It is assumed that `resolve` will always be ready to poll. |
26 | /// * It does not need a mutable reference to `self`. |
27 | /// * Since trait objects cannot make use of associated types, it requires |
28 | /// wrapping the returned `Future` and its contained `Iterator` with `Box`. |
29 | fn resolve(&self, name: Name) -> Resolving; |
30 | } |
31 | |
32 | #[derive (Clone)] |
33 | pub(crate) struct DynResolver { |
34 | resolver: Arc<dyn Resolve>, |
35 | } |
36 | |
37 | impl DynResolver { |
38 | pub(crate) fn new(resolver: Arc<dyn Resolve>) -> Self { |
39 | Self { resolver } |
40 | } |
41 | } |
42 | |
43 | impl Service<Name> for DynResolver { |
44 | type Response = Addrs; |
45 | type Error = BoxError; |
46 | type Future = Resolving; |
47 | |
48 | fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { |
49 | Poll::Ready(Ok(())) |
50 | } |
51 | |
52 | fn call(&mut self, name: Name) -> Self::Future { |
53 | self.resolver.resolve(name) |
54 | } |
55 | } |
56 | |
57 | pub(crate) struct DnsResolverWithOverrides { |
58 | dns_resolver: Arc<dyn Resolve>, |
59 | overrides: Arc<HashMap<String, Vec<SocketAddr>>>, |
60 | } |
61 | |
62 | impl DnsResolverWithOverrides { |
63 | pub(crate) fn new( |
64 | dns_resolver: Arc<dyn Resolve>, |
65 | overrides: HashMap<String, Vec<SocketAddr>>, |
66 | ) -> Self { |
67 | DnsResolverWithOverrides { |
68 | dns_resolver, |
69 | overrides: Arc::new(data:overrides), |
70 | } |
71 | } |
72 | } |
73 | |
74 | impl Resolve for DnsResolverWithOverrides { |
75 | fn resolve(&self, name: Name) -> Resolving { |
76 | match self.overrides.get(name.as_str()) { |
77 | Some(dest: &Vec) => { |
78 | let addrs: Addrs = Box::new(dest.clone().into_iter()); |
79 | Box::pin(futures_util::future::ready(Ok(addrs))) |
80 | } |
81 | None => self.dns_resolver.resolve(name), |
82 | } |
83 | } |
84 | } |
85 | |