zond_engine/core/models/port/
set.rs

1// Copyright (c) 2026 Erik Lening (hollowpointer) and Contributors
2//
3// This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
4// If a copy of the MPL was not distributed with this file, You can obtain one at
5// https://mozilla.org/MPL/2.0/.
6
7//! # Port Targeting and Range Management
8//!
9//! This module provides the [`PortSet`] model, a high-performance, thread-safe
10//! container for defining which TCP and UDP ports should be scanned.
11//!
12//! ## Overview
13//!
14//! `PortSet` is designed to be parsed once at startup and read concurrently
15//! by thousands of worker threads. It features:
16//! * **Zero-Lock Concurrency**: Immutable read access (`&self`) via eager canonicalization.
17//! * **High-Speed Lookups**: Internal binary search over collapsed `RangeInclusive` sets.
18//! * **Smart Parsing**: Human-friendly string parsing (e.g., `"80, 443, u:53, 1000-2000"`).
19
20use crate::core::models::port::Protocol; // Adjust path as necessary
21use std::{num::ParseIntError, ops::RangeInclusive, str::FromStr};
22use thiserror::Error;
23
24/// Common defaults for rapid discovery scans.
25pub const DEFAULT_PORTSET_PORTS: &str = "22, 80, 443, 445, 3389";
26
27// ══════════════════════════════════════════════════════════════════════════════
28// Error Types
29// ══════════════════════════════════════════════════════════════════════════════
30
31/// Errors that can occur when parsing a port range string.
32#[derive(Debug, Error, PartialEq, Eq)]
33pub enum PortSetParseError {
34    /// The input string could not be parsed as a 16-bit integer.
35    #[error("Failed to parse port from '{input}': {source}")]
36    InvalidPort {
37        input: String,
38        #[source]
39        source: ParseIntError,
40    },
41
42    /// The range start is higher than the end (e.g., "80-20").
43    #[error("Invalid port range: start ({start}) cannot be strictly greater than end ({end})")]
44    InvalidRange { start: u16, end: u16 },
45
46    /// The input segment did not match any known port or range format.
47    #[error("Malformed port specification, expected a single port or a range: '{0}'")]
48    MalformedSpec(String),
49}
50
51// ══════════════════════════════════════════════════════════════════════════════
52// PortSet Core Model
53// ══════════════════════════════════════════════════════════════════════════════
54
55/// A collection of TCP and UDP port ranges used for target discovery.
56///
57/// Under the hood, this stores disjoint ranges. Upon creation, all ranges
58/// are merged and sorted (canonicalized) to ensure `O(log N)` lookup times
59/// via binary search.
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub struct PortSet {
62    tcp: Vec<RangeInclusive<u16>>,
63    udp: Vec<RangeInclusive<u16>>,
64}
65
66impl PortSet {
67    /// Creates a new, empty `PortSet`.
68    pub fn new() -> Self {
69        Self {
70            tcp: Vec::new(),
71            udp: Vec::new(),
72        }
73    }
74
75    /// Returns the total number of unique port/protocol combinations.
76    ///
77    /// Note: This counts every individual port within every range.
78    pub fn len(&self) -> usize {
79        let tcp_count: usize = self
80            .tcp
81            .iter()
82            .map(|r| (r.end().saturating_sub(*r.start()) as usize).saturating_add(1))
83            .sum();
84        let udp_count: usize = self
85            .udp
86            .iter()
87            .map(|r| (r.end().saturating_sub(*r.start()) as usize).saturating_add(1))
88            .sum();
89
90        tcp_count + udp_count
91    }
92
93    /// Returns `true` if no ports are defined for either protocol.
94    pub fn is_empty(&self) -> bool {
95        self.tcp.is_empty() && self.udp.is_empty()
96    }
97
98    /// Returns an iterator over all individual ports in the set.
99    ///
100    /// Yields TCP ports first, followed by UDP ports.
101    pub fn iter(&self) -> impl Iterator<Item = (u16, Protocol)> + '_ {
102        let tcp_iter = self
103            .tcp
104            .iter()
105            .flat_map(|r| r.clone().map(|p| (p, Protocol::Tcp)));
106        let udp_iter = self
107            .udp
108            .iter()
109            .flat_map(|r| r.clone().map(|p| (p, Protocol::Udp)));
110
111        tcp_iter.chain(udp_iter)
112    }
113
114    /// Flattens the set into a vector of individual ports.
115    pub fn to_vec(&self) -> Vec<(u16, Protocol)> {
116        self.iter().collect()
117    }
118
119    /// Checks if a specific TCP port is in the target set.
120    /// Uses a highly optimized binary search over disjoint ranges.
121    pub fn has_tcp(&self, port: u16) -> bool {
122        self.tcp
123            .binary_search_by(|range| {
124                if port < *range.start() {
125                    std::cmp::Ordering::Greater
126                } else if port > *range.end() {
127                    std::cmp::Ordering::Less
128                } else {
129                    std::cmp::Ordering::Equal
130                }
131            })
132            .is_ok()
133    }
134
135    /// Checks if a specific UDP port is in the target set.
136    /// Uses a highly optimized binary search over disjoint ranges.
137    pub fn has_udp(&self, port: u16) -> bool {
138        self.udp
139            .binary_search_by(|range| {
140                if port < *range.start() {
141                    std::cmp::Ordering::Greater
142                } else if port > *range.end() {
143                    std::cmp::Ordering::Less
144                } else {
145                    std::cmp::Ordering::Equal
146                }
147            })
148            .is_ok()
149    }
150
151    // ─── Internal Utility ────────────────────────────────────────────────────
152
153    /// Sorts and merges overlapping/adjacent ranges.
154    /// Called automatically during construction.
155    fn merge_ranges(ranges: &mut Vec<RangeInclusive<u16>>) {
156        if ranges.is_empty() {
157            return;
158        }
159
160        ranges.sort_by_key(|r| *r.start());
161        let mut merged = Vec::with_capacity(ranges.len());
162        let mut it = ranges.drain(..);
163        let mut current = it.next().unwrap();
164
165        for next in it {
166            // Check for overlap or adjacency
167            if *next.start() <= (*current.end()).saturating_add(1) {
168                if *next.end() > *current.end() {
169                    current = *current.start()..=*next.end();
170                }
171            } else {
172                merged.push(current);
173                current = next;
174            }
175        }
176        merged.push(current);
177        *ranges = merged;
178    }
179}
180
181// ══════════════════════════════════════════════════════════════════════════════
182// Conversion Traits
183// ══════════════════════════════════════════════════════════════════════════════
184
185impl Default for PortSet {
186    /// Returns a default [`PortSet`] containing common discovery services.
187    fn default() -> Self {
188        Self::try_from(DEFAULT_PORTSET_PORTS).expect("Static discovery ports must be valid.")
189    }
190}
191
192impl TryFrom<&str> for PortSet {
193    type Error = PortSetParseError;
194
195    /// Parses a string into a canonicalized `PortSet`.
196    ///
197    /// ### Format Support
198    /// * **Individual**: `80`, `443`
199    /// * **Ranges**: `1000-2000`
200    /// * **Protocols**: Defaults to TCP. Use `u:` prefix for UDP (e.g., `u:53`).
201    /// * **Mixed**: `80, 443, u:53, 161-162`
202    ///
203    /// # Examples
204    ///
205    /// ```
206    /// use zond_engine::core::models::port::set::PortSet;
207    ///
208    /// let set = PortSet::try_from("80, u:53, 1000-1005").unwrap();
209    /// assert!(set.has_tcp(80));
210    /// assert!(set.has_udp(53));
211    /// assert_eq!(set.len(), 8); // 1 + 1 + 6
212    /// ```
213    fn try_from(value: &str) -> Result<Self, Self::Error> {
214        let mut tcp = Vec::new();
215        let mut udp = Vec::new();
216
217        for part in value.split([',', ' ']).filter(|s| !s.trim().is_empty()) {
218            let part = part.trim();
219
220            let (is_udp, raw_range) = if let Some(stripped) = part.strip_prefix("u:") {
221                (true, stripped)
222            } else {
223                (false, part)
224            };
225
226            let parts: Vec<&str> = raw_range.split('-').collect();
227
228            let range = match parts.as_slice() {
229                [single_port] => {
230                    let p = single_port.parse::<u16>().map_err(|source| {
231                        PortSetParseError::InvalidPort {
232                            input: single_port.to_string(),
233                            source,
234                        }
235                    })?;
236                    p..=p
237                }
238                [start_str, end_str] => {
239                    let start = start_str.parse::<u16>().map_err(|source| {
240                        PortSetParseError::InvalidPort {
241                            input: start_str.to_string(),
242                            source,
243                        }
244                    })?;
245                    let end = end_str.parse::<u16>().map_err(|source| {
246                        PortSetParseError::InvalidPort {
247                            input: end_str.to_string(),
248                            source,
249                        }
250                    })?;
251
252                    if start > end {
253                        return Err(PortSetParseError::InvalidRange { start, end });
254                    }
255
256                    start..=end
257                }
258                _ => return Err(PortSetParseError::MalformedSpec(raw_range.to_string())),
259            };
260
261            if is_udp {
262                udp.push(range);
263            } else {
264                tcp.push(range);
265            }
266        }
267
268        Self::merge_ranges(&mut tcp);
269        Self::merge_ranges(&mut udp);
270
271        Ok(Self { tcp, udp })
272    }
273}
274
275impl TryFrom<String> for PortSet {
276    type Error = PortSetParseError;
277    fn try_from(value: String) -> Result<Self, Self::Error> {
278        Self::try_from(value.as_str())
279    }
280}
281
282impl FromStr for PortSet {
283    type Err = PortSetParseError;
284    fn from_str(s: &str) -> Result<Self, Self::Err> {
285        Self::try_from(s)
286    }
287}
288
289impl FromIterator<(u16, Protocol)> for PortSet {
290    fn from_iter<T: IntoIterator<Item = (u16, Protocol)>>(iter: T) -> Self {
291        let mut tcp = Vec::new();
292        let mut udp = Vec::new();
293        let mut sctp = Vec::new();
294        for (port, proto) in iter {
295            match proto {
296                Protocol::Tcp => tcp.push(port..=port),
297                Protocol::Udp => udp.push(port..=port),
298                Protocol::Sctp => sctp.push(port..=port),
299            }
300        }
301        Self::merge_ranges(&mut tcp);
302        Self::merge_ranges(&mut udp);
303        Self { tcp, udp }
304    }
305}
306
307// ╔════════════════════════════════════════════╗
308// ║ ████████╗███████╗███████╗████████╗███████╗ ║
309// ║ ╚══██╔══╝██╔════╝██╔════╝╚══██╔══╝██╔════╝ ║
310// ║    ██║   █████╗  ███████╗   ██║   ███████╗ ║
311// ║    ██║   ██╔══╝  ╚════██║   ██║   ╚════██║ ║
312// ║    ██║   ███████╗███████║   ██║   ███████║ ║
313// ║    ╚═╝   ╚══════╝╚══════╝   ╚═╝   ╚══════╝ ║
314// ╚════════════════════════════════════════════╝
315
316#[cfg(test)]
317mod tests {
318    use super::*;
319
320    #[test]
321    fn set_try_from_str_parses_correctly() {
322        let port_set_single = PortSet::try_from("21");
323        let port_set_multiple = PortSet::try_from("21, 22 80, 800-1000, u:53 8080");
324
325        assert!(port_set_single.is_ok());
326        assert!(port_set_multiple.is_ok());
327
328        let port_set_single = port_set_single.unwrap();
329        let port_set_multiple = port_set_multiple.unwrap();
330
331        assert!(port_set_single.has_tcp(21));
332
333        assert!(port_set_multiple.has_tcp(21));
334        assert!(port_set_multiple.has_tcp(22));
335        assert!(port_set_multiple.has_tcp(80));
336        assert!(port_set_multiple.has_tcp(900));
337        assert!(port_set_multiple.has_udp(53));
338        assert!(port_set_multiple.has_tcp(8080));
339    }
340
341    #[test]
342    fn set_try_from_str_parses_udp_variants() {
343        let port_set_udp = PortSet::try_from("u:22 u:53-100, u:1024");
344
345        assert!(port_set_udp.is_ok());
346
347        let port_set_udp = port_set_udp.unwrap();
348
349        assert!(port_set_udp.has_udp(22));
350        assert!(port_set_udp.has_udp(53));
351        assert!(port_set_udp.has_udp(80));
352        assert!(port_set_udp.has_udp(100));
353        assert!(port_set_udp.has_udp(1024));
354    }
355
356    #[test]
357    fn set_empty_input() {
358        let empty = PortSet::try_from("   ");
359        assert!(empty.is_ok());
360        let set = empty.unwrap();
361        assert!(set.tcp.is_empty());
362        assert!(set.udp.is_empty());
363    }
364
365    #[test]
366    fn set_boundaries() {
367        let limits = PortSet::try_from("0, 65535, u:0-65535").unwrap();
368        assert!(limits.has_tcp(0));
369        assert!(limits.has_tcp(65535));
370        assert!(limits.has_udp(32768));
371    }
372
373    #[test]
374    fn set_messy_delimiters() {
375        let messy = PortSet::try_from(", 80, , 443 ,").unwrap();
376        assert!(messy.has_tcp(80));
377        assert!(messy.has_tcp(443));
378    }
379
380    #[test]
381    fn set_try_from_str_throws_errors() {
382        let port_set_invalid_port = PortSet::try_from("80 70000 22");
383        let port_set_invalid_range = PortSet::try_from("21 8000-80");
384        let port_set_malformed_spec = PortSet::try_from("22 60-70-80 8080");
385        let port_set_not_numeric = PortSet::try_from("u:53 abcdef 80");
386
387        assert!(matches!(
388            port_set_invalid_port,
389            Err(PortSetParseError::InvalidPort { .. })
390        ));
391
392        assert!(matches!(
393            port_set_invalid_range,
394            Err(PortSetParseError::InvalidRange {
395                start: 8000,
396                end: 80
397            })
398        ));
399
400        assert!(matches!(
401            port_set_not_numeric,
402            Err(PortSetParseError::InvalidPort { .. })
403        ));
404
405        assert!(matches!(
406            port_set_malformed_spec,
407            Err(PortSetParseError::MalformedSpec(_))
408        ));
409    }
410
411    #[test]
412    fn set_try_from_string_parses_correctly() {
413        let port_set = PortSet::try_from(String::from("21 80-100 u:5353"));
414
415        assert!(port_set.is_ok());
416
417        let port_set = port_set.unwrap();
418
419        assert!(port_set.has_tcp(21));
420        assert!(port_set.has_tcp(80));
421        assert!(port_set.has_tcp(92));
422        assert!(port_set.has_tcp(100));
423        assert!(port_set.has_udp(5353));
424    }
425
426    #[test]
427    fn set_overlap_and_adjacency_merging() {
428        // Overlap: 1-10 and 5-15 should be 1-15
429        let set = PortSet::try_from("1-10, 5-15").unwrap();
430        assert_eq!(set.len(), 15);
431        assert_eq!(set.tcp.len(), 1);
432
433        // Adjacency: 20 and 21 should be 20-21
434        let set = PortSet::try_from("20, 21").unwrap();
435        assert_eq!(set.len(), 2);
436        assert_eq!(set.tcp.len(), 1);
437
438        // Subsumption: 100-200 and 150
439        let set = PortSet::try_from("100-200, 150").unwrap();
440        assert_eq!(set.len(), 101);
441        assert_eq!(set.tcp.len(), 1);
442
443        // Mixed messy overlaps
444        let set = PortSet::try_from("u:53, u:53-53, u:50-60, u:55-65").unwrap();
445        assert_eq!(set.len(), 16); // 50 to 65
446        assert_eq!(set.udp.len(), 1);
447    }
448}
449
450#[cfg(test)]
451mod property_tests {
452    use super::*;
453    use proptest::prelude::*;
454
455    proptest::proptest! {
456        /// Verify that any single port inserted is correctly contained in the set.
457        #[test]
458        fn single_port_roundtrip(p in 0..=65535u16) {
459            let s = format!("{}", p);
460            let set = PortSet::from_str(&s).unwrap();
461            prop_assert!(set.has_tcp(p));
462            prop_assert_eq!(set.len(), 1);
463        }
464
465        /// Verify that any port range [a, b] contains all values within it.
466        #[test]
467        fn port_range_invariant(a in 0..=65535u16, b in 0..=65535u16) {
468            let (start, end) = if a < b { (a, b) } else { (b, a) };
469            let s = format!("{}-{}", start, end);
470            let set = PortSet::from_str(&s).unwrap();
471
472            prop_assert!(set.has_tcp(start));
473            prop_assert!(set.has_tcp(end));
474            prop_assert_eq!(set.len(), (end - start + 1) as usize);
475        }
476
477        /// Verify that UDP prefix 'u:' correctly assigns ports to the UDP set.
478        #[test]
479        fn udp_prefix_honored(p in 0..=65535u16) {
480            let s = format!("u:{}", p);
481            let set = PortSet::from_str(&s).unwrap();
482            prop_assert!(set.has_udp(p));
483            prop_assert!(!set.has_tcp(p));
484        }
485
486        /// Verify that comma-separated lists correctly aggregate multiple ports.
487        #[test]
488        fn multiple_ports_aggregation(p1 in 0..=1000u16, p2 in 2000..=3000u16) {
489            let s = format!("{}, {}", p1, p2);
490            let set = PortSet::from_str(&s).unwrap();
491            prop_assert!(set.has_tcp(p1));
492            prop_assert!(set.has_tcp(p2));
493            prop_assert_eq!(set.len(), 2);
494        }
495
496        /// Invariant: Normalization produces the same port count as a HashSet.
497        #[test]
498        fn normalization_invariant(ports in prop::collection::vec(0..=500u16, 1..=50)) {
499            let s = ports.iter().map(|p| p.to_string()).collect::<Vec<_>>().join(",");
500            let set = PortSet::from_str(&s).unwrap();
501
502            let unique_count = ports.into_iter().collect::<std::collections::HashSet<_>>().len();
503            prop_assert_eq!(set.len(), unique_count);
504            prop_assert!(set.tcp.len() <= unique_count);
505        }
506    }
507}
zond_engine/core/models/port/set.rs

zond_engine/core/models/port/
set.rs
