diff --git a/traffic_core/src/traffic/apply.rs b/traffic_core/src/traffic/apply.rs index d3b943c..660ccbb 100644 --- a/traffic_core/src/traffic/apply.rs +++ b/traffic_core/src/traffic/apply.rs @@ -6,7 +6,7 @@ use crate::{config::Config, netns::NsTransport, wg}; -use super::inspect::inspect_and_diff; +use super::inspect::{expected_peer_class_map, inspect_and_diff}; pub(super) fn teardown_owned_topology( transport: &dyn NsTransport, @@ -73,15 +73,11 @@ pub(super) fn apply_full_topology( ], )?; - for (idx, p) in peers.iter().enumerate() { - let idx_u16 = u16::try_from(idx) - .map_err(|_| anyhow::anyhow!("too many peers for class id allocation"))?; - let class_id = cfg - .traffic - .class_start_id - .checked_add(idx_u16) - .ok_or_else(|| anyhow::anyhow!("class id overflow: too many peers"))?; - let ip = p.ipv4.to_string(); + // The same map the inspector will verify against, so apply and verify cannot disagree + // about which class a peer belongs to. + let class_map = expected_peer_class_map(peers, cfg.traffic.class_start_id)?; + + for (ip, class_id) in &class_map { let class_egress = format!("1:{class_id}"); let class_ingress = format!("2:{class_id}"); let limit = cfg.traffic.peer_limit.as_str(); diff --git a/traffic_core/src/traffic/inspect.rs b/traffic_core/src/traffic/inspect.rs index 61d0244..b2daea1 100644 --- a/traffic_core/src/traffic/inspect.rs +++ b/traffic_core/src/traffic/inspect.rs @@ -323,3 +323,45 @@ fn inspect_cmd(transport: &dyn NsTransport, args: &[&str], what: &str) -> anyhow .run(args) .map_err(|e| anyhow::anyhow!("{what} inspect failed: {e}")) } + +#[cfg(test)] +mod tests { + use super::*; + use std::net::Ipv4Addr; + + fn peers(ips: &[&str]) -> Vec { + ips.iter() + .map(|v| wg::Peer { + ipv4: v.parse::().expect("ip"), + }) + .collect() + } + + /// Peers reach this function already sorted lexicographically by [`wg::discover_peers`], + /// so `.2` follows `.10`. Assignment must not reorder them. + #[test] + fn class_ids_follow_lexicographic_peer_order() { + let map = expected_peer_class_map(&peers(&["10.8.1.10", "10.8.1.2"]), 10).expect("map"); + assert_eq!(map.get("10.8.1.10").copied(), Some(10)); + assert_eq!(map.get("10.8.1.2").copied(), Some(11)); + } + + #[test] + fn class_ids_start_at_configured_offset() { + let map = expected_peer_class_map(&peers(&["10.8.1.2", "10.8.1.3"]), 100).expect("map"); + assert_eq!(map.get("10.8.1.2").copied(), Some(100)); + assert_eq!(map.get("10.8.1.3").copied(), Some(101)); + } + + #[test] + fn class_id_overflow_is_reported() { + let err = expected_peer_class_map(&peers(&["10.8.1.2", "10.8.1.3"]), u16::MAX) + .expect_err("overflow"); + assert!(err.to_string().contains("class id overflow")); + } + + #[test] + fn empty_peer_list_yields_empty_map() { + assert!(expected_peer_class_map(&[], 10).expect("map").is_empty()); + } +}