manycastr/worker/inbound/
mod.rs

1use log::info;
2use std::mem::MaybeUninit;
3use std::net::{Ipv6Addr, SocketAddr};
4use std::sync::Arc;
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::thread::{Builder, sleep};
7use std::time::Duration;
8use tokio::sync::mpsc::UnboundedSender;
9
10use crate::custom_module::Separated;
11use crate::custom_module::manycastr::{Address, ProtocolType, Reply, ReplyBatch};
12use crate::worker::inbound::dns::{DnsContext, parse_dns};
13use crate::worker::inbound::ping::parse_icmp;
14use crate::worker::inbound::tcp::parse_tcp;
15use crate::worker::inbound::trace::parse_trace;
16use socket2::{MaybeUninitSlice, MsgHdrMut, SockAddr, Socket};
17
18mod dns;
19mod ping;
20mod tcp;
21mod trace;
22
23/// Configuration for an inbound packet listening worker.
24///
25/// This struct holds all the parameters needed to initialize and run a worker
26/// that listens for and processes incoming measurement packets.
27#[derive(Clone)]
28pub struct InboundConfig {
29    /// The 16-bit measurement ID.
30    pub m_id: u32,
31    /// The unique ID of this specific worker.
32    pub worker_id: u16,
33    /// Protocol used
34    pub p_type: ProtocolType,
35    /// A shared signal that can be used to gracefully shut down the worker.
36    pub abort_s: Arc<AtomicBool>,
37    /// Indicates if the measurement involves traceroute.
38    pub is_traceroute: bool,
39    /// Origin ID associated with the Socket
40    pub origin_id: u32,
41    /// Source port used
42    pub sport: u16,
43    /// Source address used
44    pub src: String,
45    /// Identifies transport traceroute measurements, which require special handling
46    pub is_transport_trace: bool,
47}
48
49/// Metadata of a received reply, shared by all parse functions.
50#[derive(Clone, Copy)]
51pub struct ReplyMeta {
52    /// Source address of the packet
53    pub src: Address,
54    /// TTL / hop limit of the received packet
55    pub ttl: u32,
56    /// Kernel receive timestamp (microseconds since epoch)
57    pub rx_time: u64,
58}
59
60/// Listen for incoming packets
61/// Creates two threads, one that listens on the socket and another that forwards results to the orchestrator and shuts down the receiving socket when appropriate.
62/// Makes sure that the received packets are valid and belong to the current measurement.
63///
64/// # Arguments
65/// * `config` - configuration for the inbound worker thread
66/// * `tx` - sender to put task results in
67/// * `socket` - the socket to listen on
68///
69/// # Panics
70/// If the measurement type is invalid
71pub fn inbound(config: InboundConfig, tx: UnboundedSender<ReplyBatch>, socket: Arc<Socket>) {
72    info!(
73        "[Worker inbound] Started listener (for Origin {})",
74        config.origin_id
75    );
76    let (reply_tx, reply_rx) = std::sync::mpsc::channel::<Reply>();
77    let rx_f_c = config.abort_s.clone();
78    let dns_ctx = DnsContext {
79        is_chaos: config.p_type == ProtocolType::ChaosDns,
80        sport: config.sport,
81        m_id: config.m_id,
82        is_traceroute: config.is_transport_trace,
83    };
84    Builder::new()
85        .name("listener_thread".to_string())
86        .spawn(move || {
87            let mut received: u32 = 0;
88            // Owned by the listener so the packet slice returned by get_packet
89            // (which borrows `buf`) stays valid while we parse it.
90            let mut buf = [MaybeUninit::<u8>::uninit(); 2048];
91            let mut control_buf = [MaybeUninit::<u8>::uninit(); 128];
92            loop {
93                let (packet, meta) = match get_packet(&socket, &mut buf, &mut control_buf) {
94                    Ok(result) => result,
95                    Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
96                        if rx_f_c.load(Ordering::Relaxed) {
97                            break;
98                        }
99                        continue;
100                    }
101                    Err(e) => panic!("Socket error: {}", e),
102                };
103
104                let result = match (config.is_traceroute, config.p_type) {
105                    (true, _) => parse_trace(packet, config.m_id, meta),
106
107                    (_, ProtocolType::Icmp) => parse_icmp(packet, config.m_id, false, meta),
108
109                    (_, ProtocolType::ADns) | (_, ProtocolType::ChaosDns) => {
110                        parse_dns(packet, meta, &dns_ctx)
111                    }
112
113                    (_, ProtocolType::Tcp) => {
114                        parse_tcp(packet, config.sport, config.is_transport_trace, meta)
115                    }
116                };
117
118                if let Some(reply) = result {
119                    received += 1;
120                    let _ = reply_tx.send(reply);
121                }
122            }
123
124            if config.p_type == ProtocolType::Icmp {
125                info!(
126                    "[Worker inbound] Stopped ICMP ping listener {} (received {} packets)",
127                    config.src,
128                    received.with_separator(),
129                )
130            } else {
131                info!(
132                    "[Worker inbound] Stopped {} listener {}:{} (received {} packets)",
133                    config.p_type,
134                    config.src,
135                    config.sport,
136                    received.with_separator(),
137                );
138            }
139        })
140        .expect("Failed to spawn listener_thread");
141
142    Builder::new()
143        .name("result_sender_thread".to_string())
144        .spawn(move || {
145            handle_results(
146                &tx,
147                config.abort_s,
148                config.worker_id,
149                reply_rx,
150                config.origin_id,
151            );
152        })
153        .expect("Failed to spawn result_sender_thread");
154}
155
156/// Timestamp encodings for various measurement types
157#[derive(Clone, Copy)]
158pub(crate) enum TxEncoding {
159    /// Full 64-bit microsecond epoch (ICMP, DNS) (microseconds)
160    Micros,
161    /// 21-bit microsecond epoch (TCP probes) (microseconds)
162    Tcp21,
163    /// 14-bit millisecond epoch (traceroute hops) (microseconds)
164    Trace14,
165}
166
167/// Compute the RTT in milliseconds based on the timestamp encoding.
168pub(crate) fn rtt_ms(rx_time_us: u64, tx_time: u64, enc: TxEncoding) -> f32 {
169    match enc {
170        TxEncoding::Tcp21 => {
171            const MODULUS: u64 = 1 << 21; // 21-bit microseconds
172            const MASK: u64 = MODULUS - 1;
173            let rx = rx_time_us & MASK;
174            let tx = tx_time & MASK;
175            let us = if rx >= tx { rx - tx } else { rx + MODULUS - tx };
176            us as f32 / 1_000.0
177        }
178        TxEncoding::Trace14 => {
179            const MODULUS: u64 = 1 << 14; // 14-bit milliseconds
180            const MASK: u64 = MODULUS - 1;
181            let rx = (rx_time_us / 1_000) & MASK;
182            let ms = if rx >= tx_time {
183                rx - tx_time
184            } else {
185                rx + MODULUS - tx_time
186            };
187            ms as f32
188        }
189        // Full microsecond epoch on both ends
190        TxEncoding::Micros => (rx_time_us as i64 - tx_time as i64) as f32 / 1_000.0,
191    }
192}
193
194/// Get a packet from a socket along with its metadata (source address, TTL, and
195/// kernel receive timestamp).
196fn get_packet<'a>(
197    socket: &Socket,
198    buf: &'a mut [MaybeUninit<u8>],
199    control_buf: &'a mut [MaybeUninit<u8>],
200) -> Result<(&'a [u8], ReplyMeta), std::io::Error> {
201    let mut source_storage: SockAddr = SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 0).into();
202
203    let recv_result = {
204        let mut iov_buf = [MaybeUninitSlice::new(&mut buf[..])];
205        let mut msg = MsgHdrMut::new()
206            .with_addr(&mut source_storage)
207            .with_buffers(&mut iov_buf)
208            .with_control(&mut control_buf[..]);
209
210        socket.recvmsg(&mut msg, 0).map(|n| (n, msg.control_len()))
211    };
212
213    match recv_result {
214        Ok((bytes_read, control_len)) => {
215            let source = source_storage
216                .as_socket()
217                .ok_or_else(|| std::io::Error::other("invalid source address"))?;
218
219            let (packet_data, ancillary_data) = unsafe {
220                // TODO remove unsafe
221                let p = std::slice::from_raw_parts(buf.as_ptr() as *const u8, bytes_read);
222                let c = std::slice::from_raw_parts(control_buf.as_ptr() as *const u8, control_len);
223                (p, c)
224            };
225
226            // TODO find better approach for obtaining hop_limit/ttl
227            let hop_limit = if source.is_ipv6() {
228                // IPv6 header is never included in received data
229                parse_hop_limit(ancillary_data).unwrap_or(0)
230            } else {
231                // IPv4 raw socket includes the IP header; TTL is at byte 8
232                packet_data[8] as u32
233            };
234            // Get timestamp from the kernel, or current time if unavailable
235            let rx_time = parse_kernel_timestamp(ancillary_data).unwrap_or_else(|| {
236                std::time::SystemTime::now()
237                    .duration_since(std::time::UNIX_EPOCH)
238                    .unwrap()
239                    .as_micros() as u64
240            });
241            Ok((
242                packet_data,
243                ReplyMeta {
244                    src: source.into(),
245                    ttl: hop_limit,
246                    rx_time,
247                },
248            ))
249        }
250        Err(e) => Err(e),
251    }
252}
253
254/// Find the payload of the first control message matching (level, type) in an
255/// ancillary data buffer. Returns the buffer remainder starting at the payload.
256/// cmsghdr on 64-bit: [0..8] len, [8..12] level, [12..16] type, [16..] payload
257fn find_cmsg(data: &[u8], level: i32, type_: i32) -> Option<&[u8]> {
258    let mut pos = 0;
259    while pos + 16 <= data.len() {
260        let cmsg_len = usize::from_ne_bytes(data[pos..pos + 8].try_into().ok()?);
261        let cmsg_level = i32::from_ne_bytes(data[pos + 8..pos + 12].try_into().ok()?);
262        let cmsg_type = i32::from_ne_bytes(data[pos + 12..pos + 16].try_into().ok()?);
263
264        if cmsg_level == level && cmsg_type == type_ {
265            return Some(&data[pos + 16..]);
266        }
267
268        if cmsg_len == 0 {
269            break;
270        }
271        pos += (cmsg_len + 7) & !7; // Align to 8-byte boundary
272    }
273    None
274}
275
276/// Retrieve kernel-provided receive timestamp (SO_TIMESTAMP) from ancillary data.
277/// Returns microseconds since epoch.
278fn parse_kernel_timestamp(data: &[u8]) -> Option<u64> {
279    let payload = find_cmsg(data, libc::SOL_SOCKET, libc::SCM_TIMESTAMP)?;
280    let secs = i64::from_ne_bytes(payload.get(..8)?.try_into().ok()?);
281    let usecs = i64::from_ne_bytes(payload.get(8..16)?.try_into().ok()?);
282    Some(secs as u64 * 1_000_000 + usecs as u64)
283}
284
285/// Retrieve IPv6 hop limit from the ancillary_data buffer bytes
286fn parse_hop_limit(data: &[u8]) -> Option<u32> {
287    // IPv6 Hop Limit: Level 41 (IPPROTO_IPV6), Type 52 (IPV6_HOPLIMIT)
288    Some(*find_cmsg(data, 41, 52)?.first()? as u32)
289}
290
291/// Forward results to the Worker handler
292///
293/// # Arguments
294/// * `tx` - sender to put task results in
295/// * `rx_f` - channel that is used to signal the end of the measurement
296/// * `worker_id` - the unique worker ID of this worker
297/// * `reply_rx` - channel receiver for replies from the listener thread
298/// * `origin_id` - origin ID associated with this inbound handler
299fn handle_results(
300    tx: &UnboundedSender<ReplyBatch>,
301    rx_f: Arc<AtomicBool>,
302    worker_id: u16,
303    reply_rx: std::sync::mpsc::Receiver<Reply>,
304    origin_id: u32,
305) {
306    loop {
307        sleep(Duration::from_secs(1));
308
309        let mut rq = Vec::new();
310        while let Ok(reply) = reply_rx.try_recv() {
311            rq.push(reply);
312        }
313
314        if !rq.is_empty() {
315            tx.send(ReplyBatch {
316                rx_id: worker_id as u32,
317                results: rq,
318                origin_id,
319            })
320            .expect("Failed to send TaskResult to worker handler");
321        }
322
323        if rx_f.load(Ordering::SeqCst) {
324            tx.send(ReplyBatch::default())
325                .expect("Failed to send 'finished' signal to orchestrator");
326            break;
327        }
328    }
329}
manycastr/worker/inbound/mod.rs

manycastr/worker/inbound/
mod.rs
