From fef3a4e2a67f142ae2b4340074b8cf851dd13ec5 Mon Sep 17 00:00:00 2001
From: Litr0 <gonzalo.sobarzoa@gmail.com>
Date: Tue, 11 Jun 2024 12:57:20 -0400
Subject: [PATCH] Remove already cached answers in add

---
 src/resolver_cache.rs | 45 +++++++++++++------------------------------
 1 file changed, 13 insertions(+), 32 deletions(-)

diff --git a/src/resolver_cache.rs b/src/resolver_cache.rs
index fd99ef00..03be52f7 100644
--- a/src/resolver_cache.rs
+++ b/src/resolver_cache.rs
@@ -60,16 +60,6 @@ impl ResolverCache {
 
     /// Add an element to the answer cache.
     pub fn add_answer(&mut self, domain_name: DomainName, resource_record: ResourceRecord, qtype: Option<Qtype>, qclass: Qclass, rcode: Option<Rcode>) {
-        let key;
-        if qtype == None {
-            key = CacheKey::Secondary(qclass, domain_name.clone());
-        }
-        else {
-            key = CacheKey::Primary(qtype.unwrap(), qclass, domain_name.clone());
-        }
-        if self.cache_answer.is_cached(key.clone()) {
-            self.cache_answer.remove(domain_name.clone(), qtype, qclass);
-        }
         if resource_record.get_ttl() > 0 {
             self.cache_answer.add(domain_name, resource_record, qtype, qclass, rcode);
         }
@@ -77,16 +67,6 @@ impl ResolverCache {
 
     /// Add an element to the authority cache.
     pub fn add_authority(&mut self, domain_name: DomainName, resource_record: ResourceRecord, qtype: Option<Qtype>, qclass: Qclass, rcode: Option<Rcode>) {
-        let key;
-        if qtype == None {
-            key = CacheKey::Secondary(qclass, domain_name.clone());
-        }
-        else {
-            key = CacheKey::Primary(qtype.unwrap(), qclass, domain_name.clone());
-        }
-        if self.cache_authority.is_cached(key.clone()) {
-            self.cache_authority.remove(domain_name.clone(), qtype, qclass);
-        }
         if resource_record.get_ttl() > 0 {
             self.cache_authority.add(domain_name, resource_record, qtype, qclass, rcode);
         }
@@ -94,16 +74,6 @@ impl ResolverCache {
 
     /// Add an element to the additional cache.
     pub fn add_additional(&mut self, domain_name: DomainName, resource_record: ResourceRecord, qtype: Option<Qtype>, qclass: Qclass, rcode: Option<Rcode>) {
-        let key;
-        if qtype == None {
-            key = CacheKey::Secondary(qclass, domain_name.clone());
-        }
-        else {
-            key = CacheKey::Primary(qtype.unwrap(), qclass, domain_name.clone());
-        }
-        if self.cache_additional.is_cached(key.clone()) {
-            self.cache_additional.remove(domain_name.clone(), qtype, qclass);
-        }
         if resource_record.get_ttl() > 0 {
             if resource_record.get_rtype() != Rtype::OPT {
                 self.cache_additional.add(domain_name, resource_record, qtype, qclass, rcode);
@@ -116,13 +86,24 @@ impl ResolverCache {
         let qname = message.get_question().get_qname();
         let qtype = Some(message.get_question().get_qtype());
         let qclass = message.get_question().get_qclass();
+        let rcode = Some(message.get_header().get_rcode());
+
+        // Checks if something with the same key is already cached
+        let key;
+        if rcode == Some(Rcode::NXDOMAIN) {
+            key = CacheKey::Secondary(qclass, qname.clone());
+        }
+        else {
+            key = CacheKey::Primary(qtype.unwrap(), qclass, qname.clone());
+        }
+        if self.is_cached(key.clone()) {
+            self.remove(qname.clone(), qtype, qclass);
+        }
 
         let answers = message.get_answer();
         let authorities = message.get_authority();
         let additionals = message.get_additional();
 
-        let rcode = Some(message.get_header().get_rcode());
-
         answers.iter()
         .for_each(|rr| {
             self.add_answer(qname.clone(), rr.clone(), qtype, qclass, rcode);