Modify AEP for AEX to exit hanged enclave thread #165

intel-sgx/enclave-runner/src/tcs.rs

@@ -14,6 +14,9 @@ use std::os::raw::c_void;
 
 use sgx_isa::Enclu;
 use sgxs::loader::Tcs;
+use std::sync::Arc;
+use crate::usercalls::EnclaveState;
+use std::sync::atomic::{AtomicBool, Ordering};
 
 pub(crate) type DebugBuffer = [u8; 1024];
 
@@ -40,8 +43,9 @@ impl<T: Tcs> Usercall<T> {
         self,
         retval: (u64, u64),
         debug_buf: Option<&RefCell<DebugBuffer>>,
+        exiting: Arc<AtomicBool>
     ) -> ThreadResult<T> {
-        coenter(self.tcs, 0, retval.0, retval.1, 0, 0, debug_buf)
+        coenter(self.tcs, 0, retval.0, retval.1, 0, 0, debug_buf, exiting.clone())
     }
 
     pub fn tcs_address(&self) -> *mut c_void {
@@ -57,6 +61,7 @@ pub(crate) fn coenter<T: Tcs>(
     mut p4: u64,
     mut p5: u64,
     debug_buf: Option<&RefCell<DebugBuffer>>,
+    exiting: Arc<AtomicBool>
 ) -> ThreadResult<T> {
     /// Check if __vdso_sgx_enter_enclave exists. We're using weak linkage, so
     /// it might not.
@@ -183,13 +188,27 @@ pub(crate) fn coenter<T: Tcs>(
                 panic!("Error entering enclave (VDSO): ret = {:?}, run = {:?}", std::io::Error::from_raw_os_error(-ret), run);
             }
         } else {
+            let exiting_ptr: *mut bool = (*exiting).as_ptr();
+
             asm!("
+                    push %r14
                     lea 1f(%rip), %rcx // set SGX AEP
                     xchg {0}, %rbx
-1:                  enclu
-                    xchg %rbx, {0}
+                    jmp 2f
+
+1:                  pop %r14
+                    mov (%r14), %r15
+                    test %r15, %r15
+                    jne 3f
+                    push %r14
+
+2:                  enclu
+                    pop %r14
+3:
                 ",
                 inout(reg) tcs.address() => _, // rbx is used internally by LLVM and cannot be used as an operand for inline asm (#84658)
+                in("r14") exiting_ptr,
+                out("r15") _,
                 inout("eax") Enclu::EEnter as u32 => sgx_result,
                 out("rcx") _,
                 inout("rdx") p3,
@@ -201,6 +220,16 @@ pub(crate) fn coenter<T: Tcs>(
                 out("r11") _,
                 options(nostack, att_syntax)
             );
+
+            if sgx_result == (Enclu::EExit as u32) {
+                if crate::usercalls::abi::UsercallList::exit as u64 == p1 {
+                    (*exiting).store(true, Ordering::Relaxed);
+                }
+            } else if sgx_result == (Enclu::EResume as u32) {
+                // AEX after main thread has exited, issue #165
+                // TODO: what to return to propagate it as EnclaveAbort
+                return CoResult::Return((tcs, p2, p3))
+            }
         }
     };
 

intel-sgx/enclave-runner/src/usercalls/mod.rs

@@ -647,7 +647,7 @@ pub(crate) struct EnclaveState {
     event_queues: FnvHashMap<TcsAddress, PendingEvents>,
     fds: Mutex<FnvHashMap<Fd, Arc<AsyncFileDesc>>>,
     last_fd: AtomicUsize,
-    exiting: AtomicBool,
+    exiting: Arc<AtomicBool>,
     usercall_ext: Box<dyn UsercallExtension>,
     threads_queue: crossbeam::queue::SegQueue<StoppedTcs>,
     forward_panics: bool,
@@ -667,15 +667,15 @@ enum CoEntry {
 }
 
 impl Work {
-    fn do_work(self, io_send_queue: &mut tokio::sync::mpsc::UnboundedSender<UsercallSendData>) {
+    fn do_work(self, io_send_queue: &mut tokio::sync::mpsc::UnboundedSender<UsercallSendData>, exiting: Arc<AtomicBool>) {
         let buf = RefCell::new([0u8; 1024]);
         let usercall_send_data = match self.entry {
             CoEntry::Initial(erased_tcs, p1, p2, p3, p4, p5) => {
-                let coresult = tcs::coenter(erased_tcs, p1, p2, p3, p4, p5, Some(&buf));
+                let coresult = tcs::coenter(erased_tcs, p1, p2, p3, p4, p5, Some(&buf), exiting);
                 UsercallSendData::Sync(coresult, self.tcs, buf)
             }
             CoEntry::Resume(usercall, coresult) => {
-                let coresult = usercall.coreturn(coresult, Some(&buf));
+                let coresult = usercall.coreturn(coresult, Some(&buf), exiting);
                 UsercallSendData::Sync(coresult, self.tcs, buf)
             }
         };
@@ -739,7 +739,7 @@ impl EnclaveState {
             event_queues,
             fds: Mutex::new(fds),
             last_fd,
-            exiting: AtomicBool::new(false),
+            exiting: Arc::new(AtomicBool::new(false)),
             usercall_ext,
             threads_queue,
             forward_panics,
@@ -964,15 +964,16 @@ impl EnclaveState {
             num_of_worker_threads: usize,
             work_receiver: crossbeam::crossbeam_channel::Receiver<Work>,
             io_queue_send: tokio::sync::mpsc::UnboundedSender<UsercallSendData>,
+            exiting: Arc<AtomicBool>
         ) -> Vec<JoinHandle<()>> {
             let mut thread_handles = vec![];
             for _ in 0..num_of_worker_threads {
                 let work_receiver = work_receiver.clone();
                 let mut io_queue_send = io_queue_send.clone();
-
+                let exiting_clone = exiting.clone();
                 thread_handles.push(thread::spawn(move || {
                     while let Ok(work) = work_receiver.recv() {
-                        work.do_work(&mut io_queue_send);
+                        work.do_work(&mut io_queue_send, exiting_clone.clone());
                     }
                 }));
             }
@@ -987,7 +988,7 @@ impl EnclaveState {
             .expect("Work sender couldn't send data to receiver");
 
         let join_handlers =
-            create_worker_threads(num_of_worker_threads, work_receiver, io_queue_send.clone());
+            create_worker_threads(num_of_worker_threads, work_receiver, io_queue_send.clone(), enclave.exiting.clone());
         // main syscall polling loop
         let main_result =
             EnclaveState::syscall_loop(enclave.clone(), io_queue_receive, io_queue_send, work_sender);
@@ -1027,7 +1028,9 @@ impl EnclaveState {
             entry: CoEntry::Initial(main.tcs, argv as _, argc as _, 0, 0, 0),
         };
 
-        let num_of_worker_threads = num_cpus::get();
+        // TODO: issue #483 if only one worker thread it might be blocked
+        //    by infinite loop
+        let num_of_worker_threads = num_cpus::get() + 1;
 
         let kind = EnclaveKind::Command(Command {
             panic_reason: Mutex::new(PanicReason {