Skip to main content

manycastr/cli/commands/
start.rs

1use crate::cli::client::CliClient;
2use crate::cli::config::{
3    IpVersions, get_hitlist, get_targets, parse_configurations, resolve_workers,
4    validate_ip_versions,
5};
6use crate::cli::utils::validate_path_perms;
7use crate::custom_module::manycastr::{
8    Configuration, MeasurementType, Origin, ProtocolType, ScheduleMeasurement, TraceOptions,
9};
10use crate::custom_module::{Separated, has_anycast_origin, parse_src_address};
11use crate::{ALL_WORKERS, SINGLE_ORIGIN};
12use bimap::BiHashMap;
13use clap::ArgMatches;
14use log::{error, info, warn};
15use prettytable::{Table, format, row};
16use std::collections::HashSet;
17
18pub struct MeasurementExecutionArgs<'a> {
19    /// Determines whether results should be streamed to the command-line interface as they arrive.
20    pub is_cli: bool,
21    /// Indicates whether the results should be written in Parquet format (default: .csv.gz).
22    pub is_parquet: bool,
23    /// Specifies whether the list of targets should be shuffled before the measurement begins.
24    pub is_shuffle: bool,
25    /// The path to the file containing the list of measurement targets (the "hitlist").
26    pub hitlist_path: &'a str,
27    /// The total number of targets in the hitlist, used for estimating measurement duration.
28    pub hitlist_length: usize,
29    /// Path to write results to (may include filename and extension).
30    pub out_path: String,
31    /// A bidirectional map used to resolve worker IDs to their corresponding hostnames.
32    pub worker_map: BiHashMap<u32, String>,
33    /// If true, tags probes with session IDs for attribution (--sessions for -m feed).
34    pub is_sessions: bool,
35    /// IP versions measured, used to tag the output filename (v4/v6/mixed).
36    pub versions: IpVersions,
37}
38
39/// Handle the start command by parsing arguments and sending a measurement request to the orchestrator.
40///
41/// # Arguments
42///
43/// * `matches` - The parsed command-line arguments specific to the start command.
44/// * `grpc_client` - A mutable reference to the gRPC client used to communicate with the orchestrator.
45/// * `worker_map` - A bidirectional map of worker IDs to hostnames for selective probing.
46/// # Returns
47/// * `Result<(), Box<dyn std::error::Error>>` - Ok(()) if the measurement was successfully started, or an error if something went wrong.
48pub async fn handle(
49    matches: &ArgMatches,
50    grpc_client: &mut CliClient,
51    worker_map: BiHashMap<u32, String>,
52) -> Result<(), Box<dyn std::error::Error>> {
53    let is_responsive = matches.get_flag("responsive");
54    let is_sessions = matches.get_flag("sessions");
55    let url = matches.get_one::<String>("url");
56    let m_type = MeasurementType::from_str(matches.get_one::<String>("m_type").unwrap())
57        .expect("Invalid measurement type");
58    let is_feed = m_type.is_feed();
59
60    // Sessions only exist for live feed measurements (traceroute probes cannot carry one)
61    if is_sessions && m_type != MeasurementType::Feed {
62        // TODO enforce this in arg parsing?
63        let msg = "[CLI] --sessions requires a live feed measurement (-m feed).";
64        error!("{}", msg);
65        return Err(msg.into());
66    }
67
68    // Feed measurements read targets from stdin; hitlist-based options do not apply
69    if is_feed {
70        let invalid = [
71            (matches.contains_id("hitlist"), "--hitlist"),
72            (matches.contains_id("target"), "--target"),
73            (matches.get_flag("shuffle"), "--shuffle"),
74        ];
75        if let Some((_, flag)) = invalid.iter().find(|(is_set, _)| *is_set) {
76            let msg = format!(
77                "[CLI] {flag} cannot be combined with a feed measurement (-m {}).",
78                m_type.as_str()
79            );
80            error!("{}", msg);
81            return Err(msg.into());
82        }
83    } else if !matches.contains_id("hitlist") && !matches.contains_id("target") {
84        let msg = "[CLI] --hitlist or --target is required for hitlist-based measurements.";
85        error!("{}", msg);
86        return Err(msg.into());
87    }
88
89    // --responsive cannot be combined with trace mode (unresponsive hops) TODO enforce in arg parse
90    if m_type == MeasurementType::FeedTrace && is_responsive {
91        let msg = "[CLI] --responsive cannot be combined with feed-trace (TTL-limited probes may never reach the target).";
92        error!("{}", msg);
93        return Err(msg.into());
94    }
95
96    // Tracemap targets unresponsive prefixes; there is nothing to discover first TODO enforce in arg parse
97    if m_type == MeasurementType::Tracemap && is_responsive {
98        let msg = "[CLI] --responsive cannot be combined with tracemap (targets are assumed unresponsive).";
99        error!("{}", msg);
100        return Err(msg.into());
101    }
102
103    let configurations = if let Some(conf_path) = matches.get_one::<String>("configuration") {
104        // Use configuration set by the user
105        parse_configurations(conf_path, &worker_map)
106    } else {
107        // Create our own configuration from the arguments
108        let address = matches
109            .get_one::<String>("address")
110            .expect("--address is required unless --configuration is provided");
111        // 'unicastv4'/'unicastv6' means each worker uses its local unicast address
112        let src = parse_src_address(address);
113        let sport: u32 = *matches.get_one::<u16>("sport").unwrap() as u32;
114        let dport = *matches.get_one::<u16>("dport").unwrap() as u32;
115
116        // Get the workers that have to send out probes (worker ID, hostname, or glob like `us-*`)
117        let sender_ids: Vec<u32> = matches.get_one::<String>("selective").map_or_else(
118            || vec![ALL_WORKERS], // Default: all workers
119            |worker_entries_str| {
120                let mut ids: Vec<u32> = Vec::new();
121                for token in worker_entries_str
122                    .trim_matches(|c| c == '[' || c == ']')
123                    .split(',')
124                    .map(str::trim)
125                    .filter(|t| !t.is_empty())
126                {
127                    let matched = resolve_workers(token, &worker_map);
128                    if matched.is_empty() {
129                        warn!("'{token}' did not match any known worker ID or hostname.");
130                    }
131                    ids.extend(matched);
132                }
133                ids.sort_unstable();
134                ids.dedup(); // overlapping globs may match the same worker
135                ids
136            },
137        );
138        // Get the protocols to use (deduplicated, preserving user-specified order)
139        let p_types: Vec<ProtocolType> = {
140            let mut seen = HashSet::new();
141            matches
142                .get_many::<String>("p_type")
143                .unwrap_or_default()
144                .filter_map(|s| ProtocolType::from_str(s))
145                .filter(|p| seen.insert(*p))
146                .collect()
147        };
148
149        let num_p_types = p_types.len();
150
151        // Create the Configuration for each Worker
152        let configs: Vec<Configuration> = sender_ids
153            .iter()
154            .flat_map(|&worker_id| {
155                p_types.iter().enumerate().map(move |(idx, &p_type)| {
156                    let origin_id = if num_p_types == 1 {
157                        // Single p_type -> single Origin
158                        SINGLE_ORIGIN
159                    } else {
160                        // Multiple p_type -> iterate Origin IDs
161                        (idx + 1) as u32
162                    };
163
164                    Configuration {
165                        worker_id,
166                        origin: Some(Origin {
167                            src: Some(src),
168                            sport,
169                            dport,
170                            origin_id,
171                            p_type: p_type as i32,
172                        }),
173                    }
174                })
175            })
176            .collect();
177
178        configs
179    };
180
181    let is_shuffle = matches.get_flag("shuffle");
182    let (hitlist_path, targets, hitlist_versions, is_prefix_hitlist) = if is_feed {
183        // Streamed in live mode (reactive)
184        ("live-feed", Vec::new(), None, false)
185    } else if let Some(target_str) = matches.get_one::<String>("target") {
186        // Individual targets
187        let (targets, versions) = get_targets(target_str, is_shuffle);
188        (target_str.as_str(), targets, Some(versions), false)
189    } else {
190        // Path to hitlist
191        let path = matches.get_one::<String>("hitlist").unwrap().as_str();
192        let (targets, versions, is_prefix_hitlist) = get_hitlist(path, is_shuffle, is_responsive);
193        (path, targets, Some(versions), is_prefix_hitlist)
194    };
195
196    // --responsive is always enabled when using the multi-target hitlist
197    if is_responsive && !is_prefix_hitlist {
198        // Disallow --responsive with catchment, traceroute, and latency measurements
199        let mode = match m_type {
200            MeasurementType::Catchment => Some("catchment mappings"),
201            MeasurementType::AnycastTraceroute => Some("anycast traceroute measurements"),
202            MeasurementType::AnycastLatency if has_anycast_origin(&configurations) => {
203                Some("anycast latency measurements")
204            }
205            _ => None,
206        };
207        if let Some(mode) = mode {
208            let msg = format!(
209                "[CLI] --responsive is disabled for {mode}, except when using the ISI hitlist format."
210            );
211            error!("{}", msg);
212            return Err(msg.into());
213        }
214    }
215
216    // Validate the IP-version rules and get the measured versions
217    let versions = match validate_ip_versions(&configurations, hitlist_versions, m_type) {
218        Ok(versions) => versions,
219        Err(e) => {
220            let msg = format!("[CLI] {e}");
221            error!("{}", msg);
222            return Err(msg.into());
223        }
224    };
225
226    let dns_record = matches.get_one::<String>("query");
227    let is_cli = matches.get_flag("stream");
228    let is_parquet = matches.get_flag("parquet");
229    let worker_interval = *matches.get_one::<u32>("worker_interval").unwrap();
230    let probe_interval = *matches.get_one::<u32>("probe_interval").unwrap();
231    let probing_rate = *matches.get_one::<u32>("rate").unwrap();
232    let number_of_probes = *matches.get_one::<u32>("nprobes").unwrap();
233    let hitlist_length = targets.len();
234
235    // Get protocol and IP version
236    let ip_version = format!("({})", versions.label());
237
238    if is_feed {
239        info!(
240            "[CLI] Performing live {m_type} {ip_version} measurement using targets fed over stdin, with a rate of {}",
241            probing_rate.with_separator(),
242        );
243    } else {
244        info!(
245            "[CLI] Performing {m_type} {ip_version} measurement using targeting {} addresses, with a rate of {}, and a worker-interval of {worker_interval} seconds",
246            hitlist_length.with_separator(),
247            probing_rate.with_separator(),
248        );
249    }
250
251    // Print the origins used
252    info!("[CLI] Workers send probes using the following configurations:");
253    let mut table = Table::new();
254    table.set_format(*format::consts::FORMAT_NO_BORDER_LINE_SEPARATOR);
255    table.set_titles(
256        row![b->"Hostname", b->"Worker ID", b->"Origin ID", b->"src IP", b->"src Port", b->"Dst Port", b->"Protocol"],
257    );
258
259    for config in &configurations {
260        if let Some(origin) = &config.origin {
261            let (worker_name, worker_id_str) = if config.worker_id == ALL_WORKERS {
262                (format!("All {}", worker_map.len()), "ALL".to_string())
263            } else {
264                (
265                    worker_map
266                        .get_by_left(&config.worker_id)
267                        .cloned()
268                        .unwrap_or_else(|| "Unknown".to_string()),
269                    config.worker_id.to_string(),
270                )
271            };
272
273            // ICMP probes carry no source port (dport used as the ICMP identifier)
274            let sport = if origin.p_type() == ProtocolType::Icmp {
275                "n/a".to_string()
276            } else {
277                origin.sport.to_string()
278            };
279
280            table.add_row(row![
281                worker_name,
282                worker_id_str,
283                origin.origin_id,
284                origin.src.unwrap().to_string(),
285                sport,
286                origin.dport,
287                origin.p_type()
288            ]);
289        }
290    }
291    table.printstd();
292
293    // get optional path to write results to
294    let path = matches.get_one::<String>("out").unwrap().to_string();
295    validate_path_perms(&path)?;
296
297    let trace_options = if matches!(
298        m_type,
299        MeasurementType::AnycastTraceroute | MeasurementType::Tracemap
300    ) {
301        Some(TraceOptions {
302            max_failures: *matches.get_one::<u32>("trace_max_failures").unwrap(),
303            max_hops: *matches.get_one::<u32>("trace_max_hop").unwrap(),
304            timeout: *matches.get_one::<u32>("trace_timeout").unwrap(),
305            initial_hop: *matches.get_one::<u32>("trace_initial_hop").unwrap(),
306            star_unresponsive: *matches.get_one::<bool>("trace_star").unwrap(),
307        })
308    } else {
309        None
310    };
311
312    // Create the measurement definition and send it to the orchestrator
313    let m_definition = ScheduleMeasurement {
314        probing_rate,
315        configurations,
316        m_type: m_type.into(),
317        worker_interval,
318        is_responsive,
319        hitlist: targets,
320        record: dns_record.cloned(),
321        url: url.cloned(),
322        probe_interval,
323        number_of_probes,
324        trace_options,
325        is_prefix_hitlist,
326    };
327
328    let args = MeasurementExecutionArgs {
329        is_cli,
330        is_parquet,
331        is_shuffle,
332        hitlist_path,
333        hitlist_length,
334        out_path: path,
335        worker_map,
336        is_sessions,
337        versions,
338    };
339
340    if is_feed {
341        grpc_client
342            .do_live_measurement_to_server(m_definition, args)
343            .await
344    } else {
345        grpc_client
346            .do_measurement_to_server(m_definition, args)
347            .await
348    }
349}