From df8210eaca3f667f5339f06597c9d60231a603fc Mon Sep 17 00:00:00 2001
From: Han Xu <keepsimple@gmail.com>
Date: Sun, 3 Mar 2024 22:50:51 -0800
Subject: [PATCH] check refresh due for SRV, ADDR and TXT

---
 src/dns_parser.rs     |  9 +++++
 src/service_daemon.rs | 80 +++++++++++++++++++++++++++++++++++--------
 2 files changed, 74 insertions(+), 15 deletions(-)

diff --git a/src/dns_parser.rs b/src/dns_parser.rs
index b7bf3f93..d9f5cf6c 100644
--- a/src/dns_parser.rs
+++ b/src/dns_parser.rs
@@ -127,6 +127,15 @@ impl DnsRecord {
         self.refresh = get_expiration_time(self.created, self.ttl, 100);
     }
 
+    /// Checks if this record is due for refresh. If yes, update `refresh`.
+    pub(crate) fn refresh_maybe(&mut self, now: u64) -> bool {
+        if self.is_expired(now) || !self.refresh_due(now) {
+            return false;
+        }
+        self.refresh_no_more();
+        true
+    }
+
     /// Returns the remaining TTL in seconds
     fn get_remaining_ttl(&self, now: u64) -> u32 {
         let remaining_millis = get_expiration_time(self.created, self.ttl, 100) - now;
diff --git a/src/service_daemon.rs b/src/service_daemon.rs
index ae23d354..43e533c5 100644
--- a/src/service_daemon.rs
+++ b/src/service_daemon.rs
@@ -2127,23 +2127,73 @@ impl DnsCache {
     fn refresh_due(&mut self, ty_domain: &str) -> Vec<String> {
         let now = current_time_millis();
 
-        self.ptr
-            .get_mut(ty_domain)
-            .into_iter()
-            .flatten()
-            .filter_map(|record| {
-                let rec = record.get_record_mut();
-                if rec.is_expired(now) || !rec.refresh_due(now) {
-                    return None;
+        let mut due_list = vec![];
+        let mut not_due = vec![];
+
+        // find PTR records that are due for refresh.
+        for record in self.ptr.get_mut(ty_domain).into_iter().flatten() {
+            let is_due = record.get_record_mut().refresh_maybe(now);
+
+            if let Some(ptr) = record.any().downcast_ref::<DnsPointer>() {
+                let instance = ptr.alias.clone();
+                if is_due {
+                    due_list.push(instance);
+                } else {
+                    not_due.push(instance);
                 }
-                rec.refresh_no_more();
+            }
+        }
 
-                record
-                    .any()
-                    .downcast_ref::<DnsPointer>()
-                    .map(|dns_ptr| dns_ptr.alias.clone())
-            })
-            .collect()
+        // find other records that are due for refresh.
+        for instance in not_due {
+            if self.refresh_due_srv(&instance, now) || self.refresh_due_txt(&instance, now) {
+                due_list.push(instance);
+            }
+        }
+
+        due_list
+    }
+
+    /// Checks if any SRV records of `instance` are due to refresh.
+    /// If no, then check if any A or AAAA records are due to refresh.
+    fn refresh_due_srv(&mut self, instance: &str, now: u64) -> bool {
+        let mut not_due = vec![];
+        for record in self.srv.get_mut(instance).into_iter().flatten() {
+            if record.get_record_mut().refresh_maybe(now) {
+                return true;
+            }
+            if let Some(srv) = record.any().downcast_ref::<DnsSrv>() {
+                not_due.push(srv.host.clone());
+            }
+        }
+
+        for record in not_due {
+            if self.refresh_due_address(&record, now) {
+                return true;
+            }
+        }
+
+        false
+    }
+
+    /// Returns true if any A or AAAA records of `srv_name` are due to refresh.
+    fn refresh_due_address(&mut self, srv_name: &str, now: u64) -> bool {
+        for record in self.addr.get_mut(srv_name).into_iter().flatten() {
+            if record.get_record_mut().refresh_maybe(now) {
+                return true;
+            }
+        }
+        false
+    }
+
+    /// Returns true if any TXT records of `instance` are due to refresh.
+    fn refresh_due_txt(&mut self, instance: &str, now: u64) -> bool {
+        for record in self.txt.get_mut(instance).into_iter().flatten() {
+            if record.get_record_mut().refresh_maybe(now) {
+                return true;
+            }
+        }
+        false
     }
 }