From a542d4edf15bad99f94953ed7c1099f13720e913 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 4 Aug 2026 15:29:35 +0700 Subject: [PATCH] feat(client): add optional Tokio socket support --- Cargo.lock | 14 ++ crates/irtt-client/Cargo.toml | 7 + crates/irtt-client/src/error.rs | 4 + crates/irtt-client/src/receive.rs | 51 ++++ crates/irtt-client/src/receive/linux.rs | 70 +++++- crates/irtt-client/src/socket.rs | 281 +++++++++++++++++++++-- crates/irtt-client/src/socket_options.rs | 134 +++++++++-- 7 files changed, 528 insertions(+), 33 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 598e4fc..7c1ec3f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -713,6 +713,7 @@ dependencies = [ "libc", "socket2", "thiserror 2.0.19", + "tokio", ] [[package]] @@ -1716,6 +1717,19 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" +[[package]] +name = "tokio" +version = "1.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +dependencies = [ + "libc", + "mio", + "pin-project-lite", + "socket2", + "windows-sys", +] + [[package]] name = "typenum" version = "1.20.1" diff --git a/crates/irtt-client/Cargo.toml b/crates/irtt-client/Cargo.toml index e9ef818..9472c21 100644 --- a/crates/irtt-client/Cargo.toml +++ b/crates/irtt-client/Cargo.toml @@ -16,10 +16,17 @@ irtt-proto.workspace = true libc = { version = "0.2.189", optional = true } socket2 = { version = "0.6.5", features = ["all"] } thiserror = "2.0.19" +tokio = { + version = "1", + default-features = false, + features = ["rt", "net", "time", "sync"], + optional = true, +} [features] default = [] ancillary = ["dep:libc"] +tokio = ["dep:tokio"] [package.metadata.docs.rs] all-features = true diff --git a/crates/irtt-client/src/error.rs b/crates/irtt-client/src/error.rs index ee03b2a..4423122 100644 --- a/crates/irtt-client/src/error.rs +++ b/crates/irtt-client/src/error.rs @@ -12,6 +12,10 @@ use thiserror::Error; #[derive(Debug, Error)] #[non_exhaustive] pub enum ClientError { + /// A Tokio-backed client operation was polled without a current runtime. + #[cfg(feature = "tokio")] + #[error("Tokio AsyncClient requires a current Tokio runtime")] + NoTokioRuntime, /// The configured server address could not be resolved to a socket address. /// /// Check the host, port, address-family restrictions, and name service diff --git a/crates/irtt-client/src/receive.rs b/crates/irtt-client/src/receive.rs index f239d42..ca5fbd1 100644 --- a/crates/irtt-client/src/receive.rs +++ b/crates/irtt-client/src/receive.rs @@ -62,6 +62,32 @@ pub(crate) fn recv_datagram( linux::recv_datagram(socket, buf) } +#[cfg(all( + feature = "tokio", + not(all(target_os = "linux", feature = "ancillary")) +))] +pub(crate) fn try_recv_tokio_datagram( + socket: &tokio::net::UdpSocket, + buf: &mut [u8], +) -> Result { + let len = socket.try_recv(buf)?; + let received_at = ClientTimestamp::now(); + + Ok(ReceivedDatagram { + len, + received_at, + meta: ReceiveMeta::default(), + }) +} + +#[cfg(all(feature = "tokio", target_os = "linux", feature = "ancillary"))] +pub(crate) fn try_recv_tokio_datagram( + socket: &tokio::net::UdpSocket, + buf: &mut [u8], +) -> Result { + linux::try_recv_tokio_datagram(socket, buf) +} + #[cfg(all(target_os = "linux", feature = "ancillary"))] pub(crate) fn recv_datagram_from( socket: &UdpSocket, @@ -133,4 +159,29 @@ mod tests { assert!(datagram.received_at.mono >= before.mono); assert!(datagram.received_at.mono <= after.mono); } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_try_receive_preserves_fallback_metadata_and_would_block() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_io() + .build() + .unwrap(); + + runtime.block_on(async { + let (sender, receiver) = connected_loopback_pair(); + receiver.set_nonblocking(true).unwrap(); + let receiver = tokio::net::UdpSocket::from_std(receiver).unwrap(); + sender.send(b"tokio").unwrap(); + receiver.readable().await.unwrap(); + + let mut buf = [0_u8; 16]; + let datagram = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap(); + assert_eq!(&buf[..datagram.len], b"tokio"); + assert_eq!(datagram.meta, ReceiveMeta::default()); + + let error = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::WouldBlock); + }); + } } diff --git a/crates/irtt-client/src/receive/linux.rs b/crates/irtt-client/src/receive/linux.rs index ca0c3a4..3509058 100644 --- a/crates/irtt-client/src/receive/linux.rs +++ b/crates/irtt-client/src/receive/linux.rs @@ -2,7 +2,7 @@ use std::{ io, mem, net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, UdpSocket}, - os::fd::AsRawFd, + os::fd::{AsRawFd, RawFd}, ptr, time::{Duration, SystemTime, UNIX_EPOCH}, }; @@ -49,6 +49,24 @@ fn enable_kernel_rx_timestamps(socket: &UdpSocket) -> io::Result<()> { pub(crate) fn recv_datagram( socket: &UdpSocket, buf: &mut [u8], +) -> Result { + recv_datagram_fd(socket.as_raw_fd(), buf, 0) +} + +#[cfg(feature = "tokio")] +pub(crate) fn try_recv_tokio_datagram( + socket: &tokio::net::UdpSocket, + buf: &mut [u8], +) -> Result { + socket.try_io(tokio::io::Interest::READABLE, || { + recv_datagram_fd(socket.as_raw_fd(), buf, libc::MSG_DONTWAIT) + }) +} + +fn recv_datagram_fd( + socket_fd: RawFd, + buf: &mut [u8], + flags: libc::c_int, ) -> Result { let mut control = ControlBuffer::new(); let mut iov = libc::iovec { @@ -64,8 +82,8 @@ pub(crate) fn recv_datagram( 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) + // socket file descriptor is borrowed from a valid UDP socket. + libc::recvmsg(socket_fd, &mut msg, flags) }; if len < 0 { return Err(io::Error::last_os_error()); @@ -366,6 +384,52 @@ mod tests { assert!(datagram.received_at.mono <= after.mono); } + #[cfg(feature = "tokio")] + #[test] + fn tokio_recvmsg_would_block_clears_readiness() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_io() + .enable_time() + .build() + .unwrap(); + + runtime.block_on(async { + let (sender, receiver) = connected_ipv4_loopback_pair(); + configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap(); + receiver.set_nonblocking(true).unwrap(); + let receiver = tokio::net::UdpSocket::from_std(receiver).unwrap(); + let mut buf = [0_u8; 16]; + + sender.send(b"first").unwrap(); + tokio::time::timeout(Duration::from_secs(1), receiver.readable()) + .await + .expect("initial readability timed out") + .unwrap(); + + let first = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap(); + assert_eq!(&buf[..first.len], b"first"); + + let error = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::WouldBlock); + + assert!( + tokio::time::timeout(Duration::from_millis(20), receiver.readable()) + .await + .is_err(), + "readable completed from stale readiness after recvmsg WouldBlock" + ); + + sender.send(b"second").unwrap(); + tokio::time::timeout(Duration::from_secs(1), receiver.readable()) + .await + .expect("second readability timed out") + .unwrap(); + + let second = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap(); + assert_eq!(&buf[..second.len], b"second"); + }); + } + #[test] fn kernel_rx_timestamp_metadata_is_observed_when_kernel_provides_it() { let (sender, receiver) = connected_ipv4_loopback_pair(); diff --git a/crates/irtt-client/src/socket.rs b/crates/irtt-client/src/socket.rs index 767e825..d3ba12b 100644 --- a/crates/irtt-client/src/socket.rs +++ b/crates/irtt-client/src/socket.rs @@ -28,18 +28,57 @@ pub(crate) fn validate_open_timeouts(timeouts: &[Duration]) -> Result<(), Client } pub(crate) fn resolve_remote(config: &ClientConfig) -> Result { + #[cfg(all(test, feature = "tokio"))] + SYNC_RESOLVER_CALLS.with(|calls| calls.set(calls.get() + 1)); + let addr = normalize_server_addr(&config.server_addr); let mut addrs = addr .to_socket_addrs() .map_err(|_| ClientError::Resolve { addr: addr.clone() })?; addrs - .find(|addr| { - (!config.socket_config.ipv4_only || addr.is_ipv4()) - && (!config.socket_config.ipv6_only || addr.is_ipv6()) - }) + .find(|addr| address_family_allowed(config, *addr)) .ok_or(ClientError::Resolve { addr }) } +#[cfg(feature = "tokio")] +pub(crate) async fn resolve_remote_tokio(config: &ClientConfig) -> Result { + let addr = normalize_server_addr(&config.server_addr); + if let Ok(remote) = addr.parse::() { + return address_family_allowed(config, remote) + .then_some(remote) + .ok_or(ClientError::Resolve { addr }); + } + + #[cfg(all(test, feature = "tokio"))] + TOKIO_DNS_LOOKUPS.with(|calls| calls.set(calls.get() + 1)); + + let mut addrs = tokio::net::lookup_host(&addr) + .await + .map_err(|_| ClientError::Resolve { addr: addr.clone() })?; + addrs + .find(|remote| address_family_allowed(config, *remote)) + .ok_or_else(|| ClientError::Resolve { addr: addr.clone() }) +} + +fn address_family_allowed(config: &ClientConfig, remote: SocketAddr) -> bool { + (!config.socket_config.ipv4_only || remote.is_ipv4()) + && (!config.socket_config.ipv6_only || remote.is_ipv6()) +} + +#[cfg(all(test, feature = "tokio"))] +std::thread_local! { + static SYNC_RESOLVER_CALLS: std::cell::Cell = const { std::cell::Cell::new(0) }; + static TOKIO_DNS_LOOKUPS: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +#[cfg(all(test, feature = "tokio"))] +pub(crate) fn resolution_call_counts() -> (usize, usize) { + ( + SYNC_RESOLVER_CALLS.with(std::cell::Cell::get), + TOKIO_DNS_LOOKUPS.with(std::cell::Cell::get), + ) +} + pub(crate) fn normalize_server_addr(addr: &str) -> String { if addr.parse::().is_ok() { return addr.to_owned(); @@ -66,16 +105,16 @@ pub(crate) fn connect_udp_socket( config: &SocketConfig, remote: SocketAddr, ) -> Result { - let domain = if remote.is_ipv4() { - Domain::IPV4 - } else { - Domain::IPV6 - }; - let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?; + let socket = create_connected_udp_socket(config, remote)?; + socket.set_read_timeout(config.recv_timeout)?; + Ok(socket) +} - if config.ipv6_only && remote.is_ipv6() { - socket.set_only_v6(true)?; - } +fn create_connected_udp_socket( + config: &SocketConfig, + remote: SocketAddr, +) -> Result { + let socket = create_prebind_udp_socket(config, remote)?; let bind_addr = config.bind_addr.unwrap_or_else(|| { if remote.is_ipv4() { SocketAddr::from(([0, 0, 0, 0], 0)) @@ -95,10 +134,36 @@ pub(crate) fn connect_udp_socket( if let Some(ttl) = config.ttl { apply_ttl_to_socket(&socket, remote, ttl)?; } - socket.set_read_timeout(config.recv_timeout)?; Ok(socket) } +fn create_prebind_udp_socket( + config: &SocketConfig, + remote: SocketAddr, +) -> Result { + let domain = if remote.is_ipv4() { + Domain::IPV4 + } else { + Domain::IPV6 + }; + let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?; + + if config.ipv6_only && remote.is_ipv6() { + socket.set_only_v6(true)?; + } + Ok(socket) +} + +#[cfg(feature = "tokio")] +pub(crate) fn connect_tokio_udp_socket( + config: &SocketConfig, + remote: SocketAddr, +) -> Result { + let socket = create_connected_udp_socket(config, remote)?; + socket.set_nonblocking(true)?; + Ok(tokio::net::UdpSocket::from_std(socket)?) +} + pub(crate) fn bind_unconnected_udp_socket( config: &SocketConfig, family_remote: SocketAddr, @@ -151,4 +216,192 @@ mod tests { assert_eq!(normalize_server_addr("[::1]"), "[::1]:2112"); assert_eq!(normalize_server_addr("[::1]:1234"), "[::1]:1234"); } + + #[test] + fn address_family_filtering_remains_unchanged() { + let ipv4 = SocketAddr::from(([127, 0, 0, 1], DEFAULT_PORT)); + let ipv6 = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], DEFAULT_PORT)); + let mut config = ClientConfig::default(); + + assert!(address_family_allowed(&config, ipv4)); + assert!(address_family_allowed(&config, ipv6)); + + config.socket_config.ipv4_only = true; + assert!(address_family_allowed(&config, ipv4)); + assert!(!address_family_allowed(&config, ipv6)); + + config.socket_config.ipv4_only = false; + config.socket_config.ipv6_only = true; + assert!(!address_family_allowed(&config, ipv4)); + assert!(address_family_allowed(&config, ipv6)); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_ipv6_only_is_applied_before_bind() { + let config = SocketConfig { + ipv6_only: true, + ..SocketConfig::default() + }; + let remote = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], DEFAULT_PORT)); + + let socket = create_prebind_udp_socket(&config, remote).unwrap(); + + assert!(socket.only_v6().unwrap()); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_ipv6_socket_preserves_v6only_and_explicit_bind() { + let bind_addr = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], 0)); + let peer = match UdpSocket::bind(bind_addr) { + Ok(peer) => peer, + Err(error) if ipv6_loopback_unavailable(&error) => { + eprintln!("skipping IPv6 Tokio socket test: IPv6 loopback unavailable: {error}"); + return; + } + Err(error) => panic!("{error}"), + }; + let remote = peer.local_addr().unwrap(); + let config = SocketConfig { + bind_addr: Some(bind_addr), + ipv6_only: true, + ..SocketConfig::default() + }; + + tokio_runtime().block_on(async { + let socket = connect_tokio_udp_socket(&config, remote).unwrap(); + + assert!(socket2::SockRef::from(&socket).only_v6().unwrap()); + assert!(socket.local_addr().unwrap().is_ipv6()); + assert_eq!(socket.local_addr().unwrap().ip(), bind_addr.ip()); + assert_eq!(socket.peer_addr().unwrap(), remote); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_ipv4_socket_preserves_family_bind_peer_and_ttl() { + let peer = UdpSocket::bind(SocketAddr::from(([127, 0, 0, 1], 0))).unwrap(); + let remote = peer.local_addr().unwrap(); + let bind_addr = SocketAddr::from(([127, 0, 0, 1], 0)); + let config = SocketConfig { + bind_addr: Some(bind_addr), + ttl: Some(64), + ..SocketConfig::default() + }; + + tokio_runtime().block_on(async { + let socket = connect_tokio_udp_socket(&config, remote).unwrap(); + + assert!(socket.local_addr().unwrap().is_ipv4()); + assert_eq!(socket.local_addr().unwrap().ip(), bind_addr.ip()); + assert_eq!(socket.peer_addr().unwrap(), remote); + assert_eq!(socket2::SockRef::from(&socket).ttl_v4().unwrap(), 64); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_resolver_filters_families_without_using_blocking_resolution() { + let mut config = ClientConfig { + server_addr: "localhost:2112".to_owned(), + ..ClientConfig::default() + }; + config.socket_config.ipv4_only = true; + + tokio_runtime().block_on(async { + let before = resolution_call_counts(); + let remote = resolve_remote_tokio(&config).await.unwrap(); + let after = resolution_call_counts(); + + assert!(remote.is_ipv4()); + assert_eq!(after.0, before.0); + assert_eq!(after.1, before.1 + 1); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_resolver_selects_ipv6_for_ipv6_only_hostname() { + let mut config = ClientConfig { + server_addr: "localhost:2112".to_owned(), + ..ClientConfig::default() + }; + config.socket_config.ipv6_only = true; + + tokio_runtime().block_on(async { + let remote = resolve_remote_tokio(&config).await.unwrap(); + assert!(remote.is_ipv6()); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_resolver_literal_path_bypasses_all_name_resolution() { + let config = ClientConfig { + server_addr: "127.0.0.1:2112".to_owned(), + ..ClientConfig::default() + }; + + tokio_runtime().block_on(async { + let before = resolution_call_counts(); + let remote = resolve_remote_tokio(&config).await.unwrap(); + let after = resolution_call_counts(); + + assert_eq!(remote, SocketAddr::from(([127, 0, 0, 1], 2112))); + assert_eq!(after, before); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_resolver_rejects_disallowed_literal_family() { + let mut config = ClientConfig { + server_addr: "127.0.0.1:2112".to_owned(), + ..ClientConfig::default() + }; + config.socket_config.ipv6_only = true; + + tokio_runtime().block_on(async { + assert!(matches!( + resolve_remote_tokio(&config).await, + Err(ClientError::Resolve { .. }) + )); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_resolver_rejects_mixed_family_restrictions() { + let mut config = ClientConfig { + server_addr: "127.0.0.1:2112".to_owned(), + ..ClientConfig::default() + }; + config.socket_config.ipv4_only = true; + config.socket_config.ipv6_only = true; + + tokio_runtime().block_on(async { + assert!(matches!( + resolve_remote_tokio(&config).await, + Err(ClientError::Resolve { .. }) + )); + }); + } + + #[cfg(feature = "tokio")] + fn tokio_runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_current_thread() + .enable_io() + .build() + .unwrap() + } + + #[cfg(feature = "tokio")] + fn ipv6_loopback_unavailable(error: &std::io::Error) -> bool { + matches!( + error.kind(), + std::io::ErrorKind::AddrNotAvailable | std::io::ErrorKind::Unsupported + ) + } } diff --git a/crates/irtt-client/src/socket_options.rs b/crates/irtt-client/src/socket_options.rs index 283e87e..bc8a1d0 100644 --- a/crates/irtt-client/src/socket_options.rs +++ b/crates/irtt-client/src/socket_options.rs @@ -3,6 +3,8 @@ use std::{ net::{SocketAddr, UdpSocket}, }; +use socket2::SockRef; + use crate::{ config::{MAX_DSCP_CODEPOINT, MAX_TTL}, error::ClientError, @@ -33,6 +35,24 @@ pub(crate) fn clear_dscp_on_socket( set_socket_traffic_class(socket, remote, 0, "clear DSCP before close") } +#[cfg(feature = "tokio")] +pub(crate) fn apply_dscp_to_tokio_socket( + socket: &tokio::net::UdpSocket, + remote: SocketAddr, + dscp_codepoint: u8, +) -> Result<(), ClientError> { + let traffic_class = dscp_codepoint_to_traffic_class(dscp_codepoint)?; + set_tokio_socket_traffic_class(socket, remote, traffic_class, "set negotiated DSCP") +} + +#[cfg(feature = "tokio")] +pub(crate) fn clear_dscp_on_tokio_socket( + socket: &tokio::net::UdpSocket, + remote: SocketAddr, +) -> Result<(), ClientError> { + set_tokio_socket_traffic_class(socket, remote, 0, "clear DSCP before close") +} + pub(crate) fn validate_ttl(ttl: u32) -> Result<(), ClientError> { if ttl == 0 || ttl > MAX_TTL { return Err(ClientError::InvalidConfig { @@ -67,6 +87,20 @@ pub(crate) fn socket_traffic_class( }) } +#[cfg(all(test, feature = "tokio"))] +pub(crate) fn tokio_socket_traffic_class( + socket: &tokio::net::UdpSocket, + remote: SocketAddr, +) -> Result { + get_socket_traffic_class_ref(SockRef::from(socket), remote).map_err(|source| { + ClientError::SocketOption { + operation: "read DSCP socket option", + remote, + source, + } + }) +} + #[cfg(test)] pub(crate) fn socket_ttl(socket: &UdpSocket, remote: SocketAddr) -> Result { get_socket_ttl(socket, remote).map_err(|source| ClientError::SocketOption { @@ -77,19 +111,24 @@ pub(crate) fn socket_ttl(socket: &UdpSocket, remote: SocketAddr) -> Result io::Result<()> { + set_socket_ttl_ref(SockRef::from(socket), remote, ttl) +} + +fn set_socket_ttl_ref(socket: SockRef<'_>, remote: SocketAddr, ttl: u32) -> io::Result<()> { if remote.is_ipv4() { - socket2::SockRef::from(socket).set_ttl_v4(ttl) + socket.set_ttl_v4(ttl) } else { - socket2::SockRef::from(socket).set_unicast_hops_v6(ttl) + socket.set_unicast_hops_v6(ttl) } } #[cfg(test)] fn get_socket_ttl(socket: &UdpSocket, remote: SocketAddr) -> io::Result { + let socket = SockRef::from(socket); if remote.is_ipv4() { - socket2::SockRef::from(socket).ttl_v4() + socket.ttl_v4() } else { - socket2::SockRef::from(socket).unicast_hops_v6() + socket.unicast_hops_v6() } } @@ -112,6 +151,30 @@ fn set_socket_traffic_class_io( socket: &UdpSocket, remote: SocketAddr, traffic_class: u32, +) -> io::Result<()> { + set_socket_traffic_class_ref(SockRef::from(socket), remote, traffic_class) +} + +#[cfg(feature = "tokio")] +fn set_tokio_socket_traffic_class( + socket: &tokio::net::UdpSocket, + remote: SocketAddr, + traffic_class: u32, + operation: &'static str, +) -> Result<(), ClientError> { + set_socket_traffic_class_ref(SockRef::from(socket), remote, traffic_class).map_err(|source| { + ClientError::SocketOption { + operation, + remote, + source, + } + }) +} + +fn set_socket_traffic_class_ref( + socket: SockRef<'_>, + remote: SocketAddr, + traffic_class: u32, ) -> io::Result<()> { if remote.is_ipv4() { set_ipv4_traffic_class(socket, traffic_class) @@ -127,8 +190,8 @@ fn set_socket_traffic_class_io( target_os = "illumos", target_os = "haiku", )))] -fn set_ipv4_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result<()> { - socket2::SockRef::from(socket).set_tos_v4(traffic_class) +fn set_ipv4_traffic_class(socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> { + socket.set_tos_v4(traffic_class) } #[cfg(any( @@ -138,7 +201,7 @@ fn set_ipv4_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result< target_os = "illumos", target_os = "haiku", ))] -fn set_ipv4_traffic_class(_socket: &UdpSocket, traffic_class: u32) -> io::Result<()> { +fn set_ipv4_traffic_class(_socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> { unsupported_traffic_class(traffic_class, "IPv4 DSCP socket options") } @@ -153,8 +216,8 @@ fn set_ipv4_traffic_class(_socket: &UdpSocket, traffic_class: u32) -> io::Result target_os = "openbsd", target_os = "cygwin", ))] -fn set_ipv6_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result<()> { - socket2::SockRef::from(socket).set_tclass_v6(traffic_class) +fn set_ipv6_traffic_class(socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> { + socket.set_tclass_v6(traffic_class) } #[cfg(not(any( @@ -168,7 +231,7 @@ fn set_ipv6_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result< target_os = "openbsd", target_os = "cygwin", )))] -fn set_ipv6_traffic_class(_socket: &UdpSocket, traffic_class: u32) -> io::Result<()> { +fn set_ipv6_traffic_class(_socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> { unsupported_traffic_class(traffic_class, "IPv6 DSCP socket options") } @@ -211,8 +274,8 @@ fn unsupported_traffic_class(traffic_class: u32, feature: &'static str) -> io::R target_os = "haiku", )) ))] -fn get_ipv4_traffic_class(socket: &UdpSocket) -> io::Result { - socket2::SockRef::from(socket).tos_v4() +fn get_ipv4_traffic_class(socket: SockRef<'_>) -> io::Result { + socket.tos_v4() } #[cfg(all( @@ -225,7 +288,7 @@ fn get_ipv4_traffic_class(socket: &UdpSocket) -> io::Result { target_os = "haiku", ) ))] -fn get_ipv4_traffic_class(_socket: &UdpSocket) -> io::Result { +fn get_ipv4_traffic_class(_socket: SockRef<'_>) -> io::Result { unsupported_readback("IPv4 DSCP socket option readback") } @@ -243,8 +306,8 @@ fn get_ipv4_traffic_class(_socket: &UdpSocket) -> io::Result { target_os = "cygwin", ) ))] -fn get_ipv6_traffic_class(socket: &UdpSocket) -> io::Result { - socket2::SockRef::from(socket).tclass_v6() +fn get_ipv6_traffic_class(socket: SockRef<'_>) -> io::Result { + socket.tclass_v6() } #[cfg(all( @@ -261,12 +324,17 @@ fn get_ipv6_traffic_class(socket: &UdpSocket) -> io::Result { target_os = "cygwin", )) ))] -fn get_ipv6_traffic_class(_socket: &UdpSocket) -> io::Result { +fn get_ipv6_traffic_class(_socket: SockRef<'_>) -> io::Result { unsupported_readback("IPv6 DSCP socket option readback") } #[cfg(test)] fn get_socket_traffic_class(socket: &UdpSocket, remote: SocketAddr) -> io::Result { + get_socket_traffic_class_ref(SockRef::from(socket), remote) +} + +#[cfg(test)] +fn get_socket_traffic_class_ref(socket: SockRef<'_>, remote: SocketAddr) -> io::Result { if remote.is_ipv4() { get_ipv4_traffic_class(socket) } else { @@ -395,6 +463,40 @@ mod tests { assert_eq!(socket_traffic_class(&socket, remote).unwrap(), 0); } + #[cfg(all( + feature = "tokio", + not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + target_os = "haiku", + )) + ))] + #[test] + fn tokio_ipv4_socket_option_sets_and_clears_traffic_class() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_io() + .build() + .unwrap(); + + runtime.block_on(async { + let remote = SocketAddr::from(([127, 0, 0, 1], 9)); + let socket = UdpSocket::bind(SocketAddr::from(([127, 0, 0, 1], 0))).unwrap(); + socket.connect(remote).unwrap(); + socket.set_nonblocking(true).unwrap(); + let socket = tokio::net::UdpSocket::from_std(socket).unwrap(); + + apply_dscp_to_tokio_socket(&socket, remote, 46).unwrap(); + let traffic_class = tokio_socket_traffic_class(&socket, remote).unwrap(); + assert_eq!(traffic_class & 0xfc, 184); + assert_eq!(traffic_class & 0b11, 0); + + clear_dscp_on_tokio_socket(&socket, remote).unwrap(); + assert_eq!(tokio_socket_traffic_class(&socket, remote).unwrap(), 0); + }); + } + #[test] #[cfg(any( target_os = "android",