Detect arrival and removal of volumes

4 files changed, 112 insertions, 19 deletions

diff --git a/talpid-core/src/split_tunnel/windows/driver.rs b/talpid-core/src/split_tunnel/windows/driver.rs
index 8f7c7198f6..38bb7b8eab 100644
--- a/talpid-core/src/split_tunnel/windows/driver.rs
+++ b/talpid-core/src/split_tunnel/windows/driver.rs
@@ -316,10 +316,9 @@ impl DeviceHandle {
         for app in apps.as_ref() {
             match get_device_path(app.as_ref()) {
                 Err(error) if error.kind() == io::ErrorKind::NotFound => {
-                    log::warn!(
+                    log::debug!(
                         "{}\nPath: {}",
-                        error
-                            .display_chain_with_msg("Skipping path with non-existent drive letter"),
+                        error.display_chain_with_msg("Ignoring path on unmounted volume"),
                         app.as_ref().to_string_lossy()
                     );
                 }
diff --git a/talpid-core/src/split_tunnel/windows/mod.rs b/talpid-core/src/split_tunnel/windows/mod.rs
index 7bb2eb1bab..2dee48a2bb 100644
--- a/talpid-core/src/split_tunnel/windows/mod.rs
+++ b/talpid-core/src/split_tunnel/windows/mod.rs
@@ -1,5 +1,6 @@
 mod driver;
 mod path_monitor;
+mod volume_monitor;
 mod windows;
 
 use crate::{
@@ -329,12 +330,17 @@ impl SplitTunnel {
         let (tx, rx): (RequestTx, _) = sync_mpsc::channel();
         let (init_tx, init_rx) = sync_mpsc::channel();
 
-        let (path_monitor, path_change_rx) =
-            path_monitor::PathMonitor::spawn().map_err(Error::StartPathMonitor)?;
-
         let monitored_paths = Arc::new(Mutex::new(vec![]));
         let monitored_paths_copy = monitored_paths.clone();
 
+        let (monitor_tx, monitor_rx) = sync_mpsc::channel();
+
+        let mut volume_monitor =
+            volume_monitor::VolumeMonitor::spawn(monitor_tx.clone(), monitored_paths.clone());
+
+        let path_monitor =
+            path_monitor::PathMonitor::spawn(monitor_tx).map_err(Error::StartPathMonitor)?;
+
         std::thread::spawn(move || {
             let result = driver::DeviceHandle::new()
                 .map(Arc::new)
@@ -401,6 +407,7 @@ impl SplitTunnel {
                     error.display_chain_with_msg("Failed to shut down path monitor")
                 );
             }
+            volume_monitor.close();
         });
 
         let handle = init_rx
@@ -410,7 +417,7 @@ impl SplitTunnel {
         let handle_copy = handle.clone();
 
         std::thread::spawn(move || {
-            while let Ok(()) = path_change_rx.recv() {
+            while let Ok(()) = monitor_rx.recv() {
                 let paths = monitored_paths_copy.lock().unwrap();
                 let result = if paths.len() > 0 {
                     log::debug!("Re-resolving excluded paths");
diff --git a/talpid-core/src/split_tunnel/windows/path_monitor.rs b/talpid-core/src/split_tunnel/windows/path_monitor.rs
index 6906e574b4..e8b82ed8ff 100644
--- a/talpid-core/src/split_tunnel/windows/path_monitor.rs
+++ b/talpid-core/src/split_tunnel/windows/path_monitor.rs
@@ -472,8 +472,6 @@ impl PathMonitorHandle {
     }
 }
 
-pub type PathChangeNotifyRx = sync_mpsc::Receiver<()>;
-
 enum PathMonitorCommand {
     Shutdown,
     SetPaths(Vec<PathBuf>),
@@ -487,7 +485,7 @@ pub struct PathMonitor {
 }
 
 impl PathMonitor {
-    pub fn spawn() -> io::Result<(PathMonitorHandle, PathChangeNotifyRx)> {
+    pub fn spawn(update_notify_tx: sync_mpsc::Sender<()>) -> io::Result<PathMonitorHandle> {
         let port_handle = Arc::new(CompletionPort::create(0)?);
         let mut original_paths: Vec<PathBuf> = vec![];
 
@@ -499,7 +497,6 @@ impl PathMonitor {
         };
 
         let (cmd_tx, cmd_rx) = sync_mpsc::channel();
-        let (notify_tx, notify_rx) = sync_mpsc::channel();
 
         std::thread::spawn(move || {
             loop {
@@ -509,7 +506,7 @@ impl PathMonitor {
                 match monitor.handle_next_completion_packet() {
                     Ok(true) => match monitor.update_paths(&original_paths) {
                         Ok(true) => {
-                            let _ = notify_tx.send(());
+                            let _ = update_notify_tx.send(());
                         }
                         Ok(false) => (),
                         Err(_) => break,
@@ -526,13 +523,10 @@ impl PathMonitor {
             monitor.abort_all_requests();
         });
 
-        Ok((
-            PathMonitorHandle {
-                port_handle,
-                tx: cmd_tx,
-            },
-            notify_rx,
-        ))
+        Ok(PathMonitorHandle {
+            port_handle,
+            tx: cmd_tx,
+        })
     }
 
     fn service_commands(
diff --git a/talpid-core/src/split_tunnel/windows/volume_monitor.rs b/talpid-core/src/split_tunnel/windows/volume_monitor.rs
new file mode 100644
index 0000000000..e405a670b8
--- /dev/null
+++ b/talpid-core/src/split_tunnel/windows/volume_monitor.rs
@@ -0,0 +1,93 @@
+//! Used to monitor volume mounts and dismounts, and reapply the split
+//! tunnel config if any of the excluded paths are affected by them.
+use crate::windows::window::{create_hidden_window, WindowCloseHandle};
+use std::{
+    ffi::OsString,
+    path::{self, Path},
+    sync::{mpsc as sync_mpsc, Arc, Mutex},
+};
+use winapi::{
+    shared::minwindef::TRUE,
+    um::{
+        dbt::{
+            DBTF_NET, DBT_DEVICEARRIVAL, DBT_DEVICEREMOVECOMPLETE, DBT_DEVTYP_VOLUME,
+            DEV_BROADCAST_HDR, DEV_BROADCAST_VOLUME, WM_DEVICECHANGE,
+        },
+        winuser::DefWindowProcW,
+    },
+};
+
+pub(super) struct VolumeMonitor(());
+
+impl VolumeMonitor {
+    pub fn spawn(
+        update_tx: sync_mpsc::Sender<()>,
+        paths: Arc<Mutex<Vec<OsString>>>,
+    ) -> WindowCloseHandle {
+        create_hidden_window(move |window, message, w_param, l_param| {
+            if message != WM_DEVICECHANGE
+                || (w_param != DBT_DEVICEARRIVAL && w_param != DBT_DEVICEREMOVECOMPLETE)
+            {
+                return unsafe { DefWindowProcW(window, message, w_param, l_param) };
+            }
+
+            let paths_guard = paths.lock().unwrap();
+            let mut label_found = false;
+
+            let volumes = unsafe { parse_broadcast(&*(l_param as *const _)) };
+            for volume in volumes {
+                for path in &*paths_guard {
+                    let path = (path as &dyn AsRef<Path>).as_ref();
+                    if let Some(path::Component::Prefix(prefix)) = path.components().next() {
+                        match prefix.kind() {
+                            path::Prefix::VerbatimDisk(disk) | path::Prefix::Disk(disk) => {
+                                if disk == volume {
+                                    label_found = true;
+                                    break;
+                                }
+                            }
+                            _ => (),
+                        }
+                    }
+                }
+                if label_found {
+                    break;
+                }
+            }
+
+            if label_found {
+                // Reapply config
+                let _ = update_tx.send(());
+            }
+
+            // Always grant the request
+            TRUE as isize
+        })
+    }
+}
+
+/// Return volume labels (ASCII-encoded) affected by the device arrival or removal message, if any.
+unsafe fn parse_broadcast(broadcast: &DEV_BROADCAST_HDR) -> Vec<u8> {
+    let mut labels = vec![];
+
+    if broadcast.dbch_devicetype != DBT_DEVTYP_VOLUME {
+        return labels;
+    }
+
+    let volume_broadcast = &*(broadcast as *const _ as *const DEV_BROADCAST_VOLUME);
+    if volume_broadcast.dbcv_flags & DBTF_NET != 0 {
+        // Ignore net event
+        return labels;
+    }
+
+    // 26 = 1 + 'Z' - 'A'
+    let num_drives = 1 + 'Z' as u8 - 'A' as u8;
+    for i in 0..num_drives {
+        let is_affected = ((volume_broadcast.dbcv_unitmask >> i) & 1) != 0;
+        if is_affected {
+            labels.push('A' as u8 + i);
+        }
+    }
+
+    labels
+}