Merge branch 'linux-systemd-fix-shutdown'

2 files changed, 72 insertions, 55 deletions

diff --git a/talpid-core/src/dns/linux/systemd_resolved.rs b/talpid-core/src/dns/linux/systemd_resolved.rs
index 906efeb6be..29a8feaf85 100644
--- a/talpid-core/src/dns/linux/systemd_resolved.rs
+++ b/talpid-core/src/dns/linux/systemd_resolved.rs
@@ -1,13 +1,15 @@
 use crate::linux::{iface_index, IfaceIndexLookupError};
 use std::{
     net::IpAddr,
-    sync::{Arc, Mutex},
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
     thread,
 };
 use talpid_dbus::systemd_resolved::{DnsState, SystemdResolved as DbusInterface};
 
 pub(crate) use talpid_dbus::systemd_resolved::Error as SystemdDbusError;
-use talpid_types::ErrorExt;
 
 pub type Result<T> = std::result::Result<T, Error>;
 
@@ -22,8 +24,13 @@ pub enum Error {
 
 pub struct SystemdResolved {
     pub dbus_interface: DbusInterface,
-    state: Option<Arc<Mutex<Option<DnsState>>>>,
-    watcher_thread: Option<thread::JoinHandle<()>>,
+    state: Option<SetConfigState>,
+}
+
+struct SetConfigState {
+    dns_config: Arc<DnsState>,
+    watcher_thread: thread::JoinHandle<()>,
+    watcher_should_shutdown: Arc<AtomicBool>,
 }
 
 
@@ -34,7 +41,6 @@ impl SystemdResolved {
         let systemd_resolved = SystemdResolved {
             dbus_interface,
             state: None,
-            watcher_thread: None,
         };
 
         Ok(systemd_resolved)
@@ -43,65 +49,72 @@ impl SystemdResolved {
     pub fn set_dns(&mut self, interface_name: &str, servers: &[IpAddr]) -> Result<()> {
         let iface_index = iface_index(interface_name)?;
         let dns_state = self.dbus_interface.set_dns(iface_index, servers)?;
-        let cloned_dns_state = self.set_dns_state(dns_state);
-        let weak_dns_state = Arc::downgrade(&cloned_dns_state);
-        let dns_state_should_continue = weak_dns_state.clone();
+        let dns_config = Arc::new(dns_state);
 
-        let dbus_interface = self.dbus_interface.clone();
-        let mut applied_servers: Vec<_> = servers.iter().cloned().collect();
-        applied_servers.sort();
-        let applied_servers = Arc::new(applied_servers);
 
-        self.watcher_thread = Some(std::thread::spawn(move || {
+        let (watcher_thread, watcher_should_shutdown) =
+            self.spawn_watcher_thread(dns_config.clone());
+        self.state = Some(SetConfigState {
+            dns_config,
+            watcher_thread,
+            watcher_should_shutdown,
+        });
+
+
+        Ok(())
+    }
+
+    fn spawn_watcher_thread(
+        &mut self,
+        dns_state: Arc<DnsState>,
+    ) -> (thread::JoinHandle<()>, Arc<AtomicBool>) {
+        let dbus_interface = self.dbus_interface.clone();
+        let should_shutdown = Arc::new(AtomicBool::new(false));
+        let watch_shutdown = should_shutdown.clone();
+        let callback_shutdown = should_shutdown.clone();
+        let watcher_thread = std::thread::spawn(move || {
             let result = dbus_interface.clone().watch_dns_changes(
                 move |new_servers| {
-                    (|| {
-                        let dns_state_lock = weak_dns_state.upgrade()?;
-                        let dns_state = dns_state_lock.lock().ok()?;
-                        let dns_state_ref: &DnsState = &*dns_state.as_ref()?;
-
-                        let mut current_servers: Vec<IpAddr> = new_servers
-                                .into_iter()
-                                .filter(|server| server.iface_index == iface_index as i32)
-                                .map(|server| server.address)
-                                .collect();
-                        current_servers.sort();
-                        if current_servers != *dns_state_ref.set_servers {
-                            log::debug!("DNS config for tunnel interface changed, currently applied servers - {:?}", current_servers);
-                            if let Err(err) = dbus_interface.set_dns(iface_index, &applied_servers) {
-                                log::error!("Failed to re-apply DNS config - {}", err);
-                            }
+                    if callback_shutdown.clone().load(Ordering::Acquire) {
+                        return;
+                    }
+                    let mut current_servers: Vec<IpAddr> = new_servers
+                            .into_iter()
+                            .filter(|server| server.iface_index == dns_state.interface_index as i32)
+                            .map(|server| server.address)
+                            .collect();
+                    current_servers.sort();
+                    if current_servers != *dns_state.set_servers {
+                        log::debug!("DNS config for tunnel interface changed, currently applied servers - {:?}", current_servers);
+                        if let Err(err) = dbus_interface.set_dns(dns_state.interface_index, &dns_state.set_servers) {
+                            log::error!("Failed to re-apply DNS config - {}", err);
                         }
-                        Some(())
-                    })();
+                    }
                 },
-                || dns_state_should_continue.upgrade().is_some(),
+                move || !watch_shutdown.load(Ordering::Acquire),
             );
             if let Err(err) = result {
                 log::error!("Failed to watch DNS config updates: {}", err);
             }
-        }));
-        Ok(())
-    }
-
-    fn set_dns_state(&mut self, dns_state: DnsState) -> Arc<Mutex<Option<DnsState>>> {
-        let new_state = Arc::new(Mutex::new(Some(dns_state)));
-        self.state = Some(new_state.clone());
-        new_state
+        });
+        (watcher_thread, should_shutdown)
     }
 
     pub fn reset(&mut self) -> Result<()> {
-        if let Some(state_lock) = self.state.take() {
-            if let Some(dns_state) = state_lock.lock().expect("DNS state lock poisoned").take() {
-                if let Err(err) = self.dbus_interface.revert_link(dns_state) {
-                    log::error!("Failed to revert DNS config - {}", err.display_chain());
-                }
+        if let Some(SetConfigState {
+            dns_config,
+            watcher_thread,
+            watcher_should_shutdown,
+        }) = self.state.take()
+        {
+            watcher_should_shutdown.store(true, Ordering::Release);
+            if let Err(err) = self.dbus_interface.revert_link(&dns_config) {
+                log::error!("Failed to revert interface config: {}", err);
+            }
+
+            if watcher_thread.join().is_err() {
+                log::error!("DNS watcher thread panicked!");
             }
-        } else {
-            log::trace!("No DNS settings to reset");
-        }
-        if let Some(join_handle) = self.watcher_thread.take() {
-            let _ = join_handle.join();
         }
 
         Ok(())
diff --git a/talpid-dbus/src/systemd_resolved.rs b/talpid-dbus/src/systemd_resolved.rs
index e4e96c3b2f..2fd2f14956 100644
--- a/talpid-dbus/src/systemd_resolved.rs
+++ b/talpid-dbus/src/systemd_resolved.rs
@@ -271,8 +271,8 @@ impl SystemdResolved {
         .map_err(Error::SetDomainsError)
     }
 
-    pub fn revert_link(&mut self, dns_state: DnsState) -> std::result::Result<(), dbus::Error> {
-        let link = self.as_link_object(dns_state.interface_path);
+    pub fn revert_link(&mut self, dns_state: &DnsState) -> std::result::Result<(), dbus::Error> {
+        let link = self.as_link_object(dns_state.interface_path.clone());
 
         if let Err(error) = link.method_call::<(), _, _, _>(LINK_INTERFACE, REVERT_METHOD, ()) {
             if error.name() == Some("org.freedesktop.DBus.Error.UnknownObject") {
@@ -289,7 +289,10 @@ impl SystemdResolved {
         }
     }
 
-    pub fn watch_dns_changes<F: FnMut(Vec<DnsServer>) + Send + Sync + 'static, S: Fn() -> bool>(
+    pub fn watch_dns_changes<
+        F: FnMut(Vec<DnsServer>) + Send + Sync + 'static,
+        S: Fn() -> bool + Clone + Send + Sync + 'static,
+    >(
         &mut self,
         mut callback: F,
         should_continue: S,
@@ -298,6 +301,7 @@ impl SystemdResolved {
             MatchRule::new_signal(PropertiesPropertiesChanged::INTERFACE, DNS_SERVERS);
         match_rule.member = None;
         match_rule.path = Some(RESOLVED_MANAGER_PATH.into());
+        let should_continue_outer = should_continue.clone();
         let dns_matcher = self
             .dbus_connection
             .add_match(
@@ -319,12 +323,12 @@ impl SystemdResolved {
                             }
                         }
                     };
-                    true
+                    should_continue()
                 },
             )
             .map_err(Error::DnsUpdateMatchError)?;
 
-        while should_continue() {
+        while should_continue_outer() {
             if let Err(err) = self.dbus_connection.process(RPC_TIMEOUT) {
                 log::error!("Failed to process DBus messages: {}", err);
             }