Enable enclave exit when a TCS is blocked on I/O

This fixes issue #109

enclave-runner/src/usercalls/interface.rs

@@ -192,7 +192,7 @@ impl<'future, 'ioinput: 'future, 'tcs: 'ioinput> Usercalls<'future> for Handler<
         panic: bool,
     ) -> std::pin::Pin<Box<dyn Future<Output = (Self, EnclaveAbort<bool>)> + 'future>> {
         async move {
-            let ret = self.0.exit(panic);
+            let ret = self.0.exit(panic).await;
             return (self, ret);
         }
             .boxed_local()

enclave-runner/src/usercalls/mod.rs

@@ -540,6 +540,8 @@ pub(crate) struct EnclaveState {
     fds: Lock<FnvHashMap<Fd, Arc<AsyncFileDesc>>>,
     last_fd: AtomicUsize,
     exiting: AtomicBool,
+    exit_tx: Lock<tokio::sync::watch::Sender<bool>>,
+    exit_rx: tokio::sync::watch::Receiver<bool>,
     usercall_ext: Box<dyn UsercallExtension>,
     threads_queue: crossbeam::queue::SegQueue<StoppedTcs>,
 }
@@ -560,11 +562,11 @@ impl Work {
         let usercall_send_data = match self.entry {
             CoEntry::Initial(erased_tcs, p1, p2, p3, p4, p5) => {
                 let coresult = tcs::coenter(erased_tcs, p1, p2, p3, p4, p5, Some(&buf));
-                ((coresult, self.tcs, buf))
+                (coresult, self.tcs, buf)
             }
             CoEntry::Resume(usercall, coresult) => {
                 let coresult = usercall.coreturn(coresult, Some(&buf));
-                ((coresult, self.tcs, buf))
+                (coresult, self.tcs, buf)
             }
         };
         // if there is an error do nothing, as it means that the main thread has exited
@@ -623,12 +625,16 @@ impl EnclaveState {
             threads_queue.push(Self::event_queue_add_tcs(&mut event_queues, thread));
         }
 
+        let (exit_tx, exit_rx) = tokio::sync::watch::channel(false);
+
         Arc::new(EnclaveState {
             kind,
             event_queues,
             fds: Lock::new(fds),
             last_fd,
             exiting: AtomicBool::new(false),
+            exit_tx: Lock::new(exit_tx),
+            exit_rx,
             usercall_ext,
             threads_queue,
         })
@@ -696,9 +702,19 @@ impl EnclaveState {
                                 work_sender: &work_sender,
                             };
                             let handler = Handler(&mut input);
-                            let (_handler, result) = {
+                            let result = {
                                 let (p1, p2, p3, p4, p5) = usercall.parameters();
-                                dispatch(handler, p1, p2, p3, p4, p5).await
+                                let exit_waiter = async {
+                                    let mut stream = enclave_clone.exit_rx.clone();
+                                    while let Ok((Some(false), s)) = stream.into_future().compat().await {
+                                        stream = s;
+                                    }
+                                    Err(EnclaveAbort::Secondary)
+                                }.boxed_local();
+                                match futures::future::select(exit_waiter, dispatch(handler, p1, p2, p3, p4, p5).boxed_local()).await {
+                                    Either::Left((x, _)) => x,
+                                    Either::Right((y, _)) => y.1,
+                                }
                             };
                             let ret = match result {
                                 Ok(ret) => {
@@ -868,7 +884,7 @@ impl EnclaveState {
         };
 
         tokio::runtime::current_thread::block_on_all(async move {
-            enclave.abort_all_threads();
+            enclave.abort_all_threads().await;
             //clear the threads_queue
             while enclave.threads_queue.pop().is_ok() {}
 
@@ -949,12 +965,15 @@ impl EnclaveState {
         }
     }
 
-    fn abort_all_threads(&self) {
+    async fn abort_all_threads(&self) {
         self.exiting.store(true, Ordering::SeqCst);
         // wake other threads
         for queue in self.event_queues.values() {
             let _ = queue.unbounded_send(EV_ABORT as _);
         }
+        poll_lock_await_unwrap!(self.exit_tx)
+            .broadcast(true)
+            .expect("failed to broadcast exit");
     }
 }
 
@@ -1126,7 +1145,7 @@ impl<'tcs> IOHandlerInput<'tcs> {
     async fn write(&self, fd: Fd, buf: &[u8]) -> IoResult<usize> {
         let vec = buf.to_vec();
         let file_desc = self.lookup_fd(fd).await?;
-        return file_desc.as_stream()?.async_write(vec).await;
+        file_desc.as_stream()?.async_write(vec).await
     }
 
     #[inline(always)]
@@ -1287,8 +1306,8 @@ impl<'tcs> IOHandlerInput<'tcs> {
     }
 
     #[inline(always)]
-    fn exit(&mut self, panic: bool) -> EnclaveAbort<bool> {
-        self.enclave.abort_all_threads();
+    async fn exit(&mut self, panic: bool) -> EnclaveAbort<bool> {
+        self.enclave.abort_all_threads().await;
         EnclaveAbort::Exit { panic }
     }