perf(scanner): precompute UDP payload lookup (#886)

Build a port -> payload lookup once per UDP scan instead of cloning the whole nmap payload BTreeMap into every scanned socket and scanning it linearly. Payload selection keeps the previous "last match wins" semantics.

Adds a socket-free micro-benchmark (benches/benchmark_helpers.rs) and lookup tests (tests/udp_payload_lookup.rs).

benches/benchmark_helpers.rs

@@ -1,6 +1,10 @@
 use criterion::{criterion_group, criterion_main, Criterion};
+use rustscan::generated::get_parsed_data;
 use rustscan::input::{Opts, PortRange, ScanOrder};
 use rustscan::port_strategy::PortStrategy;
+use rustscan::scanner::build_udp_payload_lookup;
+use std::collections::BTreeMap;
+use std::hint::black_box;
 use std::net::IpAddr;
 use std::time::Duration;
 
@@ -34,6 +38,18 @@ fn bench_address_parsing() {
     let _ips = rustscan::address::parse_addresses(&opts);
 }
 
+// Replicates the old UDP payload selection behavior:
+// scan the whole UDP payload map and find the last payload whose port list contains `port`.
+fn old_payload_for_port(udp_map: &'static BTreeMap<Vec<u16>, Vec<u8>>, port: u16) -> &'static [u8] {
+    let mut payload: &'static [u8] = b"";
+    for (ports, value) in udp_map.iter() {
+        if ports.contains(&port) {
+            payload = value.as_slice();
+        }
+    }
+    payload
+}
+
 fn criterion_benchmark(c: &mut Criterion) {
     // Benching helper functions
     c.bench_function("parse address", |b| b.iter(bench_address));
@@ -46,6 +62,30 @@ fn criterion_benchmark(c: &mut Criterion) {
         b.iter(bench_address_parsing)
     });
     address_group.finish();
+
+    // UDP payload lookup micro-benchmark: compares the old linear scan of the
+    // payload map with the precomputed port -> payload lookup. No sockets.
+    let udp_map = get_parsed_data();
+    let lookup = build_udp_payload_lookup(udp_map);
+    let ports: Vec<u16> = (1..=4096).collect();
+
+    c.bench_function("udp payload lookup/old scan map 1..4096", |b| {
+        b.iter(|| {
+            for &p in ports.iter() {
+                let payload = old_payload_for_port(black_box(udp_map), black_box(p));
+                black_box(payload);
+            }
+        })
+    });
+
+    c.bench_function("udp payload lookup/new hashmap 1..4096", |b| {
+        b.iter(|| {
+            for &p in ports.iter() {
+                let payload = lookup.get(&p).copied().unwrap_or(b"");
+                black_box(payload);
+            }
+        })
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);

src/scanner/mod.rs

@@ -14,12 +14,35 @@ use colored::Colorize;
 use futures::stream::FuturesUnordered;
 use std::collections::BTreeMap;
 use std::{
-    collections::HashSet,
+    collections::{HashMap, HashSet},
     net::{IpAddr, Shutdown, SocketAddr},
     num::NonZeroU8,
+    sync::Arc,
     time::Duration,
 };
 
+/// UDP payload lookup: port -> payload bytes
+///
+/// `get_parsed_data()` returns a `&'static BTreeMap<...>`, so we can store
+/// references to the payload bytes without cloning them.
+#[doc(hidden)]
+pub type UdpPayloadLookup = HashMap<u16, &'static [u8]>;
+
+#[doc(hidden)]
+pub fn build_udp_payload_lookup(udp_map: &'static BTreeMap<Vec<u16>, Vec<u8>>) -> UdpPayloadLookup {
+    let mut lookup: UdpPayloadLookup = HashMap::new();
+
+    for (ports, payload_vec) in udp_map.iter() {
+        let payload: &'static [u8] = payload_vec.as_slice();
+        for &port in ports.iter() {
+            // Preserve existing behavior: if duplicates exist, last insert wins.
+            lookup.insert(port, payload);
+        }
+    }
+
+    lookup
+}
+
 /// The class for the scanner
 /// IP is data type IpAddr and is the IP address
 /// start & end is where the port scan starts and ends
@@ -93,11 +116,19 @@ impl Scanner {
         let mut open_sockets: Vec<SocketAddr> = Vec::new();
         let mut ftrs = FuturesUnordered::new();
         let mut errors: HashSet<String> = HashSet::new();
-        let udp_map = get_parsed_data();
+
+        // Build UDP payload lookup once (only if we are scanning UDP).
+        // This avoids cloning a big map into every spawned future and turns
+        // payload selection from O(n) to O(1).
+        let udp_payloads: Option<Arc<UdpPayloadLookup>> = if self.udp {
+            Some(Arc::new(build_udp_payload_lookup(get_parsed_data())))
+        } else {
+            None
+        };
 
         for _ in 0..self.batch_size {
             if let Some(socket) = socket_iterator.next() {
-                ftrs.push(self.scan_socket(socket, udp_map.clone()));
+                ftrs.push(self.scan_socket(socket, udp_payloads.clone()));
             } else {
                 break;
             }
@@ -111,7 +142,7 @@ impl Scanner {
 
         while let Some(result) = ftrs.next().await {
             if let Some(socket) = socket_iterator.next() {
-                ftrs.push(self.scan_socket(socket, udp_map.clone()));
+                ftrs.push(self.scan_socket(socket, udp_payloads.clone()));
             }
 
             match result {
@@ -146,10 +177,10 @@ impl Scanner {
     async fn scan_socket(
         &self,
         socket: SocketAddr,
-        udp_map: BTreeMap<Vec<u16>, Vec<u8>>,
+        udp_payloads: Option<Arc<UdpPayloadLookup>>,
     ) -> io::Result<SocketAddr> {
         if self.udp {
-            return self.scan_udp_socket(socket, udp_map).await;
+            return self.scan_udp_socket(socket, udp_payloads).await;
         }
 
         let tries = self.tries.get();
@@ -184,18 +215,16 @@ impl Scanner {
     async fn scan_udp_socket(
         &self,
         socket: SocketAddr,
-        udp_map: BTreeMap<Vec<u16>, Vec<u8>>,
+        udp_payloads: Option<Arc<UdpPayloadLookup>>,
     ) -> io::Result<SocketAddr> {
-        let mut payload: Vec<u8> = Vec::new();
-        for (key, value) in udp_map {
-            if key.contains(&socket.port()) {
-                payload = value;
-            }
-        }
+        let payload: &[u8] = udp_payloads
+            .as_ref()
+            .and_then(|m| m.get(&socket.port()).copied())
+            .unwrap_or(b"");
 
         let tries = self.tries.get();
         for _ in 1..=tries {
-            match self.udp_scan(socket, &payload, self.timeout).await {
+            match self.udp_scan(socket, payload, self.timeout).await {
                 Ok(true) => return Ok(socket),
                 Ok(false) => continue,
                 Err(e) => return Err(e),

tests/udp_payload_lookup.rs

@@ -0,0 +1,32 @@
+use rustscan::generated::get_parsed_data;
+use rustscan::scanner::build_udp_payload_lookup;
+
+#[test]
+fn udp_payload_lookup_contains_common_udp_ports() {
+    let udp_map = get_parsed_data();
+    let lookup = build_udp_payload_lookup(udp_map);
+
+    // These are common UDP services; the payload database should include them.
+    assert!(
+        lookup.contains_key(&53),
+        "expected UDP payload for DNS (53)"
+    );
+    assert!(
+        lookup.contains_key(&123),
+        "expected UDP payload for NTP (123)"
+    );
+}
+
+#[test]
+fn udp_payload_lookup_payloads_are_non_empty_for_known_ports() {
+    let udp_map = get_parsed_data();
+    let lookup = build_udp_payload_lookup(udp_map);
+
+    // Don't assert exact bytes (the generated payload set may evolve),
+    // but it should not be empty for these well-known protocols.
+    let dns = lookup.get(&53).expect("missing DNS payload");
+    assert!(!dns.is_empty(), "DNS payload should not be empty");
+
+    let ntp = lookup.get(&123).expect("missing NTP payload");
+    assert!(!ntp.is_empty(), "NTP payload should not be empty");
+}