//! Abstractions for DNS lookups.

#[cfg(feature = "trust-dns-resolver")]
mod trust_dns_resolver;

use async_trait::async_trait;
use std::{
    error::Error,
    fmt::{self, Display, Formatter},
    hash::{Hash, Hasher},
    net::{IpAddr, Ipv4Addr, Ipv6Addr},
    str::FromStr,
};

/// A result type specialised for lookup errors.
pub type LookupResult<T> = Result<T, LookupError>;

/// Errors that may occur when doing lookups.
#[derive(Debug)]
pub enum LookupError {
    /// The lookup timed out.
    Timeout,
    /// No records were found (NXDOMAIN).
    NoRecords,
    /// Any other error (I/O, encoding, protocol etc.) that causes a DNS lookup
    /// to fail, with source attached. The error cause is made available in the
    /// trace.
    Dns(Option<Box<dyn Error + Send + Sync>>),
}

impl Error for LookupError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Timeout | Self::NoRecords => None,
            Self::Dns(e) => e.as_ref().map(|e| e.as_ref() as _),
        }
    }
}

impl Display for LookupError {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Self::Timeout => write!(f, "lookup timed out"),
            Self::NoRecords => write!(f, "no records found"),
            Self::Dns(source) => match source {
                None => write!(f, "lookup returned error"),
                Some(error) => write!(f, "lookup returned error with cause: {}", error),
            },
        }
    }
}

/// A trait for entities that perform DNS resolution.
///
/// This trait uses the [`Name`] type in arguments and return values. See the
/// documentation of this type for a discussion of its design and use.
///
/// An implementation of this trait for the Trust-DNS `TokioAsyncResolver` is
/// available when the feature `trust-dns-resolver` is enabled.
///
/// **`Lookup` is not a general purpose DNS resolver.** It is a simplistic
/// ad-hoc model for performing DNS resolution as used in the SPF protocol.
#[async_trait]
pub trait Lookup: Send + Sync {
    /// Looks up the IPv4 addresses for the given name.
    ///
    /// # Errors
    ///
    /// If the lookup encounters an error, a `LookupError` variant is returned.
    async fn lookup_a<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<Ipv4Addr>>;

    /// Looks up the IPv6 addresses for the given name.
    ///
    /// # Errors
    ///
    /// If the lookup encounters an error, a `LookupError` variant is returned.
    async fn lookup_aaaa<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<Ipv6Addr>>;

    /// Looks up the mail exchanger names for the given name.
    ///
    /// Implementers may want to order the returned vector by MX preference.
    ///
    /// # Errors
    ///
    /// If the lookup encounters an error, a `LookupError` variant is returned.
    async fn lookup_mx<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<Name>>;

    /// Looks up the text records for the given name.
    ///
    /// Implementers should make sure that invalid UTF-8 does not cause the
    /// lookup to fail (eg, using a lossy transformation). A record consisting
    /// of multiple elements should be returned as one string, with elements
    /// joined without a separator.
    ///
    /// # Errors
    ///
    /// If the lookup encounters an error, a `LookupError` variant is returned.
    async fn lookup_txt<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<String>>;

    /// Looks up the names for the given address (reverse lookup).
    ///
    /// # Errors
    ///
    /// If the lookup encounters an error, a `LookupError` variant is returned.
    async fn lookup_ptr<'lookup>(&'lookup self, ip: IpAddr) -> LookupResult<Vec<Name>>;
}

#[async_trait]
impl<T: Lookup + ?Sized> Lookup for Box<T> {
    async fn lookup_a<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<Ipv4Addr>> {
        (**self).lookup_a(name).await
    }

    async fn lookup_aaaa<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<Ipv6Addr>> {
        (**self).lookup_aaaa(name).await
    }

    async fn lookup_mx<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<Name>> {
        (**self).lookup_mx(name).await
    }

    async fn lookup_txt<'lookup, 'a>(&'lookup self, name: &'a Name) -> LookupResult<Vec<String>> {
        (**self).lookup_txt(name).await
    }

    async fn lookup_ptr<'lookup>(&'lookup self, ip: IpAddr) -> LookupResult<Vec<Name>> {
        (**self).lookup_ptr(ip).await
    }
}

/// An error indicating that a domain name could not be parsed.
#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
pub struct ParseNameError;

impl Error for ParseNameError {}

impl Display for ParseNameError {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "could not parse domain name")
    }
}

/// A DNS name that can be used in DNS queries.
///
/// Use [`new`] to create a `Name` from a string. `Name::new` only checks length
/// limits and otherwise accepts all string inputs. In lookup implementations,
/// `Name` is used through `Name::new` and `Name::as_str`. `Lookup`
/// implementations should treat `Name` as a transparent wrapper of a string.
///
/// A `Name` is always fully-qualified, relative to the root. When borrowing a
/// `Name` with `as_str`, the string has a trailing dot. In the `Display`
/// presentation form the trailing dot is omitted:
///
/// ```
/// use viaspf::lookup::Name;
///
/// let name = Name::new("my.org")?;
///
/// assert_eq!(name.as_str(), "my.org.");
/// assert_eq!(name.to_string(), "my.org");
/// # Ok::<_, viaspf::lookup::ParseNameError>(())
/// ```
///
/// **`Name` is not a general purpose DNS name implementation.** It is a
/// simplistic ad-hoc model for DNS names as used in the SPF protocol. `Name`
/// plays a part in the DNS name handling of the [`Lookup`] implementation and
/// so the `Name` API is specific to its use in the evaluation logic.
///
/// [`new`]: Self::new
/// [`Lookup`]: crate::lookup::Lookup
#[derive(Clone, Debug)]
pub struct Name(String);

impl Name {
    /// Constructs a syntactically valid DNS name from a string.
    ///
    /// This only checks label length restrictions but does not apply any
    /// transformation to the given string (no case normalisation, Punycode
    /// encoding or similar).
    ///
    /// # Errors
    ///
    /// If the given string is not a valid DNS name, an error is returned.
    ///
    /// # Examples
    ///
    /// ```
    /// use viaspf::lookup::Name;
    ///
    /// assert!(Name::new("example.org").is_ok());
    /// assert!(Name::new("_spf.example.org").is_ok());
    /// ```
    pub fn new(s: &str) -> Result<Self, ParseNameError> {
        if is_valid_dns_name(s, false) {
            Ok(Self::root(s.to_owned()))
        } else {
            Err(ParseNameError)
        }
    }

    pub(crate) fn domain(s: &str) -> Result<Self, ParseNameError> {
        // §4.3: ‘Internationalized domain names MUST be encoded as A-labels’
        let s = idna::domain_to_ascii(s).map_err(|_| ParseNameError)?;
        if is_valid_dns_name(&s, true) {
            Ok(Self::root(s))
        } else {
            Err(ParseNameError)
        }
    }

    // §4.3: ‘Properly formed domains are fully qualified domains […]. That is,
    // in the DNS they are implicitly qualified relative to the root’
    fn root(mut s: String) -> Self {
        if !s.ends_with('.') {
            s.push('.');
        }
        Name(s)
    }

    /// Returns a string slice of the fully-qualified name (including the
    /// trailing dot).
    pub fn as_str(&self) -> &str {
        &self.0
    }

    /// Returns whether a name is a subdomain of another name.
    pub fn is_subdomain_of(&self, other: &Name) -> bool {
        let name = self.as_str();
        let other = other.as_str();

        name.len() > other.len() && {
            let len = name.len() - other.len();
            matches!(name.get(len..), Some(s) if s.eq_ignore_ascii_case(other))
                && matches!(name.get(..len), Some(s) if s.ends_with('.'))
        }
    }
}

impl PartialEq for Name {
    fn eq(&self, other: &Self) -> bool {
        self.0.eq_ignore_ascii_case(&other.0)
    }
}

impl Eq for Name {}

impl Hash for Name {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.to_ascii_lowercase().hash(state);
    }
}

impl Display for Name {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.0[..(self.0.len() - 1)].fmt(f)
    }
}

impl AsRef<str> for Name {
    fn as_ref(&self) -> &str {
        self.as_str()
    }
}

impl FromStr for Name {
    type Err = ParseNameError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Self::new(s)
    }
}

fn is_valid_dns_name(mut s: &str, domain: bool) -> bool {
    if let Some(sx) = s.strip_suffix('.') {
        s = sx;
    }

    if !has_valid_domain_len(s) {
        return false;
    }

    // As noted in the documentation, `Name` is a simplified domain name model
    // and does not take into account special cases that need escaping, eg
    // `"."`. See RFC 4343.
    let mut labels = s.split('.').rev().peekable();
    if domain {
        if matches!(labels.next(), Some(l) if !is_tld(l)) {
            return false;
        }
        if labels.peek().is_none() {
            return false;
        }
        labels.all(is_label)
    } else {
        labels.all(|l| has_valid_label_len(l) && !l.starts_with('-') && !l.ends_with('-'))
    }
}

// RFC 3696, §2: ‘There is an additional rule that essentially requires that
// top-level domain names not be all-numeric.’
fn is_tld(s: &str) -> bool {
    is_label(s) && !s.chars().all(|c: char| c.is_ascii_digit())
}

// ‘The LDH rule, as updated, provides that the labels […] that make up a domain
// name must consist of only the ASCII alphabetic and numeric characters, plus
// the hyphen. […] If the hyphen is used, it is not permitted to appear at
// either the beginning or end of a label.’ And: ‘A DNS label may be no more
// than 63 octets long.’
fn is_label(s: &str) -> bool {
    has_valid_label_len(s)
        && s.starts_with(|c: char| c.is_ascii_alphanumeric())
        && s.ends_with(|c: char| c.is_ascii_alphanumeric())
        && s.chars().all(|c: char| c.is_ascii_alphanumeric() || c == '-')
}

// Maximum domain length without the trailing dot. Implementation note: I have
// also seen the limit placed at 255 [254] (eg RFC 3696), but have not been able
// to find the definitive statement, so going with the more conservative 253 …
pub(crate) const MAX_DOMAIN_LENGTH: usize = 253;

fn has_valid_domain_len(s: &str) -> bool {
    matches!(s.len(), 1..=MAX_DOMAIN_LENGTH)
}

fn has_valid_label_len(s: &str) -> bool {
    matches!(s.len(), 1..=63)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn name_new_ok() {
        assert_eq!(Name::new("ch"), Ok(Name("ch.".into())));
        assert_eq!(Name::new(""), Err(ParseNameError));

        assert_eq!(Name::new("mail.gluet.ch"), Ok(Name("mail.gluet.ch.".into())));
        assert_eq!(Name::new("_dmarc.gluet.ch"), Ok(Name("_dmarc.gluet.ch.".into())));
        assert_eq!(Name::new("185.46.57.247"), Ok(Name("185.46.57.247.".into())));
        assert_eq!(Name::new("_--,"), Ok(Name("_--,.".into())));

        assert_eq!(Name::new("ch."), Ok(Name("ch.".into())));
        assert_eq!(Name::new("mail.gluet.ch."), Ok(Name("mail.gluet.ch.".into())));
        assert_eq!(Name::new("_dmarc.gluet.ch."), Ok(Name("_dmarc.gluet.ch.".into())));
        assert_eq!(Name::new("185.46.57.247."), Ok(Name("185.46.57.247.".into())));

        assert_eq!(Name::new(".mail.gluet.ch"), Err(ParseNameError));
        assert_eq!(Name::new("mail.g...t.ch"), Err(ParseNameError));
        assert_eq!(Name::new("mail.g   t.ch"), Ok(Name("mail.g   t.ch.".into())));

        assert_eq!(Name::new("www.Bücher.de"), Ok(Name("www.Bücher.de.".into())));
        assert_eq!(Name::new("www.xn--bcher-kva.de"), Ok(Name("www.xn--bcher-kva.de.".into())));
    }

    #[test]
    fn name_domain_ok() {
        assert_eq!(Name::domain("ch"), Err(ParseNameError));
        assert_eq!(Name::domain("odd-tld-lookalike"), Err(ParseNameError));
        assert_eq!(Name::domain(""), Err(ParseNameError));

        assert_eq!(Name::domain("mail.gluet.ch"), Ok(Name("mail.gluet.ch.".into())));
        assert_eq!(Name::domain("_dmarc.gluet.ch"), Err(ParseNameError));
        assert_eq!(Name::domain("185.46.57.247"), Err(ParseNameError));

        assert_eq!(Name::domain("mail.gluet.ch."), Ok(Name("mail.gluet.ch.".into())));
        assert_eq!(Name::domain("_dmarc.gluet.ch."), Err(ParseNameError));
        assert_eq!(Name::domain("185.46.57.247."), Err(ParseNameError));

        assert_eq!(Name::domain(".mail.gluet.ch"), Err(ParseNameError));
        assert_eq!(Name::domain("mail.gluet.ch.."), Err(ParseNameError));

        assert_eq!(Name::domain("myhost.local"), Ok(Name("myhost.local.".into())));

        assert_eq!(Name::domain("www.Bücher.de"), Ok(Name("www.xn--bcher-kva.de.".into())));
        assert_eq!(Name::domain("www.xn--bcher-kva.de"), Ok(Name("www.xn--bcher-kva.de.".into())));
    }

    #[test]
    fn name_normalises_case() {
        assert_eq!(
            Name::domain("www.example.org").unwrap(),
            Name::domain("wWw.ExAmPlE.OrG").unwrap()
        );
    }

    #[test]
    fn is_subdomain_of_ok() {
        let parent = Name::new("example.org").unwrap();
        assert!(Name::new("mail.example.org").unwrap().is_subdomain_of(&parent));
        assert!(!Name::new("example.org").unwrap().is_subdomain_of(&parent));
        assert!(!Name::new("myexample.org").unwrap().is_subdomain_of(&parent));
        assert!(!Name::new("example.com").unwrap().is_subdomain_of(&parent));
        assert!(!Name::new("org").unwrap().is_subdomain_of(&parent));
    }
}