Commit
+59 -7 +/-2 browse
1 | diff --git a/src/common/resolver.rs b/src/common/resolver.rs |
2 | index 246de77..85d566c 100644 |
3 | --- a/src/common/resolver.rs |
4 | +++ b/src/common/resolver.rs |
5 | @@ -27,7 +27,7 @@ use crate::{ |
6 | dmarc::Dmarc, |
7 | mta_sts::{MtaSts, TlsRpt}, |
8 | spf::{Macro, Spf}, |
9 | - Error, Resolver, Txt, MX, |
10 | + Error, IpLookupStrategy, Resolver, Txt, MX, |
11 | }; |
12 | |
13 | use super::{ |
14 | @@ -233,12 +233,48 @@ impl Resolver { |
15 | .insert(key.into_owned(), Arc::new(ips), ipv6_lookup.valid_until())) |
16 | } |
17 | |
18 | - pub async fn ip_lookup(&self, key: &str) -> crate::Result<impl Iterator<Item = IpAddr>> { |
19 | - self.resolver |
20 | - .lookup_ip(key) |
21 | - .await |
22 | - .map(|lookup| lookup.into_iter()) |
23 | - .map_err(Error::from) |
24 | + pub async fn ip_lookup( |
25 | + &self, |
26 | + key: &str, |
27 | + mut strategy: IpLookupStrategy, |
28 | + max_results: usize, |
29 | + ) -> crate::Result<Vec<IpAddr>> { |
30 | + loop { |
31 | + match strategy { |
32 | + IpLookupStrategy::Ipv4Only | IpLookupStrategy::Ipv4thenIpv6 => { |
33 | + match (self.ipv4_lookup(key).await, strategy) { |
34 | + (Ok(result), _) => { |
35 | + return Ok(result |
36 | + .iter() |
37 | + .take(max_results) |
38 | + .copied() |
39 | + .map(IpAddr::from) |
40 | + .collect()) |
41 | + } |
42 | + (Err(err), IpLookupStrategy::Ipv4Only) => return Err(err), |
43 | + _ => { |
44 | + strategy = IpLookupStrategy::Ipv6Only; |
45 | + } |
46 | + } |
47 | + } |
48 | + IpLookupStrategy::Ipv6Only | IpLookupStrategy::Ipv6thenIpv4 => { |
49 | + match (self.ipv6_lookup(key).await, strategy) { |
50 | + (Ok(result), _) => { |
51 | + return Ok(result |
52 | + .iter() |
53 | + .take(max_results) |
54 | + .copied() |
55 | + .map(IpAddr::from) |
56 | + .collect()) |
57 | + } |
58 | + (Err(err), IpLookupStrategy::Ipv6Only) => return Err(err), |
59 | + _ => { |
60 | + strategy = IpLookupStrategy::Ipv4Only; |
61 | + } |
62 | + } |
63 | + } |
64 | + } |
65 | + } |
66 | } |
67 | |
68 | pub async fn ptr_lookup<'x>(&self, addr: IpAddr) -> crate::Result<Arc<Vec<String>>> { |
69 | diff --git a/src/lib.rs b/src/lib.rs |
70 | index 5ef3194..866959e 100644 |
71 | --- a/src/lib.rs |
72 | +++ b/src/lib.rs |
73 | @@ -259,6 +259,7 @@ |
74 | |
75 | use std::{ |
76 | cell::Cell, |
77 | + default, |
78 | fmt::Display, |
79 | io, |
80 | net::{IpAddr, Ipv4Addr, Ipv6Addr}, |
81 | @@ -297,6 +298,21 @@ pub struct Resolver { |
82 | pub(crate) cache_ptr: LruCache<IpAddr, Arc<Vec<String>>>, |
83 | } |
84 | |
85 | + #[derive(Debug, Clone, Copy, Default)] |
86 | + pub enum IpLookupStrategy { |
87 | + /// Only query for A (Ipv4) records |
88 | + Ipv4Only, |
89 | + /// Only query for AAAA (Ipv6) records |
90 | + Ipv6Only, |
91 | + /// Query for A and AAAA in parallel |
92 | + //Ipv4AndIpv6, |
93 | + /// Query for Ipv6 if that fails, query for Ipv4 |
94 | + Ipv6thenIpv4, |
95 | + /// Query for Ipv4 if that fails, query for Ipv6 (default) |
96 | + #[default] |
97 | + Ipv4thenIpv6, |
98 | + } |
99 | + |
100 | #[derive(Clone)] |
101 | pub enum Txt { |
102 | Spf(Arc<Spf>), |