From 38feb11cffb54ff3dee24b78ba158216c6681146 Mon Sep 17 00:00:00 2001
From: Revertron <r@revertron.com>
Date: Sun, 21 Feb 2021 12:29:09 +0100
Subject: [PATCH] Fixed DNS-records search.

---
 src/blockchain/filter.rs | 94 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 86 insertions(+), 8 deletions(-)

diff --git a/src/blockchain/filter.rs b/src/blockchain/filter.rs
index a56c981..50a739a 100644
--- a/src/blockchain/filter.rs
+++ b/src/blockchain/filter.rs
@@ -16,29 +16,107 @@ impl BlockchainFilter {
 
 impl DnsFilter for BlockchainFilter {
     fn lookup(&self, qname: &str, qtype: QueryType) -> Option<DnsPacket> {
-        let data = self.context.lock().unwrap().blockchain.get_domain_info(qname);
+        let mut search;
+        let mut subdomain;
+        let parts: Vec<&str> = qname.rsplitn(3, ".").collect();
+        match parts.len() {
+            1 => { return None; }
+            2 => {
+                search = format!("{}.{}", parts[1], parts[0]);
+                subdomain = String::new();
+            }
+            _ => {
+                search = format!("{}.{}", parts[1], parts[0]);
+                subdomain = String::from(parts[2]);
+            }
+        }
+        debug!("Searching domain {} and record {}", &search, &subdomain);
+
+        let data = self.context.lock().unwrap().blockchain.get_domain_info(&search);
         match data {
-            None => { debug!("Not found data for domain {}", &qname); }
+            None => { debug!("Not found data for domain {}", &search); }
             Some(data) => {
-                info!("Found data for domain {}", &qname);
-                let records: Vec<DnsRecord> = match serde_json::from_str(&data) {
+                info!("Found data for domain {}", &search);
+                let mut records: Vec<DnsRecord> = match serde_json::from_str(&data) {
                     Err(_) => { return None; }
                     Ok(records) => { records }
                 };
                 let mut answers: Vec<DnsRecord> = Vec::new();
-                for mut record in records {
+                for mut record in records.iter_mut() {
                     if record.get_querytype() == qtype {
                         match &mut record {
-                            // TODO make it for all types of records
-                            DnsRecord::A { domain, .. } | DnsRecord::AAAA { domain, .. } if domain == "@" => {
+                            DnsRecord::A { domain, .. }
+                            | DnsRecord::AAAA { domain, .. }
+                            | DnsRecord::NS { domain, .. }
+                            | DnsRecord::CNAME { domain, .. }
+                            | DnsRecord::SRV { domain, .. }
+                            | DnsRecord::MX { domain, .. }
+                            | DnsRecord::UNKNOWN { domain, .. }
+                            | DnsRecord::SOA { domain, .. }
+                            | DnsRecord::TXT { domain, .. } if domain == "@" => {
                                 *domain = String::from(qname);
                             }
                             _ => ()
                         }
 
-                        answers.push(record);
+                        match record.get_domain() {
+                            None => {}
+                            Some(domain) => {
+                                if domain == search {
+                                    answers.push(record.clone());
+                                } else if domain == subdomain {
+                                    match &mut record {
+                                        DnsRecord::A { domain, .. }
+                                        | DnsRecord::AAAA { domain, .. }
+                                        | DnsRecord::NS { domain, .. }
+                                        | DnsRecord::CNAME { domain, .. }
+                                        | DnsRecord::SRV { domain, .. }
+                                        | DnsRecord::MX { domain, .. }
+                                        | DnsRecord::UNKNOWN { domain, .. }
+                                        | DnsRecord::SOA { domain, .. }
+                                        | DnsRecord::TXT { domain, .. } => {
+                                            *domain = String::from(qname);
+                                        }
+                                        _ => ()
+                                    }
+                                    answers.push(record.clone());
+                                }
+                            }
+                        }
                     }
                 }
+                if answers.is_empty() {
+                    // If there are no records found we search for *.domain.ltd record
+                    for mut record in records {
+                        if record.get_querytype() == qtype {
+                            match record.get_domain() {
+                                None => {}
+                                Some(domain) => {
+                                    if domain == search {
+                                        answers.push(record.clone());
+                                    } else if domain == "*" {
+                                        match &mut record {
+                                            DnsRecord::A { domain, .. }
+                                            | DnsRecord::AAAA { domain, .. }
+                                            | DnsRecord::NS { domain, .. }
+                                            | DnsRecord::CNAME { domain, .. }
+                                            | DnsRecord::SRV { domain, .. }
+                                            | DnsRecord::MX { domain, .. }
+                                            | DnsRecord::UNKNOWN { domain, .. }
+                                            | DnsRecord::SOA { domain, .. }
+                                            | DnsRecord::TXT { domain, .. } => {
+                                                *domain = String::from(qname);
+                                            }
+                                            _ => ()
+                                        }
+                                        answers.push(record.clone());
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+
                 if !answers.is_empty() {
                     // Create DnsPacket
                     let mut packet = DnsPacket::new();