nomilo/src/dns/dns_connector.rs

use trust_dns_proto::DnsHandle;
use trust_dns_client::client::ClientHandle;
use trust_dns_client::rr::{DNSClass, RecordType};
use trust_dns_client::op::{UpdateMessage, OpCode, MessageType, Message, Query, ResponseCode};
use trust_dns_client::error::ClientError;

use super::{Name, Record, RData};
use super::client::{ClientResponse, DnsClient};
use super::connector::{RecordConnector, ZoneConnector};


const MAX_PAYLOAD_LEN: u16 = 1232;


#[derive(Debug)]
pub enum DnsConnectorError {
    ClientError(ClientError),
    ResponceNotOk {
        code: ResponseCode,
        zone: Name,
    },
}

pub struct DnsConnectorClient {
    client: DnsClient
}

impl DnsConnectorClient {
    pub fn new(client: DnsClient) -> Self {
        DnsConnectorClient {
            client
        }
    }
}


#[async_trait]
impl RecordConnector for DnsConnectorClient {
    type Error = DnsConnectorError;

    async fn get_records(&mut self, zone: Name, class: DNSClass) -> Result<Vec<Record>, Self::Error>
    {
        let response = {
            let query = self.client.query(zone.clone(), class, RecordType::AXFR);
            query.await.map_err(|e| DnsConnectorError::ClientError(e))?
        };

        if response.response_code() != ResponseCode::NoError {
            return Err(DnsConnectorError::ResponceNotOk {
                code: response.response_code(),
                zone: zone,
            });
        }

        let answers = response.answers();
        let mut records: Vec<_> = answers.to_vec().into_iter()
            .filter(|record| !matches!(record.rdata(), RData::NULL { .. } | RData::DNSSEC(_)))
            .collect();

        // AXFR response ends with SOA, we remove it so it is not doubled in the response.
        records.pop();
        Ok(records)
    }

    async fn add_records(&mut self, zone: Name, class: DNSClass, new_records: Vec<Record>) -> Result<(), Self::Error>
    {
        // Taken from trust_dns_client::op::update_message::append
        // The original function can not be used as is because it takes a RecordSet and not a Record list

        let mut zone_query = Query::new();
        zone_query.set_name(zone.clone())
            .set_query_class(class)
            .set_query_type(RecordType::SOA);

        let mut message = Message::new();

        // TODO: set random / time based id
        message
            .set_id(0)
            .set_message_type(MessageType::Query)
            .set_op_code(OpCode::Update)
            .set_recursion_desired(false);
        message.add_zone(zone_query);
        message.add_updates(new_records);

        {
            let edns = message.edns_mut();
            edns.set_max_payload(MAX_PAYLOAD_LEN);
            edns.set_version(0);
        }

        let response = ClientResponse(self.client.send(message)).await.map_err(|e| DnsConnectorError::ClientError(e))?;

        if response.response_code() != ResponseCode::NoError {
            return Err(DnsConnectorError::ResponceNotOk {
                code: response.response_code(),
                zone: zone,
            });
        }

        Ok(())
    }

    async fn update_records(&mut self, zone: Name, class: DNSClass, old_records: Vec<Record>, new_records: Vec<Record>) -> Result<(), Self::Error>
    {
        // Taken from trust_dns_client::op::update_message::compare_and_swap
        // The original function can not be used as is because it takes a RecordSet and not a Record list

        // for updates, the query section is used for the zone
        let mut zone_query = Query::new();
        zone_query.set_name(zone.clone())
            .set_query_class(class)
            .set_query_type(RecordType::SOA);

        let mut message: Message = Message::new();

        // build the message
        // TODO: set random / time based id
        message
            .set_id(0)
            .set_message_type(MessageType::Query)
            .set_op_code(OpCode::Update)
            .set_recursion_desired(false);
        message.add_zone(zone_query);

        // make sure the record is what is expected
        let mut prerequisite = old_records.clone();
        for record in prerequisite.iter_mut() {
            record.set_ttl(0);
        }
        message.add_pre_requisites(prerequisite);

        // add the delete for the old record
        let mut delete = old_records;
        for record in delete.iter_mut() {
            // the class must be none for delete
            record.set_dns_class(DNSClass::NONE);
            // the TTL should be 0
            record.set_ttl(0);
        }
        message.add_updates(delete);

        // insert the new record...
        message.add_updates(new_records);

        // Extended dns
        {
            let edns = message.edns_mut();
            edns.set_max_payload(MAX_PAYLOAD_LEN);
            edns.set_version(0);
        }

        let response = ClientResponse(self.client.send(message)).await.map_err(|e| DnsConnectorError::ClientError(e))?;

        if response.response_code() != ResponseCode::NoError {
            return Err(DnsConnectorError::ResponceNotOk {
                code: response.response_code(),
                zone: zone,
            });
        }

        Ok(())
    }

    async fn delete_records(&mut self, zone: Name, class: DNSClass, records: Vec<Record>) -> Result<(), Self::Error>
    {
        // for updates, the query section is used for the zone
        let mut zone_query = Query::new();
        zone_query.set_name(zone.clone())
            .set_query_class(class)
            .set_query_type(RecordType::SOA);

        let mut message: Message = Message::new();

        // build the message
        // TODO: set random / time based id
        message
            .set_id(0)
            .set_message_type(MessageType::Query)
            .set_op_code(OpCode::Update)
            .set_recursion_desired(false);
        message.add_zone(zone_query);

        let mut delete = records;
        for record in delete.iter_mut() {
            // the class must be none for delete
            record.set_dns_class(DNSClass::NONE);
            // the TTL should be 0
            record.set_ttl(0);
        }
        message.add_updates(delete);

        // Extended dns
        {
            let edns = message.edns_mut();
            edns.set_max_payload(MAX_PAYLOAD_LEN);
            edns.set_version(0);
        }

        let response = ClientResponse(self.client.send(message)).await.map_err(|e| DnsConnectorError::ClientError(e))?;

        if response.response_code() != ResponseCode::NoError {
            return Err(DnsConnectorError::ResponceNotOk {
                code: response.response_code(),
                zone: zone,
            });
        }

        Ok(())

    }
}


#[async_trait]
impl ZoneConnector for DnsConnectorClient {
    type Error = DnsConnectorError;

    async fn zone_exists(&mut self, zone: Name, class: DNSClass) -> Result<(), Self::Error>
    {
        let response = {
            let query = self.client.query(zone.clone(), class, RecordType::SOA);
            query.await.map_err(|e| DnsConnectorError::ClientError(e))?
        };

        if response.response_code() != ResponseCode::NoError {
            return Err(DnsConnectorError::ResponceNotOk {
                code: response.response_code(),
                zone: zone,
            });
        }

        Ok(())
    }

}