1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83

use super::{Lookup, LookupError, LookupResult, Name};
use async_trait::async_trait;
use std::{
    error::Error,
    net::{IpAddr, Ipv4Addr, Ipv6Addr},
};
use trust_dns_resolver::{
    error::{ResolveError, ResolveErrorKind},
    Name as TrustDnsName, TokioAsyncResolver,
};

#[async_trait]
impl Lookup for TokioAsyncResolver {
    async fn lookup_a<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<Ipv4Addr>> {
        Ok(self
            .ipv4_lookup(to_trust_dns_name(name)?)
            .await
            .map_err(to_lookup_error)?
            .into_iter()
            .collect())
    }

    async fn lookup_aaaa<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<Ipv6Addr>> {
        Ok(self
            .ipv6_lookup(to_trust_dns_name(name)?)
            .await
            .map_err(to_lookup_error)?
            .into_iter()
            .collect())
    }

    async fn lookup_mx<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<Name>> {
        let mut mxs = self
            .mx_lookup(to_trust_dns_name(name)?)
            .await
            .map_err(to_lookup_error)?
            .into_iter()
            .collect::<Vec<_>>();
        mxs.sort_by_key(|mx| mx.preference());
        mxs.into_iter()
            .map(|mx| Name::new(&mx.exchange().to_ascii()).map_err(wrap_error))
            .collect()
    }

    async fn lookup_txt<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<String>> {
        Ok(self
            .txt_lookup(to_trust_dns_name(name)?)
            .await
            .map_err(to_lookup_error)?
            .into_iter()
            .map(|txt| {
                txt.iter()
                    .map(|data| String::from_utf8_lossy(data))
                    .collect()
            })
            .collect())
    }

    async fn lookup_ptr<'lookup>(&'lookup self, ip: IpAddr) -> LookupResult<Vec<Name>> {
        self.reverse_lookup(ip)
            .await
            .map_err(to_lookup_error)?
            .into_iter()
            .map(|name| Name::new(&name.to_ascii()).map_err(wrap_error))
            .collect()
    }
}

fn to_trust_dns_name(name: &Name) -> LookupResult<TrustDnsName> {
    TrustDnsName::from_ascii(name).map_err(wrap_error)
}

fn to_lookup_error(error: ResolveError) -> LookupError {
    match error.kind() {
        ResolveErrorKind::NoRecordsFound { .. } => LookupError::NoRecords,
        ResolveErrorKind::Timeout => LookupError::Timeout,
        _ => wrap_error(error),
    }
}

fn wrap_error(error: impl Error + Send + Sync + 'static) -> LookupError {
    LookupError::Dns(Some(error.into()))
}