From 1a4a8bda03379f1638900946dadc4c69ef3dbae8 Mon Sep 17 00:00:00 2001 From: al8n Date: Sun, 5 Jan 2025 15:45:05 +0800 Subject: [PATCH 1/4] Add `UdpSocket::peek` method --- tokio/src/net/udp.rs | 50 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/tokio/src/net/udp.rs b/tokio/src/net/udp.rs index 00802b036dd..edca4cb097b 100644 --- a/tokio/src/net/udp.rs +++ b/tokio/src/net/udp.rs @@ -1503,6 +1503,56 @@ impl UdpSocket { .await } + /// Receives a single datagram from the connected address without removing it from the queue. + /// On success, returns the number of bytes read from whence the data came. + /// + /// # Notes + /// + /// On Windows, if the data is larger than the buffer specified, the buffer + /// is filled with the first part of the data, and `peek_from` returns the error + /// `WSAEMSGSIZE(10040)`. The excess data is lost. + /// Make sure to always use a sufficiently large buffer to hold the + /// maximum UDP packet size, which can be up to 65536 bytes in size. + /// + /// MacOS will return an error if you pass a zero-sized buffer. + /// + /// If you're merely interested in learning the sender of the data at the head of the queue, + /// try [`peek_sender`]. + /// + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::UdpSocket; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let socket = UdpSocket::bind("127.0.0.1:8080").await?; + /// + /// let mut buf = vec![0u8; 32]; + /// let len = socket.peek(&mut buf).await?; + /// + /// println!("peeked {:?} bytes", len); + /// + /// Ok(()) + /// } + /// ``` + /// + /// [`peek_sender`]: method@Self::peek_sender + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection + pub async fn peek(&self, buf: &mut [u8]) -> io::Result { + self.io + .registration() + .async_io(Interest::READABLE, || self.io.peek(buf)) + .await + } + /// Receives data from the socket, without removing it from the input queue. /// On success, returns the number of bytes read and the address from whence /// the data came. From a40e26a50aae3c8b09423c470bd794f184f69604 Mon Sep 17 00:00:00 2001 From: al8n Date: Sun, 5 Jan 2025 20:30:07 +0800 Subject: [PATCH 2/4] Add `poll_peek` method --- tokio/src/net/udp.rs | 63 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/tokio/src/net/udp.rs b/tokio/src/net/udp.rs index edca4cb097b..a715257e2ca 100644 --- a/tokio/src/net/udp.rs +++ b/tokio/src/net/udp.rs @@ -1553,6 +1553,69 @@ impl UdpSocket { .await } + /// Receives data from the connected address, without removing it from the input queue. + /// On success, returns the sending address of the datagram. + /// + /// # Notes + /// + /// Note that on multiple calls to a `poll_*` method in the `recv` direction, only the + /// `Waker` from the `Context` passed to the most recent call will be scheduled to + /// receive a wakeup + /// + /// On Windows, if the data is larger than the buffer specified, the buffer + /// is filled with the first part of the data, and peek returns the error + /// `WSAEMSGSIZE(10040)`. The excess data is lost. + /// Make sure to always use a sufficiently large buffer to hold the + /// maximum UDP packet size, which can be up to 65536 bytes in size. + /// + /// MacOS will return an error if you pass a zero-sized buffer. + /// + /// If you're merely interested in learning the sender of the data at the head of the queue, + /// try [`poll_peek_sender`]. + /// + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// # Return value + /// + /// The function returns: + /// + /// * `Poll::Pending` if the socket is not ready to read + /// * `Poll::Ready(Ok(()))` reads data into `ReadBuf` if the socket is ready + /// * `Poll::Ready(Err(e))` if an error is encountered. + /// + /// # Errors + /// + /// This function may encounter any standard I/O error except `WouldBlock`. + /// + /// [`poll_peek_sender`]: method@Self::poll_peek_sender + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection + pub fn poll_peek( + &self, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + #[allow(clippy::blocks_in_conditions)] + let n = ready!(self.io.registration().poll_read_io(cx, || { + // Safety: will not read the maybe uninitialized bytes. + let b = unsafe { + &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit] as *mut [u8]) + }; + + self.io.peek(b) + }))?; + + // Safety: We trust `recv` to have filled up `n` bytes in the buffer. + unsafe { + buf.assume_init(n); + } + buf.advance(n); + Poll::Ready(Ok(())) + } + /// Receives data from the socket, without removing it from the input queue. /// On success, returns the number of bytes read and the address from whence /// the data came. From 60739ff28ecbda32a0fad08adf99aa0421a06f5a Mon Sep 17 00:00:00 2001 From: al8n Date: Sun, 5 Jan 2025 20:33:01 +0800 Subject: [PATCH 3/4] Fix fmt --- tokio/src/net/udp.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tokio/src/net/udp.rs b/tokio/src/net/udp.rs index a715257e2ca..638c04a2d6e 100644 --- a/tokio/src/net/udp.rs +++ b/tokio/src/net/udp.rs @@ -1593,11 +1593,7 @@ impl UdpSocket { /// /// [`poll_peek_sender`]: method@Self::poll_peek_sender /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection - pub fn poll_peek( - &self, - cx: &mut Context<'_>, - buf: &mut ReadBuf<'_>, - ) -> Poll> { + pub fn poll_peek(&self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { #[allow(clippy::blocks_in_conditions)] let n = ready!(self.io.registration().poll_read_io(cx, || { // Safety: will not read the maybe uninitialized bytes. From c0024997667d112ad6efaa8b0d042fd4a7c1459a Mon Sep 17 00:00:00 2001 From: al8n Date: Mon, 6 Jan 2025 16:56:32 +0800 Subject: [PATCH 4/4] Add `try_peek` method --- tokio/src/net/udp.rs | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tokio/src/net/udp.rs b/tokio/src/net/udp.rs index 638c04a2d6e..2690a84b602 100644 --- a/tokio/src/net/udp.rs +++ b/tokio/src/net/udp.rs @@ -1612,6 +1612,39 @@ impl UdpSocket { Poll::Ready(Ok(())) } + /// Tries to receive data on the connected address without removing it from the input queue. + /// On success, returns the number of bytes read. + /// + /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is + /// returned. This function is usually paired with `readable()`. + /// + /// # Notes + /// + /// On Windows, if the data is larger than the buffer specified, the buffer + /// is filled with the first part of the data, and peek returns the error + /// `WSAEMSGSIZE(10040)`. The excess data is lost. + /// Make sure to always use a sufficiently large buffer to hold the + /// maximum UDP packet size, which can be up to 65536 bytes in size. + /// + /// MacOS will return an error if you pass a zero-sized buffer. + /// + /// If you're merely interested in learning the sender of the data at the head of the queue, + /// try [`try_peek_sender`]. + /// + /// Note that the socket address **cannot** be implicitly trusted, because it is relatively + /// trivial to send a UDP datagram with a spoofed origin in a [packet injection attack]. + /// Because UDP is stateless and does not validate the origin of a packet, + /// the attacker does not need to be able to intercept traffic in order to interfere. + /// It is important to be aware of this when designing your application-level protocol. + /// + /// [`try_peek_sender`]: method@Self::try_peek_sender + /// [packet injection attack]: https://en.wikipedia.org/wiki/Packet_injection + pub fn try_peek(&self, buf: &mut [u8]) -> io::Result { + self.io + .registration() + .try_io(Interest::READABLE, || self.io.peek(buf)) + } + /// Receives data from the socket, without removing it from the input queue. /// On success, returns the number of bytes read and the address from whence /// the data came.