Update route manager on Windows

4 files changed, 64 insertions, 29 deletions

diff --git a/talpid-core/src/routing/windows.rs b/talpid-core/src/routing/windows.rs
index 0e75cab953..412384574b 100644
--- a/talpid-core/src/routing/windows.rs
+++ b/talpid-core/src/routing/windows.rs
@@ -5,9 +5,15 @@ use std::collections::HashSet;
 /// Windows routing errors.
 #[derive(err_derive::Error, Debug)]
 pub enum Error {
-    /// Failure to apply a route
+    /// Failure to initialize route manager
     #[error(display = "Failed to start route manager")]
     FailedToStartManager,
+    /// Failure to add routes
+    #[error(display = "Failed to add routes")]
+    AddRoutesFailed,
+    /// Failure to clear routes
+    #[error(display = "Failed to clear applied routes")]
+    ClearRoutesFailed,
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -22,27 +28,15 @@ impl RouteManager {
     /// Creates a new route manager that will apply the provided routes and ensure they exist until
     /// it's stopped.
     pub fn new(required_routes: HashSet<RequiredRoute>) -> Result<Self> {
-        let routes: Vec<_> = required_routes
-            .iter()
-            .map(|route| {
-                let destination = winnet::WinNetIpNetwork::from(route.prefix);
-                match &route.node {
-                    NetNode::DefaultNode => winnet::WinNetRoute::through_default_node(destination),
-                    NetNode::RealNode(node) => {
-                        winnet::WinNetRoute::new(winnet::WinNetNode::from(node), destination)
-                    }
-                }
-            })
-            .collect();
-
-        if !winnet::activate_routing_manager(&routes) {
+        if !winnet::activate_routing_manager() {
             return Err(Error::FailedToStartManager);
         }
-
-        Ok(Self {
+        let manager = Self {
             callback_handles: vec![],
             is_stopped: false,
-        })
+        };
+        manager.add_routes(required_routes)?;
+        Ok(manager)
     }
 
     /// Sets a callback that is called whenever the default route changes.
@@ -67,6 +61,13 @@ impl RouteManager {
         }
     }
 
+    /// Removes all routes previously applied in [`RouteManager::new`] or
+    /// [`RouteManager::add_routes`].
+    pub fn clear_default_route_callbacks(&mut self) {
+        // `WinNetCallbackHandle::drop` removes these callbacks.
+        self.callback_handles.clear();
+    }
+
     /// Stops the routing manager and invalidates the route manager - no new default route callbacks
     /// can be added
     pub fn stop(&mut self) {
@@ -76,6 +77,38 @@ impl RouteManager {
             self.is_stopped = true;
         }
     }
+
+    /// Applies the given routes until [`RouteManager::stop`] is called.
+    pub fn add_routes(&self, routes: HashSet<RequiredRoute>) -> Result<()> {
+        let routes: Vec<_> = routes
+            .iter()
+            .map(|route| {
+                let destination = winnet::WinNetIpNetwork::from(route.prefix);
+                match &route.node {
+                    NetNode::DefaultNode => winnet::WinNetRoute::through_default_node(destination),
+                    NetNode::RealNode(node) => {
+                        winnet::WinNetRoute::new(winnet::WinNetNode::from(node), destination)
+                    }
+                }
+            })
+            .collect();
+
+        if winnet::routing_manager_add_routes(&routes) {
+            Ok(())
+        } else {
+            Err(Error::AddRoutesFailed)
+        }
+    }
+
+    /// Removes all routes previously applied in [`RouteManager::new`] or
+    /// [`RouteManager::add_routes`].
+    pub fn clear_routes(&self) -> Result<()> {
+        if winnet::routing_manager_delete_applied_routes() {
+            Ok(())
+        } else {
+            Err(Error::ClearRoutesFailed)
+        }
+    }
 }
 
 impl Drop for RouteManager {
diff --git a/talpid-core/src/tunnel_state_machine/connected_state.rs b/talpid-core/src/tunnel_state_machine/connected_state.rs
index dab6e37619..03cbc41cf6 100644
--- a/talpid-core/src/tunnel_state_machine/connected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connected_state.rs
@@ -90,6 +90,8 @@ impl ConnectedState {
     }
 
     fn reset_routes(shared_values: &mut SharedTunnelStateValues) {
+        #[cfg(windows)]
+        shared_values.route_manager.clear_default_route_callbacks();
         if let Err(error) = shared_values.route_manager.clear_routes() {
             log::error!(
                 "Failed to clear routes: {:?}",
diff --git a/talpid-core/src/tunnel_state_machine/connecting_state.rs b/talpid-core/src/tunnel_state_machine/connecting_state.rs
index 4bd1345339..5d70568239 100644
--- a/talpid-core/src/tunnel_state_machine/connecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -170,6 +170,8 @@ impl ConnectingState {
     }
 
     fn reset_routes(shared_values: &mut SharedTunnelStateValues) {
+        #[cfg(windows)]
+        shared_values.route_manager.clear_default_route_callbacks();
         if let Err(error) = shared_values.route_manager.clear_routes() {
             log::error!(
                 "Failed to clear routes: {:?}",
diff --git a/talpid-core/src/winnet.rs b/talpid-core/src/winnet.rs
index a0b183d779..ab9dff5d06 100644
--- a/talpid-core/src/winnet.rs
+++ b/talpid-core/src/winnet.rs
@@ -289,17 +289,8 @@ impl Drop for WinNetRoute {
     }
 }
 
-pub fn activate_routing_manager(routes: &[WinNetRoute]) -> bool {
-    if unsafe { WinNet_ActivateRouteManager(Some(log_sink), logging_context()) } {
-        if routing_manager_add_routes(routes) {
-            true
-        } else {
-            deactivate_routing_manager();
-            false
-        }
-    } else {
-        false
-    }
+pub fn activate_routing_manager() -> bool {
+    unsafe { WinNet_ActivateRouteManager(Some(log_sink), logging_context()) }
 }
 
 pub struct WinNetCallbackHandle {
@@ -360,6 +351,10 @@ pub fn routing_manager_add_routes(routes: &[WinNetRoute]) -> bool {
     unsafe { WinNet_AddRoutes(ptr, length) }
 }
 
+pub fn routing_manager_delete_applied_routes() -> bool {
+    unsafe { WinNet_DeleteAppliedRoutes() }
+}
+
 pub fn deactivate_routing_manager() {
     unsafe { WinNet_DeactivateRouteManager() }
 }
@@ -400,6 +395,9 @@ mod api {
         // #[link_name = "WinNet_DeleteRoute"]
         // pub fn WinNet_DeleteRoute(route: *const super::WinNetRoute) -> bool;
 
+        #[link_name = "WinNet_DeleteAppliedRoutes"]
+        pub fn WinNet_DeleteAppliedRoutes() -> bool;
+
         #[link_name = "WinNet_DeactivateRouteManager"]
         pub fn WinNet_DeactivateRouteManager();