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, 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, 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, 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, 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, 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);