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(Debug)]
pub enum NewDNSCheckerError {
    InvalidResolverAddress,
    InvalidTargetCNAME,
    Unexpected(Box<dyn Error>),
}

impl<E: Error + 'static> From<E> for NewDNSCheckerError {
    fn from(e: E) -> NewDNSCheckerError {
        NewDNSCheckerError::Unexpected(Box::from(e))
    }
}

#[derive(Debug)]
pub enum CheckDomainError {
    InvalidDomainName,
    TargetCNAMENotSet,
    ChallengeTokenNotSet,
    Unexpected(Box<dyn Error>),
}

impl<E: Error + 'static> From<E> for CheckDomainError {
    fn from(e: E) -> CheckDomainError {
        CheckDomainError::Unexpected(Box::from(e))
    }
}

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

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

impl DNSChecker {
    pub fn new(target_cname: &str, resolver_addr: &str) -> Result<DNSChecker, 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)?;
        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)?;

            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")?.append_domain(&fqdn)?;
            let response = self.client.query(&fqdn, DNSClass::IN, RecordType::TXT)?;

            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(())
    }
}