From 47ce9a370c6fc1550f04d60880a690946933fa3c Mon Sep 17 00:00:00 2001
From: "Matthieu Le brazidec (r3v2d0g)" <r3v2d0g@jesus.gg>
Date: Sat, 15 Aug 2020 15:06:33 +0200
Subject: [PATCH] Add peek{,from} methods to UdpSocket

---
 src/net/udp/mod.rs | 47 ++++++++++++++++++++++++++++++++++++++++++++++
 tests/udp.rs       |  5 ++++-
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/src/net/udp/mod.rs b/src/net/udp/mod.rs
index dd47e0582..47a29facd 100644
--- a/src/net/udp/mod.rs
+++ b/src/net/udp/mod.rs
@@ -206,6 +206,29 @@ impl UdpSocket {
         self.watcher.recv_from(buf).await
     }
 
+    /// Receives data from socket without removing it from the queue.
+    ///
+    /// On success, returns the number of bytes peeked and the origin.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # fn main() -> std::io::Result<()> { async_std::task::block_on(async {
+    /// #
+    /// use async_std::net::UdpSocket;
+    ///
+    /// let socket = UdpSocket::bind("127.0.0.1:0").await?;
+    ///
+    /// let mut buf = vec![0; 1024];
+    /// let (n, peer) = socket.peek_from(&mut buf).await?;
+    /// println!("Peeked {} bytes from {}", n, peer);
+    /// #
+    /// # Ok (()) }) }
+    /// ```
+    pub async fn peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
+        self.watcher.peek_from(buf).await
+    }
+
     /// Connects the UDP socket to a remote address.
     ///
     /// When connected, methods [`send`] and [`recv`] will use the specified address for sending
@@ -301,6 +324,30 @@ impl UdpSocket {
         self.watcher.recv(buf).await
     }
 
+    /// Receives data from the socket without removing it from the queue.
+    ///
+    /// On success, returns the number of bytes peeked.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # fn main() -> std::io::Result<()> { async_std::task::block_on(async {
+    /// #
+    /// use async_std::net::UdpSocket;
+    ///
+    /// let socket = UdpSocket::bind("127.0.0.1:0").await?;
+    /// socket.connect("127.0.0.1:8080").await?;
+    ///
+    /// let mut buf = vec![0; 1024];
+    /// let n = socket.peek(&mut buf).await?;
+    /// println!("Peeked {} bytes", n);
+    /// #
+    /// # Ok(()) }) }
+    /// ```
+    pub async fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
+        self.watcher.peek(buf).await
+    }
+
     /// Gets the value of the `SO_BROADCAST` option for this socket.
     ///
     /// For more information about this option, see [`set_broadcast`].
diff --git a/tests/udp.rs b/tests/udp.rs
index 37024c478..cd119ddcd 100644
--- a/tests/udp.rs
+++ b/tests/udp.rs
@@ -12,7 +12,7 @@ const THE_MERCHANT_OF_VENICE: &[u8] = b"
 ";
 
 #[test]
-fn send_recv() -> io::Result<()> {
+fn send_recv_peek() -> io::Result<()> {
     task::block_on(async {
         let socket1 = UdpSocket::bind("127.0.0.1:0").await?;
         let socket2 = UdpSocket::bind("127.0.0.1:0").await?;
@@ -23,6 +23,9 @@ fn send_recv() -> io::Result<()> {
         socket1.send(THE_MERCHANT_OF_VENICE).await?;
 
         let mut buf = [0u8; 1024];
+        let n = socket2.peek(&mut buf).await?;
+        assert_eq!(&buf[..n], THE_MERCHANT_OF_VENICE);
+
         let n = socket2.recv(&mut buf).await?;
         assert_eq!(&buf[..n], THE_MERCHANT_OF_VENICE);