From bbb269f3fc9e5b3251832fdd568e9749733d47df Mon Sep 17 00:00:00 2001 From: meh Date: Wed, 14 Oct 2026 17:35:01 +0700 Subject: [PATCH 1/2] Port missing UDP cancellation and datagram behavior --- src/net.rs | 381 ++++++++++++++++++++++++++++++++++++------ src/reactor.rs | 2 +- src/reactor/driver.rs | 18 ++ src/reactor/socket.rs | 265 ++++++++++++++++++++++------- 4 files changed, 556 insertions(+), 110 deletions(-) diff --git a/src/net.rs b/src/net.rs index ec5dcb5..7321f4b 100644 --- a/src/net.rs +++ b/src/net.rs @@ -585,43 +585,133 @@ impl core::future::Future for Accept<'_> { } } -/// A message-oriented socket that yields instead of blocking. +pub use crate::reactor::socket::Datagram; + +/// An owned UDP socket driven by a nagoya reactor. +/// +/// Each receive consumes one datagram. Successful operations retain readiness +/// until a subsequent syscall returns `WouldBlock`, including receives of +/// empty datagrams. UDP must not use TCP's FIONREAD optimisation: zero can +/// mean that an empty datagram is queued. /// /// The receiver and sender have separate readiness slots, so one task may /// receive while another sends on the same socket. There is one waker slot per /// direction: serialize receives with each other and sends with each other. +/// Call [`Self::cancel_recv`] or [`Self::cancel_send`] when abandoning a +/// pending direct poll operation. +/// +/// Dropping a pending receive or send future releases its waiter. A pending +/// receive has consumed nothing, and a pending send has sent nothing. +/// Dropping the socket deregisters it before closing the descriptor. #[derive(Debug)] pub struct UdpSocket { - inner: DatagramSocket, + // Fields drop in declaration order. The registration must go first. registration: Registration, + inner: DatagramSocket, } impl UdpSocket { - /// Bind a non-blocking socket and register it for reads and writes. + /// Bind a non-blocking datagram socket and register it for reads and writes. pub fn bind(addr: Addr, handle: &Handle) -> Result { - Self::from_socket(DatagramSocket::bind(addr)?, handle) + Self::bind_with_reuse(addr, handle, false) } - /// Bind with `SO_REUSEADDR` enabled for IP addresses before they are - /// claimed. + /// Bind with `SO_REUSEADDR` enabled before claiming an IP address. pub fn bind_reuse_addr(addr: Addr, handle: &Handle) -> Result { - Self::from_socket(DatagramSocket::bind_reuse_addr(addr)?, handle) + Self::bind_with_reuse(addr, handle, true) + } + + /// Bind with an explicit `SO_REUSEADDR` setting applied before binding. + /// + /// IPv6 sockets are IPv6-only. A port of zero selects an ephemeral port. + pub fn bind_with_reuse(addr: Addr, handle: &Handle, reuse_addr: bool) -> Result { + Self::from_socket(DatagramSocket::bind_with_reuse(addr, reuse_addr)?, handle) } /// Adopt an existing non-blocking datagram socket. pub fn from_socket(socket: DatagramSocket, handle: &Handle) -> Result { let registration = handle.register(socket.raw(), Interest::BOTH)?; Ok(Self { - inner: socket, registration, + inner: socket, }) } - /// This socket's own address. + /// Select a peer, returning the owned socket once the operation completes. + /// + /// UDP has no handshake. The kernel filters incoming datagrams by the + /// selected peer. The peer must have the same address family as the socket. + /// + /// This future owns the socket: cancelling an in-flight connect closes + /// it, rather than exposing a socket whose peer may have changed. + pub async fn connect(self, addr: Addr) -> Result { + match self.inner.connect(addr) { + Ok(()) => { + self.registration.mark_writable(); + return Ok(self); + } + Err(error) + if error.0 == libc::EINPROGRESS + || error.0 == libc::EALREADY + || error.interrupted() => {} + Err(error) if error.0 == libc::EISCONN => { + if self.inner.peer_addr()? == addr { + return Ok(self); + } + return Err(error); + } + Err(error) => return Err(error), + } + + core::future::poll_fn(|cx| { + if !self.registration.take_writable_or_park(cx.waker()) { + return Poll::Pending; + } + loop { + match self.inner.connect_error() { + Ok(()) => break, + Err(error) if error.interrupted() => continue, + Err(error) => return Poll::Ready(Err(error)), + } + } + match self.inner.peer_addr() { + Ok(peer) if peer == addr => { + self.registration.mark_writable(); + Poll::Ready(Ok(())) + } + Ok(_) => Poll::Ready(Err(crate::reactor::error::Errno(libc::EINVAL))), + Err(error) if error.0 == libc::ENOTCONN || error.interrupted() => { + self.registration.poll_writable(cx.waker()); + Poll::Pending + } + Err(error) => Poll::Ready(Err(error)), + } + }) + .await?; + Ok(self) + } + + /// This socket's bound address and actual local port. pub fn local_addr(&self) -> Result { self.inner.local_addr() } + /// The configured peer, or `ENOTCONN` when unconnected. + pub fn peer_addr(&self) -> Result { + self.inner.peer_addr() + } + + /// Receive one datagram, returning its copied length and source. + /// + /// Use [`Self::recv_datagram`] when truncation metadata is needed. + pub fn recv_from<'a>(&'a self, buffer: &'a mut [u8]) -> RecvFrom<'a> { + RecvFrom { + socket: self, + buffer, + waiting: false, + } + } + /// Set the kernel receive-buffer capacity in bytes. pub fn set_recv_buffer(&self, size: usize) -> Result<()> { self.inner.set_recv_buffer(size) @@ -632,60 +722,170 @@ impl UdpSocket { self.inner.set_send_buffer(size) } - /// Receive one datagram directly into `buffer`. + /// Receive one datagram into an initialised slice. + /// + /// An empty buffer still consumes one datagram and reports truncation + /// when that datagram is nonempty. Zero bytes never mean end of stream. + pub fn recv_datagram<'a>(&'a self, buffer: &'a mut [u8]) -> RecvDatagram<'a> { + RecvDatagram { + socket: self, + buffer, + waiting: false, + } + } + + /// Append one datagram directly into a `BytesMut`'s spare capacity. + /// + /// No staging copy or initialisation of the spare capacity is needed. + /// The returned length counts only the appended bytes. The caller can + /// freeze the buffer to share immutable packet ownership without copying. + /// + /// An empty spare region returns `ENOBUFS` without consuming a datagram. + pub fn recv_from_buf<'a>(&'a self, buffer: &'a mut bytes::BytesMut) -> RecvFromBuf<'a> { + RecvFromBuf { + socket: self, + buffer, + waiting: false, + } + } + + /// Send one datagram without copying or retaining the supplied bytes. + /// + /// This also sends empty datagrams. An oversized packet is an error; + /// no suffix is retried as a separate datagram. + pub fn send_to<'a>(&'a self, buffer: &'a [u8], destination: Addr) -> SendTo<'a> { + SendTo { + socket: self, + buffer, + destination: Some(destination), + waiting: false, + } + } + + /// Send one datagram to the configured peer. /// - /// The receive is zero-copy at the crate boundary: the kernel writes into - /// the caller's slice, with no intermediate buffer or per-datagram - /// allocation. A larger datagram is truncated to fit. + /// Returns the kernel's error if no peer has been configured. + pub fn send<'a>(&'a self, buffer: &'a [u8]) -> SendTo<'a> { + SendTo { + socket: self, + buffer, + destination: None, + waiting: false, + } + } + + /// Poll receiving one datagram, returning its copied length and source. pub fn poll_recv_from( &self, cx: &mut Context<'_>, buffer: &mut [u8], ) -> Poll> { + self.poll_recv_datagram(cx, buffer) + .map(|result| result.map(|packet| (packet.len, packet.source))) + } + + /// Poll for one datagram, parking the reader only after `WouldBlock`. + pub fn poll_recv_datagram( + &self, + cx: &mut Context<'_>, + buffer: &mut [u8], + ) -> Poll> { + // SAFETY: the slice is live and writable for its whole length. + unsafe { self.poll_recv_into(cx, buffer.as_mut_ptr(), buffer.len()) } + } + + /// Poll a receive directly into uninitialised spare buffer capacity. + pub fn poll_recv_from_buf( + &self, + cx: &mut Context<'_>, + buffer: &mut bytes::BytesMut, + ) -> Poll> { + let filled = buffer.len(); + let spare = buffer.spare_capacity_mut(); + if spare.is_empty() { + self.cancel_recv(); + return Poll::Ready(Err(crate::reactor::error::Errno(libc::ENOBUFS))); + } + let pointer = spare.as_mut_ptr().cast::(); + let capacity = spare.len(); + // SAFETY: the spare region is allocated, writable and disjoint from + // the buffer's initialised prefix. The syscall reports copied bytes. + let result = unsafe { self.poll_recv_into(cx, pointer, capacity) }; + if let Poll::Ready(Ok(datagram)) = &result { + // SAFETY: precisely datagram.len bytes in the spare region were + // initialised, and the low-level socket clamps that to capacity. + unsafe { buffer.set_len(filled + datagram.len) }; + } + result + } + + unsafe fn poll_recv_into( + &self, + cx: &mut Context<'_>, + pointer: *mut u8, + len: usize, + ) -> Poll> { if !self.registration.take_readable_or_park(cx.waker()) { return Poll::Pending; } loop { - // SAFETY: `buffer` is live for this call, and the socket writes no - // more than its length into it. - match unsafe { self.inner.recv_from(buffer.as_mut_ptr(), buffer.len()) } { - Ok(received) => { - // One datagram does not drain the receive queue. Keep this - // registration ready so the next poll can try another; - // if there is none, it will see EWOULDBLOCK and park. - // This also handles queued zero-length datagrams, which - // FIONREAD cannot distinguish from an empty queue. + // SAFETY: the caller guarantees the writable buffer region. + match unsafe { self.inner.recv_datagram(pointer, len) } { + Ok(datagram) => { + // A single datagram never proves the receive queue empty, + // regardless of its size or whether it was truncated. self.registration.mark_readable(); - return Poll::Ready(Ok(received)); + self.cancel_recv(); + return Poll::Ready(Ok(datagram)); } Err(error) if error.would_block() => { self.registration.poll_readable(cx.waker()); return Poll::Pending; } Err(error) if error.interrupted() => continue, - Err(error) => return Poll::Ready(Err(error)), + Err(error) => { + // An asynchronous UDP error need not drain queued data. + self.registration.mark_readable(); + self.cancel_recv(); + return Poll::Ready(Err(error)); + } } } } - /// Send one datagram to `target`, parking if the send buffer is full. + /// Poll sending one datagram to an explicit destination. pub fn poll_send_to( &self, cx: &mut Context<'_>, buffer: &[u8], - target: Addr, + destination: Addr, + ) -> Poll> { + self.poll_send_datagram(cx, buffer, Some(destination)) + } + + /// Poll sending one datagram to the configured peer. + pub fn poll_send(&self, cx: &mut Context<'_>, buffer: &[u8]) -> Poll> { + self.poll_send_datagram(cx, buffer, None) + } + + fn poll_send_datagram( + &self, + cx: &mut Context<'_>, + buffer: &[u8], + destination: Option, ) -> Poll> { if !self.registration.take_writable_or_park(cx.waker()) { return Poll::Pending; } loop { - match self.inner.send_to(buffer, target) { + let result = match destination { + Some(destination) => self.inner.send_to(buffer, destination), + None => self.inner.send(buffer), + }; + match result { Ok(sent) => { - // Datagram sends are atomic: after this one fit, let the - // next send attempt the syscall. If concurrent senders - // have since filled the buffer it will park on - // EWOULDBLOCK below. self.registration.mark_writable(); + self.cancel_send(); return Poll::Ready(Ok(sent)); } Err(error) if error.would_block() => { @@ -693,34 +893,38 @@ impl UdpSocket { return Poll::Pending; } Err(error) if error.interrupted() => continue, - Err(error) => return Poll::Ready(Err(error)), + Err(error) => { + // An invalid destination or oversized packet does not + // establish that the socket's send buffer is full. + self.registration.mark_writable(); + self.cancel_send(); + return Poll::Ready(Err(error)); + } } } } - /// Receive one datagram, as a future. - pub fn recv_from<'a>(&'a self, buffer: &'a mut [u8]) -> RecvFrom<'a> { - RecvFrom { - socket: self, - buffer, - } + /// Release the waiter for an abandoned direct receive poll. + /// + /// Cached readiness is preserved. Receive futures call this on drop. + pub fn cancel_recv(&self) { + self.registration.cancel_read_wait(); } - /// Send one datagram to `target`, as a future. - pub fn send_to<'a>(&'a self, buffer: &'a [u8], target: Addr) -> SendTo<'a> { - SendTo { - socket: self, - buffer, - target, - } + /// Release the waiter for an abandoned direct send poll. + /// + /// Cached readiness is preserved. Send futures call this on drop. + pub fn cancel_send(&self) { + self.registration.cancel_write_wait(); } } -/// The future returned by [`UdpSocket::recv_from`]. +/// The cancellation-safe future returned by [`UdpSocket::recv_from`]. #[derive(Debug)] pub struct RecvFrom<'a> { socket: &'a UdpSocket, buffer: &'a mut [u8], + waiting: bool, } impl core::future::Future for RecvFrom<'_> { @@ -728,16 +932,81 @@ impl core::future::Future for RecvFrom<'_> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); - this.socket.poll_recv_from(cx, this.buffer) + let result = this.socket.poll_recv_from(cx, this.buffer); + this.waiting = result.is_pending(); + result } } -/// The future returned by [`UdpSocket::send_to`]. +impl Drop for RecvFrom<'_> { + fn drop(&mut self) { + if self.waiting { + self.socket.cancel_recv(); + } + } +} + +/// The cancellation-safe future returned by [`UdpSocket::recv_datagram`]. +#[derive(Debug)] +pub struct RecvDatagram<'a> { + socket: &'a UdpSocket, + buffer: &'a mut [u8], + waiting: bool, +} + +impl core::future::Future for RecvDatagram<'_> { + type Output = Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let result = this.socket.poll_recv_datagram(cx, this.buffer); + this.waiting = result.is_pending(); + result + } +} + +impl Drop for RecvDatagram<'_> { + fn drop(&mut self) { + if self.waiting { + self.socket.cancel_recv(); + } + } +} + +/// The cancellation-safe future returned by [`UdpSocket::recv_from_buf`]. +#[derive(Debug)] +pub struct RecvFromBuf<'a> { + socket: &'a UdpSocket, + buffer: &'a mut bytes::BytesMut, + waiting: bool, +} + +impl core::future::Future for RecvFromBuf<'_> { + type Output = Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let result = this.socket.poll_recv_from_buf(cx, this.buffer); + this.waiting = result.is_pending(); + result + } +} + +impl Drop for RecvFromBuf<'_> { + fn drop(&mut self) { + if self.waiting { + self.socket.cancel_recv(); + } + } +} + +/// The cancellation-safe future returned by UDP send operations. #[derive(Debug)] pub struct SendTo<'a> { socket: &'a UdpSocket, buffer: &'a [u8], - target: Addr, + destination: Option, + waiting: bool, } impl core::future::Future for SendTo<'_> { @@ -745,7 +1014,19 @@ impl core::future::Future for SendTo<'_> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); - this.socket.poll_send_to(cx, this.buffer, this.target) + let result = this + .socket + .poll_send_datagram(cx, this.buffer, this.destination); + this.waiting = result.is_pending(); + result + } +} + +impl Drop for SendTo<'_> { + fn drop(&mut self) { + if self.waiting { + self.socket.cancel_send(); + } } } diff --git a/src/reactor.rs b/src/reactor.rs index 93c4e52..4eb086d 100644 --- a/src/reactor.rs +++ b/src/reactor.rs @@ -83,4 +83,4 @@ pub use driver::{Handle, Reactor, Registration, Sharded}; pub use error::{Errno, Result}; pub use local::{TaskSet, block_on, block_on_with}; pub use poller::{Event, Interest, Poller}; -pub use socket::{Addr, UNIX_PATH_CAPACITY, UnixPath}; +pub use socket::{Addr, Datagram, UNIX_PATH_CAPACITY, UnixPath}; diff --git a/src/reactor/driver.rs b/src/reactor/driver.rs index 6b569b5..bd8fc11 100644 --- a/src/reactor/driver.rs +++ b/src/reactor/driver.rs @@ -362,6 +362,24 @@ impl Registration { } } + /// Release a cancelled UDP reader without consuming cached readiness. + pub(crate) fn cancel_read_wait(&self) { + let waiting = { + let mut wakers = self.slot.wakers.lock().expect("reactor slot poisoned"); + wakers.reader.take() + }; + drop(waiting); + } + + /// Release a cancelled UDP writer without consuming cached readiness. + pub(crate) fn cancel_write_wait(&self) { + let waiting = { + let mut wakers = self.slot.wakers.lock().expect("reactor slot poisoned"); + wakers.writer.take() + }; + drop(waiting); + } + /// Change what this descriptor is watched for. pub fn modify(&self, interest: Interest) -> Result<()> { self.shared.poller.modify(self.fd, self.token, interest) diff --git a/src/reactor/socket.rs b/src/reactor/socket.rs index 4d2b545..9ec23e2 100644 --- a/src/reactor/socket.rs +++ b/src/reactor/socket.rs @@ -546,33 +546,47 @@ fn stream_socket(addr: &Addr) -> Result { Ok(fd) } -/// Create a non-blocking datagram socket for `addr`'s family. +/// Create a UDP descriptor without exposing a blocking socket. fn datagram_socket(addr: &Addr) -> Result { - // SOCK_NONBLOCK and SOCK_CLOEXEC are Linux extensions to `socket`; the - // BSDs set both flags after creation, before the descriptor is published. + let protocol = if addr.family() == libc::AF_UNIX { + 0 + } else { + libc::IPPROTO_UDP + }; + #[cfg(target_os = "linux")] - let raw = check(unsafe { - libc::socket( - addr.family(), - libc::SOCK_DGRAM | libc::SOCK_NONBLOCK | libc::SOCK_CLOEXEC, - 0, - ) - })?; + let fd = { + let raw = check(unsafe { + libc::socket( + addr.family(), + libc::SOCK_DGRAM | libc::SOCK_NONBLOCK | libc::SOCK_CLOEXEC, + protocol, + ) + })?; + // SAFETY: this is a fresh descriptor owned exclusively here. + unsafe { Fd::from_raw(raw) } + }; #[cfg(not(target_os = "linux"))] - let raw = { - let raw = check(unsafe { libc::socket(addr.family(), libc::SOCK_DGRAM, 0) })?; - // SAFETY: `raw` is a descriptor the call above just returned. - unsafe { - let flags = libc::fcntl(raw, libc::F_GETFL); - libc::fcntl(raw, libc::F_SETFL, flags | libc::O_NONBLOCK); - libc::fcntl(raw, libc::F_SETFD, libc::FD_CLOEXEC); - } - raw + let fd = { + let raw = check(unsafe { libc::socket(addr.family(), libc::SOCK_DGRAM, protocol) })?; + // Take ownership before setting flags so every failure closes it. + // SAFETY: this is a fresh descriptor owned exclusively here. + let fd = unsafe { Fd::from_raw(raw) }; + let flags = check(unsafe { libc::fcntl(fd.raw(), libc::F_GETFL) })?; + check(unsafe { libc::fcntl(fd.raw(), libc::F_SETFL, flags | libc::O_NONBLOCK) })?; + let flags = check(unsafe { libc::fcntl(fd.raw(), libc::F_GETFD) })?; + check(unsafe { libc::fcntl(fd.raw(), libc::F_SETFD, flags | libc::FD_CLOEXEC) })?; + set_flag(fd.raw(), libc::SOL_SOCKET, libc::SO_NOSIGPIPE, true)?; + fd }; - // SAFETY: `raw` is a fresh descriptor owned exclusively here. - Ok(unsafe { Fd::from_raw(raw) }) + // An explicit IPv6 socket stays IPv6. Callers that need both families + // own one socket for each rather than depending on platform defaults. + if addr.family() == libc::AF_INET6 { + set_flag(fd.raw(), libc::IPPROTO_IPV6, libc::IPV6_V6ONLY, true)?; + } + Ok(fd) } /// Set a boolean socket option. @@ -877,23 +891,45 @@ impl TcpSocket { } } +/// The result of consuming one UDP datagram. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Datagram { + /// Bytes copied into the supplied buffer. + /// + /// This can be zero for an empty datagram or an empty receive buffer. + /// Zero does not indicate end of stream. + pub len: usize, + /// The sender's address, including an IPv6 scope ID when present. + pub source: Addr, + /// Whether the buffer was too small for the datagram. + /// + /// The kernel discards the remainder of that datagram. The next receive + /// starts with the next datagram. The original length is not reported + /// because platforms differ in whether they make it available. + pub truncated: bool, +} + /// A non-blocking datagram socket. #[derive(Debug)] -pub struct UdpSocket(Fd); +pub struct UdpSocket { + fd: Fd, + family: Option, +} impl UdpSocket { /// Bind a datagram socket to `addr`. pub fn bind(addr: Addr) -> Result { - Self::bind_with(addr, false) + Self::bind_with_reuse(addr, false) } /// Bind with `SO_REUSEADDR` enabled for IP addresses before they are /// claimed. pub fn bind_reuse_addr(addr: Addr) -> Result { - Self::bind_with(addr, true) + Self::bind_with_reuse(addr, true) } - fn bind_with(addr: Addr, reuse_addr: bool) -> Result { + /// Bind with an explicit reuse-address setting. + pub fn bind_with_reuse(addr: Addr, reuse_addr: bool) -> Result { let fd = datagram_socket(&addr)?; if reuse_addr && addr.family() != libc::AF_UNIX { set_flag(fd.raw(), libc::SOL_SOCKET, libc::SO_REUSEADDR, true)?; @@ -904,7 +940,10 @@ impl UdpSocket { let len = addr.write_to(&mut storage); // SAFETY: `storage` holds a valid address of `len` bytes. check(unsafe { libc::bind(fd.raw(), core::ptr::addr_of!(storage).cast(), len) })?; - Ok(Self(fd)) + Ok(Self { + fd, + family: Some(addr.family()), + }) } /// Adopt an existing non-blocking datagram descriptor. @@ -915,13 +954,16 @@ impl UdpSocket { /// will close. pub unsafe fn from_raw(fd: i32) -> Self { // SAFETY: the caller guarantees the descriptor is open and uniquely owned. - Self(unsafe { Fd::from_raw(fd) }) + Self { + fd: unsafe { Fd::from_raw(fd) }, + family: addr_of(fd, false).ok().map(|addr| addr.family()), + } } /// The descriptor, for registering with the poller. #[inline] pub fn raw(&self) -> i32 { - self.0.raw() + self.fd.raw() } /// This socket's own address. @@ -929,6 +971,48 @@ impl UdpSocket { addr_of(self.raw(), false) } + /// The configured peer, or `ENOTCONN` for an unconnected socket. + pub fn peer_addr(&self) -> Result { + addr_of(self.raw(), true) + } + + /// Select a peer and let the kernel filter incoming datagrams by it. + /// + /// UDP has no handshake. A platform that returns `EINPROGRESS` requires + /// writable readiness and [`Self::connect_error`] before use. + pub fn connect(&self, addr: Addr) -> Result<()> { + if self.family.is_some_and(|family| addr.family() != family) { + return Err(Errno(libc::EAFNOSUPPORT)); + } + // SAFETY: zeroed storage is filled for the selected address family. + let mut storage: libc::sockaddr_storage = unsafe { mem::zeroed() }; + let len = addr.write_to(&mut storage); + // SAFETY: storage holds a valid address of len bytes. + check(unsafe { libc::connect(self.raw(), core::ptr::addr_of!(storage).cast(), len) })?; + Ok(()) + } + + /// Check the result of an in-flight connect. + pub fn connect_error(&self) -> Result<()> { + let mut value: libc::c_int = 0; + let mut len = mem::size_of::() as libc::socklen_t; + // SAFETY: SO_ERROR writes one int into the live value. + check(unsafe { + libc::getsockopt( + self.raw(), + libc::SOL_SOCKET, + libc::SO_ERROR, + core::ptr::addr_of_mut!(value).cast(), + core::ptr::addr_of_mut!(len), + ) + })?; + if value == 0 { + Ok(()) + } else { + Err(Errno(value)) + } + } + /// Set the kernel receive-buffer capacity in bytes. pub fn set_recv_buffer(&self, size: usize) -> Result<()> { set_socket_int(self.raw(), libc::SO_RCVBUF, size) @@ -939,57 +1023,120 @@ impl UdpSocket { set_socket_int(self.raw(), libc::SO_SNDBUF, size) } - /// Receive one datagram directly into caller-owned memory. + /// Receive one datagram, returning its copied length and source. /// - /// A datagram larger than `len` is truncated by the kernel. The returned - /// count is the number of bytes copied into the buffer. + /// Use [`Self::recv_datagram`] when truncation metadata is needed. /// /// # Safety /// /// `pointer` must be valid for writes of `len` bytes. pub unsafe fn recv_from(&self, pointer: *mut u8, len: usize) -> Result<(usize, Addr)> { - // SAFETY: zeroed storage is valid for the kernel to fill with a peer - // address, and `address_len` describes the full storage. + // SAFETY: the caller guarantees the same buffer requirements. + unsafe { self.recv_datagram(pointer, len) }.map(|packet| (packet.len, packet.source)) + } + + /// Receive exactly one datagram into potentially uninitialised memory. + /// + /// Empty datagrams and empty receive buffers are valid. `MSG_TRUNC` + /// reports discarded bytes without consuming any following datagram. + /// + /// # Safety + /// + /// `pointer` must be valid for writes of `len` bytes. + pub unsafe fn recv_datagram(&self, pointer: *mut u8, len: usize) -> Result { + // SAFETY: zeroed storage is valid for the kernel to fill. let mut storage: libc::sockaddr_storage = unsafe { mem::zeroed() }; - let mut address_len = mem::size_of::() as libc::socklen_t; - // SAFETY: the caller guarantees the receive buffer, while the address - // storage and its length are live locals of the right types. - let read = unsafe { - libc::recvfrom( - self.raw(), - pointer.cast(), - len, - 0, - core::ptr::addr_of_mut!(storage).cast(), - core::ptr::addr_of_mut!(address_len), - ) + let mut iovec = libc::iovec { + iov_base: pointer.cast(), + iov_len: len, }; - if read < 0 { - return Err(crate::reactor::error::last()); + // SAFETY: zeroed msghdr starts with no ancillary buffers or flags. + let mut message: libc::msghdr = unsafe { mem::zeroed() }; + message.msg_name = core::ptr::addr_of_mut!(storage).cast(); + message.msg_namelen = mem::size_of::() as libc::socklen_t; + message.msg_iov = core::ptr::addr_of_mut!(iovec); + message.msg_iovlen = 1; + + #[cfg(feature = "syscall-counters")] + crate::reactor::counters::RECV.fetch_add(1, core::sync::atomic::Ordering::Relaxed); + // Do not pass MSG_TRUNC as an input flag: Linux would then return the + // original length instead of the number of initialised bytes. + // SAFETY: the caller guarantees the receive buffer; the message and + // address storage remain live throughout the syscall. + let received = unsafe { libc::recvmsg(self.raw(), core::ptr::addr_of_mut!(message), 0) }; + if received < 0 { + let error = crate::reactor::error::last(); + #[cfg(feature = "syscall-counters")] + if error.would_block() { + crate::reactor::counters::RECV_WOULD_BLOCK + .fetch_add(1, core::sync::atomic::Ordering::Relaxed); + } + return Err(error); } - let addr = Addr::read_from(&storage, address_len).ok_or(Errno(libc::EAFNOSUPPORT))?; - Ok((read as usize, addr)) + let source = + Addr::read_from(&storage, message.msg_namelen).ok_or(Errno(libc::EAFNOSUPPORT))?; + Ok(Datagram { + len: (received as usize).min(len), + source, + truncated: message.msg_flags & libc::MSG_TRUNC != 0 || received as usize > len, + }) } - /// Send one datagram from caller-owned memory to `target`. - pub fn send_to(&self, buffer: &[u8], target: Addr) -> Result { - // SAFETY: zeroed storage, then written by `write_to` for this family. + /// Send one complete datagram to an explicit destination. + pub fn send_to(&self, buffer: &[u8], destination: Addr) -> Result { + self.send_datagram(buffer, Some(destination)) + } + + /// Send one complete datagram to the configured peer. + pub fn send(&self, buffer: &[u8]) -> Result { + self.send_datagram(buffer, None) + } + + fn send_datagram(&self, buffer: &[u8], destination: Option) -> Result { + // SAFETY: zeroed storage is filled before its address is used. let mut storage: libc::sockaddr_storage = unsafe { mem::zeroed() }; - let address_len = target.write_to(&mut storage); - // SAFETY: the buffer and address storage are live for the call, and - // `address_len` describes the address written above. + let (address, len) = if let Some(destination) = destination { + if self + .family + .is_some_and(|family| destination.family() != family) + { + return Err(Errno(libc::EAFNOSUPPORT)); + } + let len = destination.write_to(&mut storage); + (core::ptr::addr_of!(storage).cast::(), len) + } else { + (core::ptr::null(), 0) + }; + #[cfg(target_os = "linux")] + let flags = libc::MSG_NOSIGNAL; + #[cfg(not(target_os = "linux"))] + let flags = 0; + + #[cfg(feature = "syscall-counters")] + crate::reactor::counters::SEND.fetch_add(1, core::sync::atomic::Ordering::Relaxed); + // SAFETY: buffer and the optional address remain live for the call. let sent = unsafe { libc::sendto( self.raw(), buffer.as_ptr().cast(), buffer.len(), - 0, - core::ptr::addr_of!(storage).cast(), - address_len, + flags, + address, + len, ) }; if sent < 0 { - Err(crate::reactor::error::last()) + let error = crate::reactor::error::last(); + #[cfg(feature = "syscall-counters")] + if error.would_block() { + crate::reactor::counters::SEND_WOULD_BLOCK + .fetch_add(1, core::sync::atomic::Ordering::Relaxed); + } + Err(error) + } else if sent as usize != buffer.len() { + // UDP writes are atomic. Never turn a short result into another + // datagram by retrying a suffix. + Err(Errno(libc::EMSGSIZE)) } else { Ok(sent as usize) } From 60d946de2cba09f0dab2f4747bc0cdcc144fd42a Mon Sep 17 00:00:00 2001 From: meh Date: Wed, 14 Oct 2026 17:35:46 +0700 Subject: [PATCH 2/2] Release 0.1.19 --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 24cc61f..1466532 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "nagoya" -version = "0.1.18" +version = "0.1.19" edition = "2024" rust-version = "1.85" license = "MIT OR Apache-2.0"