3 files changed, 14 insertions, 6 deletions
diff --git a/src/libstd/sys/sgx/mod.rs b/src/libstd/sys/sgx/mod.rs
index dff9d3b1f07..b72e4fb06f3 100644
--- a/src/libstd/sys/sgx/mod.rs
+++ b/src/libstd/sys/sgx/mod.rs
@@ -115,9 +115,15 @@ pub fn decode_error_kind(code: i32) -> ErrorKind {
 // timeouts in SGX model. The enclave runner serving usercalls may lie about
 // current time and/or ignore timeout values.
 //
+// Once the event is observed, `stop` will be used to determine whether or not
+// we should continue to wait.
+//
 // 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) {
+pub fn wait_timeout_sgx<F>(event_mask: u64, duration: crate::time::Duration, stop: F)
+where
+    F: Fn() -> bool,
+{
     use self::abi::usercalls;
     use crate::cmp;
     use crate::io::ErrorKind;
@@ -129,11 +135,13 @@ pub fn wait_timeout_sgx(event_mask: u64, duration: crate::time::Duration) {
         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);
+                if event_mask == 0 {
+                    rtabort!("expected usercalls::wait() to return Err, found Ok.");
+                }
+                rtassert!(eventset & event_mask == event_mask);
+                if stop() {
                     return;
                 }
-                rtabort!("expected usercalls::wait() to return Err, found Ok.");
             }
             Err(e) => {
                 rtassert!(e.kind() == ErrorKind::TimedOut || e.kind() == ErrorKind::WouldBlock)
diff --git a/src/libstd/sys/sgx/thread.rs b/src/libstd/sys/sgx/thread.rs
index 8ff0f1fde91..5636a6f7eab 100644
--- a/src/libstd/sys/sgx/thread.rs
+++ b/src/libstd/sys/sgx/thread.rs
@@ -76,7 +76,7 @@ impl Thread {
     }
 
     pub fn sleep(dur: Duration) {
-        wait_timeout_sgx(0, dur);
+        wait_timeout_sgx(0, dur, || true);
     }
 
     pub fn join(self) {
diff --git a/src/libstd/sys/sgx/waitqueue.rs b/src/libstd/sys/sgx/waitqueue.rs
index 71d1c22cd61..36b3f5bcc41 100644
--- a/src/libstd/sys/sgx/waitqueue.rs
+++ b/src/libstd/sys/sgx/waitqueue.rs
@@ -177,7 +177,7 @@ impl WaitQueue {
             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);
+            wait_timeout_sgx(EV_UNPARK, timeout, || entry_lock.lock().wake);
             // acquire the wait queue's lock first to avoid deadlock.
             let mut guard = lock.lock();
             let entry_guard = entry_lock.lock();