Add UdpSocket::try_recv (#255)

Also, add ::connect capabilities to UdpSocket.

Author: jgrund

src/host.rs

@@ -212,6 +212,7 @@ pub(crate) struct Udp {
 
 struct UdpBind {
     bind_addr: SocketAddr,
+    target_addr: Option<SocketAddr>,
     broadcast: bool,
     multicast_loop: bool,
     queue: mpsc::Sender<(Datagram, SocketAddr)>,
@@ -259,6 +260,7 @@ impl Udp {
         let (tx, rx) = mpsc::channel(self.capacity);
         let bind = UdpBind {
             bind_addr: addr,
+            target_addr: None,
             broadcast: DEFAULT_BROADCAST,
             multicast_loop: DEFAULT_MULTICAST_LOOP,
             queue: tx,
@@ -275,9 +277,23 @@ impl Udp {
 
         Ok(UdpSocket::new(addr, rx))
     }
+    pub(crate) fn connect(&mut self, src: SocketAddr, dst: SocketAddr) {
+        let Some(bind) = self.binds.get_mut(&src.port()) else {
+            panic!("Connect failed (no matching bind) for {src}");
+        };
+
+        bind.target_addr = Some(dst);
+    }
 
     fn receive_from_network(&mut self, src: SocketAddr, dst: SocketAddr, datagram: Datagram) {
         if let Some(bind) = self.binds.get_mut(&dst.port()) {
+            if let Some(target) = bind.target_addr {
+                if !matches(target, src) {
+                    tracing::trace!(target: TRACING_TARGET, ?src, ?dst, protocol = %Protocol::Udp(datagram), "Dropped (Connect Addr not matching)");
+                    return;
+                }
+            }
+
             if !matches(bind.bind_addr, dst) {
                 tracing::trace!(target: TRACING_TARGET, ?src, ?dst, protocol = %Protocol::Udp(datagram), "Dropped (Addr not bound)");
                 return;

src/net/udp.rs

@@ -129,7 +129,7 @@ impl Rx {
     /// will continue to return immediately until the readiness event is
     /// consumed by an attempt to read that fails with `WouldBlock` or
     /// `Poll::Pending`.
-    pub async fn readable(&mut self) -> Result<()> {
+    async fn readable(&mut self) -> Result<()> {
         if self.buffer.is_some() {
             return Ok(());
         }
@@ -156,6 +156,14 @@ impl UdpSocket {
             }),
         }
     }
+    pub async fn connect<A: ToSocketAddrs>(&self, addr: A) {
+        World::current(|world| {
+            let addr = addr.to_socket_addr(&world.dns);
+            let host = world.current_host_mut();
+
+            host.udp.connect(self.local_addr, addr);
+        })
+    }
 
     /// This function will create a new UDP socket and attempt to bind it to
     /// the `addr` provided.
@@ -292,6 +300,19 @@ impl UdpSocket {
 
         Ok((limit, origin))
     }
+    /// Tries to receive a single datagram message on the socket from the remote
+    /// address to which it is connected. On success, returns the number of
+    /// bytes read.
+    ///
+    /// This method must be called with valid byte array `buf` of sufficient size
+    /// to hold the message bytes. If a message is too long to fit in the
+    /// supplied buffer, excess bytes may be discarded.
+    ///
+    /// When there is no pending data, `Err(io::ErrorKind::WouldBlock)` is
+    /// returned. This function is usually paired with `readable()`.
+    pub fn try_recv(&self, buf: &mut [u8]) -> io::Result<usize> {
+        self.try_recv_from(buf).map(|(size, _)| size)
+    }
 
     /// Waits for the socket to become readable.
     ///

tests/udp.rs

@@ -1007,3 +1007,101 @@ fn socket_to_nonexistent_node() -> Result {
     });
     sim.run()
 }
+
+#[test]
+fn try_recv_not_connected() -> Result {
+    let mut sim = Builder::new().build();
+
+    sim.client("server1", async move {
+        let sock = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 1234)).await?;
+
+        sock.readable().await?;
+
+        let mut buf = [0; 3];
+
+        // When not connected try_recv takes any incoming Datagram.
+        let x = sock.try_recv(&mut buf)?;
+
+        assert_eq!(x, 3);
+        assert_eq!(buf, [1, 2, 3]);
+
+        Ok(())
+    });
+
+    sim.client("server2", async move {
+        let sock = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 5678)).await?;
+
+        sock.send_to(&[1, 2, 3], ("server1", 1234)).await?;
+
+        Ok(())
+    });
+
+    sim.run()
+}
+
+#[test]
+fn try_recv_connected() -> Result {
+    let mut sim = Builder::new().build();
+
+    sim.client("server1", async move {
+        let sock = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 1234)).await?;
+
+        sock.connect(("server2", 5678)).await;
+
+        sock.readable().await?;
+
+        let mut buf = [0; 3];
+
+        let x = sock.try_recv(&mut buf)?;
+
+        assert_eq!(x, 3);
+        assert_eq!(buf, [1, 2, 3]);
+
+        Ok(())
+    });
+
+    sim.client("server2", async move {
+        let sock = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 5678)).await?;
+
+        sock.send_to(&[1, 2, 3], ("server1", 1234)).await?;
+
+        Ok(())
+    });
+
+    sim.run()
+}
+
+#[test]
+fn try_recv_connected_wrong_host() -> Result {
+    let mut sim = Builder::new().build();
+
+    sim.client("server1", async move {
+        let sock = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 1234)).await?;
+
+        // Connect to server3, which doesn't send anything.
+        sock.connect(("server3", 6781)).await;
+
+        // Datagram from server2 is filtered so sock never becomes readable.
+        tokio::time::timeout(Duration::from_secs(5), sock.readable())
+            .await
+            .unwrap_err();
+
+        Ok(())
+    });
+
+    sim.client("server2", async move {
+        let sock = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 5678)).await?;
+
+        sock.send_to(&[1, 2, 9], ("server1", 1234)).await?;
+
+        Ok(())
+    });
+
+    sim.client("server3", async move {
+        _ = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 6789)).await?;
+
+        Ok(())
+    });
+
+    sim.run()
+}