From 7b79f86e173086424840c5f2e5802a17357900b5 Mon Sep 17 00:00:00 2001
From: Revertron <r@revertron.com>
Date: Sun, 12 Sep 2021 13:03:10 +0200
Subject: [PATCH] Fixed inadequate response packet lengths and made proper
 CNAME resolution.

---
 src/dns/client.rs   |  2 +-
 src/dns/protocol.rs | 25 ++++++++++++++++++++-----
 src/dns/server.rs   | 32 ++++++++++++++++----------------
 3 files changed, 37 insertions(+), 22 deletions(-)

diff --git a/src/dns/client.rs b/src/dns/client.rs
index e877f78..13b0149 100644
--- a/src/dns/client.rs
+++ b/src/dns/client.rs
@@ -433,7 +433,7 @@ impl DnsClient for HttpsDnsClient {
         packet.questions.push(DnsQuestion::new(String::from(qname), qtype));
 
         let mut req_buffer = VectorPacketBuffer::new();
-        packet.write(&mut req_buffer, 512).expect("Preparing DnsPacket failed!");
+        packet.write(&mut req_buffer, 512 - 32).expect("Preparing DnsPacket failed!");
 
         let response = self.agent
             .post(doh_url)
diff --git a/src/dns/protocol.rs b/src/dns/protocol.rs
index d795587..73a08f3 100644
--- a/src/dns/protocol.rs
+++ b/src/dns/protocol.rs
@@ -810,19 +810,34 @@ impl DnsPacket {
         None
     }
 
-    pub fn get_unresolved_cnames(&self) -> Vec<DnsRecord> {
+    pub fn get_unresolved_cnames(&self, qtype: QueryType) -> Vec<DnsRecord> {
         let mut unresolved = Vec::new();
         for answer in &self.answers {
             let mut matched = false;
             if let DnsRecord::CNAME { ref host, .. } = *answer {
                 for answer2 in &self.answers {
-                    if let DnsRecord::A { ref domain, .. } = *answer2 {
-                        if domain == host {
-                            matched = true;
-                            break;
+                    match qtype {
+                        QueryType::A => {
+                            if let DnsRecord::A { ref domain, .. } = *answer2 {
+                                if domain == host {
+                                    matched = true;
+                                    break;
+                                }
+                            }
                         }
+                        QueryType::AAAA => {
+                            if let DnsRecord::AAAA { ref domain, .. } = *answer2 {
+                                if domain == host {
+                                    matched = true;
+                                    break;
+                                }
+                            }
+                        }
+                        _ => {}
                     }
                 }
+            } else {
+                continue;
             }
 
             if !matched {
diff --git a/src/dns/server.rs b/src/dns/server.rs
index 01ca3d3..1d645d3 100644
--- a/src/dns/server.rs
+++ b/src/dns/server.rs
@@ -62,7 +62,7 @@ pub trait DnsServer {
 
 /// Utility function for resolving domains referenced in for example CNAME or SRV
 /// records. This usually spares the client from having to perform additional lookups.
-fn resolve_cnames(lookup_list: &[DnsRecord], results: &mut Vec<DnsPacket>, resolver: &mut Box<dyn DnsResolver>, depth: u16) {
+fn resolve_cnames(lookup_list: &[DnsRecord], results: &mut Vec<DnsPacket>, resolver: &mut Box<dyn DnsResolver>, qtype: QueryType, depth: u16) {
     if depth > 10 {
         return;
     }
@@ -70,17 +70,11 @@ fn resolve_cnames(lookup_list: &[DnsRecord], results: &mut Vec<DnsPacket>, resol
     for ref rec in lookup_list {
         match **rec {
             DnsRecord::CNAME { ref host, .. } | DnsRecord::SRV { ref host, .. } => {
-                if let Ok(result2) = resolver.resolve(host, QueryType::A, true) {
-                    let new_unmatched = result2.get_unresolved_cnames();
+                if let Ok(result2) = resolver.resolve(host, qtype, true) {
+                    let new_unmatched = result2.get_unresolved_cnames(qtype);
                     results.push(result2);
 
-                    resolve_cnames(&new_unmatched, results, resolver, depth + 1);
-                }
-                if let Ok(result2) = resolver.resolve(host, QueryType::AAAA, true) {
-                    let new_unmatched = result2.get_unresolved_cnames();
-                    results.push(result2);
-
-                    resolve_cnames(&new_unmatched, results, resolver, depth + 1);
+                    resolve_cnames(&new_unmatched, results, resolver, qtype, depth + 1);
                 }
             }
             _ => {}
@@ -122,10 +116,10 @@ pub fn execute_query(context: Arc<ServerContext>, request: &DnsPacket) -> DnsPac
                     packet.header.authoritative_answer = true;
                 }
 
-                let unmatched = result.get_unresolved_cnames();
+                let unmatched = result.get_unresolved_cnames(question.qtype);
                 results.push(result);
 
-                resolve_cnames(&unmatched, &mut results, &mut resolver, 0);
+                resolve_cnames(&unmatched, &mut results, &mut resolver, question.qtype, 0);
 
                 res_code
             }
@@ -139,13 +133,19 @@ pub fn execute_query(context: Arc<ServerContext>, request: &DnsPacket) -> DnsPac
 
         for result in results {
             for rec in result.answers {
-                packet.answers.push(rec);
+                if !packet.answers.contains(&rec) {
+                    packet.answers.push(rec);
+                }
             }
             for rec in result.authorities {
-                packet.authorities.push(rec);
+                if !packet.authorities.contains(&rec) {
+                    packet.authorities.push(rec);
+                }
             }
             for rec in result.resources {
-                packet.resources.push(rec);
+                if !packet.resources.contains(&rec) {
+                    packet.resources.push(rec);
+                }
             }
         }
     }
@@ -212,7 +212,7 @@ impl DnsServer for DnsUdpServer {
                         }
                     };
 
-                    let mut size_limit = 512;
+                    let mut size_limit = 512 - 32; // Minus 32 bytes for the packet header
 
                     // Check for EDNS
                     if request.resources.len() == 1 {