manycastr/cli/
config.rs

1use crate::ALL_WORKERS;
2use crate::custom_module::manycastr::{
3    Address, Configuration, MeasurementType, Origin, ProtocolType,
4};
5use crate::custom_module::parse_src_address;
6use bimap::BiHashMap;
7use bzip2::read::BzDecoder;
8use flate2::read::GzDecoder;
9use log::info;
10use rand::prelude::SliceRandom;
11use std::fs::File;
12use std::io::{BufRead, BufReader};
13use std::str::FromStr;
14
15/// Resolve a worker selector token to the matching worker IDs.
16///
17/// A token is one of:
18/// * a numeric worker ID (e.g. `1`) — matched exactly,
19/// * an exact hostname (e.g. `ams01`),
20/// * a glob with `*` wildcards (e.g. `us-*`) — matched against hostnames.
21///
22/// Returns every matching worker ID (empty if none match)
23pub fn resolve_workers(token: &str, worker_map: &BiHashMap<u32, String>) -> Vec<u32> {
24    // Numeric worker ID
25    if let Ok(id) = token.parse::<u32>() {
26        return if worker_map.contains_left(&id) {
27            vec![id]
28        } else {
29            Vec::new()
30        };
31    }
32
33    // Hostname glob
34    if token.contains('*') {
35        return worker_map
36            .iter()
37            .filter(|(_, hostname)| glob_match(token, hostname))
38            .map(|(id, _)| *id)
39            .collect();
40    }
41
42    // Exact hostname
43    worker_map
44        .get_by_right(token)
45        .map(|&id| vec![id])
46        .unwrap_or_default()
47}
48
49/// IP versions present in a set of addresses (hitlist targets or origin sources).
50#[derive(Clone, Copy, Default, PartialEq)]
51pub struct IpVersions {
52    pub has_v4: bool,
53    pub has_v6: bool,
54}
55
56impl IpVersions {
57    /// Collect the IP versions used by the (resolved) origin source addresses.
58    pub fn from_origins(configurations: &[Configuration]) -> Self {
59        let mut versions = IpVersions::default();
60        for src in configurations.iter().filter_map(|c| c.origin?.src) {
61            if src.is_v6() {
62                versions.has_v6 = true;
63            } else {
64                versions.has_v4 = true;
65            }
66        }
67        versions
68    }
69
70    /// Collect the IP versions present in a list of target addresses.
71    pub fn from_targets(targets: &[Address]) -> Self {
72        let mut versions = IpVersions::default();
73        for addr in targets {
74            if addr.is_v6() {
75                versions.has_v6 = true;
76            } else {
77                versions.has_v4 = true;
78            }
79        }
80        versions
81    }
82
83    /// Human-readable label, e.g. "IPv4" or "IPv4+IPv6".
84    pub fn label(&self) -> &'static str {
85        match (self.has_v4, self.has_v6) {
86            (true, true) => "IPv4+IPv6",
87            (false, true) => "IPv6",
88            _ => "IPv4",
89        }
90    }
91
92    /// Compact token for output filenames: "v4", "v6", or "mixed".
93    pub fn file_token(&self) -> &'static str {
94        match (self.has_v4, self.has_v6) {
95            (true, true) => "mixed",
96            (false, true) => "v6",
97            _ => "v4",
98        }
99    }
100}
101
102/// Validate the IP-version rules of a measurement and return the measured version(s).
103/// * All hitlist targets must have an origin with a matching IP version
104/// * tracemap does not support mixed IP version TODO
105///
106/// # Arguments
107/// * `configurations` - the measurement configurations
108/// * `hitlist_versions` - IP versions of the hitlist targets (`None` for a live feed)
109/// * `m_type` - the measurement type
110///
111/// # Returns
112/// The IP version(s) measured: those of the hitlist, or of the origins (live feed).
113pub fn validate_ip_versions(
114    configurations: &[Configuration],
115    hitlist_versions: Option<IpVersions>,
116    m_type: MeasurementType,
117) -> Result<IpVersions, String> {
118    let origin_versions = IpVersions::from_origins(configurations);
119    // The IP version(s) measured: those of the hitlist, or of the origins (live feed)
120    let versions = hitlist_versions.unwrap_or(origin_versions);
121
122    // Every target IP version needs at least one origin of that version
123    for (present, has_origin, label) in [
124        (versions.has_v4, origin_versions.has_v4, "IPv4"),
125        (versions.has_v6, origin_versions.has_v6, "IPv6"),
126    ] {
127        if present && !has_origin {
128            return Err(format!(
129                "The hitlist contains {label} targets but no {label} origin is configured."
130            ));
131        }
132    }
133
134    // Tracemap tasks use a single origin and cannot serve two IP versions TODO
135    if m_type == MeasurementType::Tracemap && versions.has_v4 && versions.has_v6 {
136        return Err(
137            "tracemap does not support a mixed IPv4/IPv6 hitlist (tasks use a single origin)."
138                .to_string(),
139        );
140    }
141
142    Ok(versions)
143}
144
145/// Match `text` against a `*`-wildcard `pattern`
146fn glob_match(pattern: &str, text: &str) -> bool {
147    let parts: Vec<&str> = pattern.split('*').collect();
148    if parts.len() == 1 {
149        return pattern == text; // no wildcard
150    }
151
152    let mut pos = 0;
153    // Anchor the start (unless the pattern begins with '*')
154    if !parts[0].is_empty() {
155        if !text.starts_with(parts[0]) {
156            return false;
157        }
158        pos = parts[0].len();
159    }
160    // Each interior segment must appear, in order, after the previous match
161    for part in &parts[1..parts.len() - 1] {
162        match text[pos..].find(part) {
163            Some(i) => pos += i + part.len(),
164            None => return false,
165        }
166    }
167    // Anchor the end (unless the pattern ends with '*')
168    let last = parts[parts.len() - 1];
169    text.len() >= pos + last.len() && text[pos..].ends_with(last)
170}
171
172/// Get the hitlist from a file.
173///
174/// Supports plain fsdb files (with `#fsdb` header) based on the USC/ISI ANT hitlist format.
175///
176/// # Arguments
177/// * `hitlist_path` - path to the hitlist file
178/// * `is_shuffle` - boolean whether the hitlist should be shuffled or not
179/// * `is_responsive` - whether the measurement gates probes behind a responsiveness check
180///
181/// # Returns
182/// * A tuple containing the target addresses, the IP versions present, and
183///   whether the targets are a rank-major prefix hitlist (ISI + `--responsive`).
184///
185/// # Panics
186/// * If the hitlist file cannot be opened.
187/// * If the hitlist is empty.
188pub fn get_hitlist(
189    hitlist_path: &str,
190    is_shuffle: bool,
191    is_responsive: bool,
192) -> (Vec<Address>, IpVersions, bool) {
193    let file =
194        File::open(hitlist_path).unwrap_or_else(|_| panic!("Unable to open file {hitlist_path}"));
195
196    // Create reader based on file extension
197    let reader: Box<dyn BufRead> = if hitlist_path.ends_with(".gz") {
198        let decoder = GzDecoder::new(file);
199        Box::new(BufReader::new(decoder))
200    } else if hitlist_path.ends_with(".bz2") {
201        let decoder = BzDecoder::new(file);
202        Box::new(BufReader::new(decoder))
203    } else {
204        Box::new(BufReader::new(file))
205    };
206
207    let mut lines = reader.lines().map_while(Result::ok).peekable();
208
209    if lines.peek().is_some_and(|l| l.starts_with("#fsdb")) {
210        // ISI hitlist: ranked candidate addresses per prefix (/24 for IPv4, /48 for IPv6)
211        let ranked: Vec<Vec<Address>> = lines
212            .filter(|l| !l.is_empty() && !l.starts_with('#'))
213            .filter_map(|l| parse_isi_row(&l))
214            .filter(|candidates| !candidates.is_empty())
215            .collect();
216
217        if !is_responsive {
218            // All candidates are probed; ordering is irrelevant
219            let ips = ranked.into_iter().flatten().collect();
220            let (ips, versions) = finalize_hitlist(ips, is_shuffle);
221            return (ips, versions, false);
222        }
223
224        // Sequentially (ranked) probe addresses in each prefix till one responds
225        let max_rank = ranked.iter().map(Vec::len).max().unwrap_or(0);
226        let mut ips = Vec::with_capacity(ranked.iter().map(Vec::len).sum());
227        for rank in 0..max_rank {
228            let start = ips.len();
229            ips.extend(ranked.iter().filter_map(|c| c.get(rank).copied()));
230            // Shuffle within the rank to preserve the ordering
231            if is_shuffle {
232                ips[start..].shuffle(&mut rand::rng());
233            }
234        }
235        let (ips, versions) = finalize_hitlist(ips, false);
236        return (ips, versions, true);
237    }
238
239    let ips: Vec<Address> = lines // Create a vector of addresses from the file
240        .filter(|l| !l.trim().is_empty()) // Skip empty lines
241        .map(Address::from)
242        .collect();
243
244    let (ips, versions) = finalize_hitlist(ips, is_shuffle);
245    (ips, versions, false)
246}
247
248/// Parse one USC/ISI ANT hitlist row into ranked candidate addresses.
249/// The block length selects the IP version: 8 hex digits is an IPv4 /24,
250/// 12 hex digits is an IPv6 /48.
251///
252/// IPv4 example: `01000400  01,04,09` -> 1.0.4.1, 1.0.4.4, 1.0.4.9
253/// IPv6 example: `20010db81234  1,2a3f` -> 2001:db8:1234::1, 2001:db8:1234::2a3f
254/// `-` marks a prefix with no known-responsive addresses
255fn parse_isi_row(line: &str) -> Option<Vec<Address>> {
256    let (block, suffixes) = line.split_once('\t')?;
257    let block = block.trim();
258    match block.len() {
259        // IPv4: /24 base, candidates are last-octets
260        8 => {
261            let base = u32::from_str_radix(block, 16).ok()?;
262            Some(
263                suffixes
264                    .split(',')
265                    .filter_map(|s| u8::from_str_radix(s.trim(), 16).ok())
266                    .map(|s| Address::from(base | s as u32))
267                    .collect(),
268            )
269        }
270        // IPv6: /48 base, candidates are suffixes within the 80 host bits
271        12 => {
272            let base = (u64::from_str_radix(block, 16).ok()? as u128) << 80;
273            Some(
274                suffixes
275                    .split(',')
276                    .filter_map(|s| u128::from_str_radix(s.trim(), 16).ok())
277                    .filter(|s| s >> 80 == 0)
278                    .map(|s| Address::from(base | s))
279                    .collect(),
280            )
281        }
282        _ => None,
283    }
284}
285
286/// Build a hitlist from a comma-separated list of target addresses (e.g. from the
287/// `--target` CLI flag), as an alternative to a hitlist file. The targets are
288/// treated exactly like a hitlist containing those addresses.
289///
290/// # Arguments
291/// * `targets` - comma-separated address list, e.g. "1.1.1.1" or "1.1.1.1,8.8.8.8"
292/// * `is_shuffle` - whether the resulting hitlist should be shuffled
293///
294/// # Returns
295/// * A tuple of the parsed addresses and the IP versions present.
296pub fn get_targets(targets: &str, is_shuffle: bool) -> (Vec<Address>, IpVersions) {
297    let ips: Vec<Address> = targets
298        .split(',')
299        .map(str::trim)
300        .filter(|t| !t.is_empty())
301        .map(Address::from)
302        .collect();
303
304    finalize_hitlist(ips, is_shuffle)
305}
306
307/// Validate that a parsed hitlist is non-empty, collect the IP versions it uses,
308/// and optionally shuffle it. Shared by [`get_hitlist`] (file) and
309/// [`get_targets`] (inline `--target` list).
310///
311/// # Panics
312/// * If the hitlist is empty.
313fn finalize_hitlist(mut ips: Vec<Address>, is_shuffle: bool) -> (Vec<Address>, IpVersions) {
314    if ips.is_empty() {
315        panic!("No target addresses provided (empty hitlist / target list)");
316    }
317
318    let versions = IpVersions::from_targets(&ips);
319
320    // Shuffle the hitlist, if desired
321    if is_shuffle {
322        ips.as_mut_slice().shuffle(&mut rand::rng());
323    }
324    (ips, versions)
325}
326
327/// Parse the worker configurations from a file.
328///
329/// # Arguments
330/// * `conf_file` - path to the configuration file
331/// * `worker_map` - a BiHashMap mapping worker IDs to hostnames
332///
333/// # Returns
334/// * A vector of Configuration objects parsed from the file
335///
336/// # Panics
337/// * If the configuration file cannot be opened.
338/// * If the configuration file contains invalid formats.
339/// * If no valid configurations are found in the file.
340pub fn parse_configurations(
341    conf_file: &str,
342    worker_map: &BiHashMap<u32, String>,
343) -> Vec<Configuration> {
344    info!("[CLI] Using configuration file: {conf_file}");
345    let file = File::open(conf_file)
346        .unwrap_or_else(|_| panic!("Unable to open configuration file {conf_file}"));
347    let buf_reader = BufReader::new(file);
348    let mut origin_id = 0;
349    let mut configurations: Vec<Configuration> = Vec::new();
350
351    for line in buf_reader.lines() {
352        let line = line.expect("Unable to read configuration line");
353        let line = line.trim();
354        // Skip comments and empty lines
355        if line.is_empty() || line.starts_with("#") {
356            continue;
357        }
358
359        // Worker, src_addr, src_port, dst_port, protocol
360        let parts: Vec<&str> = line.split(",").map(|s| s.trim()).collect();
361        if parts.len() != 5 {
362            panic!("Invalid configuration format: {line}");
363        }
364
365        // Get the workers for this configuration line
366        let worker_ids = if parts[0] == "ALL" {
367            vec![ALL_WORKERS]
368        } else {
369            let ids = resolve_workers(parts[0], worker_map);
370            if ids.is_empty() {
371                panic!(
372                    "'{}' did not match any known worker ID or hostname.",
373                    parts[0]
374                );
375            }
376            ids
377        };
378
379        // Parse the source address
380        let src = parse_src_address(parts[1]);
381
382        // Parse to u16 first, must fit in header
383        let sport = u16::from_str(parts[2]).expect("Unable to parse src port") as u32;
384        let dport = u16::from_str(parts[3]).expect("Unable to parse dst port") as u32;
385        let p_type = ProtocolType::from_str(parts[4]).expect("Unable to parse protocol type");
386        // Each line is one origin, shared by every worker the selector matched
387        origin_id += 1;
388
389        for worker_id in worker_ids {
390            configurations.push(Configuration {
391                worker_id,
392                origin: Some(Origin {
393                    src: Some(src),
394                    sport,
395                    dport,
396                    origin_id,
397                    p_type: p_type as i32,
398                }),
399            });
400        }
401    }
402    if configurations.is_empty() {
403        panic!("No valid configurations found in file {conf_file}");
404    }
405
406    configurations
407}
manycastr/cli/config.rs

manycastr/cli/
config.rs
