Make UdpStream block until the next non-empty msg.

mahkoh · Oct 25, 2014 · 8adfd02 · 8adfd02
1 parent 172b59a
commit 8adfd02
Showing 1 changed file with 20 additions and 13 deletions.
diff --git a/src/libstd/io/net/udp.rs b/src/libstd/io/net/udp.rs
@@ -225,13 +225,15 @@ impl UdpStream {
 }
 
 impl Reader for UdpStream {
+    /// Returns the next non-empty message from the specified address.
     fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
         let peer = self.connected_to;
         self.as_socket(|sock| {
-            match sock.recv_from(buf) {
-                Ok((_nread, src)) if src != peer => Ok(0),
-                Ok((nread, _src)) => Ok(nread),
-                Err(e) => Err(e),
+            loop {
+                let (nread, src) = try!(sock.recv_from(buf));
+                if nread > 0 && src == peer {
+                    return Ok(nread);
+                }
             }
         })
     }
@@ -337,19 +339,24 @@ mod test {
     fn stream_smoke_test_ip4() {
         let server_ip = next_test_ip4();
         let client_ip = next_test_ip4();
+        let dummy_ip = next_test_ip4();
         let (tx1, rx1) = channel();
         let (tx2, rx2) = channel();
 
         spawn(proc() {
-            match UdpSocket::bind(client_ip) {
-                Ok(client) => {
-                    let client = box client;
-                    let mut stream = client.connect(server_ip);
-                    rx1.recv();
-                    stream.write([99]).unwrap();
+            let send_as = |ip, val: &[u8]| {
+                match UdpSocket::bind(ip) {
+                    Ok(client) => {
+                        let client = box client;
+                        let mut stream = client.connect(server_ip);
+                        stream.write(val).unwrap();
+                    }
+                    Err(..) => fail!()
                 }
-                Err(..) => fail!()
-            }
+            };
+            rx1.recv();
+            send_as(dummy_ip, [98]);
+            send_as(client_ip, [99]);
             tx2.send(());
         });
 
@@ -364,7 +371,7 @@ mod test {
                         assert_eq!(nread, 1);
                         assert_eq!(buf[0], 99);
                     }
-                    Err(..) => fail!()
+                    Err(..) => fail!(),
                 }
             }
             Err(..) => fail!()