use super::{parse::TxtRecordParser, verify::DomainKey};
use crate::{
Error, IpLookupStrategy, MX, MessageAuthenticator, ResolverCache, Txt,
dkim::{Atps, DomainKeyReport},
dmarc::Dmarc,
mta_sts::{MtaSts, TlsRpt},
spf::{Macro, Spf},
};
use hickory_resolver::{
TokioResolver,
config::{CLOUDFLARE, GOOGLE, QUAD9, ResolverConfig, ResolverOpts},
net::{
DnsError, NetError,
runtime::TokioRuntimeProvider,
},
proto::{
ProtoError,
rr::{Name, RData},
},
system_conf::read_system_conf,
};
use std::{
borrow::Cow,
net::{IpAddr, Ipv4Addr, Ipv6Addr},
sync::Arc,
time::Instant,
};
pub struct DnsEntry<T> {
pub entry: T,
pub expires: Instant,
}
impl MessageAuthenticator {
pub fn new_cloudflare_tls() -> Result<Self, NetError> {
Self::new(ResolverConfig::tls(&CLOUDFLARE), ResolverOpts::default())
}
pub fn new_cloudflare() -> Result<Self, NetError> {
Self::new(
ResolverConfig::udp_and_tcp(&CLOUDFLARE),
ResolverOpts::default(),
)
}
pub fn new_google() -> Result<Self, NetError> {
Self::new(
ResolverConfig::udp_and_tcp(&GOOGLE),
ResolverOpts::default(),
)
}
pub fn new_quad9() -> Result<Self, NetError> {
Self::new(ResolverConfig::udp_and_tcp(&QUAD9), ResolverOpts::default())
}
pub fn new_quad9_tls() -> Result<Self, NetError> {
Self::new(ResolverConfig::tls(&QUAD9), ResolverOpts::default())
}
pub fn new_system_conf() -> Result<Self, NetError> {
let (config, options) = read_system_conf()?;
Self::new(config, options)
}
pub fn new(config: ResolverConfig, options: ResolverOpts) -> Result<Self, NetError> {
Ok(MessageAuthenticator(
TokioResolver::builder_with_config(config, TokioRuntimeProvider::default())
.with_options(options)
.build()?,
))
}
pub fn resolver(&self) -> &TokioResolver {
&self.0
}
pub async fn txt_raw_lookup(&self, key: impl ToFqdn) -> crate::Result<Vec<u8>> {
let mut result = vec![];
for record in self
.0
.txt_lookup(Name::from_str_relaxed::<&str>(key.to_fqdn().as_ref())?)
.await?
.answers()
{
if let RData::TXT(txt) = &record.data {
for item in &txt.txt_data {
result.extend_from_slice(item);
}
}
}
Ok(result)
}
pub async fn txt_lookup<T: TxtRecordParser + Into<Txt> + UnwrapTxtRecord>(
&self,
key: impl ToFqdn,
cache: Option<&impl ResolverCache<Box<str>, Txt>>,
) -> crate::Result<Arc<T>> {
let key = key.to_fqdn();
if let Some(value) = cache.as_ref().and_then(|c| c.get::<str>(key.as_ref())) {
return T::unwrap_txt(value);
}
#[cfg(any(test, feature = "test"))]
if true {
return mock_resolve(key.as_ref());
}
let txt_lookup = self
.0
.txt_lookup(Name::from_str_relaxed::<&str>(key.as_ref())?)
.await?;
let mut result = Err(Error::InvalidRecordType);
let records = txt_lookup.answers().iter().filter_map(|r| {
let RData::TXT(txt) = &r.data else {
return None;
};
let txt_data = &txt.txt_data;
match txt_data.len() {
1 => Some(Cow::from(txt_data[0].as_ref())),
0 => None,
_ => {
let mut entry = Vec::with_capacity(255 * txt_data.len());
for data in txt_data {
entry.extend_from_slice(data);
}
Some(Cow::from(entry))
}
}
});
for record in records {
result = T::parse(record.as_ref());
if result.is_ok() {
break;
}
}
let result: Txt = result.into();
if let Some(cache) = cache {
cache.insert(key, result.clone(), txt_lookup.valid_until());
}
T::unwrap_txt(result)
}
pub async fn mx_lookup(
&self,
key: impl ToFqdn,
cache: Option<&impl ResolverCache<Box<str>, Arc<[MX]>>>,
) -> crate::Result<Arc<[MX]>> {
let key = key.to_fqdn();
if let Some(value) = cache.as_ref().and_then(|c| c.get::<str>(key.as_ref())) {
return Ok(value);
}
#[cfg(any(test, feature = "test"))]
if true {
return mock_resolve(key.as_ref());
}
let mx_lookup = self
.0
.mx_lookup(Name::from_str_relaxed::<&str>(key.as_ref())?)
.await?;
let mx_records = mx_lookup.answers();
let mut records: Vec<(u16, Vec<Box<str>>)> = Vec::with_capacity(mx_records.len());
for mx_record in mx_records {
if let RData::MX(mx) = &mx_record.data {
let preference = mx.preference;
let exchange = mx.exchange.to_lowercase().to_string().into_boxed_str();
if let Some(record) = records.iter_mut().find(|r| r.0 == preference) {
record.1.push(exchange);
} else {
records.push((preference, vec![exchange]));
}
}
}
records.sort_unstable_by_key(|a| a.0);
let records: Arc<[MX]> = records
.into_iter()
.map(|(preference, exchanges)| MX {
preference,
exchanges: exchanges.into_boxed_slice(),
})
.collect::<Arc<[MX]>>();
if let Some(cache) = cache {
cache.insert(key, records.clone(), mx_lookup.valid_until());
}
Ok(records)
}
pub async fn ipv4_lookup(
&self,
key: impl ToFqdn,
cache: Option<&impl ResolverCache<Box<str>, Arc<[Ipv4Addr]>>>,
) -> crate::Result<Arc<[Ipv4Addr]>> {
let key = key.to_fqdn();
if let Some(value) = cache.as_ref().and_then(|c| c.get::<str>(key.as_ref())) {
return Ok(value);
}
let ipv4_lookup = self.ipv4_lookup_raw(key.as_ref()).await?;
if let Some(cache) = cache {
cache.insert(key, ipv4_lookup.entry.clone(), ipv4_lookup.expires);
}
Ok(ipv4_lookup.entry)
}
pub async fn ipv4_lookup_raw(&self, key: &str) -> crate::Result<DnsEntry<Arc<[Ipv4Addr]>>> {
#[cfg(any(test, feature = "test"))]
if true {
return mock_resolve(key);
}
let ipv4_lookup = self
.0
.ipv4_lookup(Name::from_str_relaxed::<&str>(key)?)
.await?;
let ips: Arc<[Ipv4Addr]> = ipv4_lookup
.answers()
.iter()
.filter_map(|r| {
if let RData::A(a) = &r.data {
Some(a.0)
} else {
None
}
})
.collect::<Vec<Ipv4Addr>>()
.into();
Ok(DnsEntry {
entry: ips,
expires: ipv4_lookup.valid_until(),
})
}
pub async fn ipv6_lookup(
&self,
key: impl ToFqdn,
cache: Option<&impl ResolverCache<Box<str>, Arc<[Ipv6Addr]>>>,
) -> crate::Result<Arc<[Ipv6Addr]>> {
let key = key.to_fqdn();
if let Some(value) = cache.as_ref().and_then(|c| c.get::<str>(key.as_ref())) {
return Ok(value);
}
let ipv6_lookup = self.ipv6_lookup_raw(key.as_ref()).await?;
if let Some(cache) = cache {
cache.insert(key, ipv6_lookup.entry.clone(), ipv6_lookup.expires);
}
Ok(ipv6_lookup.entry)
}
pub async fn ipv6_lookup_raw(&self, key: &str) -> crate::Result<DnsEntry<Arc<[Ipv6Addr]>>> {
#[cfg(any(test, feature = "test"))]
if true {
return mock_resolve(key);
}
let ipv6_lookup = self
.0
.ipv6_lookup(Name::from_str_relaxed::<&str>(key)?)
.await?;
let ips: Arc<[Ipv6Addr]> = ipv6_lookup
.answers()
.iter()
.filter_map(|r| {
if let RData::AAAA(aaaa) = &r.data {
Some(aaaa.0)
} else {
None
}
})
.collect::<Vec<Ipv6Addr>>()
.into();
Ok(DnsEntry {
entry: ips,
expires: ipv6_lookup.valid_until(),
})
}
pub async fn ip_lookup(
&self,
key: &str,
mut strategy: IpLookupStrategy,
max_results: usize,
cache_ipv4: Option<&impl ResolverCache<Box<str>, Arc<[Ipv4Addr]>>>,
cache_ipv6: Option<&impl ResolverCache<Box<str>, Arc<[Ipv6Addr]>>>,
) -> crate::Result<Vec<IpAddr>> {
loop {
match strategy {
IpLookupStrategy::Ipv4Only | IpLookupStrategy::Ipv4thenIpv6 => {
match (self.ipv4_lookup(key, cache_ipv4).await, strategy) {
(Ok(result), _) => {
return Ok(result
.iter()
.take(max_results)
.copied()
.map(IpAddr::from)
.collect());
}
(Err(err), IpLookupStrategy::Ipv4Only) => return Err(err),
_ => {
strategy = IpLookupStrategy::Ipv6Only;
}
}
}
IpLookupStrategy::Ipv6Only | IpLookupStrategy::Ipv6thenIpv4 => {
match (self.ipv6_lookup(key, cache_ipv6).await, strategy) {
(Ok(result), _) => {
return Ok(result
.iter()
.take(max_results)
.copied()
.map(IpAddr::from)
.collect());
}
(Err(err), IpLookupStrategy::Ipv6Only) => return Err(err),
_ => {
strategy = IpLookupStrategy::Ipv4Only;
}
}
}
}
}
}
pub async fn ptr_lookup(
&self,
addr: IpAddr,
cache: Option<&impl ResolverCache<IpAddr, Arc<[Box<str>]>>>,
) -> crate::Result<Arc<[Box<str>]>> {
if let Some(value) = cache.as_ref().and_then(|c| c.get(&addr)) {
return Ok(value);
}
#[cfg(any(test, feature = "test"))]
if true {
return mock_resolve(&addr.to_string());
}
let ptr_lookup = self.0.reverse_lookup(addr).await?;
let ptr: Arc<[Box<str>]> = ptr_lookup
.answers()
.iter()
.filter_map(|r| {
let RData::PTR(ptr) = &r.data else {
return None;
};
if !ptr.is_empty() {
Some(ptr.to_lowercase().to_string().into_boxed_str())
} else {
None
}
})
.collect::<Arc<[Box<str>]>>();
if let Some(cache) = cache {
cache.insert(addr, ptr.clone(), ptr_lookup.valid_until());
}
Ok(ptr)
}
#[cfg(any(test, feature = "test"))]
pub async fn exists(
&self,
key: impl ToFqdn,
cache_ipv4: Option<&impl ResolverCache<Box<str>, Arc<[Ipv4Addr]>>>,
cache_ipv6: Option<&impl ResolverCache<Box<str>, Arc<[Ipv6Addr]>>>,
) -> crate::Result<bool> {
let key = key.to_fqdn();
match self.ipv4_lookup(key.as_ref(), cache_ipv4).await {
Ok(_) => Ok(true),
Err(Error::DnsRecordNotFound(_)) => {
match self.ipv6_lookup(key.as_ref(), cache_ipv6).await {
Ok(_) => Ok(true),
Err(Error::DnsRecordNotFound(_)) => Ok(false),
Err(err) => Err(err),
}
}
Err(err) => Err(err),
}
}
#[cfg(not(any(test, feature = "test")))]
pub async fn exists(
&self,
key: impl ToFqdn,
cache_ipv4: Option<&impl ResolverCache<Box<str>, Arc<[Ipv4Addr]>>>,
cache_ipv6: Option<&impl ResolverCache<Box<str>, Arc<[Ipv6Addr]>>>,
) -> crate::Result<bool> {
let key = key.to_fqdn();
if cache_ipv4.is_some_and(|c| c.get::<str>(key.as_ref()).is_some())
|| cache_ipv6.is_some_and(|c| c.get::<str>(key.as_ref()).is_some())
{
return Ok(true);
}
match self
.0
.lookup_ip(Name::from_str_relaxed::<&str>(key.as_ref())?)
.await
{
Ok(result) => Ok(result.as_lookup().answers().iter().any(|r| {
matches!(
&r.data.record_type(),
hickory_resolver::proto::rr::RecordType::A
| hickory_resolver::proto::rr::RecordType::AAAA
)
})),
Err(err) if err.is_no_records_found() => Ok(false),
Err(err) => Err(err.into()),
}
}
}
impl From<ProtoError> for Error {
fn from(err: ProtoError) -> Self {
Error::DnsError(err.to_string())
}
}
impl From<NetError> for Error {
fn from(err: NetError) -> Self {
match &err {
NetError::Dns(DnsError::NoRecordsFound(no_records)) => {
Error::DnsRecordNotFound(no_records.response_code)
}
_ => Error::DnsError(err.to_string()),
}
}
}
impl From<DomainKey> for Txt {
fn from(v: DomainKey) -> Self {
Txt::DomainKey(v.into())
}
}
impl From<DomainKeyReport> for Txt {
fn from(v: DomainKeyReport) -> Self {
Txt::DomainKeyReport(v.into())
}
}
impl From<Atps> for Txt {
fn from(v: Atps) -> Self {
Txt::Atps(v.into())
}
}
impl From<Spf> for Txt {
fn from(v: Spf) -> Self {
Txt::Spf(v.into())
}
}
impl From<Macro> for Txt {
fn from(v: Macro) -> Self {
Txt::SpfMacro(v.into())
}
}
impl From<Dmarc> for Txt {
fn from(v: Dmarc) -> Self {
Txt::Dmarc(v.into())
}
}
impl From<MtaSts> for Txt {
fn from(v: MtaSts) -> Self {
Txt::MtaSts(v.into())
}
}
impl From<TlsRpt> for Txt {
fn from(v: TlsRpt) -> Self {
Txt::TlsRpt(v.into())
}
}
impl<T: Into<Txt>> From<crate::Result<T>> for Txt {
fn from(v: crate::Result<T>) -> Self {
match v {
Ok(v) => v.into(),
Err(err) => Txt::Error(err),
}
}
}
pub trait UnwrapTxtRecord: Sized {
fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>>;
}
impl UnwrapTxtRecord for DomainKey {
fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
match txt {
Txt::DomainKey(a) => Ok(a),
Txt::Error(err) => Err(err),
_ => Err(Error::Io("Invalid record type".to_string())),
}
}
}
impl UnwrapTxtRecord for DomainKeyReport {
fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
match txt {
Txt::DomainKeyReport(a) => Ok(a),
Txt::Error(err) => Err(err),
_ => Err(Error::Io("Invalid record type".to_string())),
}
}
}
impl UnwrapTxtRecord for Atps {
fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
match txt {
Txt::Atps(a) => Ok(a),
Txt::Error(err) => Err(err),
_ => Err(Error::Io("Invalid record type".to_string())),
}
}
}
impl UnwrapTxtRecord for Spf {
fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
match txt {
Txt::Spf(a) => Ok(a),
Txt::Error(err) => Err(err),
_ => Err(Error::Io("Invalid record type".to_string())),
}
}
}
impl UnwrapTxtRecord for Macro {
fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
match txt {
Txt::SpfMacro(a) => Ok(a),
Txt::Error(err) => Err(err),
_ => Err(Error::Io("Invalid record type".to_string())),
}
}
}
impl UnwrapTxtRecord for Dmarc {
fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
match txt {
Txt::Dmarc(a) => Ok(a),
Txt::Error(err) => Err(err),
_ => Err(Error::Io("Invalid record type".to_string())),
}
}
}
impl UnwrapTxtRecord for MtaSts {
fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
match txt {
Txt::MtaSts(a) => Ok(a),
Txt::Error(err) => Err(err),
_ => Err(Error::Io("Invalid record type".to_string())),
}
}
}
impl UnwrapTxtRecord for TlsRpt {
fn unwrap_txt(txt: Txt) -> crate::Result<Arc<Self>> {
match txt {
Txt::TlsRpt(a) => Ok(a),
Txt::Error(err) => Err(err),
_ => Err(Error::Io("Invalid record type".to_string())),
}
}
}
pub trait ToFqdn {
fn to_fqdn(&self) -> Box<str>;
}
impl<T: AsRef<str>> ToFqdn for T {
fn to_fqdn(&self) -> Box<str> {
let value = self.as_ref();
if value.ends_with('.') {
value.to_lowercase().into()
} else {
format!("{}.", value.to_lowercase()).into()
}
}
}
pub trait ToReverseName {
fn to_reverse_name(&self) -> String;
}
impl ToReverseName for IpAddr {
fn to_reverse_name(&self) -> String {
use std::fmt::Write;
match self {
IpAddr::V4(ip) => {
let mut segments = String::with_capacity(15);
for octet in ip.octets().iter().rev() {
if !segments.is_empty() {
segments.push('.');
}
let _ = write!(&mut segments, "{}", octet);
}
segments
}
IpAddr::V6(ip) => {
let mut segments = String::with_capacity(63);
for segment in ip.segments().iter().rev() {
for &p in format!("{segment:04x}").as_bytes().iter().rev() {
if !segments.is_empty() {
segments.push('.');
}
segments.push(char::from(p));
}
}
segments
}
}
}
}
#[cfg(any(test, feature = "test"))]
pub fn mock_resolve<T>(domain: &str) -> crate::Result<T> {
Err(if domain.contains("_parse_error.") {
Error::ParseError
} else if domain.contains("_invalid_record.") {
Error::InvalidRecordType
} else if domain.contains("_dns_error.") {
Error::DnsError("".to_string())
} else {
Error::DnsRecordNotFound(hickory_resolver::proto::op::ResponseCode::NXDomain)
})
}
#[cfg(test)]
mod test {
use std::net::IpAddr;
use crate::common::resolver::ToReverseName;
#[test]
fn reverse_lookup_addr() {
for (addr, expected) in [
("1.2.3.4", "4.3.2.1"),
(
"2001:db8::cb01",
"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",
),
(
"2a01:4f9:c011:b43c::1",
"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",
),
] {
assert_eq!(addr.parse::<IpAddr>().unwrap().to_reverse_name(), expected);
}
}
}