Giving enclaves control over their own TCSes

raoulstrackx · raoulstrackx · commit 976f41cc444a · 2020-11-06T00:22:30.000-08:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/enclave-runner/Cargo.toml b/enclave-runner/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "enclave-runner"
-version = "0.5.0"
+version = "0.6.0"
 authors = ["Fortanix, Inc."]
 license = "MPL-2.0"
 edition = "2018"
@@ -19,7 +19,7 @@ categories = ["os", "hardware-support"]
 [dependencies]
 # Project dependencies
 sgxs = { version = "0.7.2", path = "../sgxs" }
-fortanix-sgx-abi = { version = "0.4.0", path = "../fortanix-sgx-abi" }
+fortanix-sgx-abi = { version = "0.5.0", path = "../fortanix-sgx-abi" }
 sgx-isa = { version = "0.3.0", path = "../sgx-isa" }
 ipc-queue = { version = "0.1.0", path = "../ipc-queue" }
 
diff --git a/enclave-runner/src/usercalls/interface.rs b/enclave-runner/src/usercalls/interface.rs
@@ -171,9 +171,10 @@ impl<'future, 'ioinput: 'future, 'tcs: 'ioinput> Usercalls<'future> for Handler<
 
     fn launch_thread(
         self,
+        tcs: usize,
     ) -> std::pin::Pin<Box<dyn Future<Output = (Self, UsercallResult<Result>)> + 'future>> {
         async move {
-            let ret = Ok(self.0.launch_thread().to_sgx_result());
+            let ret = Ok(self.0.launch_thread(tcs).to_sgx_result());
             return (self, ret);
         }.boxed_local()
     }
diff --git a/enclave-runner/src/usercalls/mod.rs b/enclave-runner/src/usercalls/mod.rs
@@ -19,6 +19,7 @@ use std::{cmp, fmt, str};
 
 use failure::{self, bail};
 use fnv::FnvHashMap;
+use futures::executor;
 use futures::future::{poll_fn, Either, Future, FutureExt};
 use futures::lock::Mutex;
 use futures::StreamExt;
@@ -596,7 +597,7 @@ pub(crate) struct EnclaveState {
     last_fd: AtomicUsize,
     exiting: AtomicBool,
     usercall_ext: Box<dyn UsercallExtension>,
-    threads_queue: crossbeam::queue::SegQueue<StoppedTcs>,
+    available_enclave_threads: Mutex<FnvHashMap<TcsAddress, StoppedTcs>>,
     forward_panics: bool,
     // Once set to Some, the guards should not be dropped for the lifetime of the enclave.
     fifo_guards: Mutex<Option<FifoGuards>>,
@@ -677,10 +678,10 @@ impl EnclaveState {
 
         let usercall_ext = usercall_ext.unwrap_or_else(|| Box::new(UsercallExtensionDefault));
 
-        let threads_queue = crossbeam::queue::SegQueue::new();
+        let mut available_enclave_threads = FnvHashMap::default();
 
         for thread in threads_vector {
-            threads_queue.push(Self::event_queue_add_tcs(&mut event_queues, thread));
+            available_enclave_threads.insert(thread.address(), Self::event_queue_add_tcs(&mut event_queues, thread));
         }
 
         Arc::new(EnclaveState {
@@ -690,7 +691,7 @@ impl EnclaveState {
             last_fd,
             exiting: AtomicBool::new(false),
             usercall_ext,
-            threads_queue,
+            available_enclave_threads: Mutex::new(available_enclave_threads),
             forward_panics,
             fifo_guards: Mutex::new(None),
             return_queue_tx: Mutex::new(None),
@@ -860,7 +861,7 @@ impl EnclaveState {
                         let fut = async move {
                             let ret = match state.mode {
                                 EnclaveEntry::Library => {
-                                    enclave_clone.threads_queue.push(StoppedTcs {
+                                    enclave_clone.available_enclave_threads.lock().await.insert(tcs.address(), StoppedTcs {
                                         tcs,
                                         event_queue: state.event_queue,
                                     });
@@ -876,7 +877,7 @@ impl EnclaveState {
                                     // If the enclave is in the exit-state, threads are no
                                     // longer able to be launched
                                     if !enclave_clone.exiting.load(Ordering::SeqCst) {
-                                        enclave_clone.threads_queue.push(StoppedTcs {
+                                        enclave_clone.available_enclave_threads.lock().await.insert(tcs.address(), StoppedTcs {
                                             tcs,
                                             event_queue: state.event_queue,
                                         });
@@ -1008,8 +1009,7 @@ impl EnclaveState {
 
         rt.block_on(async move {
             enclave.abort_all_threads();
-            //clear the threads_queue
-            while enclave.threads_queue.pop().is_ok() {}
+            enclave.available_enclave_threads.lock().await.clear();
 
             let cmd = enclave.kind.as_command().unwrap();
             let mut cmddata = cmd.panic_reason.lock().await;
@@ -1058,7 +1058,11 @@ impl EnclaveState {
         p4: u64,
         p5: u64,
     ) -> StdResult<(u64, u64), failure::Error> {
-        let thread = enclave.threads_queue.pop().expect("threads queue empty");
+        let thread = executor::block_on(async {
+            let mut available_enclave_threads = enclave.available_enclave_threads.lock().await;
+            let k = available_enclave_threads.keys().next()?.to_owned();
+            available_enclave_threads.remove(&k)
+        }).expect("out of enclave threads");
         let work = Work {
             tcs: RunningTcs {
                 event_queue: thread.event_queue,
@@ -1389,18 +1393,32 @@ impl<'tcs> IOHandlerInput<'tcs> {
     }
 
     #[inline(always)]
-    fn launch_thread(&self) -> IoResult<()> {
+    fn launch_thread(&self, tcs_address: usize) -> IoResult<()> {
         // check if enclave is of type command
         self.enclave
             .kind
             .as_command()
             .ok_or(IoErrorKind::InvalidInput)?;
-        let new_tcs = match self.enclave.threads_queue.pop() {
-            Ok(tcs) => tcs,
-            Err(_) => {
-                return Err(IoErrorKind::WouldBlock.into());
-            }
-        };
+        let tcs_address = TcsAddress(tcs_address);
+        let new_tcs = executor::block_on(async {
+                loop {
+                    let mut guard = self.enclave.available_enclave_threads.lock().await;
+                    match guard.remove(&tcs_address) {
+                        Some(tcs) => break tcs,
+                        None => {
+                            // Release lock and try again. The enclave is in charge of its own TCS
+                            // structs. Unfortunately, there is a small issue recording terminated
+                            // TCSs; when a thread has finished with the task it was assigned, it
+                            // marks its own TCS as available for new threads, while in fact it is
+                            // still executing (see `Task::run()` in `std/src/sys/sgx/thread.rs`).
+                            // When that TCS is subsequently selected to run a new thread, it may
+                            // still not have terminated yet and thus may not yet be present in this
+                            // `available_enclave_threads`. We need to wait for it to become ready.
+                            drop(guard);
+                        }
+                    }
+                }
+            });
 
         let ret = self.work_sender.send(Work {
             tcs: RunningTcs {
@@ -1417,10 +1435,12 @@ impl<'tcs> IOHandlerInput<'tcs> {
                 let entry = e.0.entry;
                 match entry {
                     CoEntry::Initial(tcs, _, _ ,_, _, _) => {
-                        self.enclave.threads_queue.push(StoppedTcs {
-                            tcs,
-                            event_queue,
-                        });
+                        executor::block_on(async {
+                            self.enclave.available_enclave_threads.lock().await.insert(tcs.address(), StoppedTcs {
+                                tcs,
+                                event_queue,
+                            });
+                        })
                     },
                     _ => unreachable!(),
                 };
diff --git a/fortanix-sgx-abi/Cargo.toml b/fortanix-sgx-abi/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "fortanix-sgx-abi"
-version = "0.4.0"
+version = "0.5.0"
 authors = ["Fortanix, Inc."]
 license = "MPL-2.0"
 description = """
diff --git a/fortanix-sgx-abi/src/lib.rs b/fortanix-sgx-abi/src/lib.rs
@@ -488,7 +488,7 @@ impl Usercalls {
     /// [`thread_entry`]: entry/executable/fn.thread_entry.html
     /// [libraries]: entry/library/index.html
     /// [`library`]: entry/library/index.html
-    pub fn launch_thread() -> Result { unimplemented!() }
+    pub fn launch_thread(tcs: usize) -> Result { unimplemented!() }
 
     /// Signals to userspace that this enclave needs to be destroyed.
     ///
diff --git a/fortanix-sgx-tools/Cargo.toml b/fortanix-sgx-tools/Cargo.toml
@@ -21,7 +21,7 @@ edition = "2018"
 # Project dependencies
 aesm-client = { version = "0.5.0", path = "../aesm-client", features = ["sgxs"] }
 sgxs-loaders = { version = "0.3.0", path = "../sgxs-loaders" }
-enclave-runner = { version = "0.5.0", path = "../enclave-runner" }
+enclave-runner = { version = "0.6.0", path = "../enclave-runner" }
 sgxs = { version = "0.7.0", path = "../sgxs" }
 sgx-isa = { version = "0.3.0", path = "../sgx-isa" }
 
diff --git a/ipc-queue/Cargo.toml b/ipc-queue/Cargo.toml
@@ -14,7 +14,7 @@ keywords = ["sgx", "fifo", "queue", "ipc"]
 categories = ["asynchronous"]
 
 [dependencies]
-fortanix-sgx-abi = { version = "0.4.0", path = "../fortanix-sgx-abi" }
+fortanix-sgx-abi = { version = "0.5.0", path = "../fortanix-sgx-abi" }
 
 [dev-dependencies]
 static_assertions = "1.1.0"
diff --git a/report-test/Cargo.toml b/report-test/Cargo.toml
@@ -16,7 +16,7 @@ categories = ["development-tools"]
 
 [dependencies]
 # Project dependencies
-"enclave-runner" = { version = "0.5.0", path = "../enclave-runner" }
+"enclave-runner" = { version = "0.6.0", path = "../enclave-runner" }
 "sgxs" = { version = "0.7.0", path = "../sgxs" }
 "sgx-isa" = { version = "0.3.0", path = "../sgx-isa" }
 
diff --git a/sgxs-tools/Cargo.toml b/sgxs-tools/Cargo.toml
@@ -34,7 +34,7 @@ path = "src/sgx_detect/main.rs"
 "aesm-client" = { version = "0.5.0", path = "../aesm-client", features = ["sgxs"] }
 "sgx-isa" = { version = "0.3.0", path = "../sgx-isa" }
 "report-test" = { version = "0.3.1", path = "../report-test" }
-"enclave-runner" = { version = "0.5.0", path = "../enclave-runner" }
+"enclave-runner" = { version = "0.6.0", path = "../enclave-runner" }
 
 # External dependencies
 lazy_static = "1"                                # MIT/Apache-2.0
