Async usercall interface for SGX enclaves

.travis.yml

@@ -23,13 +23,15 @@ matrix:
             - libclang-3.8-dev
             - musl-tools
       rust:
-        - nightly
+        # This need to change back to `nightly` after https://github.com/fortanix/rust-sgx/issues/433 is fixed
+        - nightly-2023-01-31
       env:
         - RUST_BACKTRACE=1 LLVM_CONFIG_PATH=llvm-3.8-config
       before_script:
         - rustup target add x86_64-fortanix-unknown-sgx x86_64-unknown-linux-musl
       script:
-        - cargo test --verbose --all --exclude sgxs-loaders && [ "$(echo $(nm -D target/debug/sgx-detect|grep __vdso_sgx_enter_enclave))" = "w __vdso_sgx_enter_enclave" ]
+        - cargo test --verbose --all --exclude sgxs-loaders --exclude async-usercalls && [ "$(echo $(nm -D target/debug/sgx-detect|grep __vdso_sgx_enter_enclave))" = "w __vdso_sgx_enter_enclave" ]
+        - cargo test --verbose -p async-usercalls --target x86_64-fortanix-unknown-sgx --no-run
         - cargo test --verbose -p sgx-isa --features sgxstd --target x86_64-fortanix-unknown-sgx --no-run
         - cargo test --verbose -p sgxs-tools --features pe2sgxs --bin isgx-pe2sgx
         - cargo test --verbose -p dcap-ql --features link

Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 members = [
     "aesm-client",
+    "async-usercalls",
     "dcap-provider",
     "dcap-ql-sys",
     "dcap-ql",

aesm-client/src/imp/aesm_protobuf/mod.rs

@@ -50,7 +50,7 @@ impl AesmClient {
         let mut res_bytes = vec![0; res_len as usize];
         sock.read_exact(&mut res_bytes)?;
 
-        let res = T::Response::from_response(protobuf::parse_from_bytes(&res_bytes))?;
+        let res = T::Response::from_response(Message::parse_from_bytes(&res_bytes))?;
         Ok(res)
     }
 

async-usercalls/Cargo.toml

@@ -0,0 +1,30 @@
+[package]
+name = "async-usercalls"
+version = "0.1.0"
+authors = ["Fortanix, Inc."]
+license = "MPL-2.0"
+edition = "2018"
+description = """
+An interface for asynchronous usercalls in SGX enclaves.
+
+This is an SGX-only crate, you should compile it with the `x86_64-fortanix-unknown-sgx` target.
+"""
+repository = "https://github.com/fortanix/rust-sgx"
+documentation = "https://edp.fortanix.com/docs/api/async_usercalls/"
+homepage = "https://edp.fortanix.com/"
+keywords = ["sgx", "async", "usercall"]
+categories = ["asynchronous"]
+
+[dependencies]
+# Project dependencies
+ipc-queue = { version = "0.1", path = "../ipc-queue" }
+fortanix-sgx-abi = { version = "0.4", path = "../fortanix-sgx-abi" }
+
+# External dependencies
+lazy_static = "1.4.0"     # MIT/Apache-2.0
+crossbeam-channel = "0.4" # MIT/Apache-2.0
+fnv = "1.0"               # MIT/Apache-2.0
+
+# For cargo test --target x86_64-fortanix-unknown-sgx
+[package.metadata.fortanix-sgx]
+threads = 128

async-usercalls/src/batch_drop.rs

@@ -0,0 +1,127 @@
+use crate::hacks::Usercall;
+use crate::provider_core::ProviderCore;
+use ipc_queue::Identified;
+use std::cell::RefCell;
+use std::mem;
+use std::os::fortanix_sgx::usercalls::alloc::{User, UserSafe};
+use std::os::fortanix_sgx::usercalls::raw::UsercallNrs;
+
+pub trait BatchDroppable: private::BatchDroppable {}
+impl<T: private::BatchDroppable> BatchDroppable for T {}
+
+/// Drop the given value at some point in the future (no rush!). This is useful
+/// for freeing userspace memory when we don't particularly care about when the
+/// buffer is freed. Multiple `free` usercalls are batched together and sent to
+/// userspace asynchronously. It is also guaranteed that the memory is freed if
+/// the current thread exits before there is a large enough batch.
+///
+/// This is mainly an optimization to avoid exitting the enclave for each
+/// usercall. Note that even when sending usercalls asynchronously, if the
+/// usercall queue is empty we still need to exit the enclave to signal the
+/// userspace that the queue is not empty anymore. The batch send would send
+/// multiple usercalls and notify the userspace at most once.
+pub fn batch_drop<T: BatchDroppable>(t: T) {
+    t.batch_drop();
+}
+
+mod private {
+    use super::*;
+
+    const BATCH_SIZE: usize = 8;
+
+    struct BatchDropProvider {
+        core: ProviderCore,
+        deferred: Vec<Identified<Usercall>>,
+    }
+
+    impl BatchDropProvider {
+        pub fn new() -> Self {
+            Self {
+                core: ProviderCore::new(None),
+                deferred: Vec::with_capacity(BATCH_SIZE),
+            }
+        }
+
+        fn make_progress(&self, deferred: &[Identified<Usercall>]) -> usize {
+            let sent = self.core.try_send_multiple_usercalls(deferred);
+            if sent == 0 {
+                self.core.send_usercall(deferred[0]);
+                return 1;
+            }
+            sent
+        }
+
+        fn maybe_send_usercall(&mut self, u: Usercall) {
+            self.deferred.push(self.core.assign_id(u));
+            if self.deferred.len() < BATCH_SIZE {
+                return;
+            }
+            let sent = self.make_progress(&self.deferred);
+            let mut not_sent = self.deferred.split_off(sent);
+            self.deferred.clear();
+            self.deferred.append(&mut not_sent);
+        }
+
+        pub fn free<T: UserSafe + ?Sized>(&mut self, buf: User<T>) {
+            let ptr = buf.into_raw();
+            let size = unsafe { mem::size_of_val(&mut *ptr) };
+            let alignment = T::align_of();
+            let ptr = ptr as *mut u8;
+            let u = Usercall(UsercallNrs::free as _, ptr as _, size as _, alignment as _, 0);
+            self.maybe_send_usercall(u);
+        }
+    }
+
+    impl Drop for BatchDropProvider {
+        fn drop(&mut self) {
+            let mut sent = 0;
+            while sent < self.deferred.len() {
+                sent += self.make_progress(&self.deferred[sent..]);
+            }
+        }
+    }
+
+    std::thread_local! {
+        static PROVIDER: RefCell<BatchDropProvider> = RefCell::new(BatchDropProvider::new());
+    }
+
+    pub trait BatchDroppable {
+        fn batch_drop(self);
+    }
+
+    impl<T: UserSafe + ?Sized> BatchDroppable for User<T> {
+        fn batch_drop(self) {
+            PROVIDER.with(|p| p.borrow_mut().free(self));
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::batch_drop;
+    use std::os::fortanix_sgx::usercalls::alloc::User;
+    use std::thread;
+
+    #[test]
+    fn basic() {
+        for _ in 0..100 {
+            batch_drop(User::<[u8]>::uninitialized(100));
+        }
+    }
+
+    #[test]
+    fn multiple_threads() {
+        const THREADS: usize = 16;
+        let mut handles = Vec::with_capacity(THREADS);
+        for _ in 0..THREADS {
+            handles.push(thread::spawn(move || {
+                for _ in 0..1000 {
+                    batch_drop(User::<[u8]>::uninitialized(100));
+                }
+            }));
+        }
+        for h in handles {
+            h.join().unwrap();
+        }
+    }
+}

async-usercalls/src/callback.rs

@@ -0,0 +1,65 @@
+use crate::duplicated::{FromSgxResult, ReturnValue};
+use crate::hacks::Return;
+use fortanix_sgx_abi::{invoke_with_usercalls, Fd, Result};
+use std::io;
+
+pub struct CbFn<T>(Box<dyn FnOnce(T) + Send + 'static>);
+
+impl<T> CbFn<T> {
+    fn call(self, t: T) {
+        (self.0)(t);
+    }
+}
+
+impl<F, T> From<F> for CbFn<T>
+where
+    F: FnOnce(T) + Send + 'static,
+{
+    fn from(f: F) -> Self {
+        Self(Box::new(f))
+    }
+}
+
+macro_rules! cbfn_type {
+    ( )                      => { CbFn<()> };
+    ( -> ! )                 => { () };
+    ( -> u64 )               => { CbFn<u64> };
+    ( -> (Result, usize) )   => { CbFn<io::Result<usize>> };
+    ( -> (Result, u64) )     => { CbFn<io::Result<u64>> };
+    ( -> (Result, Fd) )      => { CbFn<io::Result<Fd>> };
+    ( -> (Result, *mut u8) ) => { CbFn<io::Result<*mut u8>> };
+    ( -> Result )            => { CbFn<io::Result<()>> };
+}
+
+macro_rules! call_cbfn {
+    ( $cb:ident, $rv:expr, )          => { let x: () = $rv; $cb.call(x); };
+    ( $cb:ident, $rv:expr, -> ! )     => { let _: ! = $rv; };
+    ( $cb:ident, $rv:expr, -> u64 )   => { let x: u64 = $rv; $cb.call(x); };
+    ( $cb:ident, $rv:expr, -> $t:ty ) => { let x: $t = $rv; $cb.call(x.from_sgx_result()); };
+}
+
+macro_rules! define_callback {
+    ($(fn $name:ident($($n:ident: $t:ty),*) $(-> $r:tt)*; )*) => {
+        #[allow(unused)]
+        #[allow(non_camel_case_types)]
+        pub(crate) enum Callback {
+            $( $name(cbfn_type! { $(-> $r)* }), )*
+        }
+
+        impl Callback {
+            pub(crate) fn call(self, ret: Return) {
+                match self {$(
+                    Callback::$name(_cb) => {
+                        call_cbfn!(
+                            _cb,
+                            ReturnValue::from_registers(stringify!($name), (ret.0, ret.1)),
+                            $(-> $r)*
+                        );
+                    }
+                )*}
+            }
+        }
+    };
+}
+
+invoke_with_usercalls!(define_callback);