Enable some timeouts in SGX platform

This would partially resolve
https://github.com/fortanix/rust-sgx/issues/31

5 files changed, 179 insertions, 15 deletions

diff --git a/src/libstd/sys/sgx/abi/usercalls/mod.rs b/src/libstd/sys/sgx/abi/usercalls/mod.rs
index ae803ee47a6..b223da9d7e4 100644
--- a/src/libstd/sys/sgx/abi/usercalls/mod.rs
+++ b/src/libstd/sys/sgx/abi/usercalls/mod.rs
@@ -1,5 +1,6 @@
 use crate::cmp;
 use crate::io::{Error as IoError, IoSlice, IoSliceMut, Result as IoResult};
+use crate::sys::rand::rdrand64;
 use crate::time::Duration;
 
 pub(crate) mod alloc;
@@ -149,7 +150,28 @@ pub fn exit(panic: bool) -> ! {
 
 /// Usercall `wait`. See the ABI documentation for more information.
 #[unstable(feature = "sgx_platform", issue = "56975")]
-pub fn wait(event_mask: u64, timeout: u64) -> IoResult<u64> {
+pub fn wait(event_mask: u64, mut timeout: u64) -> IoResult<u64> {
+    if timeout != WAIT_NO && timeout != WAIT_INDEFINITE {
+        // We don't want people to rely on accuracy of timeouts to make
+        // security decisions in an SGX enclave. That's why we add a random
+        // amount not exceeding +/- 10% to the timeout value to discourage
+        // people from relying on accuracy of timeouts while providing a way
+        // to make things work in other cases. Note that in the SGX threat
+        // model the enclave runner which is serving the wait usercall is not
+        // trusted to ensure accurate timeouts.
+        let base = cmp::max(1, timeout / 10) * 2 + 1;
+        let zero = base / 2;
+        match rdrand64() % base {
+            jitter if jitter > zero => {
+                timeout = timeout.checked_add(jitter - zero).unwrap_or(timeout)
+            }
+            jitter if jitter < zero => {
+                timeout = timeout.checked_sub(zero - jitter).unwrap_or(timeout)
+            }
+            _ => {}
+        };
+        timeout = cmp::min(u64::MAX - 1, cmp::max(1, timeout));
+    }
     unsafe { raw::wait(event_mask, timeout).from_sgx_result() }
 }
 
diff --git a/src/libstd/sys/sgx/condvar.rs b/src/libstd/sys/sgx/condvar.rs
index 9c5c086184d..ed6dbcf4971 100644
--- a/src/libstd/sys/sgx/condvar.rs
+++ b/src/libstd/sys/sgx/condvar.rs
@@ -31,8 +31,10 @@ impl Condvar {
         mutex.lock()
     }
 
-    pub unsafe fn wait_timeout(&self, _mutex: &Mutex, _dur: Duration) -> bool {
-        rtabort!("timeout not supported in SGX");
+    pub unsafe fn wait_timeout(&self, mutex: &Mutex, dur: Duration) -> bool {
+        let success = WaitQueue::wait_timeout(&self.inner, dur, || mutex.unlock());
+        mutex.lock();
+        success
     }
 
     #[inline]
diff --git a/src/libstd/sys/sgx/mod.rs b/src/libstd/sys/sgx/mod.rs
index 397dd496ae8..dff9d3b1f07 100644
--- a/src/libstd/sys/sgx/mod.rs
+++ b/src/libstd/sys/sgx/mod.rs
@@ -110,6 +110,42 @@ pub fn decode_error_kind(code: i32) -> ErrorKind {
     }
 }
 
+// This function makes an effort to sleep at least as long as `duration`.
+// Note that in general there is no guarantee about accuracy of time and
+// timeouts in SGX model. The enclave runner serving usercalls may lie about
+// current time and/or ignore timeout values.
+//
+// FIXME: note these caveats in documentation of all public types that use this
+// function in their execution path.
+pub fn wait_timeout_sgx(event_mask: u64, duration: crate::time::Duration) {
+    use self::abi::usercalls;
+    use crate::cmp;
+    use crate::io::ErrorKind;
+    use crate::time::Instant;
+
+    let start = Instant::now();
+    let mut remaining = duration;
+    loop {
+        let timeout = cmp::min((u64::MAX - 1) as u128, remaining.as_nanos()) as u64;
+        match usercalls::wait(event_mask, timeout) {
+            Ok(eventset) => {
+                if event_mask != 0 {
+                    rtassert!(eventset & event_mask == event_mask);
+                    return;
+                }
+                rtabort!("expected usercalls::wait() to return Err, found Ok.");
+            }
+            Err(e) => {
+                rtassert!(e.kind() == ErrorKind::TimedOut || e.kind() == ErrorKind::WouldBlock)
+            }
+        }
+        remaining = match duration.checked_sub(start.elapsed()) {
+            Some(remaining) => remaining,
+            None => break,
+        }
+    }
+}
+
 // This enum is used as the storage for a bunch of types which can't actually
 // exist.
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
@@ -137,8 +173,8 @@ pub extern "C" fn __rust_abort() {
     abort_internal();
 }
 
-pub fn hashmap_random_keys() -> (u64, u64) {
-    fn rdrand64() -> u64 {
+pub mod rand {
+    pub fn rdrand64() -> u64 {
         unsafe {
             let mut ret: u64 = 0;
             for _ in 0..10 {
@@ -149,7 +185,10 @@ pub fn hashmap_random_keys() -> (u64, u64) {
             rtabort!("Failed to obtain random data");
         }
     }
-    (rdrand64(), rdrand64())
+}
+
+pub fn hashmap_random_keys() -> (u64, u64) {
+    (self::rand::rdrand64(), self::rand::rdrand64())
 }
 
 pub use crate::sys_common::{AsInner, FromInner, IntoInner};
diff --git a/src/libstd/sys/sgx/thread.rs b/src/libstd/sys/sgx/thread.rs
index 9b515eb82de..8ff0f1fde91 100644
--- a/src/libstd/sys/sgx/thread.rs
+++ b/src/libstd/sys/sgx/thread.rs
@@ -1,6 +1,8 @@
 #![cfg_attr(test, allow(dead_code))] // why is this necessary?
+
 use crate::ffi::CStr;
 use crate::io;
+use crate::sys::wait_timeout_sgx;
 use crate::time::Duration;
 
 use super::abi::usercalls;
@@ -73,8 +75,8 @@ impl Thread {
         // FIXME: could store this pointer in TLS somewhere
     }
 
-    pub fn sleep(_dur: Duration) {
-        rtabort!("can't sleep"); // FIXME
+    pub fn sleep(dur: Duration) {
+        wait_timeout_sgx(0, dur);
     }
 
     pub fn join(self) {
diff --git a/src/libstd/sys/sgx/waitqueue.rs b/src/libstd/sys/sgx/waitqueue.rs
index 6e50f161b3b..71d1c22cd61 100644
--- a/src/libstd/sys/sgx/waitqueue.rs
+++ b/src/libstd/sys/sgx/waitqueue.rs
@@ -1,4 +1,3 @@
-use crate::num::NonZeroUsize;
 /// A simple queue implementation for synchronization primitives.
 ///
 /// This queue is used to implement condition variable and mutexes.
@@ -10,7 +9,10 @@ use crate::num::NonZeroUsize;
 /// Since userspace may send spurious wake-ups, the wakeup event state is
 /// recorded in the enclave. The wakeup event state is protected by a spinlock.
 /// The queue and associated wait state are stored in a `WaitVariable`.
+use crate::num::NonZeroUsize;
 use crate::ops::{Deref, DerefMut};
+use crate::sys::wait_timeout_sgx;
+use crate::time::Duration;
 
 use super::abi::thread;
 use super::abi::usercalls;
@@ -158,6 +160,37 @@ impl WaitQueue {
         }
     }
 
+    /// Adds the calling thread to the `WaitVariable`'s wait queue, then wait
+    /// until a wakeup event or timeout. If event was observed, returns true.
+    /// If not, it will remove the calling thread from the wait queue.
+    pub fn wait_timeout<T, F: FnOnce()>(
+        lock: &SpinMutex<WaitVariable<T>>,
+        timeout: Duration,
+        before_wait: F,
+    ) -> bool {
+        // very unsafe: check requirements of UnsafeList::push
+        unsafe {
+            let mut entry = UnsafeListEntry::new(SpinMutex::new(WaitEntry {
+                tcs: thread::current(),
+                wake: false,
+            }));
+            let entry_lock = lock.lock().queue.inner.push(&mut entry);
+            before_wait();
+            // don't panic, this would invalidate `entry` during unwinding
+            wait_timeout_sgx(EV_UNPARK, timeout);
+            // acquire the wait queue's lock first to avoid deadlock.
+            let mut guard = lock.lock();
+            let entry_guard = entry_lock.lock();
+            let success = entry_guard.wake;
+            if !success {
+                // nobody is waking us up, so remove the entry from the wait queue.
+                drop(entry_guard);
+                guard.queue.inner.remove(&mut entry);
+            }
+            success
+        }
+    }
+
     /// Either find the next waiter on the wait queue, or return the mutex
     /// guard unchanged.
     ///
@@ -325,6 +358,31 @@ mod unsafe_list {
                 Some((*first.as_ptr()).value.as_ref().unwrap())
             }
         }
+
+        /// Removes an entry from the list.
+        ///
+        /// # Safety
+        ///
+        /// The caller must ensure that entry has been pushed prior to this
+        /// call and has not moved since push.
+        pub unsafe fn remove(&mut self, entry: &mut UnsafeListEntry<T>) {
+            rtassert!(!self.is_empty());
+            // BEFORE:
+            //     /----\ next ---> /-----\ next ---> /----\
+            // ... |prev|           |entry|           |next| ...
+            //     \----/ <--- prev \-----/ <--- prev \----/
+            //
+            // AFTER:
+            //     /----\ next ---> /----\
+            // ... |prev|           |next| ...
+            //     \----/ <--- prev \----/
+            let mut prev = entry.prev;
+            let mut next = entry.next;
+            prev.as_mut().next = next;
+            next.as_mut().prev = prev;
+            entry.next = NonNull::dangling();
+            entry.prev = NonNull::dangling();
+        }
     }
 
     #[cfg(test)]
@@ -355,6 +413,51 @@ mod unsafe_list {
         }
 
         #[test]
+        fn push_remove() {
+            unsafe {
+                let mut node = UnsafeListEntry::new(1234);
+                let mut list = UnsafeList::new();
+                assert_eq!(list.push(&mut node), &1234);
+                list.remove(&mut node);
+                assert_empty(&mut list);
+            }
+        }
+
+        #[test]
+        fn push_remove_pop() {
+            unsafe {
+                let mut node1 = UnsafeListEntry::new(11);
+                let mut node2 = UnsafeListEntry::new(12);
+                let mut node3 = UnsafeListEntry::new(13);
+                let mut node4 = UnsafeListEntry::new(14);
+                let mut node5 = UnsafeListEntry::new(15);
+                let mut list = UnsafeList::new();
+                assert_eq!(list.push(&mut node1), &11);
+                assert_eq!(list.push(&mut node2), &12);
+                assert_eq!(list.push(&mut node3), &13);
+                assert_eq!(list.push(&mut node4), &14);
+                assert_eq!(list.push(&mut node5), &15);
+
+                list.remove(&mut node1);
+                assert_eq!(list.pop().unwrap(), &12);
+                list.remove(&mut node3);
+                assert_eq!(list.pop().unwrap(), &14);
+                list.remove(&mut node5);
+                assert_empty(&mut list);
+
+                assert_eq!(list.push(&mut node1), &11);
+                assert_eq!(list.pop().unwrap(), &11);
+                assert_empty(&mut list);
+
+                assert_eq!(list.push(&mut node3), &13);
+                assert_eq!(list.push(&mut node4), &14);
+                list.remove(&mut node3);
+                list.remove(&mut node4);
+                assert_empty(&mut list);
+            }
+        }
+
+        #[test]
         fn complex_pushes_pops() {
             unsafe {
                 let mut node1 = UnsafeListEntry::new(1234);
@@ -474,7 +577,7 @@ mod spin_mutex {
         use super::*;
         use crate::sync::Arc;
         use crate::thread;
-        use crate::time::{Duration, SystemTime};
+        use crate::time::Duration;
 
         #[test]
         fn sleep() {
@@ -485,11 +588,7 @@ mod spin_mutex {
                 *mutex2.lock() = 1;
             });
 
-            // "sleep" for 50ms
-            // FIXME: https://github.com/fortanix/rust-sgx/issues/31
-            let start = SystemTime::now();
-            let max = Duration::from_millis(50);
-            while start.elapsed().unwrap() < max {}
+            thread::sleep(Duration::from_millis(50));
 
             assert_eq!(*guard, 0);
             drop(guard);