Add Socket::peek_sender()

Signed-off-by: Austin Bonander <[email protected]>

Author: abonander

src/socket.rs

@@ -596,11 +596,36 @@ impl Socket {
     /// `peek_from` makes the same safety guarantees regarding the `buf`fer as
     /// [`recv`].
     ///
+    /// # Note: Datagram Sockets
+    /// For datagram sockets, the behavior of this method when `buf` is smaller than
+    /// the datagram at the head of the receive queue differs between Windows and
+    /// Unix-like platforms (Linux, macOS, BSDs, etc: colloquially termed "*nix").
+    ///
+    /// On *nix platforms, the datagram is truncated to the length of `buf`.
+    ///
+    /// On Windows, an error corresponding to `WSAEMSGSIZE` will be returned.
+    ///
+    /// For consistency between platforms, be sure to provide a sufficiently large buffer to avoid
+    /// truncation; the exact size required depends on the underlying protocol.
+    ///
+    /// If you just want to know the sender of the data, try [`peek_sender`].
+    ///
     /// [`recv`]: Socket::recv
+    /// [`peek_sender`]: Socket::peek_sender
     pub fn peek_from(&self, buf: &mut [MaybeUninit<u8>]) -> io::Result<(usize, SockAddr)> {
         self.recv_from_with_flags(buf, sys::MSG_PEEK)
     }
 
+    /// Retrieve the sender for the data at the head of the receive queue.
+    ///
+    /// This is equivalent to calling [`peek_from`] with a zero-sized buffer,
+    /// but suppresses the `WSAEMSGSIZE` error on Windows.
+    ///
+    /// [`peek_from`]: Socket::peek_from
+    pub fn peek_sender(&self) -> io::Result<SockAddr> {
+        sys::peek_sender(self.as_raw())
+    }
+
     /// Sends data on the socket to a connected peer.
     ///
     /// This is typically used on TCP sockets or datagram sockets which have

src/sys/unix.rs

@@ -749,6 +749,15 @@ pub(crate) fn recv_from(
     }
 }
 
+pub(crate) fn peek_sender(fd: Socket) -> io::Result<SockAddr> {
+    // Unix-like platforms simply truncate the returned data, so this implementation is trivial.
+    // However, for Windows this requires suppressing the `WSAEMSGSIZE` error,
+    // so that requires a different approach.
+    // NOTE: macOS does not populate `sockaddr` if you pass a zero-sized buffer.
+    let (_, sender) = recv_from(fd, &mut [MaybeUninit::uninit(); 8], MSG_PEEK)?;
+    Ok(sender)
+}
+
 #[cfg(not(target_os = "redox"))]
 pub(crate) fn recv_vectored(
     fd: Socket,

src/sys/windows.rs

@@ -469,6 +469,38 @@ pub(crate) fn recv_from(
     }
 }
 
+pub(crate) fn peek_sender(socket: Socket) -> io::Result<SockAddr> {
+    // Safety: `recvfrom` initialises the `SockAddr` for us.
+    let ((), sender) = unsafe {
+        SockAddr::try_init(|storage, addrlen| {
+            let res = syscall!(
+                recvfrom(
+                    socket,
+                    // Windows *appears* not to care if you pass a null pointer.
+                    ptr::null_mut(),
+                    0,
+                    MSG_PEEK,
+                    storage.cast(),
+                    addrlen,
+                ),
+                PartialEq::eq,
+                SOCKET_ERROR
+            );
+            match res {
+                Ok(_n) => Ok(()),
+                Err(e) => match e.raw_os_error() {
+                    Some(code) if code == (WSAESHUTDOWN as i32) || code == (WSAEMSGSIZE as i32) => {
+                        Ok(())
+                    }
+                    _ => Err(e),
+                },
+            }
+        })
+    }?;
+
+    Ok(sender)
+}
+
 pub(crate) fn recv_from_vectored(
     socket: Socket,
     bufs: &mut [crate::MaybeUninitSlice<'_>],

tests/socket.rs

@@ -531,6 +531,21 @@ fn out_of_band() {
     assert_eq!(unsafe { assume_init(&buf[..n]) }, DATA);
 }
 
+#[test]
+#[cfg(not(target_os = "redox"))] // cfg of `udp_pair_unconnected()`
+fn udp_peek_sender() {
+    let (socket_a, socket_b) = udp_pair_unconnected();
+
+    let socket_a_addr = socket_a.local_addr().unwrap();
+    let socket_b_addr = socket_b.local_addr().unwrap();
+
+    socket_b.send_to(b"Hello, world!", &socket_a_addr).unwrap();
+
+    let sender_addr = socket_a.peek_sender().unwrap();
+
+    assert_eq!(sender_addr.as_socket(), socket_b_addr.as_socket());
+}
+
 #[test]
 #[cfg(not(target_os = "redox"))]
 fn send_recv_vectored() {