Async usercall interface for SGX enclaves

.travis.yml

@@ -29,7 +29,8 @@ matrix:
       before_script:
         - rustup target add x86_64-fortanix-unknown-sgx x86_64-unknown-linux-musl
       script:
-        - cargo test --verbose --all
+        - cargo test --verbose --all --exclude async-usercalls
+        - cargo test --verbose -p async-usercalls --target x86_64-fortanix-unknown-sgx --no-run
         - cargo test --verbose -p sgx-isa --features sgxstd -Z package-features --target x86_64-fortanix-unknown-sgx --no-run
         - cargo test --verbose -p sgxs-tools --features pe2sgxs --bin isgx-pe2sgx -Z package-features
         - cargo test --verbose -p dcap-ql --features link -Z package-features

Cargo.toml

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

async-usercalls/Cargo.toml

@@ -0,0 +1,31 @@
+[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.3", path = "../fortanix-sgx-abi" }
+
+# External dependencies
+lazy_static = "1.4.0"     # MIT/Apache-2.0
+crossbeam-channel = "0.4" # MIT/Apache-2.0
+spin = "0.5"              # 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/rustfmt.toml

@@ -0,0 +1 @@
+max_width = 120

async-usercalls/src/alloc/allocator.rs

@@ -0,0 +1,145 @@
+use super::slab::{BufSlab, Slab, SlabAllocator, User, MAX_COUNT};
+use std::cmp;
+use std::os::fortanix_sgx::usercalls::raw::ByteBuffer;
+
+pub const MIN_BUF_SIZE: usize = 1 << 5; // 32 bytes
+pub const MAX_BUF_SIZE: usize = 1 << 16; // 64 KB
+pub const NUM_SIZES: usize = 1 + (MAX_BUF_SIZE / MIN_BUF_SIZE).trailing_zeros() as usize;
+
+pub struct SharedAllocator {
+    by_size: Vec<Vec<BufSlab>>,
+    byte_buffers: Vec<Slab<ByteBuffer>>,
+}
+
+unsafe impl Send for SharedAllocator {}
+unsafe impl Sync for SharedAllocator {}
+
+impl SharedAllocator {
+    pub fn new(buf_counts: [usize; NUM_SIZES], byte_buffer_count: usize) -> Self {
+        let mut by_size = Vec::with_capacity(NUM_SIZES);
+        for i in 0..NUM_SIZES {
+            by_size.push(make_buf_slabs(buf_counts[i], MIN_BUF_SIZE << i));
+        }
+        let byte_buffers = make_byte_buffers(byte_buffer_count);
+        Self { by_size, byte_buffers }
+    }
+
+    pub fn alloc_buf(&self, size: usize) -> Option<User<[u8]>> {
+        assert!(size > 0);
+        if size > MAX_BUF_SIZE {
+            return None;
+        }
+        let (_, index) = size_index(size);
+        self.by_size[index].alloc()
+    }
+
+    pub fn alloc_byte_buffer(&self) -> Option<User<ByteBuffer>> {
+        self.byte_buffers.alloc()
+    }
+}
+
+pub struct LocalAllocator {
+    initial_buf_counts: [usize; NUM_SIZES],
+    initial_byte_buffer_count: usize,
+    inner: SharedAllocator,
+}
+
+impl LocalAllocator {
+    pub fn new(initial_buf_counts: [usize; NUM_SIZES], initial_byte_buffer_count: usize) -> Self {
+        let mut by_size = Vec::with_capacity(NUM_SIZES);
+        by_size.resize_with(NUM_SIZES, Default::default);
+        let byte_buffers = Vec::new();
+        Self {
+            initial_buf_counts,
+            initial_byte_buffer_count,
+            inner: SharedAllocator { by_size, byte_buffers },
+        }
+    }
+
+    pub fn alloc_buf(&mut self, request_size: usize) -> User<[u8]> {
+        assert!(request_size > 0);
+        if request_size > MAX_BUF_SIZE {
+            // Always allocate very large buffers directly
+            return User::<[u8]>::uninitialized(request_size);
+        }
+        let (size, index) = size_index(request_size);
+        if let Some(buf) = self.inner.by_size[index].alloc() {
+            return buf;
+        }
+        let slabs = &mut self.inner.by_size[index];
+        if slabs.len() >= 8 {
+            // Keep the number of slabs for each size small.
+            return User::<[u8]>::uninitialized(request_size);
+        }
+        let count = slabs.last().map_or(self.initial_buf_counts[index], |s| s.count() * 2);
+        // Limit each slab's count for better worst-case performance.
+        let count = cmp::min(count, MAX_COUNT / 8);
+        slabs.push(BufSlab::new(count, size));
+        slabs.last().unwrap().alloc().expect("fresh slab failed to allocate")
+    }
+
+    pub fn alloc_byte_buffer(&mut self) -> User<ByteBuffer> {
+        let bbs = &mut self.inner.byte_buffers;
+        if let Some(byte_buffer) = bbs.alloc() {
+            return byte_buffer;
+        }
+        if bbs.len() >= 8 {
+            // Keep the number of slabs small.
+            return User::<ByteBuffer>::uninitialized();
+        }
+        let count = bbs.last().map_or(self.initial_byte_buffer_count, |s| s.count() * 2);
+        // Limit each slab's count for better worst-case performance.
+        let count = cmp::min(count, MAX_COUNT / 8);
+        bbs.push(Slab::new(count));
+        bbs.last().unwrap().alloc().expect("fresh slab failed to allocate")
+    }
+}
+
+fn make_buf_slabs(count: usize, size: usize) -> Vec<BufSlab> {
+    match count {
+        0 => Vec::new(),
+        n if n < 1024 => vec![BufSlab::new(n, size)],
+        n if n < 4 * 1024 => vec![BufSlab::new(n / 2, size), BufSlab::new(n / 2, size)],
+        n if n < 32 * 1024 => vec![
+            BufSlab::new(n / 4, size),
+            BufSlab::new(n / 4, size),
+            BufSlab::new(n / 4, size),
+            BufSlab::new(n / 4, size),
+        ],
+        n => vec![
+            BufSlab::new(n / 8, size),
+            BufSlab::new(n / 8, size),
+            BufSlab::new(n / 8, size),
+            BufSlab::new(n / 8, size),
+            BufSlab::new(n / 8, size),
+            BufSlab::new(n / 8, size),
+            BufSlab::new(n / 8, size),
+            BufSlab::new(n / 8, size),
+        ],
+    }
+}
+
+fn make_byte_buffers(count: usize) -> Vec<Slab<ByteBuffer>> {
+    match count {
+        0 => Vec::new(),
+        n if n < 1024 => vec![Slab::new(n)],
+        n if n < 4 * 1024 => vec![Slab::new(n / 2), Slab::new(n / 2)],
+        n if n < 32 * 1024 => vec![Slab::new(n / 4), Slab::new(n / 4), Slab::new(n / 4), Slab::new(n / 4)],
+        n => vec![
+            Slab::new(n / 8),
+            Slab::new(n / 8),
+            Slab::new(n / 8),
+            Slab::new(n / 8),
+            Slab::new(n / 8),
+            Slab::new(n / 8),
+            Slab::new(n / 8),
+            Slab::new(n / 8),
+        ],
+    }
+}
+
+fn size_index(request_size: usize) -> (usize, usize) {
+    let size = cmp::max(MIN_BUF_SIZE, request_size.next_power_of_two());
+    let index = (size / MIN_BUF_SIZE).trailing_zeros() as usize;
+    (size, index)
+}

async-usercalls/src/alloc/bitmap.rs

@@ -0,0 +1,156 @@
+use spin::Mutex;
+use std::sync::atomic::*;
+
+pub struct OptionalBitmap(BitmapKind);
+
+struct LargeBitmap(Mutex<LargeBitmapInner>);
+
+struct LargeBitmapInner {
+    bits: Box<[u64]>,
+    unset_count: usize, // optimization
+}
+
+enum BitmapKind {
+    None,
+    V1(AtomicU8),
+    V2(AtomicU16),
+    V3(AtomicU32),
+    V4(AtomicU64),
+    V5(LargeBitmap),
+}
+
+impl OptionalBitmap {
+    pub fn none() -> Self {
+        Self(BitmapKind::None)
+    }
+
+    /// `bit_count` must be >= 8 and a power of two
+    pub fn new(bit_count: usize) -> Self {
+        Self(match bit_count {
+            8 => BitmapKind::V1(AtomicU8::new(0)),
+            16 => BitmapKind::V2(AtomicU16::new(0)),
+            32 => BitmapKind::V3(AtomicU32::new(0)),
+            64 => BitmapKind::V4(AtomicU64::new(0)),
+            n if n > 0 && n % 64 == 0 => {
+                let bits = vec![0u64; n / 64].into_boxed_slice();
+                BitmapKind::V5(LargeBitmap(Mutex::new(LargeBitmapInner {
+                    bits,
+                    unset_count: bit_count,
+                })))
+            }
+            _ => panic!("bit_count must be >= 8 and a power of two"),
+        })
+    }
+
+    /// set the bit at given index to 0 and panic if the old value was not 1.
+    pub fn unset(&self, index: usize) {
+        match self.0 {
+            BitmapKind::None => {}
+            BitmapKind::V1(ref a) => a.unset(index),
+            BitmapKind::V2(ref b) => b.unset(index),
+            BitmapKind::V3(ref c) => c.unset(index),
+            BitmapKind::V4(ref d) => d.unset(index),
+            BitmapKind::V5(ref e) => e.unset(index),
+        }
+    }
+
+    /// return the index of a previously unset bit and set that bit to 1.
+    pub fn reserve(&self) -> Option<usize> {
+        match self.0 {
+            BitmapKind::None => None,
+            BitmapKind::V1(ref a) => a.reserve(),
+            BitmapKind::V2(ref b) => b.reserve(),
+            BitmapKind::V3(ref c) => c.reserve(),
+            BitmapKind::V4(ref d) => d.reserve(),
+            BitmapKind::V5(ref e) => e.reserve(),
+        }
+    }
+}
+
+trait BitmapOps {
+    fn unset(&self, index: usize);
+    fn reserve(&self) -> Option<usize>;
+}
+
+macro_rules! impl_bitmap_ops {
+    ( $( $t:ty ),* $(,)? ) => {$(
+        impl BitmapOps for $t {
+            fn unset(&self, index: usize) {
+                let bit = 1 << index;
+                let old = self.fetch_and(!bit, Ordering::Release) & bit;
+                assert!(old != 0);
+            }
+
+            fn reserve(&self) -> Option<usize> {
+                let initial = self.load(Ordering::Relaxed);
+                let unset_count = initial.count_zeros();
+                let (mut index, mut bit) = match unset_count {
+                    0 => return None,
+                    _ => (0, 1),
+                };
+                for _ in 0..unset_count {
+                    // find the next unset bit
+                    while bit & initial != 0 {
+                        index += 1;
+                        bit = bit << 1;
+                    }
+                    let old = self.fetch_or(bit, Ordering::Acquire) & bit;
+                    if old == 0 {
+                        return Some(index);
+                    }
+                    index += 1;
+                    bit = bit << 1;
+                }
+                None
+            }
+        }
+    )*};
+}
+
+impl_bitmap_ops!(AtomicU8, AtomicU16, AtomicU32, AtomicU64);
+
+impl BitmapOps for LargeBitmap {
+    fn unset(&self, index: usize) {
+        let mut inner = self.0.lock();
+        let array = &mut inner.bits;
+        assert!(index < array.len() * 64);
+        let slot = index / 64;
+        let offset = index % 64;
+        let element = &mut array[slot];
+
+        let bit = 1 << offset;
+        let old = *element & bit;
+        *element = *element & !bit;
+        inner.unset_count += 1;
+        assert!(old != 0);
+    }
+
+    fn reserve(&self) -> Option<usize> {
+        let mut inner = self.0.lock();
+        if inner.unset_count == 0 {
+            return None;
+        }
+        let array = &mut inner.bits;
+        for slot in 0..array.len() {
+            if let (Some(offset), val) = reserve_u64(array[slot]) {
+                array[slot] = val;
+                inner.unset_count -= 1;
+                return Some(slot * 64 + offset);
+            }
+        }
+        unreachable!()
+    }
+}
+
+fn reserve_u64(element: u64) -> (Option<usize>, u64) {
+    let (mut index, mut bit) = match element.count_zeros() {
+        0 => return (None, element),
+        _ => (0, 1),
+    };
+    // find the first unset bit
+    while bit & element != 0 {
+        index += 1;
+        bit = bit << 1;
+    }
+    (Some(index), element | bit)
+}