domani/src/domain/checker.rs

use std::error::Error;
use std::str::FromStr;

use mockall::automock;
use trust_dns_client::client::{Client, SyncClient};
use trust_dns_client::rr::{DNSClass, Name, RData, RecordType};
use trust_dns_client::udp::UdpClientConnection;

#[derive(thiserror::Error, Debug)]
pub enum NewDNSCheckerError {
    #[error("invalid resolver address")]
    InvalidResolverAddress,

    #[error("invalid target CNAME")]
    InvalidTargetCNAME,

    #[error(transparent)]
    Unexpected(Box<dyn Error>),
}

#[derive(thiserror::Error, Debug)]
pub enum CheckDomainError {
    #[error("invalid domain name")]
    InvalidDomainName,

    #[error("target CNAME not set")]
    TargetCNAMENotSet,

    #[error("challenge token not set")]
    ChallengeTokenNotSet,

    #[error(transparent)]
    Unexpected(Box<dyn Error>),
}

#[automock]
pub trait Checker {
    fn check_domain(&self, domain: &str, challenge_token: &str) -> Result<(), CheckDomainError>;
}

pub struct DNSChecker {
    target_cname: Name,
    client: SyncClient<UdpClientConnection>,
}

pub fn new(target_cname: &str, resolver_addr: &str) -> Result<impl Checker, NewDNSCheckerError> {
    let target_cname =
        Name::from_str(target_cname).map_err(|_| NewDNSCheckerError::InvalidTargetCNAME)?;

    let resolver_addr = resolver_addr
        .parse()
        .map_err(|_| NewDNSCheckerError::InvalidResolverAddress)?;

    let conn = UdpClientConnection::new(resolver_addr)
        .map_err(|e| NewDNSCheckerError::Unexpected(Box::from(e)))?;

    let client = SyncClient::new(conn);

    Ok(DNSChecker {
        target_cname,
        client,
    })
}

impl Checker for DNSChecker {
    fn check_domain(&self, domain: &str, challenge_token: &str) -> Result<(), CheckDomainError> {
        let mut fqdn = Name::from_str(domain).map_err(|_| CheckDomainError::InvalidDomainName)?;
        fqdn.set_fqdn(true);

        // check that the CNAME is installed correctly on the domain
        {
            let response = self
                .client
                .query(&fqdn, DNSClass::IN, RecordType::CNAME)
                .map_err(|e| CheckDomainError::Unexpected(Box::from(e)))?;

            let records = response.answers();

            if records.len() != 1 {
                return Err(CheckDomainError::TargetCNAMENotSet);
            }

            // if the single record isn't a CNAME, or it's not the target CNAME, then return
            // TargetCNAMENotSet
            match records[0].data() {
                Some(RData::CNAME(remote_cname)) if remote_cname == &self.target_cname => (),
                _ => return Err(CheckDomainError::TargetCNAMENotSet),
            }
        }

        // check that the TXT record with the challenge token is correctly installed on the domain
        {
            let fqdn = Name::from_str("_gateway")
                .map_err(|e| CheckDomainError::Unexpected(Box::from(e)))?
                .append_domain(&fqdn)
                .map_err(|e| CheckDomainError::Unexpected(Box::from(e)))?;

            let response = self
                .client
                .query(&fqdn, DNSClass::IN, RecordType::TXT)
                .map_err(|e| CheckDomainError::Unexpected(Box::from(e)))?;

            let records = response.answers();

            if !records.iter().any(|record| -> bool {
                match record.data() {
                    Some(RData::TXT(txt)) => txt.to_string().contains(challenge_token),
                    _ => false,
                }
            }) {
                return Err(CheckDomainError::ChallengeTokenNotSet);
            }
        }

        Ok(())
    }
}