From 3ff6a072252ab3747afe2eb45b2b6c4e4e704b6f Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Thu, 7 May 2026 17:58:40 +0700 Subject: [PATCH] feat(client): add Linux ancillary receive metadata --- Cargo.lock | 1 + crates/irtt-client/Cargo.toml | 2 + .../src/client/tests/probes_replies.rs | 23 +- crates/irtt-client/src/lib.rs | 5 +- crates/irtt-client/src/receive.rs | 30 +- crates/irtt-client/src/receive/linux.rs | 278 ++++++++++++++++++ crates/irtt-client/src/socket.rs | 6 + 7 files changed, 339 insertions(+), 6 deletions(-) create mode 100644 crates/irtt-client/src/receive/linux.rs diff --git a/Cargo.lock b/Cargo.lock index e3bda24..e55760a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -174,6 +174,7 @@ name = "irtt-client" version = "0.1.0" dependencies = [ "irtt-proto", + "libc", "socket2", "thiserror", ] diff --git a/crates/irtt-client/Cargo.toml b/crates/irtt-client/Cargo.toml index 245c5e0..71296a9 100644 --- a/crates/irtt-client/Cargo.toml +++ b/crates/irtt-client/Cargo.toml @@ -7,8 +7,10 @@ rust-version.workspace = true [dependencies] irtt-proto = { path = "../irtt-proto" } +libc = { version = "0.2", optional = true } socket2 = { version = "0.5", features = ["all"] } thiserror = "2" [features] default = [] +ancillary = ["dep:libc"] diff --git a/crates/irtt-client/src/client/tests/probes_replies.rs b/crates/irtt-client/src/client/tests/probes_replies.rs index c89c31f..c3cc14a 100644 --- a/crates/irtt-client/src/client/tests/probes_replies.rs +++ b/crates/irtt-client/src/client/tests/probes_replies.rs @@ -299,10 +299,25 @@ fn recv_once_decodes_echo_reply_and_emits_event() { assert!(server_timing.is_some()); let st = server_timing.as_ref().unwrap(); assert!(st.processing.is_some()); - assert_eq!(packet_meta.traffic_class, None); - assert_eq!(packet_meta.dscp, None); - assert_eq!(packet_meta.ecn, None); - assert_eq!(packet_meta.kernel_rx_timestamp, None); + #[cfg(not(all(target_os = "linux", feature = "ancillary")))] + { + assert_eq!(packet_meta.traffic_class, None); + assert_eq!(packet_meta.dscp, None); + assert_eq!(packet_meta.ecn, None); + assert_eq!(packet_meta.kernel_rx_timestamp, None); + } + #[cfg(all(target_os = "linux", feature = "ancillary"))] + { + assert_eq!( + packet_meta.dscp, + packet_meta.traffic_class.map(|tc| tc >> 2) + ); + assert_eq!( + packet_meta.ecn, + packet_meta.traffic_class.map(|tc| tc & 0b11) + ); + assert_eq!(packet_meta.kernel_rx_timestamp, None); + } } other => panic!("expected EchoReply, got {other:?}"), } diff --git a/crates/irtt-client/src/lib.rs b/crates/irtt-client/src/lib.rs index 06a36d6..f418e8e 100644 --- a/crates/irtt-client/src/lib.rs +++ b/crates/irtt-client/src/lib.rs @@ -1,4 +1,7 @@ -#![forbid(unsafe_code)] +#![cfg_attr( + not(all(target_os = "linux", feature = "ancillary")), + forbid(unsafe_code) +)] mod client; mod config; diff --git a/crates/irtt-client/src/receive.rs b/crates/irtt-client/src/receive.rs index fbfd39d..9bb64d3 100644 --- a/crates/irtt-client/src/receive.rs +++ b/crates/irtt-client/src/receive.rs @@ -1,7 +1,13 @@ -use std::{io, net::UdpSocket}; +use std::{ + io, + net::{SocketAddr, UdpSocket}, +}; use crate::{metadata::ReceiveMeta, timing::ClientTimestamp}; +#[cfg(all(target_os = "linux", feature = "ancillary"))] +mod linux; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct ReceivedDatagram { pub(crate) len: usize, @@ -9,6 +15,7 @@ pub(crate) struct ReceivedDatagram { pub(crate) meta: ReceiveMeta, } +#[cfg(not(all(target_os = "linux", feature = "ancillary")))] pub(crate) fn recv_datagram( socket: &UdpSocket, buf: &mut [u8], @@ -23,6 +30,27 @@ pub(crate) fn recv_datagram( }) } +#[cfg(all(target_os = "linux", feature = "ancillary"))] +pub(crate) fn recv_datagram( + socket: &UdpSocket, + buf: &mut [u8], +) -> Result { + linux::recv_datagram(socket, buf) +} + +#[cfg(not(all(target_os = "linux", feature = "ancillary")))] +pub(crate) fn configure_receive_metadata( + _socket: &UdpSocket, + _remote: SocketAddr, +) -> io::Result<()> { + Ok(()) +} + +#[cfg(all(target_os = "linux", feature = "ancillary"))] +pub(crate) fn configure_receive_metadata(socket: &UdpSocket, remote: SocketAddr) -> io::Result<()> { + linux::configure_receive_metadata(socket, remote) +} + #[cfg(test)] mod tests { use std::net::UdpSocket; diff --git a/crates/irtt-client/src/receive/linux.rs b/crates/irtt-client/src/receive/linux.rs new file mode 100644 index 0000000..16a424b --- /dev/null +++ b/crates/irtt-client/src/receive/linux.rs @@ -0,0 +1,278 @@ +use std::{ + io, mem, + net::{SocketAddr, UdpSocket}, + os::fd::AsRawFd, + ptr, +}; + +use crate::{metadata::ReceiveMeta, receive::ReceivedDatagram, timing::ClientTimestamp}; + +const CONTROL_LEN: usize = 64; + +pub(crate) fn configure_receive_metadata(socket: &UdpSocket, remote: SocketAddr) -> io::Result<()> { + let socket = socket2::SockRef::from(socket); + if remote.is_ipv4() { + socket.set_recv_tos(true) + } else { + socket.set_recv_tclass_v6(true) + } +} + +pub(crate) fn recv_datagram( + socket: &UdpSocket, + buf: &mut [u8], +) -> Result { + let mut control = ControlBuffer::new(); + let mut iov = libc::iovec { + iov_base: buf.as_mut_ptr().cast(), + iov_len: buf.len(), + }; + let mut msg: libc::msghdr = unsafe { mem::zeroed() }; + msg.msg_iov = &mut iov; + msg.msg_iovlen = 1; + msg.msg_control = control.as_mut_ptr().cast(); + msg.msg_controllen = control.len(); + + let len = unsafe { + // SAFETY: `msg` points to one writable iovec backed by `buf`, and the + // control buffer is writable and lives until `recvmsg` returns. The + // socket file descriptor is borrowed from a valid `UdpSocket`. + libc::recvmsg(socket.as_raw_fd(), &mut msg, 0) + }; + if len < 0 { + return Err(io::Error::last_os_error()); + } + let received_at = ClientTimestamp::now(); + let meta = unsafe { + // SAFETY: `msg` was initialized by a successful `recvmsg` call. The + // parser only reads cmsghdr entries within `msg_controllen` and ignores + // short or unrelated control messages. + parse_receive_meta(&msg) + }; + + Ok(ReceivedDatagram { + len: len as usize, + received_at, + meta, + }) +} + +unsafe fn parse_receive_meta(msg: &libc::msghdr) -> ReceiveMeta { + let mut meta = ReceiveMeta::default(); + let mut cmsg = libc::CMSG_FIRSTHDR(msg); + while !cmsg.is_null() { + let cmsg_ref = &*cmsg; + if cmsg_ref.cmsg_len >= libc::CMSG_LEN(0) as usize { + let data_len = cmsg_ref.cmsg_len - libc::CMSG_LEN(0) as usize; + if is_traffic_class_cmsg(cmsg_ref) { + meta.traffic_class = read_int_cmsg_low_byte(cmsg, data_len); + } + } + cmsg = libc::CMSG_NXTHDR(msg, cmsg); + } + meta +} + +fn is_traffic_class_cmsg(cmsg: &libc::cmsghdr) -> bool { + (cmsg.cmsg_level == libc::IPPROTO_IP && cmsg.cmsg_type == libc::IP_TOS) + || (cmsg.cmsg_level == libc::IPPROTO_IPV6 && cmsg.cmsg_type == libc::IPV6_TCLASS) +} + +unsafe fn read_u8_cmsg(cmsg: *mut libc::cmsghdr, data_len: usize) -> Option { + if data_len >= mem::size_of::() { + Some(*libc::CMSG_DATA(cmsg).cast::()) + } else { + None + } +} + +unsafe fn read_int_cmsg_low_byte(cmsg: *mut libc::cmsghdr, data_len: usize) -> Option { + if data_len >= mem::size_of::() { + let value = ptr::read_unaligned(libc::CMSG_DATA(cmsg).cast::()); + Some((value & 0xff) as u8) + } else { + read_u8_cmsg(cmsg, data_len) + } +} + +#[repr(align(8))] +struct ControlBuffer([u8; CONTROL_LEN]); + +impl ControlBuffer { + fn new() -> Self { + Self([0; CONTROL_LEN]) + } + + fn as_mut_ptr(&mut self) -> *mut u8 { + self.0.as_mut_ptr() + } + + fn len(&self) -> usize { + self.0.len() + } +} + +#[cfg(test)] +mod tests { + use std::{ + io, + net::{SocketAddr, UdpSocket}, + time::Duration, + }; + + use crate::{ + event::PacketMeta, + receive::{configure_receive_metadata, recv_datagram}, + socket_options::apply_dscp_to_socket, + timing::ClientTimestamp, + }; + + fn connected_ipv4_loopback_pair() -> (UdpSocket, UdpSocket) { + let a = UdpSocket::bind(SocketAddr::from(([127, 0, 0, 1], 0))).unwrap(); + let b = UdpSocket::bind(SocketAddr::from(([127, 0, 0, 1], 0))).unwrap(); + a.connect(b.local_addr().unwrap()).unwrap(); + b.connect(a.local_addr().unwrap()).unwrap(); + a.set_read_timeout(Some(Duration::from_millis(200))) + .unwrap(); + b.set_read_timeout(Some(Duration::from_millis(200))) + .unwrap(); + (a, b) + } + + fn connected_ipv6_loopback_pair() -> Option<(UdpSocket, UdpSocket)> { + let bind_addr = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], 0)); + let a = match UdpSocket::bind(bind_addr) { + Ok(socket) => socket, + Err(error) if is_unavailable_ipv6_loopback(&error) => { + eprintln!( + "skipping IPv6 ancillary receive test: IPv6 loopback unavailable: {error}" + ); + return None; + } + Err(error) => panic!("{error}"), + }; + let b = match UdpSocket::bind(bind_addr) { + Ok(socket) => socket, + Err(error) if is_unavailable_ipv6_loopback(&error) => { + eprintln!( + "skipping IPv6 ancillary receive test: IPv6 loopback unavailable: {error}" + ); + return None; + } + Err(error) => panic!("{error}"), + }; + match a.connect(b.local_addr().unwrap()) { + Ok(()) => {} + Err(error) if is_unavailable_ipv6_loopback(&error) => { + eprintln!( + "skipping IPv6 ancillary receive test: IPv6 loopback unavailable: {error}" + ); + return None; + } + Err(error) => panic!("{error}"), + } + match b.connect(a.local_addr().unwrap()) { + Ok(()) => {} + Err(error) if is_unavailable_ipv6_loopback(&error) => { + eprintln!( + "skipping IPv6 ancillary receive test: IPv6 loopback unavailable: {error}" + ); + return None; + } + Err(error) => panic!("{error}"), + } + a.set_read_timeout(Some(Duration::from_millis(200))) + .unwrap(); + b.set_read_timeout(Some(Duration::from_millis(200))) + .unwrap(); + Some((a, b)) + } + + #[test] + fn recvmsg_receive_returns_length_payload_and_timestamp() { + let (sender, receiver) = connected_ipv4_loopback_pair(); + configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap(); + sender.send(b"hello").unwrap(); + + let before = ClientTimestamp::now(); + let mut buf = [0_u8; 16]; + let datagram = recv_datagram(&receiver, &mut buf).unwrap(); + let after = ClientTimestamp::now(); + + assert_eq!(datagram.len, 5); + assert_eq!(&buf[..datagram.len], b"hello"); + assert!(datagram.received_at.mono >= before.mono); + assert!(datagram.received_at.mono <= after.mono); + } + + #[test] + fn ipv4_traffic_class_metadata_is_observed_when_kernel_provides_it() { + let (sender, receiver) = connected_ipv4_loopback_pair(); + configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap(); + apply_dscp_to_socket(&sender, receiver.local_addr().unwrap(), 46).unwrap(); + sender.send(b"dscp").unwrap(); + + let mut buf = [0_u8; 16]; + let datagram = recv_datagram(&receiver, &mut buf).unwrap(); + let Some(traffic_class) = datagram.meta.traffic_class else { + eprintln!("skipping IPv4 ancillary metadata assertion: kernel did not provide IP_TOS"); + return; + }; + + let packet_meta = PacketMeta::from(datagram.meta); + assert_eq!(traffic_class & 0xfc, 184); + assert_eq!(packet_meta.dscp, Some(46)); + assert_eq!(packet_meta.ecn, Some(0)); + } + + #[test] + fn ipv6_traffic_class_metadata_is_observed_when_kernel_provides_it() { + let Some((sender, receiver)) = connected_ipv6_loopback_pair() else { + return; + }; + configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap(); + apply_dscp_to_socket(&sender, receiver.local_addr().unwrap(), 46).unwrap(); + sender.send(b"dscp").unwrap(); + + let mut buf = [0_u8; 16]; + let datagram = recv_datagram(&receiver, &mut buf).unwrap(); + let Some(traffic_class) = datagram.meta.traffic_class else { + eprintln!( + "skipping IPv6 ancillary metadata assertion: kernel did not provide IPV6_TCLASS" + ); + return; + }; + + let packet_meta = PacketMeta::from(datagram.meta); + assert_eq!(traffic_class & 0xfc, 184); + assert_eq!(packet_meta.dscp, Some(46)); + assert_eq!(packet_meta.ecn, Some(0)); + } + + #[test] + fn observed_zero_traffic_class_preserves_some_zero() { + let (sender, receiver) = connected_ipv4_loopback_pair(); + configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap(); + sender.send(b"zero").unwrap(); + + let mut buf = [0_u8; 16]; + let datagram = recv_datagram(&receiver, &mut buf).unwrap(); + let Some(traffic_class) = datagram.meta.traffic_class else { + eprintln!("skipping observed-zero assertion: kernel did not provide IP_TOS"); + return; + }; + + let packet_meta = PacketMeta::from(datagram.meta); + assert_eq!(traffic_class, 0); + assert_eq!(packet_meta.traffic_class, Some(0)); + assert_eq!(packet_meta.dscp, Some(0)); + assert_eq!(packet_meta.ecn, Some(0)); + } + + fn is_unavailable_ipv6_loopback(error: &io::Error) -> bool { + matches!( + error.kind(), + io::ErrorKind::AddrNotAvailable | io::ErrorKind::Unsupported + ) + } +} diff --git a/crates/irtt-client/src/socket.rs b/crates/irtt-client/src/socket.rs index 672885a..6b8c134 100644 --- a/crates/irtt-client/src/socket.rs +++ b/crates/irtt-client/src/socket.rs @@ -8,6 +8,7 @@ use socket2::{Domain, Protocol, Socket, Type}; use crate::{ config::{ClientConfig, SocketConfig, DEFAULT_PORT, MIN_OPEN_TIMEOUT}, error::ClientError, + receive::configure_receive_metadata, socket_options::apply_ttl_to_socket, }; @@ -86,6 +87,11 @@ pub(crate) fn connect_udp_socket( socket.connect(&remote.into())?; let socket: UdpSocket = socket.into(); + configure_receive_metadata(&socket, remote).map_err(|source| ClientError::SocketOption { + operation: "enable receive metadata", + remote, + source, + })?; if let Some(ttl) = config.ttl { apply_ttl_to_socket(&socket, remote, ttl)?; }