From edef97f26a033d7930a173e123cad7fcb84a486c Mon Sep 17 00:00:00 2001
From: Litr0 <gonzalo.sobarzoa@gmail.com>
Date: Tue, 11 Jun 2024 10:42:00 -0400
Subject: [PATCH] refactor remove in cache

---
 src/dns_cache.rs      | 10 +++++++---
 src/resolver_cache.rs | 16 ++++++++--------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/src/dns_cache.rs b/src/dns_cache.rs
index 31e8722b..a0d92758 100644
--- a/src/dns_cache.rs
+++ b/src/dns_cache.rs
@@ -109,9 +109,13 @@ impl DnsCache {
     }
 
     /// Removes an element from cache
-    pub fn remove(&mut self, domain_name: DomainName, qtype: Qtype, qclass: Qclass) {
+    pub fn remove(&mut self, domain_name: DomainName, qtype: Option<Qtype>, qclass: Qclass) {
         let mut cache_data = self.get_cache();
-        let _extracted = cache_data.pop(&CacheKey::Primary(qtype, qclass, domain_name));
+        if qtype != None {
+            let _extracted = cache_data.pop(&CacheKey::Primary(qtype.unwrap(), qclass, domain_name));
+        } else {
+            let _extracted = cache_data.pop(&CacheKey::Secondary(qclass, domain_name));
+        }
         self.set_cache(cache_data); 
     }
 
@@ -458,7 +462,7 @@ mod dns_cache_test {
 
         cache.add(domain_name.clone(), resource_record.clone(), Some(Qtype::A), Qclass::IN, None);
 
-        cache.remove(domain_name.clone(), Qtype::A, Qclass::IN);
+        cache.remove(domain_name.clone(), Some(Qtype::A), Qclass::IN);
 
         let rr_cache_vec = cache.get(domain_name.clone(), Qtype::A, Qclass::IN);
 
diff --git a/src/resolver_cache.rs b/src/resolver_cache.rs
index c0478baa..0995fb3f 100644
--- a/src/resolver_cache.rs
+++ b/src/resolver_cache.rs
@@ -199,22 +199,22 @@ impl ResolverCache {
     }
 
     /// Removes an element from the answer cache.
-    pub fn remove_answer(&mut self, domain_name: DomainName, qtype: Qtype, qclass: Qclass) {
+    pub fn remove_answer(&mut self, domain_name: DomainName, qtype: Option<Qtype>, qclass: Qclass) {
         self.cache_answer.remove(domain_name, qtype, qclass);
     }
 
     /// Removes an element from the authority cache.
-    pub fn remove_authority(&mut self, domain_name: DomainName, qtype: Qtype, qclass: Qclass) {
+    pub fn remove_authority(&mut self, domain_name: DomainName, qtype: Option<Qtype>, qclass: Qclass) {
         self.cache_authority.remove(domain_name, qtype, qclass);
     }
 
     /// Removes an element from the additional cache.
-    pub fn remove_additional(&mut self, domain_name: DomainName, qtype: Qtype, qclass: Qclass) {
+    pub fn remove_additional(&mut self, domain_name: DomainName, qtype: Option<Qtype>, qclass: Qclass) {
         self.cache_additional.remove(domain_name, qtype, qclass);
     }
 
     /// Removes an element from the cache.
-    pub fn remove(&mut self, domain_name: DomainName, qtype: Qtype, qclass: Qclass) {
+    pub fn remove(&mut self, domain_name: DomainName, qtype: Option<Qtype>, qclass: Qclass) {
         self.remove_answer(domain_name.clone(), qtype, qclass);
         self.remove_authority(domain_name.clone(), qtype, qclass);
         self.remove_additional(domain_name.clone(), qtype, qclass);
@@ -750,7 +750,7 @@ mod resolver_cache_test{
         resolver_cache.add_answer(domain_name.clone(), resource_record_2.clone(), Some(Qtype::A), Qclass::IN, None);
         resolver_cache.add_answer(domain_name.clone(), resource_record_3.clone(), Some(Qtype::A), Qclass::IN, None);
 
-        resolver_cache.remove_answer(domain_name.clone(), Qtype::A, Qclass::IN);
+        resolver_cache.remove_answer(domain_name.clone(), Some(Qtype::A), Qclass::IN);
 
         let rr = resolver_cache.get_answer(domain_name.clone(), Qtype::A, Qclass::IN);
 
@@ -801,7 +801,7 @@ mod resolver_cache_test{
         resolver_cache.add_authority(domain_name.clone(), resource_record_2.clone(), Some(Qtype::A), Qclass::IN, None);
         resolver_cache.add_authority(domain_name.clone(), resource_record_3.clone(), Some(Qtype::A), Qclass::IN, None);
 
-        resolver_cache.remove_authority(domain_name.clone(), Qtype::A, Qclass::IN);
+        resolver_cache.remove_authority(domain_name.clone(), Some(Qtype::A), Qclass::IN);
 
         let rr = resolver_cache.get_authority(domain_name.clone(), Qtype::A, Qclass::IN);
 
@@ -852,7 +852,7 @@ mod resolver_cache_test{
         resolver_cache.add_additional(domain_name.clone(), resource_record_2.clone(), Some(Qtype::A), Qclass::IN, None);
         resolver_cache.add_additional(domain_name.clone(), resource_record_3.clone(), Some(Qtype::A), Qclass::IN, None);
         
-        resolver_cache.remove_additional(domain_name.clone(), Qtype::A, Qclass::IN);
+        resolver_cache.remove_additional(domain_name.clone(), Some(Qtype::A), Qclass::IN);
 
         let rr = resolver_cache.get_additional(domain_name.clone(), Qtype::A, Qclass::IN);
 
@@ -912,7 +912,7 @@ mod resolver_cache_test{
         
         let query = DnsMessage::new_query_message(qname.clone(), qtype.clone(), qclass.clone(), op_code.clone(), rd.clone(), id.clone());
 
-        resolver_cache.remove(domain_name.clone(), Qtype::A, Qclass::IN);
+        resolver_cache.remove(domain_name.clone(), Some(Qtype::A), Qclass::IN);
 
         let message = resolver_cache.get(query);