manycastr/cli/
feed.rs

1//! Live feed support: reading NDJSON targets from stdin for feed-based measurements.
2
3use crate::cli::config::resolve_workers;
4use crate::custom_module::manycastr::{
5    Address, CliMessage, Configuration, LiveTarget, ProtocolType, TargetBatch, cli_message,
6};
7use crate::{ALL_ORIGINS, ALL_WORKERS};
8use bimap::BiHashMap;
9use futures_core::Stream;
10use log::warn;
11use std::collections::HashMap;
12use std::io::BufRead;
13use std::pin::Pin;
14use std::task::{Context, Poll};
15use tokio::sync::mpsc;
16
17/// Size of the bounded stdin-to-gRPC feed channel (blocks stdin when full).
18pub const FEED_CHANNEL_SIZE: usize = 1024;
19
20/// Per-origin properties the feed parser needs to validate targets.
21pub struct FeedOrigin {
22    /// IP version of the origin's source address
23    pub is_v6: bool,
24    /// Protocol the origin probes with
25    pub p_type: ProtocolType,
26}
27
28impl FeedOrigin {
29    /// TCP and DNS CHAOS do not support session IDs.
30    fn supports_sessions(&self) -> bool {
31        matches!(self.p_type, ProtocolType::Icmp | ProtocolType::ADns)
32    }
33}
34
35/// The configured origins available to feed targets, and the default origins for each IP version.
36pub struct FeedOrigins {
37    /// Origin ID -> IP version and protocol, to match feed targets with compatible origins
38    by_id: HashMap<u32, FeedOrigin>,
39    /// First configured IPv4 origin (default for IPv4 targets without an `origin` field)
40    default_v4: Option<u32>,
41    /// First configured IPv6 origin (default for IPv6 targets without an `origin` field)
42    default_v6: Option<u32>,
43}
44
45impl FeedOrigins {
46    /// Collect the unique origins of a measurement.
47    /// Set the first origin of each IP version as that version's default.
48    pub fn new(configurations: &[Configuration]) -> Self {
49        let mut by_id = HashMap::new();
50        let mut default_v4 = None;
51        let mut default_v6 = None;
52        for origin in configurations.iter().filter_map(|c| c.origin) {
53            by_id.entry(origin.origin_id).or_insert(FeedOrigin {
54                is_v6: origin.is_v6(),
55                p_type: origin.p_type(),
56            });
57            let default = if origin.is_v6() {
58                &mut default_v6
59            } else {
60                &mut default_v4
61            };
62            default.get_or_insert(origin.origin_id);
63        }
64        Self {
65            by_id,
66            default_v4,
67            default_v6,
68        }
69    }
70
71    /// The default origin for a target of the given IP version.
72    /// Returns `None` if no origin of that version is configured.
73    fn default_for(&self, is_v6: bool) -> Option<u32> {
74        let default = if is_v6 {
75            self.default_v6
76        } else {
77            self.default_v4
78        };
79        if default.is_none() {
80            warn!(
81                "[CLI] No {} origin is configured.",
82                if is_v6 { "IPv6" } else { "IPv4" }
83            );
84        }
85        default
86    }
87
88    /// Whether an origin of the given IP version is configured.
89    fn has_version(&self, is_v6: bool) -> bool {
90        self.by_id.values().any(|origin| origin.is_v6 == is_v6)
91    }
92}
93
94/// Wraps the feed receiver as a `Stream` so it can be used as a gRPC streaming request.
95pub struct FeedStream {
96    pub(crate) inner: mpsc::Receiver<CliMessage>,
97}
98
99/// Implement `Stream` to enable for async message feeding (rate-limited stream at Orc).
100impl Stream for FeedStream {
101    type Item = CliMessage;
102
103    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
104        self.inner.poll_recv(cx)
105    }
106}
107
108/// Read target lines from stdin and forward them to the live feed.
109///
110/// Each line is a JSON object (e.g., `{"dst":"1.1.1.1","worker":"ams01","origin":2}`),
111/// or a bare address (e.g., `1.1.1.1`).
112/// The optional `worker` field selects the probing worker(s): a worker ID, a hostname, a glob
113/// (e.g. `us-*` — probes the target from every matched worker, spaced by the worker interval),
114/// `"all"` (probe from all workers), or `"any"` (round-robin, default).
115/// The optional `origin` field selects the origin to send from: an origin ID, or
116/// `"all"` (all configured origins). Defaults to the first configured origin of
117/// the target's IP version.
118/// The optional `nprobes` field sets how many measurement probes are sent to
119/// the target (default 1), spaced by the measurement's probe interval.
120/// The optional `ttl` field (feed-trace only) sets the probe TTL (default 255)
121/// The optional `session` field (`--sessions` only) tags the target with a
122/// session; replies carry it back for attribution (ICMP/DNS-A only).
123/// Blocks when the feed channel is full (rate-limiting set by Orchestrator).
124/// Runs on a dedicated thread; dropping the sender (at EOF) signals the end of the feed.
125pub fn read_stdin_feed(
126    feed_tx: mpsc::Sender<CliMessage>,
127    worker_map: BiHashMap<u32, String>,
128    origins: FeedOrigins,
129    is_trace: bool,    // feed-trace measurement (enables the per-target `ttl` field)
130    is_sessions: bool, // feed sessions enabled (--sessions; enables the per-target `session` field)
131) {
132    let stdin = std::io::stdin();
133    for line in stdin.lock().lines() {
134        let Ok(line) = line else {
135            break;
136        };
137        let line = line.trim();
138        if line.is_empty() {
139            continue;
140        }
141
142        let Some(target) = parse_feed_line(line, &worker_map, &origins, is_trace, is_sessions)
143        else {
144            warn!("[CLI] Skipping invalid feed line: {line}");
145            continue;
146        };
147
148        let addr = target.dst.expect("parsed target always has a dst");
149        // Skip origin:all targets when no origin with the same IP version exists
150        if !origins.has_version(addr.is_v6()) {
151            warn!(
152                "[CLI] Skipping target {addr}: no {} origin is configured",
153                if addr.is_v6() { "IPv6" } else { "IPv4" }
154            );
155            continue;
156        }
157
158        let msg = CliMessage {
159            message: Some(cli_message::Message::Targets(TargetBatch {
160                targets: vec![target],
161            })),
162        };
163        if feed_tx.blocking_send(msg).is_err() {
164            break; // Feed closed (measurement ended)
165        }
166    }
167}
168
169/// Parse a single feed line into a live target: an NDJSON object
170/// (e.g., `{"dst":"1.1.1.1","worker":"ams01","origin":2}`) or a bare address (e.g., `1.1.1.1`).
171///
172/// Returns `None` when the line was invalid or matched no worker.
173fn parse_feed_line(
174    line: &str,
175    worker_map: &BiHashMap<u32, String>,
176    origins: &FeedOrigins,
177    is_trace: bool,
178    is_sessions: bool,
179) -> Option<LiveTarget> {
180    // Parse a bare address (with default configs)
181    if !line.starts_with('{') {
182        let dst = line.parse::<Address>().ok()?;
183        return Some(LiveTarget {
184            dst: Some(dst),
185            worker_ids: Vec::new(), // any worker (round-robin)
186            origin_id: origins.default_for(dst.is_v6())?,
187            nprobes: 1,    // TODO use unset value of 0 (defaulting to 1)
188            ttl: 0,        // unset (default 255 for feed-trace)
189            session_id: 0, // no session
190        });
191    }
192
193    // parse NDJSON format
194    parse_feed_object(line, worker_map, origins, is_trace, is_sessions)
195}
196
197/// Parse an NDJSON feed object into a live target carrying its worker selection.
198/// Returns `None` on a malformed object, an unknown worker/origin, or an invalid nprobes/ttl.
199fn parse_feed_object(
200    line: &str,
201    worker_map: &BiHashMap<u32, String>,
202    origins: &FeedOrigins,
203    is_trace: bool,
204    is_sessions: bool,
205) -> Option<LiveTarget> {
206    let value: serde_json::Value = serde_json::from_str(line).ok()?;
207    let dst = value.get("dst")?.as_str()?.parse::<Address>().ok()?;
208    let worker_ids = match value.get("worker") {
209        None => Vec::new(), // any worker (round-robin)
210        Some(worker) => parse_worker(worker, worker_map)?,
211    };
212    let origin_id = match value.get("origin") {
213        None => origins.default_for(dst.is_v6())?,
214        Some(origin) => parse_origin(origin, origins, dst.is_v6())?,
215    };
216    let nprobes = match value.get("nprobes") {
217        None => 1,
218        Some(nprobes) => parse_nprobes(nprobes)?,
219    };
220    let ttl = match value.get("ttl") {
221        None => 0, // unset (default 255 for feed-trace)
222        Some(_) if !is_trace => {
223            warn!("[CLI] The 'ttl' field requires a feed-trace measurement (-m feed-trace).");
224            return None;
225        }
226        Some(ttl) => parse_ttl(ttl)?,
227    };
228    let session_id = match value.get("session") {
229        None => 0, // no session
230        Some(_) if !is_sessions => {
231            warn!("[CLI] Ignoring 'session': sessions are not enabled (start with --sessions).");
232            0
233        }
234        Some(session) => parse_session(session, origin_id, origins, dst.is_v6())?,
235    };
236
237    Some(LiveTarget {
238        dst: Some(dst),
239        worker_ids,
240        origin_id,
241        nprobes,
242        ttl,
243        session_id,
244    })
245}
246
247/// Parse `session`, a 16-bit session ID (0-65535) tagging this target; replies echo it
248/// back so they can be attributed to the submitting session.
249/// When set to 0, the target is not attributed to any session (replies are reported with session 0).
250///
251/// Warns when setting a session for TCP/CHAOS probes (does not support session encoding).
252fn parse_session(
253    session: &serde_json::Value,
254    origin_id: u32,
255    origins: &FeedOrigins,
256    dst_is_v6: bool,
257) -> Option<u32> {
258    let n = match session {
259        // Session as JSON number (e.g., "session":7)
260        serde_json::Value::Number(n) => n.as_u64()?,
261        // Session as numeric string (e.g., "session":"7")
262        serde_json::Value::String(s) => s.parse::<u64>().ok()?,
263        _ => return None,
264    };
265
266    if n > u16::MAX as u64 {
267        warn!("[CLI] '{n}' is not a valid session value (0-65535).");
268        return None;
269    }
270    if n == 0 {
271        return Some(0); // Explicit "no session"
272    }
273
274    // Warn when using CHAOS/TCP that cannot encode sessions in probes
275    let unattributable = match origin_id {
276        ALL_ORIGINS => origins
277            .by_id
278            .values()
279            .any(|origin| origin.is_v6 == dst_is_v6 && !origin.supports_sessions()),
280        id => origins
281            .by_id
282            .get(&id)
283            .is_some_and(|o| !o.supports_sessions()),
284    };
285    if unattributable {
286        warn!(
287            "[CLI] Session {n}: TCP/CHAOS replies cannot echo the session ID; their rows will report session 0."
288        );
289    }
290
291    Some(n as u32)
292}
293
294/// Parse `ttl`, the probe TTL used for this target (feed-trace only, 1-255).
295fn parse_ttl(ttl: &serde_json::Value) -> Option<u32> {
296    let n = match ttl {
297        // TTL as JSON number (e.g., "ttl":12)
298        serde_json::Value::Number(n) => n.as_u64()?,
299        // TTL as numeric string (e.g., "ttl":"12")
300        serde_json::Value::String(s) => s.parse::<u64>().ok()?,
301        _ => return None,
302    };
303
304    if (1..=u8::MAX as u64).contains(&n) {
305        Some(n as u32)
306    } else {
307        warn!("[CLI] '{n}' is not a valid ttl value (1-255).");
308        None
309    }
310}
311
312/// Parse `nprobes`, that specifies the number of probes to send to this target.
313fn parse_nprobes(nprobes: &serde_json::Value) -> Option<u32> {
314    let n = match nprobes {
315        // Probe count as JSON number (e.g., "nprobes":3)
316        serde_json::Value::Number(n) => n.as_u64()?,
317        // Probe count as numeric string (e.g., "nprobes":"3")
318        serde_json::Value::String(s) => s.parse::<u64>().ok()?,
319        _ => return None,
320    };
321
322    if (1..=u8::MAX as u64).contains(&n) {
323        Some(n as u32)
324    } else {
325        warn!("[CLI] '{n}' is not a valid nprobes value (1-255).");
326        None
327    }
328}
329
330/// Resolve a feed line's `origin` value to an origin ID:
331/// an origin ID (number or numeric string) of a configured origin, or `"all"`.
332/// A specific origin must match the target's IP version.
333fn parse_origin(origin: &serde_json::Value, origins: &FeedOrigins, dst_is_v6: bool) -> Option<u32> {
334    let id = match origin {
335        // Origin ID as JSON number (e.g., "origin":2)
336        serde_json::Value::Number(n) => u32::try_from(n.as_u64()?).ok()?,
337        serde_json::Value::String(s) if s == "all" => return Some(ALL_ORIGINS),
338        // Origin ID as numeric string (e.g., "origin":"2")
339        serde_json::Value::String(s) => match s.parse::<u32>() {
340            Ok(id) => id,
341            Err(_) => {
342                warn!("[CLI] '{s}' is not a valid origin ID.");
343                return None;
344            }
345        },
346        _ => return None,
347    };
348
349    // IP version of origin must match the target address
350    match origins.by_id.get(&id) {
351        Some(origin) if origin.is_v6 == dst_is_v6 => Some(id),
352        Some(origin) => {
353            warn!(
354                "[CLI] Origin {id} is {} but the target is {}.",
355                if origin.is_v6 { "IPv6" } else { "IPv4" },
356                if dst_is_v6 { "IPv6" } else { "IPv4" }
357            );
358            None
359        }
360        None => {
361            warn!("[CLI] Origin ID '{id}' is not a configured origin.");
362            None
363        }
364    }
365}
366
367/// Resolve a feed line's `worker` value to the target's `worker_ids`:
368/// `"any"` (round-robin, empty list) or `"all"` (broadcast, `[ALL_WORKERS]`) sentinels,
369/// or a worker ID, hostname, or glob (e.g. `us-*`) resolved via [`resolve_workers`] —
370/// a glob yields every matched worker (probed staggered by the worker interval).
371/// Returns `None` (skip the line) on an unknown worker or a glob that matched nothing.
372fn parse_worker(
373    worker: &serde_json::Value,
374    worker_map: &BiHashMap<u32, String>,
375) -> Option<Vec<u32>> {
376    // Worker ID as JSON number (e.g., "worker":1)
377    if let Some(id) = worker.as_u64() {
378        let id = u32::try_from(id).ok()?;
379        if worker_map.contains_left(&id) {
380            return Some(vec![id]);
381        }
382        warn!("[CLI] Worker ID '{id}' is not a known worker.");
383        return None;
384    }
385
386    let worker = worker.as_str()?;
387    match worker {
388        "any" => Some(Vec::new()),
389        "all" => Some(vec![ALL_WORKERS]),
390        // Worker ID, hostname, or glob
391        _ => {
392            let ids = resolve_workers(worker, worker_map);
393            if ids.is_empty() {
394                warn!("[CLI] '{worker}' did not match any known worker ID or hostname.");
395                return None;
396            }
397            Some(ids)
398        }
399    }
400}
manycastr/cli/feed.rs

manycastr/cli/
feed.rs
