Force caller to deal with default route callback handle

6 files changed, 26 insertions, 35 deletions

diff --git a/talpid-core/src/routing/windows.rs b/talpid-core/src/routing/windows.rs
index cba05cb936..be812ff99a 100644
--- a/talpid-core/src/routing/windows.rs
+++ b/talpid-core/src/routing/windows.rs
@@ -24,6 +24,9 @@ pub enum Error {
     /// Failure to clear routes
     #[error(display = "Failed to clear applied routes")]
     ClearRoutesFailed,
+    /// WinNet returned an error while adding default route callback
+    #[error(display = "Failed to set callback for default route")]
+    FailedToAddDefaultRouteCallback,
     /// Attempt to use route manager that has been dropped
     #[error(display = "Cannot send message to route manager since it is down")]
     RouteManagerDown,
@@ -33,7 +36,6 @@ pub type Result<T> = std::result::Result<T, Error>;
 
 /// Manages routes by calling into WinNet
 pub struct RouteManager {
-    callback_handles: Vec<winnet::WinNetCallbackHandle>,
     runtime: tokio::runtime::Handle,
     manage_tx: Option<UnboundedSender<RouteManagerCommand>>,
 }
@@ -73,7 +75,6 @@ impl RouteManager {
         }
         let (manage_tx, manage_rx) = mpsc::unbounded();
         let manager = Self {
-            callback_handles: vec![],
             runtime: runtime.clone(),
             manage_tx: Some(manage_tx),
         };
@@ -129,32 +130,16 @@ impl RouteManager {
     }
 
     /// Sets a callback that is called whenever the default route changes.
-    #[cfg(target_os = "windows")]
     pub fn add_default_route_callback<T: 'static>(
         &mut self,
         callback: Option<winnet::DefaultRouteChangedCallback>,
         context: T,
-    ) {
+    ) -> Result<winnet::WinNetCallbackHandle> {
         if self.manage_tx.is_none() {
-            return;
+            return Err(Error::RouteManagerDown);
         }
-
-        match winnet::add_default_route_change_callback(callback, context) {
-            Err(_e) => {
-                // not sure if this should panic
-                log::error!("Failed to add callback!");
-            }
-            Ok(handle) => {
-                self.callback_handles.push(handle);
-            }
-        }
-    }
-
-    /// 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();
+        winnet::add_default_route_change_callback(callback, context)
+            .map_err(|_| Error::FailedToAddDefaultRouteCallback)
     }
 
     /// Stops the routing manager and invalidates the route manager - no new default route callbacks
@@ -165,7 +150,6 @@ impl RouteManager {
                 log::error!("RouteManager channel already down or thread panicked");
             }
 
-            self.callback_handles.clear();
             winnet::deactivate_routing_manager();
         }
     }
diff --git a/talpid-core/src/tunnel/wireguard/mod.rs b/talpid-core/src/tunnel/wireguard/mod.rs
index 899589dbc4..47c86e4e5f 100644
--- a/talpid-core/src/tunnel/wireguard/mod.rs
+++ b/talpid-core/src/tunnel/wireguard/mod.rs
@@ -90,6 +90,8 @@ pub struct WireguardMonitor {
     stop_setup_tx: Option<futures::channel::oneshot::Sender<()>>,
     pinger_stop_sender: mpsc::Sender<()>,
     _tcp_proxies: Vec<TcpProxy>,
+    #[cfg(target_os = "windows")]
+    _callback_handle: Option<crate::winnet::WinNetCallbackHandle>,
 }
 
 #[cfg(target_os = "linux")]
@@ -191,8 +193,13 @@ impl WireguardMonitor {
         runtime.block_on((on_event)(TunnelEvent::InterfaceUp(metadata.clone())));
 
         #[cfg(target_os = "windows")]
-        route_manager
-            .add_default_route_callback(Some(WgGoTunnel::default_route_changed_callback), ());
+        let callback_handle = route_manager
+            .add_default_route_callback(Some(WgGoTunnel::default_route_changed_callback), ())
+            .ok();
+        #[cfg(target_os = "windows")]
+        if callback_handle.is_none() {
+            log::warn!("Failed to register default route callback");
+        }
 
         let event_callback = Box::new(on_event.clone());
         let (close_msg_sender, close_msg_receiver) = mpsc::channel();
@@ -209,6 +216,8 @@ impl WireguardMonitor {
             stop_setup_tx: Some(stop_setup_tx),
             pinger_stop_sender: pinger_tx,
             _tcp_proxies: tcp_proxies,
+            #[cfg(target_os = "windows")]
+            _callback_handle: callback_handle,
         };
 
         let gateway = config.ipv4_gateway;
diff --git a/talpid-core/src/tunnel_state_machine/connected_state.rs b/talpid-core/src/tunnel_state_machine/connected_state.rs
index 2511713fc5..68766ec23a 100644
--- a/talpid-core/src/tunnel_state_machine/connected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connected_state.rs
@@ -139,8 +139,6 @@ 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!("{}", error.display_chain_with_msg("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 cb337c02ab..9d7f768207 100644
--- a/talpid-core/src/tunnel_state_machine/connecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/connecting_state.rs
@@ -216,8 +216,6 @@ 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!("{}", error.display_chain_with_msg("Failed to clear routes"));
         }
diff --git a/talpid-core/src/tunnel_state_machine/windows.rs b/talpid-core/src/tunnel_state_machine/windows.rs
index 3ea26bf819..52b81055af 100644
--- a/talpid-core/src/tunnel_state_machine/windows.rs
+++ b/talpid-core/src/tunnel_state_machine/windows.rs
@@ -69,13 +69,11 @@ pub(super) fn update_split_tunnel_addresses(
         .register_ips(tunnel_ipv4, tunnel_ipv6, internet_ipv4, internet_ipv6)
         .map_err(BoxedError::new)?;
 
-    // FIXME: Do this via the route manager
     shared_values.st_route_handler = Some(
-        crate::winnet::add_default_route_change_callback(
-            Some(split_tunnel_default_route_change_handler),
-            context,
-        )
-        .map_err(BoxedError::new)?,
+        shared_values
+            .route_manager
+            .add_default_route_callback(Some(split_tunnel_default_route_change_handler), context)
+            .map_err(BoxedError::new)?,
     );
 
     Ok(())
diff --git a/windows/winnet/src/winnet/winnet.cpp b/windows/winnet/src/winnet/winnet.cpp
index e82345ebee..9dda5cbbf2 100644
--- a/windows/winnet/src/winnet/winnet.cpp
+++ b/windows/winnet/src/winnet/winnet.cpp
@@ -524,6 +524,10 @@ WinNet_UnregisterDefaultRouteChangedCallback(
 

 	if (nullptr == g_RouteManager)

 	{

+		if (nullptr != g_RouteManagerLogSink)

+		{

+			g_RouteManagerLogSink->debug("Cannot remove default route callback for deactivated route manager");

+		}

 		return;

 	}