1use super::{parse::TxtRecordParser, verify::DomainKey};
8use crate::{
9 Error, IpLookupStrategy, MX, MessageAuthenticator, ResolverCache, Txt,
10 dkim::{Atps, DomainKeyReport},
11 dmarc::Dmarc,
12 mta_sts::{MtaSts, TlsRpt},
13 spf::{Macro, Spf},
14};
15use hickory_resolver::{
16 Name, TokioResolver,
17 config::{ResolverConfig, ResolverOpts},
18 name_server::TokioConnectionProvider,
19 proto::{ProtoError, ProtoErrorKind},
20 system_conf::read_system_conf,
21};
22use std::{
23 borrow::Cow,
24 net::{IpAddr, Ipv4Addr, Ipv6Addr},
25 sync::Arc,
26 time::Instant,
27};
28
29pub struct DnsEntry<T> {
30 pub entry: T,
31 pub expires: Instant,
32}
33
34impl MessageAuthenticator {
35 pub fn new_cloudflare_tls() -> Result<Self, ProtoError> {
36 Self::new(ResolverConfig::cloudflare_tls(), ResolverOpts::default())
37 }
38
39 pub fn new_cloudflare() -> Result<Self, ProtoError> {
40 Self::new(ResolverConfig::cloudflare(), ResolverOpts::default())
41 }
42
43 pub fn new_google() -> Result<Self, ProtoError> {
44 Self::new(ResolverConfig::google(), ResolverOpts::default())
45 }
46
47 pub fn new_quad9() -> Result<Self, ProtoError> {
48 Self::new(ResolverConfig::quad9(), ResolverOpts::default())
49 }
50
51 pub fn new_quad9_tls() -> Result<Self, ProtoError> {
52 Self::new(ResolverConfig::quad9_tls(), ResolverOpts::default())
53 }
54
55 pub fn new_system_conf() -> Result<Self, ProtoError> {
56 let (config, options) = read_system_conf()?;
57 Self::new(config, options)
58 }
59
60 pub fn new(config: ResolverConfig, options: ResolverOpts) -> Result<Self, ProtoError> {
61 Ok(MessageAuthenticator(
62 TokioResolver::builder_with_config(config, TokioConnectionProvider::default())
63 .with_options(options)
64 .build(),
65 ))
66 }
67
68 pub fn resolver(&self) -> &TokioResolver {
69 &self.0
70 }
71
72 pub async fn txt_raw_lookup(&self, key: impl IntoFqdn<'_>) -> crate::Result<Vec<u8>> {
73 let mut result = vec![];
74 for record in self
75 .0
76 .txt_lookup(Name::from_str_relaxed(key.into_fqdn().as_ref())?)
77 .await?
78 .as_lookup()
79 .record_iter()
80 {
81 if let Some(txt_data) = record.data().as_txt() {
82 for item in txt_data.txt_data() {
83 result.extend_from_slice(item);
84 }
85 }
86 }
87
88 Ok(result)
89 }
90
91 pub async fn txt_lookup<'x, T: TxtRecordParser + Into<Txt> + UnwrapTxtRecord>(
92 &self,
93 key: impl IntoFqdn<'x>,
94 cache: Option<&impl ResolverCache<String, Txt>>,
95 ) -> crate::Result<Arc<T>> {
96 let key = key.into_fqdn();
97 if let Some(value) = cache.as_ref().and_then(|c| c.get(key.as_ref())) {
98 return T::unwrap_txt(value);
99 }
100
101 #[cfg(any(test, feature = "test"))]
102 if true {
103 return mock_resolve(key.as_ref());
104 }
105
106 let txt_lookup = self
107 .0
108 .txt_lookup(Name::from_str_relaxed(key.as_ref())?)
109 .await?;
110 let mut result = Err(Error::InvalidRecordType);
111 let records = txt_lookup.as_lookup().record_iter().filter_map(|r| {
112 let txt_data = r.data().as_txt()?.txt_data();
113 match txt_data.len() {
114 1 => Some(Cow::from(txt_data[0].as_ref())),
115 0 => None,
116 _ => {
117 let mut entry = Vec::with_capacity(255 * txt_data.len());
118 for data in txt_data {
119 entry.extend_from_slice(data);
120 }
121 Some(Cow::from(entry))
122 }
123 }
124 });
125
126 for record in records {
127 result = T::parse(record.as_ref());
128 if result.is_ok() {
129 break;
130 }
131 }
132
133 let result: Txt = result.into();
134
135 if let Some(cache) = cache {
136 cache.insert(key.into_owned(), result.clone(), txt_lookup.valid_until());
137 }
138
139 T::unwrap_txt(result)
140 }
141
142 pub async fn mx_lookup<'x>(
143 &self,
144 key: impl IntoFqdn<'x>,
145 cache: Option<&impl ResolverCache<String, Arc<Vec<MX>>>>,
146 ) -> crate::Result<Arc<Vec<MX>>> {
147 let key = key.into_fqdn();
148 if let Some(value) = cache.as_ref().and_then(|c| c.get(key.as_ref())) {
149 return Ok(value);
150 }
151
152 #[cfg(any(test, feature = "test"))]
153 if true {
154 return mock_resolve(key.as_ref());
155 }
156
157 let mx_lookup = self
158 .0
159 .mx_lookup(Name::from_str_relaxed(key.as_ref())?)
160 .await?;
161 let mx_records = mx_lookup.as_lookup().records();
162 let mut records: Vec<MX> = Vec::with_capacity(mx_records.len());
163 for mx_record in mx_records {
164 if let Some(mx) = mx_record.data().as_mx() {
165 let preference = mx.preference();
166 let exchange = mx.exchange().to_lowercase().to_string();
167
168 if let Some(record) = records.iter_mut().find(|r| r.preference == preference) {
169 record.exchanges.push(exchange);
170 } else {
171 records.push(MX {
172 exchanges: vec![exchange],
173 preference,
174 });
175 }
176 }
177 }
178
179 records.sort_unstable_by(|a, b| a.preference.cmp(&b.preference));
180 let records = Arc::new(records);
181
182 if let Some(cache) = cache {
183 cache.insert(key.into_owned(), records.clone(), mx_lookup.valid_until());
184 }
185
186 Ok(records)
187 }
188
189 pub async fn ipv4_lookup<'x>(
190 &self,
191 key: impl IntoFqdn<'x>,
192 cache: Option<&impl ResolverCache<String, Arc<Vec<Ipv4Addr>>>>,
193 ) -> crate::Result<Arc<Vec<Ipv4Addr>>> {
194 let key = key.into_fqdn();
195 if let Some(value) = cache.as_ref().and_then(|c| c.get(key.as_ref())) {
196 return Ok(value);
197 }
198
199 let ipv4_lookup = self.ipv4_lookup_raw(key.as_ref()).await?;
200
201 if let Some(cache) = cache {
202 cache.insert(
203 key.into_owned(),
204 ipv4_lookup.entry.clone(),
205 ipv4_lookup.expires,
206 );
207 }
208
209 Ok(ipv4_lookup.entry)
210 }
211
212 pub async fn ipv4_lookup_raw(&self, key: &str) -> crate::Result<DnsEntry<Arc<Vec<Ipv4Addr>>>> {
213 #[cfg(any(test, feature = "test"))]
214 if true {
215 return mock_resolve(key);
216 }
217
218 let ipv4_lookup = self.0.ipv4_lookup(Name::from_str_relaxed(key)?).await?;
219 let ips: Arc<Vec<Ipv4Addr>> = ipv4_lookup
220 .as_lookup()
221 .record_iter()
222 .filter_map(|r| r.data().as_a()?.0.into())
223 .collect::<Vec<Ipv4Addr>>()
224 .into();
225
226 Ok(DnsEntry {
227 entry: ips,
228 expires: ipv4_lookup.valid_until(),
229 })
230 }
231
232 pub async fn ipv6_lookup<'x>(
233 &self,
234 key: impl IntoFqdn<'x>,
235 cache: Option<&impl ResolverCache<String, Arc<Vec<Ipv6Addr>>>>,
236 ) -> crate::Result<Arc<Vec<Ipv6Addr>>> {
237 let key = key.into_fqdn();
238 if let Some(value) = cache.as_ref().and_then(|c| c.get(key.as_ref())) {
239 return Ok(value);
240 }
241
242 let ipv6_lookup = self.ipv6_lookup_raw(key.as_ref()).await?;
243
244 if let Some(cache) = cache {
245 cache.insert(
246 key.into_owned(),
247 ipv6_lookup.entry.clone(),
248 ipv6_lookup.expires,
249 );
250 }
251
252 Ok(ipv6_lookup.entry)
253 }
254
255 pub async fn ipv6_lookup_raw(&self, key: &str) -> crate::Result<DnsEntry<Arc<Vec<Ipv6Addr>>>> {
256 #[cfg(any(test, feature = "test"))]
257 if true {
258 return mock_resolve(key);
259 }
260
261 let ipv6_lookup = self.0.ipv6_lookup(Name::from_str_relaxed(key)?).await?;
262 let ips: Arc<Vec<Ipv6Addr>> = ipv6_lookup
263 .as_lookup()
264 .record_iter()
265 .filter_map(|r| r.data().as_aaaa()?.0.into())
266 .collect::<Vec<Ipv6Addr>>()
267 .into();
268
269 Ok(DnsEntry {
270 entry: ips,
271 expires: ipv6_lookup.valid_until(),
272 })
273 }
274
275 pub async fn ip_lookup(
276 &self,
277 key: &str,
278 mut strategy: IpLookupStrategy,
279 max_results: usize,
280 cache_ipv4: Option<&impl ResolverCache<String, Arc<Vec<Ipv4Addr>>>>,
281 cache_ipv6: Option<&impl ResolverCache<String, Arc<Vec<Ipv6Addr>>>>,
282 ) -> crate::Result<Vec<IpAddr>> {
283 loop {
284 match strategy {
285 IpLookupStrategy::Ipv4Only | IpLookupStrategy::Ipv4thenIpv6 => {
286 match (self.ipv4_lookup(key, cache_ipv4).await, strategy) {
287 (Ok(result), _) => {
288 return Ok(result
289 .iter()
290 .take(max_results)
291 .copied()
292 .map(IpAddr::from)
293 .collect());
294 }
295 (Err(err), IpLookupStrategy::Ipv4Only) => return Err(err),
296 _ => {
297 strategy = IpLookupStrategy::Ipv6Only;
298 }
299 }
300 }
301 IpLookupStrategy::Ipv6Only | IpLookupStrategy::Ipv6thenIpv4 => {
302 match (self.ipv6_lookup(key, cache_ipv6).await, strategy) {
303 (Ok(result), _) => {
304 return Ok(result
305 .iter()
306 .take(max_results)
307 .copied()
308 .map(IpAddr::from)
309 .collect());
310 }
311 (Err(err), IpLookupStrategy::Ipv6Only) => return Err(err),
312 _ => {
313 strategy = IpLookupStrategy::Ipv4Only;
314 }
315 }
316 }
317 }
318 }
319 }
320
321 pub async fn ptr_lookup(
322 &self,
323 addr: IpAddr,
324 cache: Option<&impl ResolverCache<IpAddr, Arc<Vec<String>>>>,
325 ) -> crate::Result<Arc<Vec<String>>> {
326 if let Some(value) = cache.as_ref().and_then(|c| c.get(&addr)) {
327 return Ok(value);
328 }
329
330 #[cfg(any(test, feature = "test"))]
331 if true {
332 return mock_resolve(&addr.to_string());
333 }
334
335 let ptr_lookup = self.0.reverse_lookup(addr).await?;
336 let ptr: Arc<Vec<String>> = ptr_lookup
337 .as_lookup()
338 .record_iter()
339 .filter_map(|r| {
340 let r = r.data().as_ptr()?;
341 if !r.is_empty() {
342 r.to_lowercase().to_string().into()
343 } else {
344 None
345 }
346 })
347 .collect::<Vec<String>>()
348 .into();
349
350 if let Some(cache) = cache {
351 cache.insert(addr, ptr.clone(), ptr_lookup.valid_until());
352 }
353
354 Ok(ptr)
355 }
356
357 #[cfg(any(test, feature = "test"))]
358 pub async fn exists<'x>(
359 &self,
360 key: impl IntoFqdn<'x>,
361 cache_ipv4: Option<&impl ResolverCache<String, Arc<Vec<Ipv4Addr>>>>,
362 cache_ipv6: Option<&impl ResolverCache<String, Arc<Vec<Ipv6Addr>>>>,
363 ) -> crate::Result<bool> {
364 let key = key.into_fqdn().into_owned();
365 match self.ipv4_lookup(key.as_str(), cache_ipv4).await {
366 Ok(_) => Ok(true),
367 Err(Error::DnsRecordNotFound(_)) => {
368 match self.ipv6_lookup(key.as_str(), cache_ipv6).await {
369 Ok(_) => Ok(true),
370 Err(Error::DnsRecordNotFound(_)) => Ok(false),
371 Err(err) => Err(err),
372 }
373 }
374 Err(err) => Err(err),
375 }
376 }
377
378 #[cfg(not(any(test, feature = "test")))]
379 pub async fn exists<'x>(
380 &self,
381 key: impl IntoFqdn<'x>,
382 cache_ipv4: Option<&impl ResolverCache<String, Arc<Vec<Ipv4Addr>>>>,
383 cache_ipv6: Option<&impl ResolverCache<String, Arc<Vec<Ipv6Addr>>>>,
384 ) -> crate::Result<bool> {
385 let key = key.into_fqdn();
386
387 if cache_ipv4.is_some_and(|c| c.get(key.as_ref()).is_some())
388 || cache_ipv6.is_some_and(|c| c.get(key.as_ref()).is_some())
389 {
390 return Ok(true);
391 }
392
393 match self
394 .0
395 .lookup_ip(Name::from_str_relaxed(key.as_ref())?)
396 .await
397 {
398 Ok(result) => Ok(result.as_lookup().record_iter().any(|r| {
399 matches!(
400 r.data().record_type(),
401 hickory_resolver::proto::rr::RecordType::A
402 | hickory_resolver::proto::rr::RecordType::AAAA
403 )
404 })),
405 Err(err) => match err.kind() {
406 ProtoErrorKind::NoRecordsFound { .. } => Ok(false),
407 _ => Err(err.into()),
408 },
409 }
410 }
411}
412
413impl From<ProtoError> for Error {
414 fn from(err: ProtoError) -> Self {
415 match err.kind() {
416 ProtoErrorKind::NoRecordsFound(response_code) => {
417 Error::DnsRecordNotFound(response_code.response_code)
418 }
419 _ => Error::DnsError(err.to_string()),
420 }
421 }
422}
423
424impl From<DomainKey> for Txt {
425 fn from(v: DomainKey) -> Self {
426 Txt::DomainKey(v.into())
427 }
428}
429
430impl From<DomainKeyReport> for Txt {
431 fn from(v: DomainKeyReport) -> Self {
432 Txt::DomainKeyReport(v.into())
433 }
434}
435
436impl From<Atps> for Txt {
437 fn from(v: Atps) -> Self {
438 Txt::Atps(v.into())
439 }
440}
441
442impl From<Spf> for Txt {
443 fn from(v: Spf) -> Self {
444 Txt::Spf(v.into())
445 }
446}
447
448impl From<Macro> for Txt {
449 fn from(v: Macro) -> Self {
450 Txt::SpfMacro(v.into())
451 }
452}
453
454impl From<Dmarc> for Txt {
455 fn from(v: Dmarc) -> Self {
456 Txt::Dmarc(v.into())
457 }
458}
459
460impl From<MtaSts> for Txt {
461 fn from(v: MtaSts) -> Self {
462 Txt::MtaSts(v.into())
463 }
464}
465
466impl From<TlsRpt> for Txt {
467 fn from(v: TlsRpt) -> Self {
468 Txt::TlsRpt(v.into())
469 }
470}
471
472impl<T: Into<Txt>> From<crate::Result<T>> for Txt {
473 fn from(v: crate::Result<T>) -> Self {
474 match v {
475 Ok(v) => v.into(),
476 Err(err) => Txt::Error(err),
477 }
478 }
479}
480
481pub trait UnwrapTxtRecord: Sized {
482 fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>>;
483}
484
485impl UnwrapTxtRecord for DomainKey {
486 fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
487 match txt {
488 Txt::DomainKey(a) => Ok(a),
489 Txt::Error(err) => Err(err),
490 _ => Err(Error::Io("Invalid record type".to_string())),
491 }
492 }
493}
494
495impl UnwrapTxtRecord for DomainKeyReport {
496 fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
497 match txt {
498 Txt::DomainKeyReport(a) => Ok(a),
499 Txt::Error(err) => Err(err),
500 _ => Err(Error::Io("Invalid record type".to_string())),
501 }
502 }
503}
504
505impl UnwrapTxtRecord for Atps {
506 fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
507 match txt {
508 Txt::Atps(a) => Ok(a),
509 Txt::Error(err) => Err(err),
510 _ => Err(Error::Io("Invalid record type".to_string())),
511 }
512 }
513}
514
515impl UnwrapTxtRecord for Spf {
516 fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
517 match txt {
518 Txt::Spf(a) => Ok(a),
519 Txt::Error(err) => Err(err),
520 _ => Err(Error::Io("Invalid record type".to_string())),
521 }
522 }
523}
524
525impl UnwrapTxtRecord for Macro {
526 fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
527 match txt {
528 Txt::SpfMacro(a) => Ok(a),
529 Txt::Error(err) => Err(err),
530 _ => Err(Error::Io("Invalid record type".to_string())),
531 }
532 }
533}
534
535impl UnwrapTxtRecord for Dmarc {
536 fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
537 match txt {
538 Txt::Dmarc(a) => Ok(a),
539 Txt::Error(err) => Err(err),
540 _ => Err(Error::Io("Invalid record type".to_string())),
541 }
542 }
543}
544
545impl UnwrapTxtRecord for MtaSts {
546 fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
547 match txt {
548 Txt::MtaSts(a) => Ok(a),
549 Txt::Error(err) => Err(err),
550 _ => Err(Error::Io("Invalid record type".to_string())),
551 }
552 }
553}
554
555impl UnwrapTxtRecord for TlsRpt {
556 fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
557 match txt {
558 Txt::TlsRpt(a) => Ok(a),
559 Txt::Error(err) => Err(err),
560 _ => Err(Error::Io("Invalid record type".to_string())),
561 }
562 }
563}
564
565pub trait IntoFqdn<'x> {
566 fn into_fqdn(self) -> Cow<'x, str>;
567}
568
569impl<'x> IntoFqdn<'x> for String {
570 fn into_fqdn(self) -> Cow<'x, str> {
571 if self.ends_with('.') {
572 self.to_lowercase().into()
573 } else {
574 format!("{}.", self.to_lowercase()).into()
575 }
576 }
577}
578
579impl<'x> IntoFqdn<'x> for &'x str {
580 fn into_fqdn(self) -> Cow<'x, str> {
581 if self.ends_with('.') {
582 self.to_lowercase().into()
583 } else {
584 format!("{}.", self.to_lowercase()).into()
585 }
586 }
587}
588
589impl<'x> IntoFqdn<'x> for &String {
590 fn into_fqdn(self) -> Cow<'x, str> {
591 if self.ends_with('.') {
592 self.to_lowercase().into()
593 } else {
594 format!("{}.", self.to_lowercase()).into()
595 }
596 }
597}
598
599pub trait ToReverseName {
600 fn to_reverse_name(&self) -> String;
601}
602
603impl ToReverseName for IpAddr {
604 fn to_reverse_name(&self) -> String {
605 use std::fmt::Write;
606
607 match self {
608 IpAddr::V4(ip) => {
609 let mut segments = String::with_capacity(15);
610 for octet in ip.octets().iter().rev() {
611 if !segments.is_empty() {
612 segments.push('.');
613 }
614 let _ = write!(&mut segments, "{}", octet);
615 }
616 segments
617 }
618 IpAddr::V6(ip) => {
619 let mut segments = String::with_capacity(63);
620 for segment in ip.segments().iter().rev() {
621 for &p in format!("{segment:04x}").as_bytes().iter().rev() {
622 if !segments.is_empty() {
623 segments.push('.');
624 }
625 segments.push(char::from(p));
626 }
627 }
628 segments
629 }
630 }
631 }
632}
633
634#[cfg(any(test, feature = "test"))]
635pub fn mock_resolve<T>(domain: &str) -> crate::Result<T> {
636 Err(if domain.contains("_parse_error.") {
637 Error::ParseError
638 } else if domain.contains("_invalid_record.") {
639 Error::InvalidRecordType
640 } else if domain.contains("_dns_error.") {
641 Error::DnsError("".to_string())
642 } else {
643 Error::DnsRecordNotFound(hickory_resolver::proto::op::ResponseCode::NXDomain)
644 })
645}
646
647#[cfg(test)]
648mod test {
649 use std::net::IpAddr;
650
651 use crate::common::resolver::ToReverseName;
652
653 #[test]
654 fn reverse_lookup_addr() {
655 for (addr, expected) in [
656 ("1.2.3.4", "4.3.2.1"),
657 (
658 "2001:db8::cb01",
659 "1.0.b.c.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.8.b.d.0.1.0.0.2",
660 ),
661 (
662 "2a01:4f9:c011:b43c::1",
663 "1.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.c.3.4.b.1.1.0.c.9.f.4.0.1.0.a.2",
664 ),
665 ] {
666 assert_eq!(addr.parse::<IpAddr>().unwrap().to_reverse_name(), expected);
667 }
668 }
669}