nomilo/src/dns/dns_api.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::api::{RecordApi, ZoneApi};


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

pub struct DnsApiClient {
	client: DnsClient
}

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


#[async_trait]
impl RecordApi for DnsApiClient {
	type Error = DnsApiError;

	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| DnsApiError::ClientError(e))?
		};

	    if response.response_code() != ResponseCode::NoError {
	    	return Err(DnsApiError::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(1232);
	        edns.set_version(0);
	    }

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

        if response.response_code() != ResponseCode::NoError {
	    	return Err(DnsApiError::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 = 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;
	    // the class must be none for delete
	    for record in delete.iter_mut() {
	    	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(1232);
	        edns.set_version(0);
	    }

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

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

        Ok(())
	}
}


#[async_trait]
impl ZoneApi for DnsApiClient {
	type Error = DnsApiError;

	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| DnsApiError::ClientError(e))?
		};

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

	    Ok(())
	}

}