Skip to content

Commit

Permalink
Cancel queue
Browse files Browse the repository at this point in the history
  • Loading branch information
@mzohreva
mzohreva committed Sep 18, 2020
1 parent 00055a3 commit 5e0efc2
Show file tree
Hide file tree
Showing 5 changed files with 186 additions and 33 deletions.
6 changes: 3 additions & 3 deletions enclave-runner/src/usercalls/abi.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand All @@ -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;
}

Expand All @@ -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,)*
}
Expand Down
7 changes: 4 additions & 3 deletions enclave-runner/src/usercalls/interface.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()))
Expand Down Expand Up @@ -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;
}

Expand Down
148 changes: 129 additions & 19 deletions enclave-runner/src/usercalls/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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]>),
Expand All @@ -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>>>;
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -587,6 +588,7 @@ impl EnclaveKind {
struct FifoDescriptors {
usercall_queue: FifoDescriptor<Usercall>,
return_queue: FifoDescriptor<Return>,
cancel_queue: FifoDescriptor<Cancel>,
}

pub(crate) struct EnclaveState {
Expand Down Expand Up @@ -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>>,
Expand Down Expand Up @@ -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) => {
Expand All @@ -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,
Expand All @@ -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())
}
Expand Down Expand Up @@ -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);
Expand All @@ -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);
}
Expand Down Expand Up @@ -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());
}
Expand Down Expand Up @@ -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)),
Expand All @@ -1578,6 +1678,7 @@ impl<'tcs> IOHandlerInput<'tcs> {
enum Queue {
Usercall,
Return,
Cancel,
}

struct QueueSynchronizer {
Expand All @@ -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.
Expand All @@ -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)
}
}

Expand All @@ -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"),
_ => {}
}
Expand All @@ -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 _);
Expand Down
Loading

0 comments on commit 5e0efc2

Please sign in to comment.