Fix cargo test

examples/udp-hole-punch.rs

@@ -133,8 +133,13 @@ fn main() {
     // Now we use the socket, our private rendezvous info and their public rendezvous info to
     // complete the connection.
     let punched_socket = match PunchedUdpSocket::punch_hole(socket, our_priv_info, their_pub_info) {
-        Ok(punched_socket) => punched_socket,
-        Err(e) => {
+        WOk(punched_socket, warnings) => {
+            for warning in warnings {
+                println!("Warning when punching hole: {}", warning);
+            }
+            punched_socket
+        },
+        WErr(e) => {
             println!("IO error punching udp socket: {}", e);
             println!("Exiting.");
             return;

src/lib.rs

@@ -69,7 +69,7 @@ pub use rendezvous_info::{PrivRendezvousInfo, PubRendezvousInfo,
                          gen_rendezvous_info};
 pub use mapped_udp_socket::{MappedUdpSocket, MappedUdpSocketMapError,
                             MappedUdpSocketMapWarning, MappedUdpSocketNewError};
-pub use punched_udp_socket::PunchedUdpSocket;
+pub use punched_udp_socket::{PunchedUdpSocket, filter_udp_hole_punch_packet};
 pub use mapped_tcp_socket::{MappedTcpSocket, tcp_punch_hole};
 pub use simple_udp_hole_punch_server::{SimpleUdpHolePunchServer, SimpleUdpHolePunchServerNewError};
 

src/punched_udp_socket.rs

@@ -281,6 +281,21 @@ impl PunchedUdpSocket {
     }
 }
 
+/// Returns `None` if `data` looks like a hole punching message. Otherwise returns the data it was
+/// given.
+///
+/// Punching a hole with a udp socket involves packets being sent and received on the socket. After
+/// hole punching succeeds it's possible that more hole punching packets sent by the remote peer
+/// may yet arrive on the socket. This function can be used to filter out those packets.
+pub fn filter_udp_hole_punch_packet(data: &[u8]) -> Option<&[u8]> {
+    match ::cbor::Decoder::from_reader(data)
+                 .decode::<HolePunch>().next()
+    {
+        Some(Ok(_)) => None,
+        _ => Some(data),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::mpsc;
@@ -290,9 +305,10 @@ mod tests {
 
     use mapping_context::MappingContext;
     use mapped_udp_socket::MappedUdpSocket;
-    use punched_udp_socket::PunchedUdpSocket;
+    use punched_udp_socket::{PunchedUdpSocket, filter_udp_hole_punch_packet};
     use rendezvous_info::gen_rendezvous_info;
 
+    #[test]
     fn two_peers_punch_hole_over_loopback() {
         let mapping_context = unwrap_result!(MappingContext::new().result_discard());
         let mapped_socket_0 = unwrap_result!(MappedUdpSocket::new(&mapping_context).result_discard());
@@ -310,36 +326,52 @@ mod tests {
             let res = PunchedUdpSocket::punch_hole(socket_0,
                                                    priv_info_0,
                                                    pub_info_1);
-            tx_0.send(res);
+            unwrap_result!(tx_0.send(res));
         });
         let jh_1 = thread!("two_peers_hole_punch_over_loopback punch socket 1", move || {
             let res = PunchedUdpSocket::punch_hole(socket_1,
                                                    priv_info_1,
                                                    pub_info_0);
-            tx_1.send(res);
+            unwrap_result!(tx_1.send(res));
         });
 
         thread::sleep(Duration::from_millis(500));
-        let punched_socket_0 = unwrap_result!(unwrap_result!(rx_0.try_recv()));
-        let punched_socket_1 = unwrap_result!(unwrap_result!(rx_1.try_recv()));
+        let punched_socket_0 = unwrap_result!(unwrap_result!(rx_0.try_recv()).result_discard());
+        let punched_socket_1 = unwrap_result!(unwrap_result!(rx_1.try_recv()).result_discard());
 
         const DATA_LEN: usize = 8;
         let data_send: [u8; DATA_LEN] = rand::random();
         let mut data_recv;
 
         // Send data from 0 to 1
-        data_recv = [0u8; DATA_LEN];
-        punched_socket_0.socket.send_to(&data_send[..], &*punched_socket_0.peer_addr);
-        let (n, _) = unwrap_result!(punched_socket_1.socket.recv_from(&mut data_recv[..]));
+        data_recv = [0u8; 1024];
+        let n = unwrap_result!(punched_socket_0.socket.send_to(&data_send[..], &*punched_socket_0.peer_addr));
         assert_eq!(n, DATA_LEN);
-        assert_eq!(data_send, data_recv);
+        loop {
+            let (n, _) = unwrap_result!(punched_socket_1.socket.recv_from(&mut data_recv[..]));
+            match filter_udp_hole_punch_packet(&data_recv[..n]) {
+                Some(d) => {
+                    assert_eq!(data_send, d);
+                    break;
+                },
+                None => continue,
+            }
+        }
 
         // Send data from 1 to 0
-        data_recv = [0u8; DATA_LEN];
-        punched_socket_1.socket.send_to(&data_send[..], &*punched_socket_1.peer_addr);
-        let (n, _) = unwrap_result!(punched_socket_0.socket.recv_from(&mut data_recv[..]));
+        data_recv = [0u8; 1024];
+        let n = unwrap_result!(punched_socket_1.socket.send_to(&data_send[..], &*punched_socket_1.peer_addr));
         assert_eq!(n, DATA_LEN);
-        assert_eq!(data_send, data_recv);
+        loop {
+            let (n, _) = unwrap_result!(punched_socket_0.socket.recv_from(&mut data_recv[..]));
+            match filter_udp_hole_punch_packet(&data_recv[..n]) {
+                Some(d) => {
+                    assert_eq!(data_send, d);
+                    break;
+                },
+                None => continue,
+            }
+        }
 
         unwrap_result!(jh_0.join());
         unwrap_result!(jh_1.join());