More tests and fixes

CHANGELOG.md

@@ -1,4 +1,8 @@
 # Changelog
+## v0.1.4
+* Another fix for the routes with different protocols
+* Added more tests to detect more possible regressions
+
 ## v0.1.3
 * Fix issue with new route overwriting existing one on same address but different protocol
 

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "iptables-proxy"
-version = "0.1.1"
+version = "0.1.4"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

src/lib.rs

@@ -3,6 +3,48 @@ use std::borrow::Cow;
 
 mod rule_parser;
 
+pub struct Routes {
+    routes: Vec<ForwardingRoute>,
+}
+
+impl Routes {
+    pub fn new() -> Self {
+        Self { routes: Vec::new() }
+    }
+
+    pub fn add(&mut self, route: ForwardingRoute) -> Option<ForwardingRoute> {
+        if let Some((idx, _)) = self.routes.iter().enumerate().find(|(_, r)| {
+            r.public_ip() == route.public_ip()
+                && r.public_port() == route.public_port()
+                && r.protocol() == route.protocol()
+        }) {
+            let previous = self.routes.swap_remove(idx);
+
+            self.routes.push(route);
+
+            Some(previous)
+        } else {
+            self.routes.push(route);
+
+            None
+        }
+    }
+
+    pub fn remove(
+        &mut self,
+        public_ip: &str,
+        public_port: u16,
+        protocol: &Protocol,
+    ) -> Option<ForwardingRoute> {
+        match self.routes.iter().enumerate().find(|(_, r)| {
+            r.public_ip() == public_ip && r.public_port() == public_port && r.protocol() == protocol
+        }) {
+            Some((i, _)) => Some(self.routes.swap_remove(i)),
+            None => None,
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub struct ForwardingRoute {
     pub_ip: String,
@@ -165,6 +207,39 @@ mod tests {
     use super::*;
 
     #[test]
-    #[ignore = "This is not yet implemented properly"]
-    fn parse_rules() {}
+    fn add_route_to_empty_routes() {
+        let mut routes = Routes::new();
+
+        let prev = routes.add(ForwardingRoute {
+            pub_ip: "".into(),
+            pub_port: 1234,
+            dest_ip: "".into(),
+            dest_port: 1234,
+            protocol: Protocol::Tcp,
+        });
+        assert_eq!(None, prev);
+    }
+
+    #[test]
+    fn add_route_with_different_protocol() {
+        let mut routes = Routes::new();
+
+        let prev = routes.add(ForwardingRoute {
+            pub_ip: "".into(),
+            pub_port: 1234,
+            dest_ip: "".into(),
+            dest_port: 1234,
+            protocol: Protocol::Tcp,
+        });
+        assert_eq!(None, prev);
+
+        let prev = routes.add(ForwardingRoute {
+            pub_ip: "".into(),
+            pub_port: 1234,
+            dest_ip: "".into(),
+            dest_port: 1234,
+            protocol: Protocol::Udp,
+        });
+        assert_eq!(None, prev);
+    }
 }

src/main.rs

@@ -36,7 +36,7 @@ impl<S> Filter<S> for CrateFilter {
 }
 
 struct AppState {
-    routes: Mutex<Vec<ForwardingRoute>>,
+    routes: Mutex<iptables_proxy::Routes>,
     dry_run: bool,
     public_ip: String,
 }
@@ -45,36 +45,13 @@ impl AppState {
     pub fn add(&self, route: ForwardingRoute) -> Option<ForwardingRoute> {
         let mut existing_routes = self.routes.lock().unwrap();
 
-        if let Some((idx, _)) = existing_routes.iter().enumerate().find(|(_, r)| {
-            r.public_ip() == route.public_ip()
-                && r.public_port() == route.public_port()
-                && r.protocol() == route.protocol()
-        }) {
-            tracing::error!("Route already exists, replacing existing route");
-            let old_route = existing_routes.remove(idx);
-
-            existing_routes.push(route.clone());
-            drop(existing_routes);
-
-            Some(old_route)
-        } else {
-            existing_routes.push(route.clone());
-
-            None
-        }
+        existing_routes.add(route)
     }
 
     pub fn remove(&self, req: &RemoveRequest) -> Option<ForwardingRoute> {
         let mut existing_routes = self.routes.lock().unwrap();
 
-        match existing_routes
-            .iter()
-            .enumerate()
-            .find(|(_, r)| r.public_ip() == self.public_ip && r.public_port() == req.public_port)
-        {
-            Some((i, _)) => Some(existing_routes.remove(i)),
-            None => None,
-        }
+        existing_routes.remove(&self.public_ip, req.public_port, &req.protocol)
     }
 }
 
@@ -96,7 +73,7 @@ async fn main() {
     }
 
     let app_state = Arc::new(AppState {
-        routes: Mutex::new(Vec::new()),
+        routes: Mutex::new(iptables_proxy::Routes::new()),
         dry_run: args.dry_run,
         public_ip: format!("{}", args.public_ip),
     });
@@ -166,6 +143,7 @@ async fn create(
 #[derive(Debug, Deserialize)]
 struct RemoveRequest {
     public_port: u16,
+    protocol: Protocol,
 }
 
 #[tracing::instrument(skip(state, payload))]