domani/src/domain/checker.rs

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

use crate::domain;

use trust_dns_client::client::{AsyncClient, ClientHandle};
use trust_dns_client::rr::{DNSClass, Name, RData, RecordType};
use trust_dns_client::udp;

#[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("target CNAME not set")]
    TargetCNAMENotSet,

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

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

pub struct DNSChecker {
    target_cname: Name,

    // TODO we should use some kind of connection pool here, I suppose
    client: tokio::sync::Mutex<AsyncClient>,
}

pub fn new(
    tokio_runtime: sync::Arc<tokio::runtime::Runtime>,
    target_cname: domain::Name,
    resolver_addr: &str,
) -> Result<DNSChecker, NewDNSCheckerError> {
    let resolver_addr = resolver_addr
        .parse()
        .map_err(|_| NewDNSCheckerError::InvalidResolverAddress)?;

    let stream = udp::UdpClientStream::<tokio::net::UdpSocket>::new(resolver_addr);

    let (client, bg) = tokio_runtime
        .block_on(async { AsyncClient::connect(stream).await })
        .map_err(|e| NewDNSCheckerError::Unexpected(Box::from(e)))?;

    tokio_runtime.spawn(bg);

    Ok(DNSChecker {
        target_cname: target_cname.inner,
        client: tokio::sync::Mutex::new(client),
    })
}

impl DNSChecker {
    pub async fn check_domain(
        &self,
        domain: &domain::Name,
        challenge_token: &str,
    ) -> Result<(), CheckDomainError> {
        let domain = &domain.inner;

        // check that the CNAME is installed correctly on the domain
        {
            let response = match self
                .client
                .lock()
                .await
                .query(domain.clone(), DNSClass::IN, RecordType::CNAME)
                .await
            {
                Ok(res) => res,
                Err(e) => return Err(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 domain = Name::from_str("_domiply_challenge")
                .map_err(|e| CheckDomainError::Unexpected(Box::from(e)))?
                .append_domain(&domain)
                .map_err(|e| CheckDomainError::Unexpected(Box::from(e)))?;

            let response = match self
                .client
                .lock()
                .await
                .query(domain, DNSClass::IN, RecordType::TXT)
                .await
            {
                Ok(res) => res,
                Err(e) => return Err(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(())
    }
}