From 094539d376eeeeea19af12696460ebf01af32001 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ga=C3=ABl=20Berthaud-M=C3=BCller?= <blacksponge@tuta.io>
Date: Fri, 4 Mar 2022 13:50:57 +0100
Subject: [PATCH] move record fetching in message module

---
 src/dns/message.rs  | 36 ++++++++++++++++++++++++++++++++----
 src/routes/zones.rs | 38 +++++++++++++++-----------------------
 2 files changed, 47 insertions(+), 27 deletions(-)

diff --git a/src/dns/message.rs b/src/dns/message.rs
index 1e1e62f..f2e9c65 100644
--- a/src/dns/message.rs
+++ b/src/dns/message.rs
@@ -1,24 +1,52 @@
 use trust_dns_proto::DnsHandle;
 use trust_dns_client::rr::{DNSClass, RecordType};
-use trust_dns_client::op::{UpdateMessage, OpCode, MessageType, Message, Query};
+use trust_dns_client::op::{UpdateMessage, OpCode, MessageType, Message, Query, ResponseCode};
+use trust_dns_client::error::ClientError;
 use trust_dns_proto::error::ProtoError;
+use trust_dns_client::proto::xfer::{DnsRequestOptions};
 
-use super::trust_dns_types::{Name, Record};
+use super::trust_dns_types::{Name, Record, RData};
 use super::client::{ClientResponse};
 
 
+#[derive(Debug)]
 pub enum MessageError {
 	RecordNotInZone {
 		zone: Name,
 		class: DNSClass,
 		mismatched_class: Vec<Record>,
 		mismatched_zone: Vec<Record>,
-	}
+	},
+	ClientError(ClientError),
+	ResponceNotOk(ResponseCode)
 }
 
 
+#[async_trait]
 pub trait DnsMessage: DnsHandle<Error = ProtoError> + Send {
-	fn add_records(&mut self, zone: Name, class: DNSClass, new_records: Vec<Record>) -> Result<ClientResponse<Self::Response>, MessageError>
+	async fn get_records(&mut self, zone: Name, class: DNSClass) -> Result<Vec<Record>, MessageError>
+	{
+		let response = {
+			let mut query = Query::query(zone, RecordType::AXFR);
+			query.set_query_class(class);
+			ClientResponse(self.lookup(query, DnsRequestOptions::default())).await.map_err(|e| MessageError::ClientError(e))?
+	    };
+
+	    if response.response_code() != ResponseCode::NoError {
+	    	return Err(MessageError::ResponceNotOk(response.response_code()));
+	    }
+
+	    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)
+	}
+
+	fn add_records(&mut self, zone: Name, class: DNSClass, new_records: Vec<Record>) -> Result<ClientResponse<<Self as DnsHandle>::Response>, MessageError>
 	{
 		let mut mismatched_class = Vec::new();
 		let mut mismatched_zone = Vec::new();
diff --git a/src/routes/zones.rs b/src/routes/zones.rs
index d251543..b41e0ea 100644
--- a/src/routes/zones.rs
+++ b/src/routes/zones.rs
@@ -41,30 +41,21 @@ pub async fn get_zone_records(
         }
     }).await?;
 
-    let response = {
-        let query = client.query(zone.clone(), DNSClass::IN, RecordType::AXFR);
-        query.await.map_err(make_500)?
+    let records: Vec<_> = match client.get_records(zone.clone(), DNSClass::IN).await {
+        Ok(records) => records.into_iter().map(dns::record::Record::from).collect(),
+
+        Err(MessageError::ResponceNotOk(code)) => {
+            println!("Querrying AXFR of zone {} failed with code {}", *zone, code);
+            return ErrorResponse::new(
+                Status::NotFound,
+                "Zone could not be found".into()
+            ).with_details(json!({
+                "zone_name": zone.to_utf8()
+            })).err();
+        },
+        Err(err) => { return make_500(err).err(); },
     };
 
-    if response.response_code() != ResponseCode::NoError {
-        println!("Querrying AXFR of zone {} failed with code {}", *zone, response.response_code());
-        return ErrorResponse::new(
-            Status::NotFound,
-            "Zone could not be found".into()
-        ).with_details(json!({
-            "zone_name": zone.to_utf8()
-        })).err();
-    }
-
-    let answers = response.answers();
-    let mut records: Vec<_> = answers.to_vec().into_iter()
-        .map(dns::record::Record::from)
-        .filter(|record| !matches!(record.rdata, dns::rdata::RData::NULL { .. } | dns::rdata::RData::DNSSEC(_)))
-        .collect();
-
-    // AXFR response ends with SOA, we remove it so it is not doubled in the response.
-    records.pop();
-
     Ok(Json(records))
 }
 
@@ -127,7 +118,8 @@ pub async fn create_zone_records(
                     "mismatched_zone": mismatched_zone.into_iter().map(|r| r.clone().into()).collect::<Vec<dns::record::Record>>(),
                 })
             ).err();
-        }
+        },
+        Err(e) => return make_500(e).err()
     };
 
     // TODO: better error handling