Merge pull request #787 from mikhail5555/improve-excluded_nets-filter-performance

Improve excluded nets filter performance dractically

Author: PsypherPunk

benches/benchmark_portscan.rs

@@ -1,8 +1,9 @@
 use async_std::task::block_on;
-use criterion::{black_box, criterion_group, criterion_main, Criterion};
-use rustscan::input::{PortRange, ScanOrder};
+use criterion::{criterion_group, criterion_main, Criterion};
+use rustscan::input::{Opts, PortRange, ScanOrder};
 use rustscan::port_strategy::PortStrategy;
 use rustscan::scanner::Scanner;
+use std::hint::black_box;
 use std::net::IpAddr;
 use std::time::Duration;
 
@@ -26,6 +27,24 @@ fn bench_port_strategy() {
     let _strategy = PortStrategy::pick(&Some(range.clone()), None, ScanOrder::Serial);
 }
 
+fn bench_address_parsing() {
+    let opts = Opts {
+        addresses: vec![
+            "127.0.0.1".to_owned(),
+            "10.2.0.1".to_owned(),
+            "192.168.0.0/24".to_owned(),
+        ],
+        exclude_addresses: Some(vec![
+            "10.0.0.0/8".to_owned(),
+            "172.16.0.0/12".to_owned(),
+            "192.168.0.0/16".to_owned(),
+            "172.16.0.1".to_owned(),
+        ]),
+        ..Default::default()
+    };
+    let _ips = rustscan::address::parse_addresses(&opts);
+}
+
 fn criterion_benchmark(c: &mut Criterion) {
     let addrs = vec!["127.0.0.1".parse::<IpAddr>().unwrap()];
     let range = PortRange {
@@ -74,6 +93,13 @@ fn criterion_benchmark(c: &mut Criterion) {
     c.bench_function("parse address", |b| b.iter(bench_address));
 
     c.bench_function("port strategy", |b| b.iter(bench_port_strategy));
+
+    let mut address_group = c.benchmark_group("address parsing");
+    address_group.measurement_time(Duration::from_secs(10));
+    address_group.bench_function("parse addresses with exclusions", |b| {
+        b.iter(|| bench_address_parsing())
+    });
+    address_group.finish();
 }
 
 criterion_group!(benches, criterion_benchmark);

build.rs

@@ -141,7 +141,7 @@ fn ports_v(fp_map: &BTreeMap<i32, String>) -> BTreeMap<i32, Vec<u16>> {
                 } else if !segment.is_empty() {
                     match segment.parse::<u16>() {
                         Ok(port) => port_list.push(port),
-                        Err(_) => println!("Error parsing port: {}", segment),
+                        Err(_) => println!("Error parsing port: {segment}"),
                     }
                 }
             }

src/address.rs

@@ -69,17 +69,11 @@ pub fn parse_addresses(input: &Opts) -> Vec<IpAddr> {
         }
     }
 
-    // Finally, craft a list of addresses to be excluded from the scan.
-    let mut excluded_ips: BTreeSet<IpAddr> = BTreeSet::new();
-    if let Some(exclude_addresses) = &input.exclude_addresses {
-        for addr in exclude_addresses {
-            excluded_ips.extend(parse_address(addr, &backup_resolver));
-        }
-    }
+    let excluded_cidrs = parse_excluded_networks(&input.exclude_addresses, &backup_resolver);
 
     // Remove duplicated/excluded IPs.
     let mut seen = BTreeSet::new();
-    ips.retain(|ip| seen.insert(*ip) && !excluded_ips.contains(ip));
+    ips.retain(|ip| seen.insert(*ip) && !excluded_cidrs.iter().any(|cidr| cidr.contains(ip)));
 
     ips
 }
@@ -125,6 +119,46 @@ fn resolve_ips_from_host(source: &str, backup_resolver: &Resolver) -> Vec<IpAddr
     ips
 }
 
+/// Parses excluded networks from a list of addresses.
+///
+/// This function handles three types of inputs:
+/// 1. CIDR notation (e.g. "192.168.0.0/24")
+/// 2. Single IP addresses (e.g. "192.168.0.1")
+/// 3. Hostnames that need to be resolved (e.g. "example.com")
+///
+/// ```rust
+/// # use rustscan::address::parse_excluded_networks;
+/// # use hickory_resolver::Resolver;
+/// let resolver = Resolver::default().unwrap();
+/// let excluded = parse_excluded_networks(&Some(vec!["192.168.0.0/24".to_owned()]), &resolver);
+/// ```
+pub fn parse_excluded_networks(
+    exclude_addresses: &Option<Vec<String>>,
+    resolver: &Resolver,
+) -> Vec<IpCidr> {
+    exclude_addresses
+        .iter()
+        .flatten()
+        .flat_map(|addr| parse_single_excluded_address(addr, resolver))
+        .collect()
+}
+
+/// Parses a single address into an IpCidr, handling CIDR notation, IP addresses, and hostnames.
+fn parse_single_excluded_address(addr: &str, resolver: &Resolver) -> Vec<IpCidr> {
+    if let Ok(cidr) = IpCidr::from_str(addr) {
+        return vec![cidr];
+    }
+
+    if let Ok(ip) = IpAddr::from_str(addr) {
+        return vec![IpCidr::new_host(ip)];
+    }
+
+    resolve_ips_from_host(addr, resolver)
+        .into_iter()
+        .map(IpCidr::new_host)
+        .collect()
+}
+
 /// Derive a DNS resolver.
 ///
 /// 1. if the `resolver` parameter has been set:

src/main.rs

@@ -46,7 +46,7 @@ fn main() {
     let config = Config::read(opts.config_path.clone());
     opts.merge(&config);
 
-    debug!("Main() `opts` arguments are {:?}", opts);
+    debug!("Main() `opts` arguments are {opts:?}");
 
     let scripts_to_run: Vec<ScriptFile> = match init_scripts(&opts.scripts) {
         Ok(scripts_to_run) => scripts_to_run,
@@ -94,7 +94,7 @@ fn main() {
         opts.exclude_ports.unwrap_or_default(),
         opts.udp,
     );
-    debug!("Scanner finished building: {:?}", scanner);
+    debug!("Scanner finished building: {scanner:?}");
 
     let mut portscan_bench = NamedTimer::start("Portscan");
     let scan_result = block_on(scanner.run());
@@ -146,7 +146,7 @@ fn main() {
             // This part allows us to add commandline arguments to the Script call_format, appending them to the end of the command.
             if !opts.command.is_empty() {
                 let user_extra_args = &opts.command.join(" ");
-                debug!("Extra args vec {:?}", user_extra_args);
+                debug!("Extra args vec {user_extra_args:?}");
                 if script_f.call_format.is_some() {
                     let mut call_f = script_f.call_format.unwrap();
                     call_f.push(' ');
@@ -156,7 +156,7 @@ fn main() {
                         opts.greppable,
                         opts.accessible
                     );
-                    debug!("Call format {}", call_f);
+                    debug!("Call format {call_f}");
                     script_f.call_format = Some(call_f);
                 }
             }
@@ -187,7 +187,7 @@ fn main() {
     benchmarks.push(script_bench);
     rustscan_bench.end();
     benchmarks.push(rustscan_bench);
-    debug!("Benchmarks raw {:?}", benchmarks);
+    debug!("Benchmarks raw {benchmarks:?}");
     info!("{}", benchmarks.summary());
 }
 

src/scanner/mod.rs

@@ -112,7 +112,7 @@ impl Scanner {
                 }
             }
         }
-        debug!("Typical socket connection errors {:?}", errors);
+        debug!("Typical socket connection errors {errors:?}");
         debug!("Open Sockets found: {:?}", &open_sockets);
         open_sockets
     }
@@ -153,7 +153,7 @@ impl Scanner {
                     }
                     self.fmt_ports(socket);
 
-                    debug!("Return Ok after {} tries", nr_try);
+                    debug!("Return Ok after {nr_try} tries");
                     return Ok(socket);
                 }
                 Err(e) => {
@@ -164,7 +164,7 @@ impl Scanner {
                     if nr_try == tries {
                         error_string.push(' ');
                         error_string.push_str(&socket.ip().to_string());
-                        return Err(io::Error::new(io::ErrorKind::Other, error_string));
+                        return Err(io::Error::other(error_string));
                     }
                 }
             };
@@ -193,10 +193,9 @@ impl Scanner {
             }
         }
 
-        Err(io::Error::new(
-            io::ErrorKind::Other,
-            format!("UDP scan timed-out for all tries on socket {}", socket),
-        ))
+        Err(io::Error::other(format!(
+            "UDP scan timed-out for all tries on socket {socket}"
+        )))
     }
 
     /// Performs the connection to the socket with timeout
@@ -275,7 +274,7 @@ impl Scanner {
 
                 match io::timeout(wait, udp_socket.recv(&mut buf)).await {
                     Ok(size) => {
-                        debug!("Received {} bytes", size);
+                        debug!("Received {size} bytes");
                         self.fmt_ports(socket);
                         Ok(true)
                     }
@@ -289,7 +288,7 @@ impl Scanner {
                 }
             }
             Err(e) => {
-                println!("Err E binding sock {:?}", e);
+                println!("Err E binding sock {e:?}");
                 Err(e)
             }
         }

src/scripts/mod.rs

@@ -109,7 +109,7 @@ pub fn init_scripts(scripts: &ScriptsRequired) -> Result<Vec<ScriptFile>> {
         }
         ScriptsRequired::Custom => {
             let script_config = ScriptConfig::read_config()?;
-            debug!("Script config \n{:?}", script_config);
+            debug!("Script config \n{script_config:?}");
 
             let script_dir_base = if let Some(config_directory) = &script_config.directory {
                 PathBuf::from(config_directory)
@@ -118,10 +118,10 @@ pub fn init_scripts(scripts: &ScriptsRequired) -> Result<Vec<ScriptFile>> {
             };
 
             let script_paths = find_scripts(script_dir_base)?;
-            debug!("Scripts paths \n{:?}", script_paths);
+            debug!("Scripts paths \n{script_paths:?}");
 
             let parsed_scripts = parse_scripts(script_paths);
-            debug!("Scripts parsed \n{:?}", parsed_scripts);
+            debug!("Scripts parsed \n{parsed_scripts:?}");
 
             // Only Scripts that contain all the tags found in ScriptConfig will be selected.
             if let Some(config_hashset) = script_config.tags {
@@ -142,7 +142,7 @@ pub fn init_scripts(scripts: &ScriptsRequired) -> Result<Vec<ScriptFile>> {
                     }
                 }
             }
-            debug!("\nScript(s) to run {:?}", scripts_to_run);
+            debug!("\nScript(s) to run {scripts_to_run:?}");
         }
     }
 
@@ -269,14 +269,14 @@ impl Script {
             };
             to_run = default_template.fill_with_struct(&exec_parts)?;
         }
-        debug!("\nScript format to run {}", to_run);
+        debug!("\nScript format to run {to_run}");
         execute_script(&to_run)
     }
 }
 
 #[cfg(not(tarpaulin_include))]
 fn execute_script(script: &str) -> Result<String> {
-    debug!("\nScript arguments {}", script);
+    debug!("\nScript arguments {script}");
 
     let (cmd, arg) = if cfg!(unix) {
         ("sh", "-c")
@@ -314,7 +314,7 @@ fn execute_script(script: &str) -> Result<String> {
             Ok(String::from_utf8_lossy(&output.stdout).into_owned())
         }
         Err(error) => {
-            debug!("Command error {}", error.to_string());
+            debug!("Command error {error}",);
             Err(anyhow!(error.to_string()))
         }
     }
@@ -373,7 +373,7 @@ impl ScriptFile {
                 Some(parsed)
             }
             Err(e) => {
-                debug!("Failed to parse ScriptFile headers {}", e.to_string());
+                debug!("Failed to parse ScriptFile headers {e}");
                 None
             }
         }