manycastr/worker/
measurement.rs

1use crate::custom_module::manycastr::{
2    Finished, MeasurementType, Origin, ProtocolType, ReplyBatch, Start,
3};
4use crate::dns_identifier;
5use crate::worker::bpf::{
6    attach_dns_filter, attach_icmp_filter, attach_tcp_filter, attach_traceroute_filter,
7};
8use crate::worker::config::{Worker, set_unicast_origins};
9use crate::worker::inbound::{InboundConfig, inbound};
10use crate::worker::outbound::{OutboundConfig, outbound};
11use log::{error, info, warn};
12use socket2::{Domain, Protocol, SockAddr, Socket, Type};
13use std::error::Error;
14use std::net::{IpAddr, SocketAddr};
15use std::os::fd::AsRawFd;
16use std::sync::Arc;
17use std::sync::atomic::{AtomicBool, Ordering};
18
19impl Worker {
20    /// Initialize a new measurement by creating outbound and inbound threads, and ensures task results are sent back to the orchestrator.
21    ///
22    /// Extracts the protocol type from the measurement definition, and determines which source address to use.
23    /// Creates a socket to send out probes and receive replies with, calls the appropriate inbound & outbound functions.
24    /// Creates an additional thread that forwards task results to the orchestrator.
25    ///
26    /// # Arguments
27    /// * `start` - Definition of the new measurement
28    /// * `worker_id` - the unique ID of this worker
29    /// * `abort_outbound` - Forcefully signal the outbound thread to stop sending probes
30    pub(crate) fn init(
31        &mut self,
32        start: Start,
33        worker_id: u16,
34        abort_outbound: Arc<AtomicBool>,
35    ) -> Result<(), Box<dyn Error>> {
36        let m_id = start.m_id;
37        let m_type = start.m_type();
38
39        // Channel for sending from inbound to the orchestrator forwarder thread
40        let (inbound_tx, mut inbound_rx) = tokio::sync::mpsc::unbounded_channel();
41
42        // Replace unicast placeholder addresses in rx_origins, tx_origins with local addresses
43        let rx_origins = set_unicast_origins(start.rx_origins);
44        let tx_origins = set_unicast_origins(start.tx_origins);
45        let tx_origin_ids: std::collections::HashSet<_> =
46            tx_origins.iter().map(|o| o.origin_id).collect();
47
48        // Traceroute mode: raw-only sockets, ICMP Time Exceeded BPF filter, trace reply parsing
49        let is_traceroute = matches!(
50            m_type,
51            MeasurementType::AnycastTraceroute
52                | MeasurementType::Tracemap
53                | MeasurementType::FeedTrace
54        );
55
56        // Start inbound/outbound threads for each origin
57        for rx_origin in rx_origins {
58            // The IP version is a per-origin property (mixed-version measurements)
59            let is_ipv6 = rx_origin.src.expect("no src").is_v6();
60            let is_transport_traceroute =
61                is_traceroute && !matches!(rx_origin.p_type(), ProtocolType::Icmp);
62
63            // UDP/TCP (Paris) traceroute uses two sockets (ICMP and UDP/TCP)
64            let (rx_socket, tx_socket) = if is_transport_traceroute {
65                let rx = Self::get_socket(
66                    is_ipv6,
67                    ProtocolType::Icmp,
68                    rx_origin,
69                    true, // Attaches Time Exceeded + Dest Unreachable BPF filter
70                    m_id,
71                );
72                let tx = Self::get_socket(
73                    is_ipv6,
74                    rx_origin.p_type(),
75                    rx_origin,
76                    true, // TTL + checksum control
77                    m_id,
78                );
79                (rx, tx)
80            } else {
81                let socket =
82                    Self::get_socket(is_ipv6, rx_origin.p_type(), rx_origin, is_traceroute, m_id);
83                (socket.clone(), socket)
84            };
85
86            let inbound_config = InboundConfig {
87                m_id,
88                worker_id,
89                p_type: rx_origin.p_type(),
90                abort_s: self.abort_inbound.clone(),
91                is_traceroute,
92                origin_id: rx_origin.origin_id,
93                sport: rx_origin.sport as u16,
94                src: rx_origin.src.expect("no src").to_string(),
95                is_transport_trace: false,
96            };
97
98            // For transport traceroute, listen on the raw transport socket for discovery replies
99            if is_transport_traceroute {
100                inbound(
101                    InboundConfig {
102                        is_traceroute: false, // parse as normal DNS/TCP discovery replies
103                        is_transport_trace: true,
104                        ..inbound_config.clone()
105                    },
106                    inbound_tx.clone(),
107                    tx_socket.clone(),
108                );
109            }
110
111            // Primary listener (ICMP trace replies for transport traceroute)
112            inbound(inbound_config, inbound_tx.clone(), rx_socket);
113
114            // See if this origin_id is in tx_origins
115            if tx_origin_ids.contains(&rx_origin.origin_id) {
116                self.log_probe_details(&rx_origin);
117
118                // Channel for forwarding tasks to outbound
119                let (outbound_tx, outbound_rx) = tokio::sync::mpsc::channel(1000);
120                self.outbound_txs.push((rx_origin.origin_id, outbound_tx));
121
122                let outbound_handle = outbound(
123                    OutboundConfig {
124                        worker_id,
125                        abort_outbound: abort_outbound.clone(),
126                        m_id,
127                        p_type: rx_origin.p_type(),
128                        qname: start.record.clone(),
129                        info_url: start.url.clone(),
130                        probing_rate: start.rate / tx_origins.len() as u32, // Adjust probing rate for multiple origins
131                        src: rx_origin.src.unwrap(),
132                        sport: rx_origin.sport as u16,
133                        dport: rx_origin.dport as u16,
134                        origin_id: rx_origin.origin_id,
135                    },
136                    outbound_rx,
137                    tx_socket,
138                );
139                self.outbound_handles.push(outbound_handle);
140            }
141        }
142
143        // Spawn thread to forward reply batches to the CLI
144        let is_busy = self.is_busy.clone();
145        let mut grpc_client_clone = self.grpc_client.clone();
146        tokio::spawn(async move {
147            while let Some(batch) = inbound_rx.recv().await {
148                if batch == ReplyBatch::default() {
149                    // Mark the worker as idle (no active measurement)
150                    is_busy.store(false, Ordering::SeqCst);
151                    info!(
152                        "[Worker] Letting the orchestrator know that this worker finished the measurement"
153                    );
154                    let _ = grpc_client_clone
155                        .measurement_finished(Finished {
156                            m_id,
157                            worker_id: worker_id.into(),
158                        })
159                        .await;
160                    break;
161                }
162
163                // TODO retry with backoff instead of breaking
164                if let Err(e) = grpc_client_clone.send_result(batch).await {
165                    error!("[Worker] Failed to forward batch: {e}");
166                    break;
167                }
168            }
169            inbound_rx.close();
170        });
171
172        Ok(())
173    }
174
175    /// Print the Origins (i.e., source address and port values) used for this measurement
176    ///
177    /// # Arguments
178    /// * `p_type` - Protocol used
179    /// * `origins` - Sending origins used by this Worker
180    fn log_probe_details(&self, origin: &Origin) {
181        match origin.p_type() {
182            ProtocolType::Icmp => info!(
183                "[Worker] Sending {} on: {} using ICMP ID {}",
184                origin.p_type(),
185                origin.src.unwrap(),
186                origin.dport
187            ),
188            _ => info!(
189                "[Worker] Sending {} on: {}, {}:{}",
190                origin.p_type(),
191                origin.src.unwrap(),
192                origin.sport,
193                origin.dport
194            ),
195        }
196    }
197
198    /// Obtain a raw socket for the given IP version and protocol.
199    ///
200    /// # Arguments
201    /// * `is_ipv6` - IP version used (true: IPv6)
202    /// * `p_type` - Protocol type used (ICMP, UDP, or TCP)
203    /// * `origin` - Origin used in this measurement (anycast or local unicast address)
204    /// * `is_traceroute` - Whether this is a traceroute measurement
205    ///
206    /// # Returns
207    /// `Arc<Socket>` containing the raw socket to send/receive from
208    fn get_socket(
209        is_ipv6: bool,
210        p_type: ProtocolType,
211        origin: Origin,
212        is_traceroute: bool,
213        m_id: u32,
214    ) -> Arc<Socket> {
215        let domain = if is_ipv6 { Domain::IPV6 } else { Domain::IPV4 };
216
217        let protocol = match p_type {
218            ProtocolType::Icmp => {
219                if is_ipv6 {
220                    Protocol::ICMPV6
221                } else {
222                    Protocol::ICMPV4
223                }
224            }
225            ProtocolType::Tcp => Protocol::TCP,
226            ProtocolType::ADns | ProtocolType::ChaosDns => Protocol::UDP,
227        };
228
229        let addr: IpAddr = (origin.src.as_ref().expect("no src")).into();
230
231        let socket = Self::try_raw_socket(domain, protocol, is_ipv6)
232            .expect("Failed to create raw socket. sudo or CAP_NET_RAW required.");
233
234        let sock_addr = SockAddr::from(SocketAddr::new(addr, origin.sport as u16));
235        socket
236            .bind(&sock_addr)
237            .expect("Failed to bind socket to address.");
238
239        // Attach a cBPF filter so the kernel drops non-matching packets
240        let filter = match p_type {
241            ProtocolType::Icmp if is_traceroute => (
242                // Time-Exceeded + Echo-Reply by type (identifier is per-probe dynamic)
243                attach_traceroute_filter(&socket, is_ipv6),
244                "ICMP traceroute".to_string(),
245            ),
246            ProtocolType::Icmp => (
247                // Echo replies with id == dport
248                attach_icmp_filter(&socket, origin.dport as u16, is_ipv6),
249                format!("ICMP (id {})", origin.dport),
250            ),
251            ProtocolType::Tcp => (
252                // RST flag + sport filtering
253                attach_tcp_filter(&socket, origin.sport as u16, is_ipv6),
254                format!("TCP RST (sport {})", origin.sport),
255            ),
256            // DNS Identifier + sport filtering
257            ProtocolType::ADns | ProtocolType::ChaosDns => (
258                attach_dns_filter(&socket, origin.sport as u16, dns_identifier(m_id), is_ipv6),
259                format!("DNS (sport {})", origin.sport),
260            ),
261        };
262        match filter {
263            (Ok(()), desc) => info!("[Worker] Attached {desc} BPF filter"),
264            (Err(e), desc) => warn!("[Worker] Failed to attach {desc} BPF filter: {e}"),
265        }
266
267        socket.set_send_buffer_size(4 * 1024 * 1024).ok(); // 4 MB buffer for sending
268        socket.set_recv_buffer_size(16 * 1024 * 1024).ok(); // 16 MB for receiving
269
270        // enable SO_TIMESTAMP (get kernel timestamp when packet is received)
271        let ts_ret = unsafe {
272            let val: libc::c_int = 1;
273            libc::setsockopt(
274                socket.as_raw_fd(),
275                libc::SOL_SOCKET,
276                libc::SO_TIMESTAMP,
277                &val as *const _ as *const libc::c_void,
278                size_of::<libc::c_int>() as libc::socklen_t,
279            )
280        };
281        if ts_ret != 0 {
282            warn!(
283                "[Worker] Failed to enable SO_TIMESTAMP: {}",
284                std::io::Error::last_os_error()
285            );
286        }
287
288        socket
289            .set_read_timeout(Some(std::time::Duration::from_millis(1)))
290            .expect("Failed to set read timeout");
291
292        Arc::new(socket)
293    }
294
295    /// Try to create a RAW socket with IP_HDRINCL and appropriate options.
296    /// Returns None if creation or setup fails (e.g. missing CAP_NET_RAW).
297    fn try_raw_socket(domain: Domain, protocol: Protocol, is_ipv6: bool) -> Option<Socket> {
298        let socket = match Socket::new(domain, Type::RAW, Some(protocol)) {
299            Ok(s) => s,
300            Err(e) => {
301                warn!("[Worker] RAW socket creation failed: {e}");
302                return None;
303            }
304        };
305
306        let res = if is_ipv6 {
307            // Always request the received hop limit as ancillary data.
308            let r = socket.set_recv_hoplimit_v6(true);
309
310            if protocol == Protocol::ICMPV6 {
311                r
312            } else {
313                r.and_then(|_| socket.set_header_included_v6(true))
314            }
315        } else {
316            socket.set_header_included_v4(true)
317        };
318        if let Err(e) = res {
319            warn!("[Worker] RAW socket option failed: {e}");
320            return None;
321        }
322
323        Some(socket)
324    }
325}
manycastr/worker/measurement.rs

manycastr/worker/
measurement.rs
