Selecting TCS used when new thread launches

Author: raoulstrackx

Diff for: Cargo.lock

@@ -1190,9 +1190,8 @@ dependencies = [
 
 [[package]]
 name = "fortanix-sgx-abi"
-version = "0.3.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c56c422ef86062869b2d57ae87270608dc5929969dd130a6e248979cf4fb6ca6"
+version = "0.5.0"
+source = "git+https://github.com/fortanix/rust-sgx.git?branch=raoul/tcs_control#495f16cb6cd0d8f6a7b25ade5196a96a0d22fe35"
 dependencies = [
  "compiler_builtins",
  "rustc-std-workspace-core",

Diff for: Cargo.toml

@@ -101,6 +101,8 @@ cargo-util = { path = "src/tools/cargo/crates/cargo-util" }
 rustfmt-nightly = { path = "src/tools/rustfmt" }
 
 [patch.crates-io]
+fortanix-sgx-abi = { git = "https://github.com/fortanix/rust-sgx.git", branch = "raoul/tcs_control" }
+
 # See comments in `src/tools/rustc-workspace-hack/README.md` for what's going on
 # here
 rustc-workspace-hack = { path = 'src/tools/rustc-workspace-hack' }

Diff for: library/std/Cargo.toml

@@ -40,7 +40,7 @@ rand = "0.7"
 dlmalloc = { version = "0.2.1", features = ['rustc-dep-of-std'] }
 
 [target.x86_64-fortanix-unknown-sgx.dependencies]
-fortanix-sgx-abi = { version = "0.3.2", features = ['rustc-dep-of-std'] }
+fortanix-sgx-abi = { version = "0.5.0", features = ['rustc-dep-of-std'] }
 
 [target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), target_os = "hermit"))'.dependencies]
 hermit-abi = { version = "0.1.17", features = ['rustc-dep-of-std'] }

Diff for: library/std/src/sys/sgx/abi/entry.S

@@ -11,7 +11,7 @@ IMAGE_BASE:
     .long 1                    /* type = NT_VERSION */
 0:  .asciz "toolchain-version" /* name */
 1:  .align 4
-2:  .long 1                    /* desc - toolchain version number, 32-bit LE */
+2:  .long 2                    /* desc - toolchain version number, 32-bit LE */
 3:  .align 4
 
 .section .rodata
@@ -90,15 +90,17 @@ IMAGE_BASE:
 .equ tcsls_last_rsp,            0x10 /*  initialized by loader to 0 */
 .equ tcsls_panic_last_rsp,      0x18 /*  initialized by loader to 0 */
 .equ tcsls_debug_panic_buf_ptr, 0x20 /*  initialized by loader to 0 */
-.equ tcsls_user_rsp,            0x28
-.equ tcsls_user_retip,          0x30
-.equ tcsls_user_rbp,            0x38
-.equ tcsls_user_r12,            0x40
-.equ tcsls_user_r13,            0x48
-.equ tcsls_user_r14,            0x50
-.equ tcsls_user_r15,            0x58
-.equ tcsls_tls_ptr,             0x60
-.equ tcsls_tcs_addr,            0x68
+.equ tcsls_static_tcs_addr,     0x28 /*  initialized by loader to *offset* from image base to static TCS */
+.equ tcsls_clist_next,          0x30 /*  initialized by loader to *offset* from image base to next TCLS, circular linked list */
+.equ tcsls_user_rsp,            0x38
+.equ tcsls_user_retip,          0x40
+.equ tcsls_user_rbp,            0x48
+.equ tcsls_user_r12,            0x50
+.equ tcsls_user_r13,            0x58
+.equ tcsls_user_r14,            0x60
+.equ tcsls_user_r15,            0x68
+.equ tcsls_tls_ptr,             0x70
+.equ tcsls_tcs_addr,            0x78
 
 .macro load_tcsls_flag_secondary_bool reg:req comments:vararg
     .ifne tcsls_flag_secondary /* to convert to a bool, must be the first bit */
@@ -370,3 +372,24 @@ take_debug_panic_buf_ptr:
     pop %r11
     lfence
     jmp *%r11
+
+.global next_tcsls
+next_tcsls:
+    mov %gs:tcsls_clist_next,%rax
+    pop %r11
+    lfence
+    jmp *%r11
+
+.global static_tcs_offset
+static_tcs_offset:
+    mov $tcsls_static_tcs_addr, %rax
+    pop %r11
+    lfence
+    jmp *%r11
+
+.global clist_next_offset
+clist_next_offset:
+    mov $tcsls_clist_next, %rax
+    pop %r11
+    lfence
+    jmp *%r11

Diff for: library/std/src/sys/sgx/abi/mem.rs

@@ -89,3 +89,37 @@ pub fn is_user_range(p: *const u8, len: usize) -> bool {
     let base = image_base() as usize;
     end < base || start > base + (unsafe { ENCLAVE_SIZE } - 1) // unsafe ok: link-time constant
 }
+
+#[repr(C, packed)]
+#[derive(Default)]
+struct TcslsTcsListItem {
+    tcs_offset: u64,
+    next_offset: u64,
+}
+
+extern "C" {
+    fn next_tcsls() -> *const u8;
+    fn static_tcs_offset() -> u64;
+    fn clist_next_offset() -> u64;
+}
+
+/// Returns the location of all TCSes available at compile time in the enclave
+#[unstable(feature = "sgx_platform", issue = "56975")]
+pub fn static_tcses() -> Vec<*const u8> {
+    unsafe {
+        let mut tcsls = next_tcsls();
+        let mut tcses = Vec::new();
+
+        loop {
+            let tcs_addr = rel_ptr(*rel_ptr::<u64>(tcsls as u64 + static_tcs_offset()));
+            tcsls = *(rel_ptr::<*const u8>(tcsls as u64 + clist_next_offset()));
+
+            if tcses.first() != Some(&tcs_addr) {
+                tcses.push(tcs_addr);
+            } else {
+                break;
+            }
+        }
+        tcses
+    }
+}

Diff for: library/std/src/sys/sgx/abi/usercalls/mod.rs

@@ -139,9 +139,9 @@ pub fn connect_stream(addr: &str) -> IoResult<(Fd, String, String)> {
 
 /// Usercall `launch_thread`. See the ABI documentation for more information.
 #[unstable(feature = "sgx_platform", issue = "56975")]
-pub unsafe fn launch_thread() -> IoResult<()> {
+pub unsafe fn launch_thread(tcs: Option<Tcs>) -> IoResult<()> {
     // SAFETY: The caller must uphold the safety contract for `launch_thread`.
-    unsafe { raw::launch_thread().from_sgx_result() }
+    unsafe { raw::launch_thread(tcs).from_sgx_result() }
 }
 
 /// Usercall `exit`. See the ABI documentation for more information.

Diff for: library/std/src/sys/sgx/thread.rs

@@ -5,6 +5,7 @@ use crate::io;
 use crate::num::NonZeroUsize;
 use crate::time::Duration;
 
+use super::abi::thread;
 use super::abi::usercalls;
 
 pub struct Thread(task_queue::JoinHandle);
@@ -13,8 +14,65 @@ pub const DEFAULT_MIN_STACK_SIZE: usize = 4096;
 
 pub use self::task_queue::JoinNotifier;
 
+mod tcs_queue {
+    use super::super::abi::mem as sgx_mem;
+    use super::super::abi::thread;
+    use crate::ptr::NonNull;
+    use crate::sync::{Mutex, MutexGuard, Once};
+
+    #[derive(Clone, PartialEq, Eq, Debug)]
+    pub(super) struct Tcs {
+        address: NonNull<u8>,
+    }
+
+    impl Tcs {
+        fn new(address: NonNull<u8>) -> Tcs {
+            Tcs { address }
+        }
+
+        pub(super) fn address(&self) -> &NonNull<u8> {
+            &self.address
+        }
+    }
+
+    /// A queue of not running TCS structs
+    static mut TCS_QUEUE: Option<Mutex<Vec<Tcs>>> = None;
+    static TCS_QUEUE_INIT: Once = Once::new();
+
+    fn init_tcs_queue() -> Vec<Tcs> {
+        sgx_mem::static_tcses()
+            .iter()
+            .filter_map(|addr| if NonNull::new(*addr as _) != Some(thread::current()) {
+                    Some(Tcs::new(NonNull::new(*addr as _).expect("Compile-time value unexpected NULL")))
+                } else {
+                    None
+                })
+            .collect()
+    }
+
+    fn lock() -> MutexGuard<'static, Vec<Tcs>> {
+        unsafe {
+            TCS_QUEUE_INIT.call_once(|| TCS_QUEUE = Some(Mutex::new(init_tcs_queue())));
+            TCS_QUEUE.as_ref().unwrap().lock().unwrap()
+        }
+    }
+
+    pub(super) fn take_tcs() -> Option<Tcs> {
+        let mut tcs_queue = lock();
+        if let Some(tcs) = tcs_queue.pop() { Some(tcs) } else { None }
+    }
+
+    pub(super) fn add_tcs(tcs: Tcs) {
+        let mut tcs_queue = lock();
+        tcs_queue.insert(0, tcs);
+    }
+}
+
 mod task_queue {
+    use super::tcs_queue::{self, Tcs};
     use super::wait_notify;
+    use crate::ptr::NonNull;
+    use crate::sync::mpsc;
     use crate::sync::{Mutex, MutexGuard, Once};
 
     pub type JoinHandle = wait_notify::Waiter;
@@ -30,18 +88,24 @@ mod task_queue {
     pub(super) struct Task {
         p: Box<dyn FnOnce()>,
         done: JoinNotifier,
+        tcs: Tcs,
     }
 
     impl Task {
-        pub(super) fn new(p: Box<dyn FnOnce()>) -> (Task, JoinHandle) {
+        pub(super) fn new(tcs: Tcs, p: Box<dyn FnOnce()>) -> (Task, JoinHandle) {
             let (done, recv) = wait_notify::new();
             let done = JoinNotifier(Some(done));
-            (Task { p, done }, recv)
+            let task = Task { p, done, tcs };
+            (task, recv)
         }
 
         pub(super) fn run(self) -> JoinNotifier {
-            (self.p)();
-            self.done
+            let Task { tcs, p, done } = self;
+
+            p();
+
+            tcs_queue::add_tcs(tcs);
+            done
         }
     }
 
@@ -58,6 +122,13 @@ mod task_queue {
             TASK_QUEUE.as_ref().unwrap().lock().unwrap()
         }
     }
+
+    pub(super) fn take_task(tcs: NonNull<u8>) -> Option<Task> {
+        let mut tasks = lock();
+        let (i, _) = tasks.iter().enumerate().find(|(_i, task)| *task.tcs.address() == tcs)?;
+        let task = tasks.remove(i);
+        Some(task)
+    }
 }
 
 /// This module provides a synchronization primitive that does not use thread
@@ -105,17 +176,17 @@ pub mod wait_notify {
 impl Thread {
     // unsafe: see thread::Builder::spawn_unchecked for safety requirements
     pub unsafe fn new(_stack: usize, p: Box<dyn FnOnce()>) -> io::Result<Thread> {
-        let mut queue_lock = task_queue::lock();
-        unsafe { usercalls::launch_thread()? };
-        let (task, handle) = task_queue::Task::new(p);
-        queue_lock.push(task);
+        let tcs = tcs_queue::take_tcs().ok_or(io::Error::from(io::ErrorKind::WouldBlock))?;
+        let mut tasks = task_queue::lock();
+        unsafe { usercalls::launch_thread(Some(*tcs.address()))? };
+        let (task, handle) = task_queue::Task::new(tcs, p);
+        tasks.push(task);
         Ok(Thread(handle))
     }
 
     pub(super) fn entry() -> JoinNotifier {
-        let mut pending_tasks = task_queue::lock();
-        let task = rtunwrap!(Some, pending_tasks.pop());
-        drop(pending_tasks); // make sure to not hold the task queue lock longer than necessary
+        let task = task_queue::take_task(thread::current())
+            .expect("enclave entered through TCS unexpectedly");
         task.run()
     }