Merge branch 'macos-fix-custom-dns-reset'

2 files changed, 277 insertions, 33 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index af15f7ddf1..d37ca3025a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,10 @@ Line wrap the file at 100 chars.                                              Th
 ### Added
 - Add custom bridge settings in GUI.
 
+### Fixed
+#### macOS
+- DNS was not properly restored in some cases when using custom DNS.
+
 
 ## [2024.2-beta1] - 2024-04-15
 ### Added
diff --git a/talpid-core/src/dns/macos.rs b/talpid-core/src/dns/macos.rs
index b125bbaf29..e0535fedf7 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);
             }
         },
     );
@@ -447,3 +491,199 @@ fn state_to_setup_path(state_path: &str) -> Option<String> {
         None
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::{DnsSettings, State};
+    use std::collections::{BTreeSet, HashMap};
+
+    /// The initial backup should equal whatever the first provided state is.
+    #[test]
+    fn test_backup_new_dns_config() {
+        let prev_state = HashMap::new();
+
+        let new_state = HashMap::from([
+            ("a".to_owned(), None),
+            (
+                "b".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["1.2.3.4".to_owned()],
+                    "iface_b".to_owned(),
+                )),
+            ),
+            // One of our states already equals the desired state. It should be stored regardless.
+            (
+                "c".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["10.64.0.1".to_owned()],
+                    "iface_c".to_owned(),
+                )),
+            ),
+        ]);
+
+        let desired_addresses: BTreeSet<String> = ["10.64.0.1".to_owned()].into();
+
+        let merged_state = State::merge_states(&new_state, prev_state, desired_addresses);
+
+        assert_eq!(merged_state, new_state);
+    }
+
+    /// Any changes equal to the desired state should be ignored. Other changes should be recorded.
+    #[test]
+    fn test_backup_ignore_desired_state() {
+        let prev_state = HashMap::from([
+            ("a".to_owned(), None),
+            (
+                "b".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["1.2.3.4".to_owned()],
+                    "iface_b".to_owned(),
+                )),
+            ),
+            (
+                "c".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["10.64.0.1".to_owned()],
+                    "iface_c".to_owned(),
+                )),
+            ),
+            (
+                "d".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["1.3.3.7".to_owned()],
+                    "iface_d".to_owned(),
+                )),
+            ),
+        ]);
+        let new_state = HashMap::from([
+            // This change should be ignored
+            (
+                "a".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["10.64.0.1".to_owned()],
+                    "iface_a".to_owned(),
+                )),
+            ),
+            // This change should be ignored
+            (
+                "b".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["10.64.0.1".to_owned()],
+                    "iface_b".to_owned(),
+                )),
+            ),
+            // This change should be ignored
+            (
+                "c".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["4.3.2.1".to_owned()],
+                    "iface_c".to_owned(),
+                )),
+            ),
+            // This change should NOT be ignored
+            (
+                "d".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["4.3.2.1".to_owned()],
+                    "iface_d".to_owned(),
+                )),
+            ),
+        ]);
+        let expect_state = HashMap::from([
+            ("a".to_owned(), None),
+            (
+                "b".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["1.2.3.4".to_owned()],
+                    "iface_b".to_owned(),
+                )),
+            ),
+            (
+                "c".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["4.3.2.1".to_owned()],
+                    "iface_c".to_owned(),
+                )),
+            ),
+            (
+                "d".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["4.3.2.1".to_owned()],
+                    "iface_d".to_owned(),
+                )),
+            ),
+        ]);
+
+        let desired_addresses: BTreeSet<String> = ["10.64.0.1".to_owned()].into();
+
+        let merged_state = State::merge_states(&new_state, prev_state, desired_addresses);
+
+        assert_eq!(merged_state, expect_state);
+    }
+
+    /// Services not specified in the new state should be removed from the backed up state
+    #[test]
+    fn test_backup_remove_dns_config() {
+        let prev_state = HashMap::from([
+            (
+                "a".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["10.64.0.1".to_owned()],
+                    "iface_a".to_owned(),
+                )),
+            ),
+            (
+                "b".to_owned(),
+                Some(DnsSettings::from_server_addresses(
+                    &["1.2.3.4".to_owned()],
+                    "iface_b".to_owned(),
+                )),
+            ),
+            ("c".to_owned(), None),
+        ]);
+        let new_state = HashMap::from([("c".to_owned(), None)]);
+        let expected_state = new_state.clone();
+
+        let desired_addresses: BTreeSet<String> = ["10.64.0.1".to_owned()].into();
+
+        let merged_state = State::merge_states(&new_state, prev_state, desired_addresses);
+
+        assert_eq!(merged_state, expected_state);
+    }
+
+    /// If DHCP provides an IP identical to our desired state, the tracked state will not reflect
+    /// this. This is a known limitation.
+    // TODO: This should actually succeed. If we happen to switch to a network whose IP equals
+    //       the "desired IP", we should still back up the result.
+    #[test]
+    #[should_panic]
+    fn test_backup_change_equals_desired_state() {
+        let prev_state = HashMap::from([(
+            "a".to_owned(),
+            Some(DnsSettings::from_server_addresses(
+                &["192.168.100.1".to_owned()],
+                "iface_a".to_owned(),
+            )),
+        )]);
+        let new_state = HashMap::from([(
+            "a".to_owned(),
+            Some(DnsSettings::from_server_addresses(
+                &["192.168.1.1".to_owned()],
+                "iface_a".to_owned(),
+            )),
+        )]);
+        let expect_state = HashMap::from([(
+            "a".to_owned(),
+            Some(DnsSettings::from_server_addresses(
+                &["192.168.1.1".to_owned()],
+                "iface_a".to_owned(),
+            )),
+        )]);
+
+        let desired_addresses: BTreeSet<String> = ["192.168.1.1".to_owned()].into();
+
+        let merged_state = State::merge_states(&new_state, prev_state, desired_addresses);
+
+        assert_eq!(merged_state, expect_state);
+    }
+}