Add timeout

src/main.rs

@@ -7,6 +7,7 @@ use rand::SeedableRng;
 use snow::Builder;
 use std::io;
 use std::net::UdpSocket;
+use std::time::Duration;
 use x25519_dalek::PublicKey;
 use x25519_dalek::StaticSecret;
 use zerocopy::byteorder::U32;
@@ -101,6 +102,7 @@ fn perform_handshake() -> Result<u128> {
     initiation[116..132].copy_from_slice(&hasher.finalize().into_bytes());
 
     let socket = UdpSocket::bind("0.0.0.0:0")?;
+    socket.set_read_timeout(Some(Duration::from_millis(100)))?;
 
     let start = Instant::now();
 
@@ -127,12 +129,19 @@ struct Record {
 
 fn main() -> Result<()> {
     let mut wtr = csv::Writer::from_writer(io::stdout());
+    let mut i = 0;
 
-    for _ in 0..SAMPLE_SIZE {
+    while i < SAMPLE_SIZE {
+        match perform_handshake() {
+            Ok(duration) => {
                 wtr.serialize(Record {
                     time: SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(),
-            duration: perform_handshake()?,
+                    duration,
                 })?;
+                i += 1;
+            }
+            _ => continue,
+        }
     }
 
     wtr.flush()?;