Cancel queue

Author: mzohreva

enclave-runner/src/usercalls/abi.rs

@@ -19,7 +19,7 @@ use futures::future::Future;
 
 type Register = u64;
 
-trait RegisterArgument {
+pub(super) trait RegisterArgument {
     fn from_register(_: Register) -> Self;
     fn into_register(self) -> Register;
 }
@@ -29,7 +29,7 @@ type EnclaveAbort = super::EnclaveAbort<bool>;
 pub(crate) type UsercallResult<T> = ::std::result::Result<T, EnclaveAbort>;
 pub(crate) type DispatchResult = UsercallResult<(Register, Register)>;
 
-trait ReturnValue {
+pub(super) trait ReturnValue {
     fn into_registers(self) -> DispatchResult;
 }
 
@@ -38,7 +38,7 @@ macro_rules! define_usercalls {
     ($(fn $f:ident($($n:ident: $t:ty),*) $(-> $r:tt)*; )*) => {
         #[repr(C)]
         #[allow(non_camel_case_types)]
-        enum UsercallList {
+        pub(super) enum UsercallList {
             __enclave_usercalls_invalid,
             $($f,)*
         }

enclave-runner/src/usercalls/interface.rs

@@ -252,12 +252,13 @@ impl<'future, 'ioinput: 'future, 'tcs: 'ioinput> Usercalls<'future> for Handler<
         self,
         usercall_queue: *mut FifoDescriptor<Usercall>,
         return_queue: *mut FifoDescriptor<Return>,
+        cancel_queue: *mut FifoDescriptor<Cancel>,
     ) -> std::pin::Pin<Box<dyn Future<Output = (Self, UsercallResult<Result>)> + 'future>> {
         async move {
             unsafe {
                 let ret = match (usercall_queue.as_mut(), return_queue.as_mut()) {
                     (Some(usercall_queue), Some(return_queue)) => {
-                        self.0.async_queues(usercall_queue, return_queue).await.map(Ok)
+                        self.0.async_queues(usercall_queue, return_queue, cancel_queue.as_mut()).await.map(Ok)
                     },
                     _ => {
                         Ok(Err(IoErrorKind::InvalidInput.into()))
@@ -321,13 +322,13 @@ fn result_from_io_error(err: IoError) -> Result {
     ret as _
 }
 
-trait ToSgxResult {
+pub(super) trait ToSgxResult {
     type Return;
 
     fn to_sgx_result(self) -> Self::Return;
 }
 
-trait SgxReturn {
+pub(super) trait SgxReturn {
     fn on_error() -> Self;
 }
 

enclave-runner/src/usercalls/mod.rs

@@ -24,7 +24,7 @@ use ipc_queue::{self, QueueEvent};
 use sgxs::loader::Tcs as SgxsTcs;
 use std::alloc::{GlobalAlloc, Layout, System};
 use std::cell::RefCell;
-use std::collections::VecDeque;
+use std::collections::{HashMap, VecDeque};
 use std::io::{self, ErrorKind as IoErrorKind, Read, Result as IoResult};
 use std::pin::Pin;
 use std::result::Result as StdResult;
@@ -46,16 +46,17 @@ lazy_static! {
 pub(crate) mod abi;
 mod interface;
 
-use self::abi::dispatch;
+use self::abi::{dispatch, UsercallList};
 use self::interface::{Handler, OutputBuffer};
 
-const EV_ABORT: u64 = 0b0000_0000_0000_1000;
+const EV_ABORT: u64 = 0b0000_0000_0001_0000;
 
 // Experiments show that tha actual size of these queues is less important than
 // the ratio between them. It appears that a much larger return queue performs
 // much better when multiple enclave threads send usercalls.
 const USERCALL_QUEUE_SIZE: usize = 16;
 const RETURN_QUEUE_SIZE: usize = 1024;
+const CANCEL_QUEUE_SIZE: usize = USERCALL_QUEUE_SIZE * 2;
 
 enum UsercallSendData {
     Sync(ThreadResult<ErasedTcs>, RunningTcs, RefCell<[u8; 1024]>),
@@ -64,7 +65,7 @@ enum UsercallSendData {
 
 enum UsercallHandleData {
     Sync(tcs::Usercall<ErasedTcs>, RunningTcs, RefCell<[u8; 1024]>),
-    Async(Usercall),
+    Async(Usercall, Option<async_mpsc::UnboundedSender<UsercallEvent>>),
 }
 
 type EnclaveResult = StdResult<(u64, u64), EnclaveAbort<Option<EnclavePanic>>>;
@@ -513,11 +514,11 @@ struct PendingEvents {
 }
 
 impl PendingEvents {
-    const MAX_EVENT: usize = 8;
+    const MAX_EVENT: usize = 16;
 
     fn new() -> Self {
         // sanity check to ensure we update MAX_EVENT if new events are added in the future
-        const EV_ALL: u64 = EV_USERCALLQ_NOT_FULL | EV_RETURNQ_NOT_EMPTY | EV_UNPARK;
+        const EV_ALL: u64 = EV_USERCALLQ_NOT_FULL | EV_RETURNQ_NOT_EMPTY | EV_CANCELQ_NOT_FULL | EV_UNPARK;
         assert!(EV_ALL < Self::MAX_EVENT as u64);
         assert!(Self::MAX_EVENT <= 1usize + u8::max_value() as usize);
 
@@ -587,6 +588,7 @@ impl EnclaveKind {
 struct FifoDescriptors {
     usercall_queue: FifoDescriptor<Usercall>,
     return_queue: FifoDescriptor<Return>,
+    cancel_queue: FifoDescriptor<Cancel>,
 }
 
 pub(crate) struct EnclaveState {
@@ -630,6 +632,31 @@ impl Work {
     }
 }
 
+enum UsercallEvent {
+    Started(u64, tokio::sync::oneshot::Sender<()>),
+    Finished(u64),
+    Cancelled(u64, Instant),
+}
+
+fn ignore_cancel_impl(usercall_nr: u64) -> bool {
+    usercall_nr != UsercallList::read as u64 &&
+    usercall_nr != UsercallList::read_alloc as u64 &&
+    usercall_nr != UsercallList::write as u64 &&
+    usercall_nr != UsercallList::accept_stream as u64 &&
+    usercall_nr != UsercallList::connect_stream as u64 &&
+    usercall_nr != UsercallList::wait as u64
+}
+
+trait IgnoreCancel {
+    fn ignore_cancel(&self) -> bool;
+}
+impl IgnoreCancel for Usercall {
+    fn ignore_cancel(&self) -> bool { ignore_cancel_impl(self.args.0) }
+}
+impl IgnoreCancel for Cancel {
+    fn ignore_cancel(&self) -> bool { ignore_cancel_impl(self.usercall_nr) }
+}
+
 impl EnclaveState {
     fn event_queue_add_tcs(
         event_queues: &mut FnvHashMap<TcsAddress, futures::channel::mpsc::UnboundedSender<u8>>,
@@ -702,15 +729,41 @@ impl EnclaveState {
         tx_return_channel: tokio::sync::mpsc::UnboundedSender<(EnclaveResult, EnclaveEntry)>,
         mut handle_data: UsercallHandleData,
     ) {
+        let notifier_rx = match handle_data {
+            UsercallHandleData::Async(ref usercall, Some(ref usercall_event_tx)) => {
+                let (notifier_tx, notifier_rx) = tokio::sync::oneshot::channel();
+                usercall_event_tx.send(UsercallEvent::Started(usercall.id, notifier_tx)).ok()
+                    .expect("failed to send usercall event");
+                Some(notifier_rx)
+            },
+            _ => None,
+        };
         let (parameters, mode, tcs) = match handle_data {
             UsercallHandleData::Sync(ref usercall, ref mut tcs, _) => (usercall.parameters(), tcs.mode, Some(tcs)),
-            UsercallHandleData::Async(ref usercall)                => (usercall.args, EnclaveEntry::ExecutableNonMain, None),
+            UsercallHandleData::Async(ref usercall, _)             => (usercall.args, EnclaveEntry::ExecutableNonMain, None),
         };
         let mut input = IOHandlerInput { enclave: enclave.clone(), tcs, work_sender: &work_sender };
         let handler = Handler(&mut input);
-        let (_handler, result) = {
+        let result = {
+            use self::interface::ToSgxResult;
+            use self::abi::ReturnValue;
+
             let (p1, p2, p3, p4, p5) = parameters;
-            dispatch(handler, p1, p2, p3, p4, p5).await
+            match notifier_rx {
+                None => dispatch(handler, p1, p2, p3, p4, p5).await.1,
+                Some(notifier_rx) => {
+                    let a = dispatch(handler, p1, p2, p3, p4, p5).boxed_local();
+                    let b = notifier_rx;
+                    match futures::future::select(a, b).await {
+                        Either::Left((ret, _)) => ret.1,
+                        Either::Right((Ok(()), _)) => {
+                            let result: IoResult<usize> = Err(IoErrorKind::Interrupted.into());
+                            ReturnValue::into_registers(Ok(result.to_sgx_result()))
+                        },
+                        Either::Right((Err(_), _)) => panic!("notifier channel closed unexpectedly"),
+                    }
+                },
+            }
         };
         let ret = match result {
             Ok(ret) => {
@@ -721,7 +774,11 @@ impl EnclaveState {
                             entry: CoEntry::Resume(usercall, ret),
                         }).expect("Work sender couldn't send data to receiver");
                     }
-                    UsercallHandleData::Async(usercall) => {
+                    UsercallHandleData::Async(usercall, usercall_event_tx) => {
+                        if let Some(usercall_event_tx) = usercall_event_tx {
+                            usercall_event_tx.send(UsercallEvent::Finished(usercall.id)).ok()
+                                .expect("failed to send usercall event");
+                        }
                         let return_queue_tx = enclave.return_queue_tx.lock().await.clone().expect("return_queue_tx not initialized");
                         let ret = Return {
                             id: usercall.id,
@@ -740,7 +797,7 @@ impl EnclaveState {
                         trap_attached_debugger(usercall.tcs_address() as _).await;
                         EnclavePanic::from(debug_buf.into_inner())
                     }
-                    UsercallHandleData::Async(_) => {
+                    UsercallHandleData::Async(_, _) => {
                         // FIXME: find a better panic message
                         EnclavePanic::DebugStr("async exit with a panic".to_owned())
                     }
@@ -815,14 +872,16 @@ impl EnclaveState {
         };
         let enclave_clone = enclave.clone();
         let io_future = async move {
-            let (usercall_queue_synchronizer, return_queue_synchronizer, sync_usercall_tx) = QueueSynchronizer::new(enclave_clone.clone());
+            let (uqs, rqs, cqs, sync_usercall_tx) = QueueSynchronizer::new(enclave_clone.clone());
 
-            let (usercall_queue_tx, usercall_queue_rx) = ipc_queue::bounded_async(USERCALL_QUEUE_SIZE, usercall_queue_synchronizer);
-            let (return_queue_tx, return_queue_rx) = ipc_queue::bounded_async(RETURN_QUEUE_SIZE, return_queue_synchronizer);
+            let (usercall_queue_tx, usercall_queue_rx) = ipc_queue::bounded_async(USERCALL_QUEUE_SIZE, uqs);
+            let (return_queue_tx, return_queue_rx) = ipc_queue::bounded_async(RETURN_QUEUE_SIZE, rqs);
+            let (cancel_queue_tx, cancel_queue_rx) = ipc_queue::bounded_async(CANCEL_QUEUE_SIZE, cqs);
 
             let fifo_descriptors = FifoDescriptors {
                 usercall_queue: usercall_queue_tx.into_descriptor(),
                 return_queue: return_queue_rx.into_descriptor(),
+                cancel_queue: cancel_queue_tx.into_descriptor(),
             };
 
             *enclave_clone.fifo_descriptors.lock().await = Some(fifo_descriptors);
@@ -834,14 +893,51 @@ impl EnclaveState {
                 }
             });
 
+            let (usercall_event_tx, mut usercall_event_rx) = async_mpsc::unbounded_channel();
+            let usercall_event_tx_clone = usercall_event_tx.clone();
+            tokio::task::spawn_local(async move {
+                while let Ok(c) = cancel_queue_rx.recv().await {
+                    if !c.ignore_cancel() {
+                        let _ = usercall_event_tx_clone.send(UsercallEvent::Cancelled(c.id, Instant::now()));
+                    }
+                }
+            });
+
+            tokio::task::spawn_local(async move {
+                let mut notifiers = HashMap::new();
+                let mut cancels: HashMap<u64, Instant> = HashMap::new();
+                // This should be greater than the amount of time it takes for the enclave runner
+                // to start executing a usercall after the enclave sends it on the usercall_queue.
+                const CANCEL_EXPIRY: Duration = Duration::from_millis(100);
+                loop {
+                    match usercall_event_rx.recv().await.expect("usercall_event channel closed unexpectedly") {
+                        UsercallEvent::Started(id, notifier) => match cancels.remove(&id) {
+                            Some(t) if t.elapsed() < CANCEL_EXPIRY => { let _ = notifier.send(()); },
+                            _ => { notifiers.insert(id, notifier); },
+                        },
+                        UsercallEvent::Finished(id) => { notifiers.remove(&id); },
+                        UsercallEvent::Cancelled(id, t) => if t.elapsed() < CANCEL_EXPIRY {
+                            match notifiers.remove(&id) {
+                                Some(notifier) => { let _ = notifier.send(()); },
+                                None => { cancels.insert(id, t); },
+                            }
+                        },
+                    }
+                    // cleanup expired cancels
+                    let now = Instant::now();
+                    cancels.retain(|_id, &mut t| now - t < CANCEL_EXPIRY);
+                }
+            });
+
             let mut recv_queue = io_queue_receive.into_future();
             while let (Some(work), stream) = recv_queue.await {
                 recv_queue = stream.into_future();
                 let enclave_clone = enclave_clone.clone();
                 let tx_return_channel = tx_return_channel.clone();
                 match work {
                     UsercallSendData::Async(usercall) => {
-                        let uchd = UsercallHandleData::Async(usercall);
+                        let usercall_event_tx = if usercall.ignore_cancel() { None } else { Some(usercall_event_tx.clone()) };
+                        let uchd = UsercallHandleData::Async(usercall, usercall_event_tx);
                         let fut = Self::handle_usercall(enclave_clone, work_sender.clone(), tx_return_channel, uchd);
                         tokio::task::spawn_local(fut);
                     }
@@ -1416,7 +1512,7 @@ impl<'tcs> IOHandlerInput<'tcs> {
     }
 
     fn check_event_set(set: u64) -> IoResult<u8> {
-        const EV_ALL: u64 = EV_USERCALLQ_NOT_FULL | EV_RETURNQ_NOT_EMPTY | EV_UNPARK;
+        const EV_ALL: u64 = EV_USERCALLQ_NOT_FULL | EV_RETURNQ_NOT_EMPTY | EV_CANCELQ_NOT_FULL | EV_UNPARK;
         if (set & !EV_ALL) != 0 {
             return Err(IoErrorKind::InvalidInput.into());
         }
@@ -1561,12 +1657,16 @@ impl<'tcs> IOHandlerInput<'tcs> {
         &mut self,
         usercall_queue: &mut FifoDescriptor<Usercall>,
         return_queue: &mut FifoDescriptor<Return>,
+        cancel_queue: Option<&mut FifoDescriptor<Cancel>>,
     ) -> StdResult<(), EnclaveAbort<bool>> {
         let fifo_descriptors = self.enclave.fifo_descriptors.lock().await.take();
         match fifo_descriptors {
             Some(fifo_descriptors) => {
                 *usercall_queue = fifo_descriptors.usercall_queue;
                 *return_queue = fifo_descriptors.return_queue;
+                if let Some(cancel_queue) = cancel_queue {
+                    *cancel_queue = fifo_descriptors.cancel_queue;
+                }
                 Ok(())
             }
             None => Err(self.exit(true)),
@@ -1578,6 +1678,7 @@ impl<'tcs> IOHandlerInput<'tcs> {
 enum Queue {
     Usercall,
     Return,
+    Cancel,
 }
 
 struct QueueSynchronizer {
@@ -1590,7 +1691,7 @@ struct QueueSynchronizer {
 }
 
 impl QueueSynchronizer {
-    fn new(enclave: Arc<EnclaveState>) -> (Self, Self, async_pubsub::Sender<()>) {
+    fn new(enclave: Arc<EnclaveState>) -> (Self, Self, Self, async_pubsub::Sender<()>) {
         // This broadcast channel is used to notify enclave-runner of any
         // synchronous usercalls made by the enclave for the purpose of
         // synchronizing access to usercall and return queues.
@@ -1605,11 +1706,17 @@ impl QueueSynchronizer {
         };
         let return_queue_synchronizer = QueueSynchronizer {
             queue: Queue::Return,
+            enclave: enclave.clone(),
+            subscription: Mutex::new(tx.subscribe()),
+            subscription_maker: tx.clone(),
+        };
+        let cancel_queue_synchronizer = QueueSynchronizer {
+            queue: Queue::Cancel,
             enclave,
             subscription: Mutex::new(tx.subscribe()),
             subscription_maker: tx.clone(),
         };
-        (usercall_queue_synchronizer, return_queue_synchronizer, tx)
+        (usercall_queue_synchronizer, return_queue_synchronizer, cancel_queue_synchronizer, tx)
     }
 }
 
@@ -1628,6 +1735,7 @@ impl ipc_queue::AsyncSynchronizer for QueueSynchronizer {
     fn wait(&self, event: QueueEvent) -> Pin<Box<dyn Future<Output = StdResult<(), ipc_queue::SynchronizationError>> + '_>> {
         match (self.queue, event) {
             (Queue::Usercall, QueueEvent::NotFull) => panic!("enclave runner should not send on the usercall queue"),
+            (Queue::Cancel, QueueEvent::NotFull)   => panic!("enclave runner should not send on the cancel queue"),
             (Queue::Return, QueueEvent::NotEmpty)  => panic!("enclave runner should not receive on the return queue"),
             _ => {}
         }
@@ -1646,12 +1754,14 @@ impl ipc_queue::AsyncSynchronizer for QueueSynchronizer {
     fn notify(&self, event: QueueEvent) {
         let ev = match (self.queue, event) {
             (Queue::Usercall, QueueEvent::NotEmpty) => panic!("enclave runner should not send on the usercall queue"),
+            (Queue::Cancel, QueueEvent::NotEmpty)   => panic!("enclave runner should not send on the cancel queue"),
             (Queue::Return, QueueEvent::NotFull)    => panic!("enclave runner should not receive on the return queue"),
             (Queue::Usercall, QueueEvent::NotFull)  => EV_USERCALLQ_NOT_FULL,
+            (Queue::Cancel, QueueEvent::NotFull)    => EV_CANCELQ_NOT_FULL,
             (Queue::Return, QueueEvent::NotEmpty)   => EV_RETURNQ_NOT_EMPTY,
         };
         // When the enclave needs to wait on a queue, it executes the wait() usercall synchronously,
-        // specifying EV_USERCALLQ_NOT_FULL, EV_RETURNQ_NOT_EMPTY, or both in the event_mask.
+        // specifying EV_USERCALLQ_NOT_FULL, EV_RETURNQ_NOT_EMPTY or EV_CANCELQ_NOT_FULL in the event_mask.
         // Userspace will wake any or all threads waiting on the appropriate event when it is triggered.
         for queue in self.enclave.event_queues.values() {
             let _ = queue.unbounded_send(ev as _);