Return a handle for fetching excluded processes instead of using a tunnel command

7 files changed, 48 insertions, 61 deletions

diff --git a/talpid-core/src/split_tunnel/windows/mod.rs b/talpid-core/src/split_tunnel/windows/mod.rs
index 8b679106ea..fc75e80331 100644
--- a/talpid-core/src/split_tunnel/windows/mod.rs
+++ b/talpid-core/src/split_tunnel/windows/mod.rs
@@ -86,8 +86,8 @@ pub enum Error {
     RequestThreadStuck,
 
     /// The request handling thread is down
-    #[error(display = "The ST request thread is down")]
-    RequestThreadDown,
+    #[error(display = "The split tunnel monitor is down")]
+    SplitTunnelDown,
 
     /// Failed to start the NTFS reparse point monitor
     #[error(display = "Failed to start path monitor")]
@@ -104,7 +104,7 @@ pub struct SplitTunnel {
     request_tx: RequestTx,
     event_thread: Option<std::thread::JoinHandle<()>>,
     quit_event: Arc<QuitEvent>,
-    excluded_pids: Arc<RwLock<HashMap<usize, ExcludedProcess>>>,
+    excluded_processes: Arc<RwLock<HashMap<usize, ExcludedProcess>>>,
     _route_change_callback: Option<WinNetCallbackHandle>,
     daemon_tx: Weak<mpsc::UnboundedSender<TunnelCommand>>,
     async_path_update_in_progress: Arc<AtomicBool>,
@@ -163,11 +163,30 @@ pub struct ExcludedProcess {
     pub inherited: bool,
 }
 
+/// Cloneable handle for interacting with the split tunnel module.
+#[derive(Debug, Clone)]
+pub struct SplitTunnelHandle {
+    excluded_processes: Weak<RwLock<HashMap<usize, ExcludedProcess>>>,
+}
+
+impl SplitTunnelHandle {
+    /// Return processes that are currently being excluded, including
+    /// their pids, paths, and reason for being excluded.
+    pub fn get_processes(&self) -> Result<Vec<ExcludedProcess>, Error> {
+        let processes = self
+            .excluded_processes
+            .upgrade()
+            .ok_or(Error::SplitTunnelDown)?;
+        let processes = processes.read().unwrap();
+        Ok(processes.values().cloned().collect())
+    }
+}
+
 struct EventThreadContext {
     handle: Arc<driver::DeviceHandle>,
     event_overlapped: OVERLAPPED,
     quit_event: Arc<QuitEvent>,
-    excluded_pids: Arc<RwLock<HashMap<usize, ExcludedProcess>>>,
+    excluded_processes: Arc<RwLock<HashMap<usize, ExcludedProcess>>>,
 }
 unsafe impl Send for EventThreadContext {}
 
@@ -188,13 +207,13 @@ impl SplitTunnel {
         }
 
         let quit_event = Arc::new(QuitEvent::new());
-        let excluded_pids = Arc::new(RwLock::new(HashMap::new()));
+        let excluded_processes = Arc::new(RwLock::new(HashMap::new()));
 
         let event_context = EventThreadContext {
             handle: handle.clone(),
             event_overlapped,
             quit_event: quit_event.clone(),
-            excluded_pids: excluded_pids.clone(),
+            excluded_processes: excluded_processes.clone(),
         };
 
         let event_thread = std::thread::spawn(move || {
@@ -309,7 +328,7 @@ impl SplitTunnel {
                         reason,
                         image,
                     } => {
-                        let mut pids = event_context.excluded_pids.write().unwrap();
+                        let mut pids = event_context.excluded_processes.write().unwrap();
                         match event_id {
                             EventId::StartSplittingProcess => {
                                 if let Some(prev_entry) = pids.get(&process_id) {
@@ -372,7 +391,7 @@ impl SplitTunnel {
             _route_change_callback: None,
             daemon_tx,
             async_path_update_in_progress: Arc::new(AtomicBool::new(false)),
-            excluded_pids,
+            excluded_processes,
         })
     }
 
@@ -511,7 +530,7 @@ impl SplitTunnel {
 
         request_tx
             .send((request, response_tx))
-            .map_err(|_| Error::RequestThreadDown)?;
+            .map_err(|_| Error::SplitTunnelDown)?;
 
         response_rx
             .recv_timeout(REQUEST_TIMEOUT)
@@ -553,8 +572,8 @@ impl SplitTunnel {
         let wait_task = move || {
             request_tx
                 .send((request, response_tx))
-                .map_err(|_| Error::RequestThreadDown)?;
-            response_rx.recv().map_err(|_| Error::RequestThreadDown)?
+                .map_err(|_| Error::SplitTunnelDown)?;
+            response_rx.recv().map_err(|_| Error::SplitTunnelDown)?
         };
         let in_progress = self.async_path_update_in_progress.clone();
         self.runtime.spawn_blocking(move || {
@@ -610,15 +629,11 @@ impl SplitTunnel {
         self.send_request(Request::RegisterIps(InterfaceAddresses::default()))
     }
 
-    /// Return processes that are currently being excluded.
-    pub fn get_processes(&self) -> Result<Vec<ExcludedProcess>, Error> {
-        Ok(self
-            .excluded_pids
-            .read()
-            .unwrap()
-            .values()
-            .cloned()
-            .collect())
+    /// Returns a handle used for interacting with the split tunnel module.
+    pub fn handle(&self) -> SplitTunnelHandle {
+        SplitTunnelHandle {
+            excluded_processes: Arc::downgrade(&self.excluded_processes),
+        }
     }
 }
 
diff --git a/talpid-core/src/tunnel_state_machine/connected_state.rs b/talpid-core/src/tunnel_state_machine/connected_state.rs
index beb29590c1..80e5e28957 100644
--- a/talpid-core/src/tunnel_state_machine/connected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connected_state.rs
@@ -277,11 +277,6 @@ impl ConnectedState {
                 shared_values.split_tunnel.set_paths(&paths, result_tx);
                 SameState(self.into())
             }
-            #[cfg(windows)]
-            Some(TunnelCommand::GetExcludedProcesses(result_tx)) => {
-                let _ = result_tx.send(shared_values.split_tunnel.get_processes());
-                SameState(self.into())
-            }
         }
     }
 
diff --git a/talpid-core/src/tunnel_state_machine/connecting_state.rs b/talpid-core/src/tunnel_state_machine/connecting_state.rs
index 64d21086cd..7536b26b09 100644
--- a/talpid-core/src/tunnel_state_machine/connecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -390,11 +390,6 @@ impl ConnectingState {
                 shared_values.split_tunnel.set_paths(&paths, result_tx);
                 SameState(self.into())
             }
-            #[cfg(windows)]
-            Some(TunnelCommand::GetExcludedProcesses(result_tx)) => {
-                let _ = result_tx.send(shared_values.split_tunnel.get_processes());
-                SameState(self.into())
-            }
         }
     }
 
diff --git a/talpid-core/src/tunnel_state_machine/disconnected_state.rs b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
index 281b22b8f2..fa91284f6b 100644
--- a/talpid-core/src/tunnel_state_machine/disconnected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
@@ -208,11 +208,6 @@ impl TunnelState for DisconnectedState {
                 shared_values.split_tunnel.set_paths(&paths, result_tx);
                 SameState(self.into())
             }
-            #[cfg(windows)]
-            Some(TunnelCommand::GetExcludedProcesses(result_tx)) => {
-                let _ = result_tx.send(shared_values.split_tunnel.get_processes());
-                SameState(self.into())
-            }
             None => {
                 Self::reset_dns(shared_values);
                 Finished
diff --git a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
index 2c42a1a864..2a14d881d6 100644
--- a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
@@ -57,11 +57,6 @@ impl DisconnectingState {
                     shared_values.split_tunnel.set_paths(&paths, result_tx);
                     AfterDisconnect::Nothing
                 }
-                #[cfg(windows)]
-                Some(TunnelCommand::GetExcludedProcesses(result_tx)) => {
-                    let _ = result_tx.send(shared_values.split_tunnel.get_processes());
-                    AfterDisconnect::Nothing
-                }
             },
             AfterDisconnect::Block(reason) => match command {
                 Some(TunnelCommand::AllowLan(allow_lan)) => {
@@ -102,11 +97,6 @@ impl DisconnectingState {
                     shared_values.split_tunnel.set_paths(&paths, result_tx);
                     AfterDisconnect::Block(reason)
                 }
-                #[cfg(windows)]
-                Some(TunnelCommand::GetExcludedProcesses(result_tx)) => {
-                    let _ = result_tx.send(shared_values.split_tunnel.get_processes());
-                    AfterDisconnect::Block(reason)
-                }
                 None => AfterDisconnect::Block(reason),
             },
             AfterDisconnect::Reconnect(retry_attempt) => match command {
@@ -148,11 +138,6 @@ impl DisconnectingState {
                     shared_values.split_tunnel.set_paths(&paths, result_tx);
                     AfterDisconnect::Reconnect(retry_attempt)
                 }
-                #[cfg(windows)]
-                Some(TunnelCommand::GetExcludedProcesses(result_tx)) => {
-                    let _ = result_tx.send(shared_values.split_tunnel.get_processes());
-                    AfterDisconnect::Reconnect(retry_attempt)
-                }
             },
         };
 
diff --git a/talpid-core/src/tunnel_state_machine/error_state.rs b/talpid-core/src/tunnel_state_machine/error_state.rs
index d64393a432..a11e49d859 100644
--- a/talpid-core/src/tunnel_state_machine/error_state.rs
+++ b/talpid-core/src/tunnel_state_machine/error_state.rs
@@ -209,11 +209,6 @@ impl TunnelState for ErrorState {
                 shared_values.split_tunnel.set_paths(&paths, result_tx);
                 SameState(self.into())
             }
-            #[cfg(windows)]
-            Some(TunnelCommand::GetExcludedProcesses(result_tx)) => {
-                let _ = result_tx.send(shared_values.split_tunnel.get_processes());
-                SameState(self.into())
-            }
         }
     }
 }
diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs
index 16d6e6d55b..e55cc9c0d5 100644
--- a/talpid-core/src/tunnel_state_machine/mod.rs
+++ b/talpid-core/src/tunnel_state_machine/mod.rs
@@ -150,6 +150,9 @@ pub async fn spawn(
     )
     .await?;
 
+    #[cfg(windows)]
+    let split_tunnel = state_machine.shared_values.split_tunnel.handle();
+
     tokio::task::spawn_blocking(move || {
         state_machine.run(state_change_listener);
         if shutdown_tx.send(()).is_err() {
@@ -160,6 +163,8 @@ pub async fn spawn(
     Ok(TunnelStateMachineHandle {
         command_tx,
         shutdown_rx,
+        #[cfg(windows)]
+        split_tunnel,
     })
 }
 
@@ -192,12 +197,6 @@ pub enum TunnelCommand {
         oneshot::Sender<Result<(), split_tunnel::Error>>,
         Vec<OsString>,
     ),
-    /// Return a list of processes that are currently being split, as well as their
-    /// paths.
-    #[cfg(windows)]
-    GetExcludedProcesses(
-        oneshot::Sender<Result<Vec<split_tunnel::ExcludedProcess>, split_tunnel::Error>>,
-    ),
 }
 
 type TunnelCommandReceiver = stream::Fuse<mpsc::UnboundedReceiver<TunnelCommand>>;
@@ -603,6 +602,8 @@ state_wrapper! {
 pub struct TunnelStateMachineHandle {
     command_tx: Arc<mpsc::UnboundedSender<TunnelCommand>>,
     shutdown_rx: oneshot::Receiver<()>,
+    #[cfg(windows)]
+    split_tunnel: split_tunnel::SplitTunnelHandle,
 }
 
 impl TunnelStateMachineHandle {
@@ -621,4 +622,10 @@ impl TunnelStateMachineHandle {
     pub fn command_tx(&self) -> &Arc<mpsc::UnboundedSender<TunnelCommand>> {
         &self.command_tx
     }
+
+    /// Returns split tunnel object handle.
+    #[cfg(windows)]
+    pub fn split_tunnel(&self) -> &split_tunnel::SplitTunnelHandle {
+        &self.split_tunnel
+    }
 }