diff --git a/src/main.rs b/src/main.rs index 8e992ee..55e09aa 100644 --- a/src/main.rs +++ b/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 { 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 { - wtr.serialize(Record { - time: SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(), - duration: perform_handshake()?, - })?; + while i < SAMPLE_SIZE { + match perform_handshake() { + Ok(duration) => { + wtr.serialize(Record { + time: SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(), + duration, + })?; + i += 1; + } + _ => continue, + } } wtr.flush()?;