From c364b029fd7a46ae5d82c3000f42b6a190fe9b09 Mon Sep 17 00:00:00 2001
From: Litr0 <gonzalo.sobarzoa@gmail.com>
Date: Tue, 11 Jun 2024 12:07:29 -0400
Subject: [PATCH] refactor DNS cache implementation to use `CacheKey` enum

---
 src/resolver_cache.rs | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/resolver_cache.rs b/src/resolver_cache.rs
index 0995fb3f..5a69d432 100644
--- a/src/resolver_cache.rs
+++ b/src/resolver_cache.rs
@@ -1,4 +1,4 @@
-use crate::dns_cache::DnsCache;
+use crate::dns_cache::{CacheKey, DnsCache};
 use crate::domain_name::DomainName;
 use crate::message::resource_record::ResourceRecord;
 use crate::message::type_qtype::Qtype;
@@ -41,6 +41,23 @@ impl ResolverCache {
         }
     }
 
+    /// Set the maximum size of the cache.
+    pub fn set_max_size(&mut self, size: NonZeroUsize) {
+        self.cache_answer.set_max_size(size);
+        self.cache_authority.set_max_size(size);
+        self.cache_additional.set_max_size(size);
+    }
+
+    /// See if the cache is empty.
+    pub fn is_empty(&self) -> bool {
+        self.cache_answer.is_empty() && self.cache_authority.is_empty() && self.cache_additional.is_empty()
+    }
+
+    /// See if an element is in the cache.
+    pub fn is_cached(&self, cache_key: CacheKey) -> bool {
+        self.cache_answer.is_cached(cache_key.clone()) || self.cache_authority.is_cached(cache_key.clone()) || self.cache_additional.is_cached(cache_key.clone())
+    }
+
     /// 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>) {
         if resource_record.get_ttl() > 0 {