refactor remove in cache

niclabs · Jun 11, 2024 · edef97f · edef97f
1 parent 641b4a1
commit edef97f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 11 deletions.
diff --git 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
@@ -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);