Allow arbitrary timeouts for wait usercall

This should partially address issue #31

enclave-runner/src/usercalls/mod.rs

@@ -21,7 +21,7 @@ use std::pin::Pin;
 
 use std::sync::Arc;
 use std::thread;
-use std::time;
+use std::time::{self, Duration, Instant};
 use std::task::{Poll, Context, Waker};
 
 use failure;
@@ -1317,12 +1317,6 @@ impl<'tcs> IOHandlerInput<'tcs> {
 
     #[inline(always)]
     async fn wait(&mut self, event_mask: u64, timeout: u64) -> IoResult<u64> {
-        let wait = match timeout {
-            WAIT_NO => false,
-            WAIT_INDEFINITE => true,
-            _ => return Err(IoErrorKind::InvalidInput.into()),
-        };
-
         let event_mask = Self::check_event_set(event_mask)?;
 
         let mut ret = None;
@@ -1340,17 +1334,29 @@ impl<'tcs> IOHandlerInput<'tcs> {
         }
 
         if ret.is_none() {
+            let start = Instant::now();
             loop {
-                let ev = if wait {
-                    Ok(self.tcs.event_queue.next().await.unwrap())
-                } else {
-                    match self.tcs.event_queue.try_next() {
+                let ev = match timeout {
+                    WAIT_INDEFINITE => self.tcs.event_queue.next().await.ok_or(()),
+                    WAIT_NO => match self.tcs.event_queue.try_next() {
                         Ok(Some(ev)) => Ok(ev),
-                        Err(_) => break,
                         Ok(None) => Err(()),
+                        Err(_) => break,
+                    },
+                    timeout => {
+                        let remaining = match Duration::from_nanos(timeout).checked_sub(start.elapsed()) {
+                            None => break,
+                            Some(ref duration) if duration.as_nanos() == 0 => break,
+                            Some(duration) => duration,
+                        };
+                        match tokio::time::timeout(remaining, self.tcs.event_queue.next()).await {
+                            Ok(Some(ev)) => Ok(ev),
+                            Ok(None) => Err(()),
+                            Err(_) => break, // timed out
+                        }
                     }
-                }
-                .expect("TCS event queue disconnected unexpectedly");
+                }.expect("TCS event queue disconnected unexpectedly");
+
                 if (ev & (EV_ABORT as u8)) != 0 {
                     // dispatch will make sure this is not returned to enclave
                     return Err(IoErrorKind::Other.into());
@@ -1369,7 +1375,7 @@ impl<'tcs> IOHandlerInput<'tcs> {
         if let Some(ret) = ret {
             Ok(ret.into())
         } else {
-            Err(IoErrorKind::WouldBlock.into())
+            Err(if timeout == WAIT_NO { IoErrorKind::WouldBlock } else { IoErrorKind::TimedOut }.into())
         }
     }
 

fortanix-sgx-abi/src/lib.rs

@@ -503,14 +503,18 @@ impl Usercalls {
 
     /// Wait for an event to occur, or check if an event is currently pending.
     ///
-    /// `timeout` must be [`WAIT_NO`] or [`WAIT_INDEFINITE`]. If it is another
-    /// value, userspace will return an error.
+    /// `timeout` must be [`WAIT_NO`] or [`WAIT_INDEFINITE`] or a positive
+    /// value smaller than u64::MAX specifying number of nanoseconds to wait.
     ///
     /// If `timeout` is [`WAIT_INDEFINITE`], this call will block and return
     /// once a matching event is queued on this TCS. If `timeout` is
     /// [`WAIT_NO`], this call will return immediately, and the return value
     /// will indicate if an event was pending. If it was, it has been dequeued.
-    /// If not, the [`WouldBlock`] error value will be returned.
+    /// If not, the [`WouldBlock`] error value will be returned. If `timeout`
+    /// is a value other than [`WAIT_NO`] and [`WAIT_INDEFINITE`], this call
+    /// will block until either a matching event is queued in which case the
+    /// return value will indicate the event, or the timeout is reached in
+    /// which case the [`TimedOut`] error value will be returned.
     ///
     /// A matching event is one whose bits are equal to or a subset of
     /// `event_mask`. If `event_mask` is `0`, this call will never return due
@@ -534,6 +538,7 @@ impl Usercalls {
     /// [`WAIT_INDEFINITE`]: constant.WAIT_INDEFINITE.html
     /// [`EV_RETURNQ_NOT_EMPTY`]: constant.EV_RETURNQ_NOT_EMPTY.html
     /// [`WouldBlock`]: enum.Error.html#variant.WouldBlock
+    /// [`TimedOut`]: enum.Error.html#variant.TimedOut
     pub fn wait(event_mask: u64, timeout: u64) -> (Result, u64) { unimplemented!() }
 
     /// Send an event to one or all TCSes.