Add tests for RIB insertion/removal

Signed-off-by: Trey Aspelund <[email protected]>

rdb/Cargo.toml

@@ -12,6 +12,8 @@ schemars.workspace = true
 serde_json.workspace = true
 thiserror.workspace = true
 slog.workspace = true
+slog-bunyan.workspace = true
+slog-async.workspace = true
 mg-common.workspace = true
 itertools.workspace = true
 chrono.workspace = true

rdb/src/db.rs

@@ -379,6 +379,15 @@ impl Db {
         }
     }
 
+    pub fn get_selected_prefix_paths(&self, prefix: &Prefix) -> Vec<Path> {
+        let rib = lock!(self.rib_loc);
+        let paths = rib.get(prefix);
+        match paths {
+            None => Vec::new(),
+            Some(p) => p.iter().cloned().collect(),
+        }
+    }
+
     pub fn update_loc_rib(rib_in: &Rib, rib_loc: &mut Rib, prefix: Prefix) {
         let bp = bestpaths(prefix, rib_in, BESTPATH_FANOUT);
         match bp {
@@ -701,3 +710,190 @@ impl Reaper {
             });
     }
 }
+
+#[cfg(test)]
+mod test {
+    use crate::StaticRouteKey;
+    use crate::{
+        db::Db, BgpPathProperties, Path, Prefix, Prefix4,
+        DEFAULT_RIB_PRIORITY_BGP, DEFAULT_RIB_PRIORITY_STATIC,
+    };
+    use slog::{Drain, Logger};
+    use std::fs::File;
+    use std::io::Write;
+
+    pub fn init_file_logger(filename: &str) -> Logger {
+        build_logger(File::create(filename).expect("build logger"))
+    }
+
+    pub fn build_logger<W: Write + Send + 'static>(w: W) -> Logger {
+        let drain = slog_bunyan::new(w).build().fuse();
+        let drain = slog_async::Async::new(drain)
+            .chan_size(0x8000)
+            .build()
+            .fuse();
+        slog::Logger::root(drain, slog::o!())
+    }
+
+    #[test]
+    fn test_rib() {
+        // init test vars
+        let p0 = Prefix::from("192.168.0.0/24".parse::<Prefix4>().unwrap());
+        let p1 = Prefix::from("192.168.1.0/24".parse::<Prefix4>().unwrap());
+        let p2 = Prefix::from("192.168.2.0/24".parse::<Prefix4>().unwrap());
+        let nh0 = "203.0.113.0";
+        let nh1 = "203.0.113.1";
+        let bgp_path0 = Path {
+            nexthop: nh0.parse().unwrap(),
+            rib_priority: DEFAULT_RIB_PRIORITY_BGP,
+            shutdown: false,
+            bgp: Some(BgpPathProperties {
+                origin_as: 1111,
+                id: 1111,
+                med: Some(1111),
+                local_pref: Some(1111),
+                as_path: vec![1111, 1111, 1111],
+                stale: None,
+            }),
+            vlan_id: None,
+        };
+        let bgp_path1 = Path {
+            nexthop: nh1.parse().unwrap(),
+            rib_priority: DEFAULT_RIB_PRIORITY_BGP,
+            shutdown: false,
+            bgp: Some(BgpPathProperties {
+                origin_as: 2222,
+                id: 2222,
+                med: Some(2222),
+                local_pref: Some(2222),
+                as_path: vec![2222, 2222, 2222],
+                stale: None,
+            }),
+            vlan_id: None,
+        };
+        let static_key0 = StaticRouteKey {
+            prefix: p0,
+            nexthop: nh0.parse().unwrap(),
+            vlan_id: None,
+            rib_priority: DEFAULT_RIB_PRIORITY_STATIC,
+        };
+        let static_path0 = Path::from(static_key0);
+        let static_key1 = StaticRouteKey {
+            prefix: p0,
+            nexthop: nh0.parse().unwrap(),
+            vlan_id: None,
+            rib_priority: DEFAULT_RIB_PRIORITY_STATIC + 10,
+        };
+        let static_path1 = Path::from(static_key1);
+
+        // setup
+        let log = init_file_logger(&"/tmp/rib.log");
+        let db_path = "/tmp/rb.db".to_string();
+        let _ = std::fs::remove_dir_all(&db_path);
+        let db = Db::new(&db_path, log.clone()).expect("create db");
+
+        // Start test cases
+
+        // start from empty rib
+        assert!(db.full_rib().is_empty());
+        assert!(db.loc_rib().is_empty());
+
+        // both paths have the same next-hop, but not all fields
+        // from StaticRouteKey match (rib_priority is different).
+        db.add_static_routes(&vec![static_key0, static_key1])
+            .expect(
+                "add_static_routes failed for {static_key0} and {static_key1}",
+            );
+
+        // both paths should get installed in rib_in, but only
+        // static_path0 should get into rib_loc.
+        assert_eq!(
+            db.get_prefix_paths(&p0),
+            vec![static_path0.clone(), static_path1.clone()]
+        );
+        assert_eq!(
+            db.get_selected_prefix_paths(&p0),
+            vec![static_path0.clone()]
+        );
+
+        // rib_priority differs, so removal of static_key0
+        // should not affect path from static_key1
+        db.remove_static_routes(vec![static_key0])
+            .expect("remove_static_routes_failed for {static_key0}");
+        assert_eq!(db.get_prefix_paths(&p0), vec![static_path1.clone()]);
+        assert_eq!(
+            db.get_selected_prefix_paths(&p0),
+            vec![static_path1.clone()]
+        );
+
+        // install bgp routes
+        db.add_bgp_prefixes(vec![p0, p1], bgp_path0.clone());
+        db.add_bgp_prefixes(vec![p1, p2], bgp_path1.clone());
+
+        // p0 via paths static_path0 and bgp_path0 in rib_in
+        assert_eq!(
+            db.get_prefix_paths(&p0),
+            vec![static_path1.clone(), bgp_path0.clone()]
+        );
+        // p0 via path static_path0 in rib_loc
+        assert_eq!(
+            db.get_selected_prefix_paths(&p0),
+            vec![static_path1.clone()]
+        );
+
+        // p1 via both bgp_path0 and bgp_path1 in rib_in
+        assert_eq!(
+            db.get_prefix_paths(&p1),
+            vec![bgp_path0.clone(), bgp_path1.clone()]
+        );
+        // p1 via bgp_path1 in rib_loc (highest local_pref)
+        assert_eq!(db.get_selected_prefix_paths(&p1), vec![bgp_path1.clone()]);
+
+        // p2 via bgp_path1 in rib_in
+        assert_eq!(db.get_prefix_paths(&p2), vec![bgp_path1.clone()]);
+        // p2 via bgp_path1 in rib_loc
+        assert_eq!(db.get_selected_prefix_paths(&p2), vec![bgp_path1.clone()]);
+
+        // withdrawal of p2 via bgp_path1
+        db.remove_bgp_prefixes(vec![p2], bgp_path1.bgp.clone().unwrap().id);
+        assert!(db.get_prefix_paths(&p2).is_empty());
+        assert!(db.get_selected_prefix_paths(&p2).is_empty());
+        // p1 is unaffected
+        assert_eq!(
+            db.get_prefix_paths(&p1),
+            vec![bgp_path0.clone(), bgp_path1.clone()]
+        );
+        assert_eq!(db.get_selected_prefix_paths(&p1), vec![bgp_path1.clone()]);
+
+        // yank all routes from bgp_path0, simulating peer shutdown
+        db.remove_bgp_peer_prefixes(bgp_path0.bgp.clone().unwrap().id);
+        // p0 via static_path1 in rib_in and rib_loc
+        assert_eq!(db.get_prefix_paths(&p0), vec![static_path1.clone()]);
+        assert_eq!(
+            db.get_selected_prefix_paths(&p0),
+            vec![static_path1.clone()]
+        );
+        // p1 via bgp_path1 in rib_in and rib_loc
+        assert_eq!(db.get_prefix_paths(&p1), vec![bgp_path1.clone()]);
+        assert_eq!(db.get_selected_prefix_paths(&p1), vec![bgp_path1.clone()]);
+        // p2 not present, test that key is missing
+        assert!(db.get_prefix_paths(&p2).is_empty());
+        assert!(db.get_selected_prefix_paths(&p2).is_empty());
+
+        // withdrawal of p1 via bgp_path1
+        db.remove_bgp_prefixes(vec![p1], bgp_path1.bgp.clone().unwrap().id);
+        assert!(db.get_prefix_paths(&p1).is_empty());
+        assert!(db.get_selected_prefix_paths(&p1).is_empty());
+
+        // removal of final static route (from static_key1) should result
+        // in the prefix being completely deleted
+        db.remove_static_routes(vec![static_key1])
+            .expect("remove_static_routes_failed for {static_key1}");
+        assert!(db.get_prefix_paths(&p0).is_empty());
+        assert!(db.get_selected_prefix_paths(&p0).is_empty());
+
+        // rib should be empty again
+        assert!(db.full_rib().is_empty());
+        assert!(db.loc_rib().is_empty());
+    }
+}