Add RouteManager::add_routes

2 files changed, 64 insertions, 35 deletions

diff --git a/talpid-core/src/routing/linux.rs b/talpid-core/src/routing/linux.rs
index f128c71fc9..b1087bfdc0 100644
--- a/talpid-core/src/routing/linux.rs
+++ b/talpid-core/src/routing/linux.rs
@@ -155,31 +155,12 @@ impl RouteManagerImplInner {
 
         let iface_map = Self::initialize_link_map(&handle).await?;
 
-
-        let mut required_normal_routes = HashSet::new();
-        let mut required_default_routes = HashSet::new();
-
-        for route in required_routes {
-            match route.node {
-                NetNode::RealNode(node) => {
-                    required_normal_routes
-                        .insert(Route::new(node, route.prefix).table(route.table_id));
-                }
-                NetNode::DefaultNode => {
-                    required_default_routes.insert(RequiredDefaultRoute {
-                        table_id: route.table_id,
-                        destination: route.prefix,
-                    });
-                }
-            }
-        }
-
         let mut monitor = Self {
             iface_map,
             handle,
             messages,
 
-            required_default_routes,
+            required_default_routes: HashSet::new(),
             added_routes: HashSet::new(),
 
             default_routes: HashSet::new(),
@@ -193,24 +174,49 @@ impl RouteManagerImplInner {
         monitor.best_default_node_v6 =
             Self::pick_best_default_node(&monitor.default_routes, IpVersion::V6);
 
+        monitor.add_required_routes(required_routes).await?;
+
+        Ok(monitor)
+    }
+
+    async fn add_required_routes(&mut self, required_routes: HashSet<RequiredRoute>) -> Result<()> {
+        let mut required_normal_routes = HashSet::new();
+        let mut required_default_routes = HashSet::new();
+
+        for route in required_routes {
+            match route.node {
+                NetNode::RealNode(node) => {
+                    required_normal_routes
+                        .insert(Route::new(node, route.prefix).table(route.table_id));
+                }
+                NetNode::DefaultNode => {
+                    required_default_routes.insert(RequiredDefaultRoute {
+                        table_id: route.table_id,
+                        destination: route.prefix,
+                    });
+                }
+            }
+        }
 
         for normal_route in required_normal_routes.into_iter() {
-            monitor.add_route(normal_route).await?;
+            self.add_route(normal_route).await?;
         }
 
-        for route in monitor.required_default_routes.clone().into_iter() {
+        for route in required_default_routes.into_iter() {
+            self.required_default_routes.insert(route);
             if let (false, _, Some(default_node)) | (true, Some(default_node), _) = (
                 route.destination.is_ipv4(),
-                &monitor.best_default_node_v4,
-                &monitor.best_default_node_v6,
+                &self.best_default_node_v4,
+                &self.best_default_node_v6,
             ) {
                 // best to pick a single node identifier rather than device + ip
                 let new_route =
                     Route::new(default_node.clone(), route.destination).table(route.table_id);
-                monitor.add_route(new_route).await?;
+                self.add_route(new_route).await?;
             }
         }
-        Ok(monitor)
+
+        Ok(())
     }
 
     async fn get_default_routes(&self) -> Result<HashSet<Route>> {
@@ -417,15 +423,7 @@ impl RouteManagerImplInner {
         loop {
             futures::select! {
                 command = manage_rx.select_next_some().fuse() => {
-                    match command {
-                        RouteManagerCommand::Shutdown(shutdown_signal) => {
-                            log::trace!("Shutting down route manager");
-                            self.cleanup_routes().await;
-                            log::trace!("Route manager done");
-                            let _ = shutdown_signal.send(());
-                            return Ok(());
-                        }
-                    }
+                    self.process_command(command).await?;
                 },
                 (route_change, socket) = self.messages.select_next_some().fuse() => {
                     self.process_netlink_message(route_change).await?;
@@ -434,6 +432,23 @@ impl RouteManagerImplInner {
         }
     }
 
+    async fn process_command(&mut self, command: RouteManagerCommand) -> Result<()> {
+        match command {
+            RouteManagerCommand::Shutdown(shutdown_signal) => {
+                log::trace!("Shutting down route manager");
+                self.cleanup_routes().await;
+                log::trace!("Route manager done");
+                let _ = shutdown_signal.send(());
+                return Ok(());
+            }
+            RouteManagerCommand::AddRoutes(routes) => {
+                log::debug!("Adding routes: {:?}", routes);
+                self.add_required_routes(routes).await?;
+            }
+        }
+        Ok(())
+    }
+
     async fn process_netlink_message(&mut self, msg: NetlinkMessage<RtnlMessage>) -> Result<()> {
         match msg.payload {
             NetlinkPayload::InnerMessage(RtnlMessage::NewLink(new_link)) => {
diff --git a/talpid-core/src/routing/unix.rs b/talpid-core/src/routing/unix.rs
index 5b123d180f..f15877f494 100644
--- a/talpid-core/src/routing/unix.rs
+++ b/talpid-core/src/routing/unix.rs
@@ -37,10 +37,14 @@ pub enum Error {
     /// Failed to spawn route manager future
     #[error(display = "Failed to spawn route manager on the provided executor")]
     FailedToSpawnManager,
+    /// Attempt to use route manager that has been dropped
+    #[error(display = "Cannot send message to route manager since it is down")]
+    RouteManagerDown,
 }
 
 #[derive(Debug)]
 pub enum RouteManagerCommand {
+    AddRoutes(HashSet<RequiredRoute>),
     Shutdown(oneshot::Sender<()>),
 }
 
@@ -99,6 +103,16 @@ impl RouteManager {
             }
         }
     }
+
+    /// Applies the given routes until [`RouteManager::stop`] is called.
+    pub fn add_routes(&mut self, routes: HashSet<RequiredRoute>) -> Result<(), Error> {
+        if let Some(tx) = &self.manage_tx {
+            let _ = tx.unbounded_send(RouteManagerCommand::AddRoutes(routes));
+            Ok(())
+        } else {
+            Err(Error::RouteManagerDown)
+        }
+    }
 }
 
 impl Drop for RouteManager {