Merge branch 'windows-st-monitor-drives'

7 files changed, 274 insertions, 149 deletions

diff --git a/talpid-core/src/offline/windows.rs b/talpid-core/src/offline/windows.rs
index c04277aea5..15b5d47169 100644
--- a/talpid-core/src/offline/windows.rs
+++ b/talpid-core/src/offline/windows.rs
@@ -1,41 +1,21 @@
-use crate::winnet;
+use crate::{
+    windows::window::{create_hidden_window, WindowCloseHandle},
+    winnet,
+};
 use futures::channel::mpsc::UnboundedSender;
 use parking_lot::Mutex;
 use std::{
     ffi::c_void,
     io,
-    mem::zeroed,
-    os::windows::io::{IntoRawHandle, RawHandle},
-    ptr,
     sync::{Arc, Weak},
     thread,
     time::Duration,
 };
 use talpid_types::ErrorExt;
-use winapi::{
-    shared::{
-        basetsd::LONG_PTR,
-        minwindef::{DWORD, LPARAM, LRESULT, UINT, WPARAM},
-        windef::HWND,
-    },
-    um::{
-        handleapi::CloseHandle,
-        libloaderapi::GetModuleHandleW,
-        processthreadsapi::GetThreadId,
-        synchapi::WaitForSingleObject,
-        winbase::INFINITE,
-        winuser::{
-            CreateWindowExW, DefWindowProcW, DestroyWindow, DispatchMessageW, GetMessageW,
-            GetWindowLongPtrW, PostQuitMessage, PostThreadMessageW, SetWindowLongPtrW,
-            GWLP_USERDATA, GWLP_WNDPROC, PBT_APMRESUMEAUTOMATIC, PBT_APMSUSPEND, WM_DESTROY,
-            WM_POWERBROADCAST, WM_USER,
-        },
-    },
+use winapi::um::winuser::{
+    DefWindowProcW, PBT_APMRESUMEAUTOMATIC, PBT_APMSUSPEND, WM_POWERBROADCAST,
 };
 
-const CLASS_NAME: &[u8] = b"S\0T\0A\0T\0I\0C\0\0\0";
-const REQUEST_THREAD_SHUTDOWN: UINT = WM_USER + 1;
-
 #[derive(err_derive::Error, Debug)]
 pub enum Error {
     #[error(display = "Unable to create listener thread")]
@@ -45,8 +25,7 @@ pub enum Error {
 }
 
 pub struct BroadcastListener {
-    thread_handle: RawHandle,
-    thread_id: DWORD,
+    window: WindowCloseHandle,
     system_state: Arc<Mutex<SystemState>>,
     _callback_handle: winnet::WinNetCallbackHandle,
     _notify_tx: Arc<UnboundedSender<bool>>,
@@ -67,7 +46,7 @@ impl BroadcastListener {
 
         let power_broadcast_state_ref = system_state.clone();
 
-        let power_broadcast_callback = move |message: UINT, wparam: WPARAM, _lparam: LPARAM| {
+        let power_broadcast_callback = move |window, message, wparam, lparam| {
             let state = power_broadcast_state_ref.clone();
             if message == WM_POWERBROADCAST {
                 if wparam == PBT_APMSUSPEND {
@@ -83,22 +62,16 @@ impl BroadcastListener {
                     });
                 }
             }
+            unsafe { DefWindowProcW(window, message, wparam, lparam) }
         };
 
-        let join_handle = thread::Builder::new()
-            .spawn(move || unsafe {
-                Self::message_pump(power_broadcast_callback);
-            })
-            .map_err(Error::ThreadCreationError)?;
-
-        let real_handle = join_handle.into_raw_handle();
+        let window = create_hidden_window(power_broadcast_callback);
 
         let callback_handle =
             unsafe { Self::setup_network_connectivity_listener(system_state.clone())? };
 
         Ok(BroadcastListener {
-            thread_handle: real_handle,
-            thread_id: unsafe { GetThreadId(real_handle) },
+            window,
             system_state,
             _callback_handle: callback_handle,
             _notify_tx: notify_tx,
@@ -131,88 +104,6 @@ impl BroadcastListener {
         (v4_connectivity, v6_connectivity)
     }
 
-    unsafe fn message_pump<F>(client_callback: F)
-    where
-        F: Fn(UINT, WPARAM, LPARAM),
-    {
-        let dummy_window = CreateWindowExW(
-            0,
-            CLASS_NAME.as_ptr() as *const u16,
-            ptr::null_mut(),
-            0,
-            0,
-            0,
-            0,
-            0,
-            ptr::null_mut(),
-            ptr::null_mut(),
-            GetModuleHandleW(ptr::null_mut()),
-            ptr::null_mut(),
-        );
-
-        // Move callback information to the heap.
-        // This enables us to reach the callback through a "thin pointer".
-        let boxed_callback = Box::new(client_callback);
-
-        // Detach callback from Box.
-        let raw_callback = Box::into_raw(boxed_callback) as *mut c_void;
-
-        SetWindowLongPtrW(dummy_window, GWLP_USERDATA, raw_callback as LONG_PTR);
-        SetWindowLongPtrW(
-            dummy_window,
-            GWLP_WNDPROC,
-            Self::window_procedure::<F> as LONG_PTR,
-        );
-
-        let mut msg = zeroed();
-
-        loop {
-            let status = GetMessageW(&mut msg, 0 as HWND, 0, 0);
-
-            if status < 0 {
-                continue;
-            }
-            if status == 0 {
-                break;
-            }
-
-            if msg.hwnd.is_null() {
-                if msg.message == REQUEST_THREAD_SHUTDOWN {
-                    DestroyWindow(dummy_window);
-                }
-            } else {
-                DispatchMessageW(&mut msg);
-            }
-        }
-
-        // Reattach callback to Box for proper clean-up.
-        let _ = Box::from_raw(raw_callback as *mut F);
-    }
-
-    unsafe extern "system" fn window_procedure<F>(
-        window: HWND,
-        message: UINT,
-        wparam: WPARAM,
-        lparam: LPARAM,
-    ) -> LRESULT
-    where
-        F: Fn(UINT, WPARAM, LPARAM),
-    {
-        let raw_callback = GetWindowLongPtrW(window, GWLP_USERDATA);
-
-        if raw_callback != 0 {
-            let typed_callback = &mut *(raw_callback as *mut F);
-            typed_callback(message, wparam, lparam);
-        }
-
-        if message == WM_DESTROY {
-            PostQuitMessage(0);
-            return 0;
-        }
-
-        DefWindowProcW(window, message, wparam, lparam)
-    }
-
     /// The caller must make sure the `system_state` reference is valid
     /// until after `WinNet_DeactivateConnectivityMonitor` has been called.
     unsafe fn setup_network_connectivity_listener(
@@ -252,11 +143,7 @@ impl BroadcastListener {
 
 impl Drop for BroadcastListener {
     fn drop(&mut self) {
-        unsafe {
-            PostThreadMessageW(self.thread_id, REQUEST_THREAD_SHUTDOWN, 0, 0);
-            WaitForSingleObject(self.thread_handle, INFINITE);
-            CloseHandle(self.thread_handle);
-        }
+        self.window.close();
     }
 }
 
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..34c9a67864 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,19 @@ 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 path_monitor = path_monitor::PathMonitor::spawn(monitor_tx.clone())
+            .map_err(Error::StartPathMonitor)?;
+        let mut volume_monitor = volume_monitor::VolumeMonitor::spawn(
+            path_monitor.clone(),
+            monitor_tx,
+            monitored_paths.clone(),
+        );
+
         std::thread::spawn(move || {
             let result = driver::DeviceHandle::new()
                 .map(Arc::new)
@@ -395,6 +403,7 @@ impl SplitTunnel {
                 }
             }
 
+            volume_monitor.close();
             if let Err(error) = path_monitor.shutdown() {
                 log::error!(
                     "{}",
@@ -410,7 +419,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..b27deb1e21 100644
--- a/talpid-core/src/split_tunnel/windows/path_monitor.rs
+++ b/talpid-core/src/split_tunnel/windows/path_monitor.rs
@@ -445,6 +445,7 @@ impl StrippedPath {
     }
 }
 
+#[derive(Clone)]
 pub struct PathMonitorHandle {
     port_handle: Arc<CompletionPort>,
     tx: sync_mpsc::Sender<PathMonitorCommand>,
@@ -458,6 +459,11 @@ impl PathMonitorHandle {
         self.notify_monitor()
     }
 
+    pub fn refresh(&self) -> io::Result<()> {
+        let _ = self.tx.send(PathMonitorCommand::Refresh);
+        self.notify_monitor()
+    }
+
     pub fn shutdown(&self) -> io::Result<()> {
         let _ = self.tx.send(PathMonitorCommand::Shutdown);
         self.notify_monitor()
@@ -472,11 +478,10 @@ impl PathMonitorHandle {
     }
 }
 
-pub type PathChangeNotifyRx = sync_mpsc::Receiver<()>;
-
 enum PathMonitorCommand {
-    Shutdown,
     SetPaths(Vec<PathBuf>),
+    Refresh,
+    Shutdown,
 }
 
 pub struct PathMonitor {
@@ -487,7 +492,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 +504,6 @@ impl PathMonitor {
         };
 
         let (cmd_tx, cmd_rx) = sync_mpsc::channel();
-        let (notify_tx, notify_rx) = sync_mpsc::channel();
 
         std::thread::spawn(move || {
             loop {
@@ -507,9 +511,9 @@ impl PathMonitor {
                     break;
                 }
                 match monitor.handle_next_completion_packet() {
-                    Ok(true) => match monitor.update_paths(&original_paths) {
+                    Ok(true) => match monitor.update_paths(&original_paths, false) {
                         Ok(true) => {
-                            let _ = notify_tx.send(());
+                            let _ = update_notify_tx.send(());
                         }
                         Ok(false) => (),
                         Err(_) => break,
@@ -526,13 +530,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(
@@ -547,20 +548,23 @@ impl PathMonitor {
                 }
                 PathMonitorCommand::SetPaths(new_paths) => {
                     *original_paths = new_paths;
-                    return !self.update_paths(&original_paths).is_err();
+                    return !self.update_paths(&original_paths, false).is_err();
+                }
+                PathMonitorCommand::Refresh => {
+                    return !self.update_paths(&original_paths, true).is_err();
                 }
             }
         }
         true
     }
 
-    fn update_paths(&mut self, unresolved_paths: &[PathBuf]) -> Result<bool, ()> {
+    fn update_paths(&mut self, unresolved_paths: &[PathBuf], force: bool) -> Result<bool, ()> {
         let resolved_paths = resolve_all_links_multiple(unresolved_paths);
         let new_stripped_paths = resolved_paths
             .iter()
             .filter_map(|p| StrippedPath::new(p).ok())
             .collect();
-        if new_stripped_paths != self.stripped_paths {
+        if force || new_stripped_paths != self.stripped_paths {
             self.stripped_paths = new_stripped_paths;
             if let Err(error) = self.update_directory_contexts() {
                 log::error!("Failed to open new directory handles: {}", error);
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..e9060e5528
--- /dev/null
+++ b/talpid-core/src/split_tunnel/windows/volume_monitor.rs
@@ -0,0 +1,96 @@
+//! Used to monitor volume mounts and dismounts, and reapply the split
+//! tunnel config if any of the excluded paths are affected by them.
+use super::path_monitor::PathMonitorHandle;
+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(
+        path_monitor: PathMonitorHandle,
+        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(());
+                let _ = path_monitor.refresh();
+            }
+
+            // 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
+}
diff --git a/talpid-core/src/windows.rs b/talpid-core/src/windows/mod.rs
index 1febc5bc21..4ec1fe19f3 100644
--- a/talpid-core/src/windows.rs
+++ b/talpid-core/src/windows/mod.rs
@@ -31,6 +31,8 @@ use winapi::shared::{
     ws2ipdef::{SOCKADDR_IN6_LH as sockaddr_in6, SOCKADDR_INET},
 };
 
+pub mod window;
+
 /// Result type for this module.
 pub type Result<T> = std::result::Result<T, Error>;
 
diff --git a/talpid-core/src/windows/window.rs b/talpid-core/src/windows/window.rs
new file mode 100644
index 0000000000..05e457e5d6
--- /dev/null
+++ b/talpid-core/src/windows/window.rs
@@ -0,0 +1,128 @@
+//! Utilities for windows.
+
+use std::{os::windows::io::AsRawHandle, ptr, thread};
+use winapi::{
+    shared::{
+        basetsd::LONG_PTR,
+        minwindef::{LPARAM, LRESULT, UINT, WPARAM},
+        windef::HWND,
+    },
+    um::{
+        libloaderapi::GetModuleHandleW,
+        processthreadsapi::GetThreadId,
+        winuser::{
+            CreateWindowExW, DefWindowProcW, DestroyWindow, DispatchMessageW, GetMessageW,
+            GetWindowLongPtrW, PostQuitMessage, PostThreadMessageW, SetWindowLongPtrW,
+            TranslateMessage, GWLP_USERDATA, GWLP_WNDPROC, WM_DESTROY, WM_USER,
+        },
+    },
+};
+
+const CLASS_NAME: &[u8] = b"S\0T\0A\0T\0I\0C\0\0\0";
+const REQUEST_THREAD_SHUTDOWN: UINT = WM_USER + 1;
+
+/// Handle for closing an associated window.
+/// The window is not destroyed when this is dropped.
+pub struct WindowCloseHandle {
+    thread: Option<std::thread::JoinHandle<()>>,
+}
+
+impl WindowCloseHandle {
+    /// Close the window and wait for the thread.
+    pub fn close(&mut self) {
+        if let Some(thread) = self.thread.take() {
+            let thread_id = unsafe { GetThreadId(thread.as_raw_handle()) };
+            unsafe { PostThreadMessageW(thread_id, REQUEST_THREAD_SHUTDOWN, 0, 0) };
+            let _ = thread.join();
+        }
+    }
+}
+
+/// Creates a dummy window whose messages are handled by `wnd_proc`.
+pub fn create_hidden_window<F: (Fn(HWND, UINT, WPARAM, LPARAM) -> LRESULT) + Send + 'static>(
+    wnd_proc: F,
+) -> WindowCloseHandle {
+    let join_handle = thread::spawn(move || {
+        let dummy_window = unsafe {
+            CreateWindowExW(
+                0,
+                CLASS_NAME.as_ptr() as *const u16,
+                ptr::null_mut(),
+                0,
+                0,
+                0,
+                0,
+                0,
+                ptr::null_mut(),
+                ptr::null_mut(),
+                GetModuleHandleW(ptr::null_mut()),
+                ptr::null_mut(),
+            )
+        };
+
+        // Move callback information to the heap.
+        // This enables us to reach the callback through a "thin pointer".
+        let raw_callback = Box::into_raw(Box::new(wnd_proc));
+
+        unsafe {
+            SetWindowLongPtrW(dummy_window, GWLP_USERDATA, raw_callback as LONG_PTR);
+            SetWindowLongPtrW(
+                dummy_window,
+                GWLP_WNDPROC,
+                window_procedure::<F> as LONG_PTR,
+            );
+        }
+
+        let mut msg = unsafe { std::mem::zeroed() };
+
+        loop {
+            let status = unsafe { GetMessageW(&mut msg, ptr::null_mut(), 0, 0) };
+
+            if status < 0 {
+                continue;
+            }
+            if status == 0 {
+                break;
+            }
+
+            if msg.hwnd.is_null() {
+                if msg.message == REQUEST_THREAD_SHUTDOWN {
+                    unsafe { DestroyWindow(dummy_window) };
+                }
+            } else {
+                unsafe {
+                    TranslateMessage(&mut msg);
+                    DispatchMessageW(&mut msg);
+                }
+            }
+        }
+
+        // Free callback.
+        let _ = unsafe { Box::from_raw(raw_callback) };
+    });
+
+    WindowCloseHandle {
+        thread: Some(join_handle),
+    }
+}
+
+unsafe extern "system" fn window_procedure<F>(
+    window: HWND,
+    message: UINT,
+    wparam: WPARAM,
+    lparam: LPARAM,
+) -> LRESULT
+where
+    F: Fn(HWND, UINT, WPARAM, LPARAM) -> LRESULT,
+{
+    if message == WM_DESTROY {
+        PostQuitMessage(0);
+        return 0;
+    }
+    let raw_callback = GetWindowLongPtrW(window, GWLP_USERDATA);
+    if raw_callback != 0 {
+        let typed_callback = &mut *(raw_callback as *mut F);
+        return typed_callback(window, message, wparam, lparam);
+    }
+    DefWindowProcW(window, message, wparam, lparam)
+}