Read server socket until `WouldBlock`

Using the `-n` feature of @Aaron1011's excellent client I was able to
find and fix a nagging bug.

Post conversion from the hand-rolled event loop to a `mio` based one,
the server would intermittently fail to respond to client requests.

The server's UDP socket is edge-triggered and the `mio` docs clearly
state that [operations must be performed until `WouldBlock`](https://carllerche.github.io/mio/mio/struct.Poll.html#edge-triggered-and-level-triggered)
otherwise `poll` might not return subsequent events.

Fix is to loop inside the MESSAGE token match until `recv_from`
returns `WouldBlock`.

The resulting code is excessively nested; a cleanup will follow.

1 files changed, 25 insertions, 13 deletions

diff --git a/src/bin/server.rs b/src/bin/server.rs
index 9a0897e..31e9392 100644
--- a/src/bin/server.rs
+++ b/src/bin/server.rs
@@ -57,6 +57,7 @@ use std::env;
 use std::process;
 use std::fs::File;
 use std::io::Read;
+use std::io;
 use std::time::Duration;
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -282,19 +283,30 @@ fn polling_loop(addr: &SocketAddr, mut ephemeral_key: &mut Signer, cert_bytes: &
         for event in events.iter() {
             match event.token() {
                 MESSAGE => {
-                    let (num_bytes, src_addr) = socket.recv_from(&mut buf).expect("recv_from failed");
-
-                    if let Ok(nonce) = nonce_from_request(&buf, num_bytes) {
-                        let resp = make_response(&mut ephemeral_key, cert_bytes, nonce);
-                        let resp_bytes = resp.encode().unwrap();
-
-                        let bytes_sent = socket.send_to(&resp_bytes, &src_addr).expect("send_to failed");
-
-                        num_responses += 1;
-                        info!("Responded {} bytes to {} for '{}..' (resp #{})", bytes_sent, src_addr, hex::encode(&nonce[0..4]), num_responses);
-                    } else {
-                        num_bad_requests += 1;
-                        info!("Invalid request ({} bytes) from {} (resp #{})", num_bytes, src_addr, num_responses);
+                    loop {
+                        match socket.recv_from(&mut buf) {
+                            Ok((num_bytes, src_addr)) => {
+                                if let Ok(nonce) = nonce_from_request(&buf, num_bytes) {
+                                    let resp = make_response(&mut ephemeral_key, cert_bytes, nonce);
+                                    let resp_bytes = resp.encode().unwrap();
+
+                                    let bytes_sent = socket.send_to(&resp_bytes, &src_addr).expect("send_to failed");
+
+                                    num_responses += 1;
+                                    info!("Responded {} bytes to {} for '{}..' (resp #{})", bytes_sent, src_addr, hex::encode(&nonce[0..4]), num_responses);
+                                } else {
+                                    num_bad_requests += 1;
+                                    info!("Invalid request ({} bytes) from {} (resp #{})", num_bytes, src_addr, num_responses);
+                                }
+                            }
+                            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
+                                break
+                            }
+                            Err(ref e) => {
+                                error!("Error {:?}: {:?}", e.kind(), e);
+                                break
+                            }
+                        }
                     }
                 }