Merge pull request #11 from BugenZhao/bz-list-input

tracerouter: allow ip list file as input

Author: BugenZhao

README.md

@@ -38,6 +38,10 @@ Compared to the original implementation, the main features of *flashroute.rs* ar
     ```shell
     cargo run --release -- 192.168.1.1/32 --grain 0
     ```
+- Probe all hosts listed in a file line by line
+    ```shell
+    cargo run --release -- path/to/file --grain 8
+    ```
 
 Listening on ICMP socket requires superuser permission, the *flashroute.rs* may automatically restart in sudo mode.
 

res/test.in

@@ -0,0 +1,3 @@
+8.8.8.8
+115.159.1.233
+39.156.27.1

src/error.rs

@@ -16,6 +16,10 @@ pub enum Error {
     UnexpectedIcmpType(pnet::packet::icmp::IcmpType, pnet::packet::icmp::IcmpCode),
     #[error("")]
     BadGrainOrNet(u8, ipnet::Ipv4Net),
+    #[error("")]
+    InvalidIpv4Addr(String),
+    #[error("")]
+    CannotResolveTargets(String),
     #[error("io error: {0}")]
     IoError(#[from] std::io::Error), // thus io::Error can implicitly `into` IoError
 }

src/main.rs

@@ -41,6 +41,8 @@ async fn main() -> Result<()> {
     #[cfg(unix)]
     utils::ensure_su();
 
+    log::info!("{:?}", *OPT);
+
     #[cfg(debug_assertions)]
     log::warn!(
         "{} is built in DEBUG mode, thus may perform quite poorly.",
@@ -49,13 +51,12 @@ async fn main() -> Result<()> {
 
     let tr = Arc::new(Tracerouter::new()?);
     let r = tr.clone();
-    let int_handler = tokio::spawn(async move {
+    tokio::spawn(async move {
         tokio::signal::ctrl_c().await.unwrap();
         r.stop();
     });
 
     let topo = tr.run().await?;
-    int_handler.abort();
     process_topo(topo).await?;
 
     #[cfg(windows)]

src/opt.rs

@@ -1,6 +1,7 @@
 use std::path::PathBuf;
 use structopt::StructOpt;
 
+use crate::error::*;
 use crate::utils;
 
 #[derive(Debug, StructOpt)]
@@ -60,8 +61,8 @@ pub struct Opt {
     // Target
     #[structopt(short, long, default_value = "8")]
     pub grain: u8,
-    #[structopt()]
-    pub targets: ipnet::Ipv4Net,
+    #[structopt(parse(try_from_str = parse_targets))]
+    pub targets: Targets,
     #[structopt(long)]
     pub global_only: bool,
     #[structopt(long)]
@@ -72,6 +73,22 @@ pub struct Opt {
     pub local_addr: std::net::Ipv4Addr,
 }
 
+#[derive(Debug, Clone)]
+pub enum Targets {
+    Net(ipnet::Ipv4Net),
+    List(PathBuf),
+}
+
+pub fn parse_targets(arg: &str) -> Result<Targets> {
+    if let Ok(net) = arg.parse() {
+        Ok(Targets::Net(net))
+    } else if let Ok(path) = arg.parse() {
+        Ok(Targets::List(path))
+    } else {
+        Err(Error::CannotResolveTargets(arg.to_owned()))
+    }
+}
+
 pub fn get_opt() -> Opt {
     let mut opt: Opt = Opt::from_args();
     opt.local_addr = crate::utils::get_interface_ipv4_addr(&opt.interface).unwrap();

src/tracerouter.rs

@@ -17,6 +17,7 @@ use crate::{
     dcb::DstCtrlBlock,
     error::*,
     network::NetworkManager,
+    opt::Targets,
     prober::ProbePhase,
     prober::ProbeResult,
     prober::Prober,
@@ -46,31 +47,14 @@ pub struct Tracerouter {
 
 impl Tracerouter {
     pub fn new() -> Result<Self> {
-        if OPT.grain > (OPT.targets.max_prefix_len() - OPT.targets.prefix_len()) {
-            return Err(Error::BadGrainOrNet(OPT.grain, OPT.targets));
-        }
-
         log::info!(
             "Using interface `{}` ({})",
             OPT.interface.name,
             crate::utils::get_interface_ipv4_addr(&OPT.interface).unwrap()
         );
 
-        log::info!("Generating targets...");
-        let all_count = Self::targets_count();
-        let mut targets = DcbMap::with_capacity(all_count);
-        for addr in Self::random_targets() {
-            targets.insert(
-                Self::addr_to_key(addr),
-                DstCtrlBlock::new(addr, OPT.default_ttl),
-            );
-        }
-        let filtered_count = targets.len();
-        log::info!(
-            "Generated {} targets, {} removed",
-            filtered_count,
-            all_count - filtered_count
-        );
+        log::info!("Initializing targets...");
+        let targets = Self::generate_targets()?;
 
         Ok(Self {
             targets: Arc::new(targets),
@@ -83,28 +67,67 @@ impl Tracerouter {
         (u >> (OPT.grain)) as AddrKey
     }
 
-    fn targets_count() -> usize {
-        1 << ((OPT.targets.max_prefix_len() - OPT.targets.prefix_len()) - OPT.grain)
-    }
+    fn generate_targets() -> Result<DcbMap> {
+        match OPT.targets.clone() {
+            Targets::Net(net) => {
+                if OPT.grain > (net.max_prefix_len() - net.prefix_len()) {
+                    return Err(Error::BadGrainOrNet(OPT.grain, net));
+                }
 
-    fn random_targets() -> impl Iterator<Item = Ipv4Addr> {
-        let mut rng = StdRng::seed_from_u64(OPT.seed);
-        let subnets = OPT
-            .targets
-            .subnets(OPT.targets.max_prefix_len() - OPT.grain)
-            .unwrap();
-
-        subnets
-            .map(move |net| net.addr().saturating_add(rng.gen_range(0, 1 << OPT.grain)))
-            .filter(|addr| {
-                if OPT.global_only && OPT.allow_private {
-                    addr.is_bz_global() || addr.is_private()
-                } else if OPT.global_only {
-                    addr.is_bz_global()
-                } else {
-                    true
+                let mut rng = StdRng::seed_from_u64(OPT.seed);
+                let subnets = net.subnets(net.max_prefix_len() - OPT.grain).unwrap();
+
+                let iter = subnets
+                    .map(move |net| net.addr().saturating_add(rng.gen_range(0, 1 << OPT.grain)))
+                    .filter(|addr| {
+                        if OPT.global_only && OPT.allow_private {
+                            addr.is_bz_global() || addr.is_private()
+                        } else if OPT.global_only {
+                            addr.is_bz_global()
+                        } else {
+                            true
+                        }
+                    });
+
+                let all_count = 1 << ((net.max_prefix_len() - net.prefix_len()) - OPT.grain);
+                let mut generated_targets = DcbMap::with_capacity(all_count);
+                for addr in iter {
+                    generated_targets.insert(
+                        Self::addr_to_key(addr),
+                        DstCtrlBlock::new(addr, OPT.default_ttl),
+                    );
                 }
-            })
+                let filtered_count = generated_targets.len();
+                log::info!(
+                    "Generated {} targets, {} removed",
+                    filtered_count,
+                    all_count - filtered_count
+                );
+
+                Ok(generated_targets)
+            }
+
+            Targets::List(path) => {
+                let mut generated_targets = DcbMap::new();
+
+                let content = std::fs::read_to_string(path)?;
+                for line in content.lines() {
+                    if line.is_empty() {
+                        continue;
+                    }
+                    let addr = line
+                        .parse()
+                        .or(Err(Error::InvalidIpv4Addr(line.to_owned())))?;
+                    generated_targets.insert(
+                        Self::addr_to_key(addr),
+                        DstCtrlBlock::new(addr, OPT.default_ttl),
+                    );
+                }
+                log::info!("Imported {} targets from file", generated_targets.len());
+
+                Ok(generated_targets)
+            }
+        }
     }
 }
 
@@ -349,13 +372,14 @@ mod test {
     #[test]
     fn test_generation() {
         let tr = Tracerouter::new().unwrap();
-        assert_eq!(
-            tr.targets.len(),
-            1 << (32 - OPT.targets.prefix_len() - OPT.grain)
-        );
-        assert!(tr
-            .targets
-            .values()
-            .all(|dcb| OPT.targets.contains(&dcb.addr)));
+        if let Targets::Net(targets) = OPT.targets {
+            assert_eq!(
+                tr.targets.len(),
+                1 << (32 - targets.prefix_len() - OPT.grain)
+            );
+            assert!(tr.targets.values().all(|dcb| targets.contains(&dcb.addr)));
+        } else {
+            panic!();
+        }
     }
 }