Fix custom DNS address being restored incorrectly

1 files changed, 77 insertions, 33 deletions

diff --git a/talpid-core/src/dns/macos.rs b/talpid-core/src/dns/macos.rs
index b125bbaf29..ea70951b33 100644
--- a/talpid-core/src/dns/macos.rs
+++ b/talpid-core/src/dns/macos.rs
@@ -84,7 +84,7 @@ impl State {
         match &self.dns_settings {
             None => {
                 self.dns_settings = Some(new_settings);
-                self.update_known_state(store);
+                self.update_and_apply_state(store);
             }
             Some(old_settings) => {
                 if new_settings.address_set() != old_settings.address_set() {
@@ -99,54 +99,98 @@ impl State {
         Ok(())
     }
 
-    /// Apply the desired DNS settings to all interfaces, and save the original state. The operation
-    /// is idempotent.
-    fn update_known_state(&mut self, store: &SCDynamicStore) {
-        let Some(expected_settings) = &self.dns_settings else {
+    /// Store changes to the DNS config, ignoring any changes that we have applied. Then apply our
+    /// desired state to any services to which it has not already been applied.
+    fn update_and_apply_state(&mut self, store: &SCDynamicStore) {
+        let actual_state = read_all_dns(store);
+        self.update_backup_state(&actual_state);
+        self.apply_desired_state(store, &actual_state);
+    }
+
+    /// Store changes to the DNS config, ignoring any changes that we have applied. The operation is
+    /// idempotent.
+    fn update_backup_state(&mut self, actual_state: &HashMap<ServicePath, Option<DnsSettings>>) {
+        let Some(ref desired_settings) = self.dns_settings else {
             return;
         };
 
-        let new_settings = read_all_dns(store);
-        let mut prev_settings = mem::take(&mut self.backup);
+        let prev_state = mem::take(&mut self.backup);
+        let desired_set = desired_settings.address_set();
 
-        for (path, settings) in new_settings {
-            let old_entry = prev_settings.remove(&path).flatten();
+        self.backup = Self::merge_states(actual_state, prev_state, desired_set);
+    }
 
-            let should_set_dns = match settings {
-                Some(settings) => {
-                    if settings.address_set() != expected_settings.address_set() {
-                        let servers = settings.server_addresses().join(",");
-                        log::debug!("Saving DNS settings [{}] for {}", servers, path);
-                        self.backup.insert(path.to_owned(), Some(settings));
-                        true
-                    } else {
-                        self.backup.insert(path.to_owned(), old_entry);
-                        false
-                    }
-                }
-                None => {
-                    self.backup.insert(path.to_owned(), None);
-                    true
-                }
-            };
+    /// Merge `new_state` set by the OS with a previous `prev_state`, but ignore any service whose
+    /// addresses are `ignore_addresses`.
+    fn merge_states(
+        new_state: &HashMap<ServicePath, Option<DnsSettings>>,
+        mut prev_state: HashMap<ServicePath, Option<DnsSettings>>,
+        ignore_addresses: BTreeSet<String>,
+    ) -> HashMap<ServicePath, Option<DnsSettings>> {
+        let mut modified_state = HashMap::new();
 
-            if should_set_dns {
-                let path_cf = CFString::new(&path);
-                if let Err(e) = expected_settings.save(store, path_cf) {
-                    log::error!("Failed changing DNS for {}: {}", path, e);
+        for (path, settings) in new_state {
+            let old_entry = prev_state.remove(path);
+            match settings {
+                // If the service is using the desired addresses, don't save changes
+                Some(settings) if settings.address_set() == ignore_addresses => {
+                    let settings = old_entry.unwrap_or_else(|| Some(settings.to_owned()));
+                    modified_state.insert(path.to_owned(), settings);
+                }
+                // Otherwise, save the new settings
+                settings => {
+                    let servers = settings
+                        .as_ref()
+                        .map(|settings| settings.server_addresses().join(","))
+                        .unwrap_or_default();
+                    log::debug!("Saving DNS settings [{}] for {}", servers, path);
+                    modified_state.insert(path.to_owned(), settings.to_owned());
                 }
             }
         }
 
-        for path in prev_settings.keys() {
+        for path in prev_state.keys() {
             log::debug!("DNS removed for {path}");
         }
+
+        modified_state
+    }
+
+    /// Apply the desired addresses to all network services. The operation is idempotent.
+    fn apply_desired_state(
+        &mut self,
+        store: &SCDynamicStore,
+        actual_state: &HashMap<ServicePath, Option<DnsSettings>>,
+    ) {
+        let Some(ref desired_settings) = self.dns_settings else {
+            return;
+        };
+        let desired_set = desired_settings.address_set();
+
+        for (path, settings) in actual_state {
+            match settings {
+                // Do nothing if the state is already what we want
+                Some(settings) if settings.address_set() == desired_set => (),
+                // Apply desired state to service
+                _ => {
+                    let path_cf = CFString::new(path);
+                    if let Err(e) = desired_settings.save(store, path_cf) {
+                        log::error!("Failed changing DNS for {}: {}", path, e);
+                    }
+                }
+            }
+        }
     }
 
     fn reset(&mut self, store: &SCDynamicStore) -> Result<()> {
         log::trace!("Restoring DNS settings to: {:#?}", self.backup);
-        let old_backup = std::mem::take(&mut self.backup);
+
+        let actual_state = read_all_dns(store);
+        self.update_backup_state(&actual_state);
         self.dns_settings.take();
+
+        let old_backup = std::mem::take(&mut self.backup);
+
         for (service_path, settings) in old_backup {
             if let Some(settings) = settings {
                 settings.save(store, service_path.as_str())?;
@@ -359,7 +403,7 @@ fn create_dynamic_store(state: Arc<Mutex<State>>) -> Result<SCDynamicStore> {
         BURST_LONGEST_BUFFER_PERIOD,
         move || {
             if let Some(store) = &*store_container.read().unwrap() {
-                state.lock().update_known_state(&store.store);
+                state.lock().update_and_apply_state(&store.store);
             }
         },
     );