diff --git a/Cargo.lock b/Cargo.lock index 4eafc528..1850a21f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -130,9 +130,9 @@ checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" [[package]] name = "chacha20" -version = "0.10.1" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" dependencies = [ "cfg-if", "cpufeatures", diff --git a/netwatch/src/lib.rs b/netwatch/src/lib.rs index ff172974..29866433 100644 --- a/netwatch/src/lib.rs +++ b/netwatch/src/lib.rs @@ -9,4 +9,4 @@ mod udp; pub use self::ip_family::IpFamily; #[cfg(not(wasm_browser))] -pub use self::udp::{UdpSender, UdpSocket}; +pub use self::udp::{BindOptions, SocketRef, UdpSender, UdpSocket}; diff --git a/netwatch/src/udp.rs b/netwatch/src/udp.rs index be5ce69f..725bf910 100644 --- a/netwatch/src/udp.rs +++ b/netwatch/src/udp.rs @@ -1,3 +1,7 @@ +#[cfg(unix)] +use std::os::fd::{AsFd, BorrowedFd}; +#[cfg(windows)] +use std::os::windows::io::{AsSocket, BorrowedSocket}; use std::{ future::Future, io, @@ -28,6 +32,82 @@ pub struct UdpSocket { /// UDP socket read/write buffer size (7MB). The value of 7MB is chosen as it /// is the max supported by a default configuration of macOS. Some platforms will silently clamp the value. const SOCKET_BUFFER_SIZE: usize = 7 << 20; + +/// A socket that is about to be bound, handed to the hook set with +/// [`BindOptions::configure_socket`]. +/// +/// It implements [`AsFd`] on unix and [`AsSocket`] on Windows, which is what +/// socket wrappers take, so the hook can set options with the socket crate of +/// its choice: `socket2::SockRef::from(&socket)`, or plain `libc::setsockopt` +/// on the raw fd. +#[derive(Debug)] +pub struct SocketRef<'a>(&'a socket2::Socket); + +#[cfg(unix)] +impl AsFd for SocketRef<'_> { + fn as_fd(&self) -> BorrowedFd<'_> { + self.0.as_fd() + } +} + +#[cfg(windows)] +impl AsSocket for SocketRef<'_> { + fn as_socket(&self) -> BorrowedSocket<'_> { + self.0.as_socket() + } +} + +/// The hook set with [`BindOptions::configure_socket`]. +type Configurator = Arc, IpFamily) -> io::Result<()> + Send + Sync>; + +/// Options to bind a [`UdpSocket`] with. +/// +/// Used by [`UdpSocket::bind_with`]. The default options match what the other +/// `bind_*` constructors use. +#[derive(derive_more::Debug, Default, Clone)] +pub struct BindOptions { + #[debug(skip)] + configure: Option, +} + +impl BindOptions { + /// Creates the default options. + pub fn new() -> Self { + Self::default() + } + + /// Sets a hook to run on the socket just before it is bound. + /// + /// This is the escape hatch for socket options this crate does not model + /// itself, like `SO_MARK` on Linux or `IP_BOUND_IF` on Apple platforms. The + /// hook runs on every bind, including the rebinds [`UdpSocket`] does to + /// recover from network changes, since each of those creates a new socket. + /// That also lets the hook pick up state that changed in between, like the + /// interface the default route now points at. + /// + /// An error from the hook fails the bind, rather than leaving a socket that + /// silently missed its configuration. + /// + /// ```no_run + /// use std::net::{Ipv4Addr, SocketAddr}; + /// + /// use netwatch::{BindOptions, UdpSocket}; + /// + /// let opts = BindOptions::new().configure_socket(|socket, _family| { + /// socket2::SockRef::from(&socket).set_recv_buffer_size(1 << 20) + /// }); + /// let socket = UdpSocket::bind_with(SocketAddr::from((Ipv4Addr::UNSPECIFIED, 0)), opts)?; + /// # Ok::<(), std::io::Error>(()) + /// ``` + pub fn configure_socket( + mut self, + configure: impl Fn(SocketRef<'_>, IpFamily) -> io::Result<()> + Send + Sync + 'static, + ) -> Self { + self.configure = Some(Arc::new(configure)); + self + } +} + impl UdpSocket { /// Bind only Ipv4 on any interface. pub fn bind_v4(port: u16) -> io::Result { @@ -52,18 +132,30 @@ impl UdpSocket { /// Bind to the given port only on localhost. pub fn bind_local(network: IpFamily, port: u16) -> io::Result { let addr = SocketAddr::new(network.local_addr(), port); - Self::bind_raw(addr) + Self::bind_with(addr, BindOptions::default()) } /// Bind to the given port and listen on all interfaces. pub fn bind(network: IpFamily, port: u16) -> io::Result { let addr = SocketAddr::new(network.unspecified_addr(), port); - Self::bind_raw(addr) + Self::bind_with(addr, BindOptions::default()) } /// Bind to any provided [`SocketAddr`]. pub fn bind_full(addr: impl Into) -> io::Result { - Self::bind_raw(addr) + Self::bind_with(addr, BindOptions::default()) + } + + /// Bind to any provided [`SocketAddr`], using the given [`BindOptions`]. + pub fn bind_with(addr: impl Into, opts: BindOptions) -> io::Result { + let socket = SocketState::bind(addr.into(), opts.configure)?; + + Ok(UdpSocket { + socket: RwLock::new(socket), + recv_waker: AtomicWaker::default(), + send_waker: AtomicWaker::default(), + is_broken: AtomicBool::new(false), + }) } /// Is the socket broken and needs a rebind? @@ -96,17 +188,6 @@ impl UdpSocket { Ok(()) } - fn bind_raw(addr: impl Into) -> io::Result { - let socket = SocketState::bind(addr.into())?; - - Ok(UdpSocket { - socket: RwLock::new(socket), - recv_waker: AtomicWaker::default(), - send_waker: AtomicWaker::default(), - is_broken: AtomicBool::new(false), - }) - } - /// Receives a single datagram message on the socket from the remote address /// to which it is connected. On success, returns the number of bytes read. /// @@ -729,17 +810,23 @@ impl Future for SendToFut<'_, '_> { } } -#[derive(Debug)] +#[derive(derive_more::Debug)] enum SocketState { Connected { socket: tokio::net::UdpSocket, state: noq_udp::UdpSocketState, /// The addr we are binding to. addr: SocketAddr, + /// The hook to rerun when rebinding, if any. + #[debug(skip)] + configure: Option, }, Closed { /// The addr to rebind to when recovering. addr: SocketAddr, + /// The hook to rerun when rebinding, if any. + #[debug(skip)] + configure: Option, last_max_gso_segments: NonZeroUsize, last_gro_segments: NonZeroUsize, last_may_fragment: bool, @@ -753,6 +840,7 @@ impl SocketState { socket, state, addr: _, + configure: _, } => Ok((socket, state)), Self::Closed { .. } => { warn!("socket closed"); @@ -761,7 +849,7 @@ impl SocketState { } } - fn bind(addr: SocketAddr) -> io::Result { + fn bind(addr: SocketAddr, configure: Option) -> io::Result { let network = IpFamily::from(addr.ip()); let socket = socket2::Socket::new( network.into(), @@ -786,6 +874,13 @@ impl SocketState { socket.set_only_v6(true)?; } + // Let the caller configure the socket before it is bound. An error here + // fails the bind: a socket that silently missed its configuration would + // send traffic where the caller did not want it. + if let Some(configure) = &configure { + configure(SocketRef(&socket), network)?; + } + // Binding must happen before calling noq, otherwise `local_addr` // is not yet available on all OSes. socket.bind(&addr.into())?; @@ -814,13 +909,18 @@ impl SocketState { socket, state: socket_state, addr: local_addr, + configure, }) } fn rebind(&mut self) -> io::Result<()> { - let addr = match self { - Self::Connected { addr, .. } => *addr, - Self::Closed { addr, .. } => *addr, + let (addr, configure) = match self { + Self::Connected { + addr, configure, .. + } + | Self::Closed { + addr, configure, .. + } => (*addr, configure.clone()), }; debug!("rebinding {}", addr); @@ -829,13 +929,14 @@ impl SocketState { if let Self::Connected { state, .. } = self { *self = SocketState::Closed { addr, + configure: configure.clone(), last_max_gso_segments: state.max_gso_segments(), last_gro_segments: state.gro_segments(), last_may_fragment: state.may_fragment(), }; } - match Self::bind(addr) { + match Self::bind(addr, configure) { Ok(new_state) => { *self = new_state; Ok(()) @@ -854,9 +955,15 @@ impl SocketState { fn close(&mut self) -> Option<(tokio::net::UdpSocket, noq_udp::UdpSocketState)> { match self { - Self::Connected { state, addr, .. } => { + Self::Connected { + state, + addr, + configure, + .. + } => { let s = SocketState::Closed { addr: *addr, + configure: configure.clone(), last_max_gso_segments: state.max_gso_segments(), last_gro_segments: state.gro_segments(), last_may_fragment: state.may_fragment(), @@ -1068,6 +1175,11 @@ impl Future for SendFutNoq<'_, '_> { #[cfg(test)] mod tests { + use std::{ + net::Ipv4Addr, + sync::atomic::{AtomicUsize, Ordering}, + }; + use testresult::TestResult; use super::*; @@ -1129,6 +1241,44 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_configure_socket_runs_on_every_bind() -> TestResult { + let calls = Arc::new(AtomicUsize::new(0)); + let seen = calls.clone(); + let opts = BindOptions::new().configure_socket(move |socket, family| { + assert_eq!(family, IpFamily::V4); + // The hook gets a real socket, before it is bound. + socket2::SockRef::from(&socket).set_reuse_address(true)?; + seen.fetch_add(1, Ordering::SeqCst); + Ok(()) + }); + + let socket = UdpSocket::bind_with(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)), opts)?; + assert_eq!(calls.load(Ordering::SeqCst), 1, "hook did not run on bind"); + + socket.rebind()?; + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "hook did not run again on rebind" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_configure_socket_error_fails_the_bind() -> TestResult { + let opts = + BindOptions::new().configure_socket(|_socket, _family| Err(io::Error::other("nope"))); + + assert!( + UdpSocket::bind_with(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)), opts).is_err(), + "a failing hook must fail the bind" + ); + + Ok(()) + } + #[tokio::test] async fn test_udp_mark_broken() -> TestResult { let socket_a = UdpSocket::bind_local(IpFamily::V4, 0)?;