Do not update td state in the case of direct return

Signed-off-by: Ming-Wei Shih <[email protected]>
samuel-lee-msft · Dec 1, 2021 · b1b9356 · b1b9356
1 parent 09c7cb8
commit b1b9356
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 21 deletions.
diff --git a/enclave/core/sgx/exit.S b/enclave/core/sgx/exit.S
@@ -58,8 +58,9 @@ oe_asm_exit:
     mov %rdx, %r11
 
 .determine_exit_type:
-    // Check if the argument direct_return is set
-    cmp $1, %rcx
+    // Save the direct_return argument in r13 and check if it is set
+    mov %rcx, %r13
+    cmp $1, %r13
     je .return
 
     // Check the depth of the ECALL stack (zero for clean exit)
@@ -117,23 +118,33 @@ oe_asm_exit:
     mov SGX_SSA_URSP_OFFSET(%r12), %rsp
     mov SGX_SSA_URBP_OFFSET(%r12), %rbp
 
+    // Bypass update_td_state if direct_return is set
+    cmp $1, %r13
+    je .execute_eexit
+
+.update_td_state:
+    lfence
+
     // Do not update the state if the enclave exits in the middle
     // the exception handling (e.g., exiting from the first-level
-    // exception handler) or exits after an illegal instruction
-    // emulation
+    // exception handler), exits after an illegal instruction
+    // emulation, or the enclave is aborted
     cmpq $TD_STATE_SECOND_LEVEL_EXCEPTION_HANDLING, td_state(%r11)
     je .execute_eexit
     cmpq $TD_STATE_FIRST_LEVEL_EXCEPTION_HANDLING, td_previous_state(%r11)
     je .execute_eexit
+    cmpq $TD_STATE_ABORTED, td_state(%r11)
+    je .execute_eexit
+
+    lfence
 
     // Update the state to indicate that the enclave is returning
     // to the host
     movq $TD_STATE_EXITED, td_state(%r11)
 
+.execute_eexit:
     oe_cleanup_registers
 
-.execute_eexit:
-    lfence
     // Check oe_sgx_td_t.simulate flag
     // simulate-flag-check.
     mov td_simulate(%r11), %rax

diff --git a/host/sgx/enclavemanager.c b/host/sgx/enclavemanager.c
@@ -201,7 +201,7 @@ oe_enclave_t* oe_query_enclave_instance(void* tcs)
     }
 
     if (!ret)
-        OE_TRACE_ERROR("tcs=0x%x\n", tcs);
+        OE_TRACE_ERROR("tcs=0x%lx\n", tcs);
 
     return ret;
 }
diff --git a/tests/sgx/thread_interrupt/enc/enc.c b/tests/sgx/thread_interrupt/enc/enc.c
@@ -232,6 +232,26 @@ void enc_thread_interrupt_nonblocking(void)
     host_join_thread();
 }
 
+uint64_t thread_terminate_handler(oe_exception_record_t* exception_record)
+{
+    int self_tid = 0;
+
+    OE_TEST(exception_record->code == OE_EXCEPTION_UNKNOWN);
+
+    host_get_tid(&self_tid);
+    OE_TEST(_thread_info_blocking.tid == self_tid);
+
+    printf("(tid=%d) thread is interrupted...\n", self_tid);
+
+    OE_TEST(
+        _thread_info_blocking.td->state ==
+        OE_TD_STATE_SECOND_LEVEL_EXCEPTION_HANDLING);
+
+    __atomic_store_n(&_thread_info_blocking.lock, 2, __ATOMIC_RELEASE);
+
+    return OE_EXCEPTION_CONTINUE_EXECUTION;
+}
+
 void enc_run_thread_blocking(int tid)
 {
     oe_result_t result = OE_OK;
@@ -242,7 +262,7 @@ void enc_run_thread_blocking(int tid)
     printf("(tid=%d) blocking thread is running...\n", self_tid);
 
     OE_CHECK(
-        oe_add_vectored_exception_handler(false, thread_interrupt_handler));
+        oe_add_vectored_exception_handler(false, thread_terminate_handler));
 
     _thread_info_blocking.tid = tid;
     _thread_info_blocking.td = oe_sgx_get_td();
@@ -268,7 +288,7 @@ void enc_run_thread_blocking(int tid)
     return;
 }
 
-void enc_thread_interrupt_blocking(void)
+void enc_thread_interrupt_blocking(int* ret)
 {
     oe_result_t result;
     int tid = 0;
@@ -295,9 +315,9 @@ void enc_thread_interrupt_blocking(void)
     _handler_entered = 0;
     while (!_handler_entered)
     {
-        OE_TEST(
-            oe_sgx_td_register_host_signal(_thread_info_blocking.td, SIGUSR1) ==
-            true);
+        if (oe_sgx_td_register_host_signal(_thread_info_blocking.td, SIGUSR1) !=
+            true)
+            break;
 
         printf(
             "(tid=%d) Sending registered signal (SIGUSR1) to (td=0x%lx, "
@@ -307,6 +327,9 @@ void enc_thread_interrupt_blocking(void)
             _thread_info_blocking.tid,
             ++retry);
 
+        if (_thread_info_blocking.td->state != OE_TD_STATE_RUNNING)
+            break;
+
         host_send_interrupt(_thread_info_blocking.tid, SIGUSR1);
 
         if (retry == 10)
@@ -320,7 +343,8 @@ void enc_thread_interrupt_blocking(void)
         host_sleep_msec(30);
     }
 
-    OE_TEST(retry == 10);
+    if (retry != 10)
+        goto done;
 
     retry = 0;
     _handler_entered = 0;
@@ -339,6 +363,9 @@ void enc_thread_interrupt_blocking(void)
             _thread_info_blocking.tid,
             ++retry);
 
+        if (_thread_info_blocking.td->state != OE_TD_STATE_RUNNING)
+            break;
+
         host_send_interrupt(_thread_info_blocking.tid, SIGUSR2);
 
         if (retry == 10)
@@ -352,9 +379,27 @@ void enc_thread_interrupt_blocking(void)
         host_sleep_msec(30);
     }
 
-    OE_TEST(retry == 10);
+    if (retry != 10)
+        goto done;
 
-    oe_abort();
+    *ret = 1;
+
+done:
+    /* unblock and shutdown the thread */
+    OE_TEST(
+        oe_sgx_td_register_host_signal(_thread_info_blocking.td, SIGUSR2) ==
+        true);
+
+    printf(
+        "(tid=%d) Shutting down (td=0x%lx, "
+        "tid=%d)\n",
+        tid,
+        (uint64_t)_thread_info_blocking.td,
+        _thread_info_blocking.tid);
+
+    host_send_interrupt(_thread_info_blocking.tid, SIGUSR2);
+
+    host_join_thread();
 }
 
 OE_SET_ENCLAVE_SGX(

diff --git a/tests/sgx/thread_interrupt/host/host.c b/tests/sgx/thread_interrupt/host/host.c
@@ -98,11 +98,12 @@ int main(int argc, const char* argv[])
     }
     else if (!strcmp(argv[2], "blocking"))
     {
-        result = enc_thread_interrupt_blocking(enclave);
-        OE_TEST(result == OE_ENCLAVE_ABORTING);
-        /* Expcet a non-OE_OK result. The error code may be different
-         * between debug and release build. */
-        OE_TEST(oe_terminate_enclave(enclave) != OE_OK);
+        int ret = 0;
+        result = enc_thread_interrupt_blocking(enclave, &ret);
+        OE_TEST(ret == 1);
+        if (result != OE_OK)
+            oe_put_err("oe_call_enclave() failed: result=%u", result);
+        OE_TEST(oe_terminate_enclave(enclave) == OE_OK);
     }
     else
     {

diff --git a/tests/sgx/thread_interrupt/thread_interrupt.edl b/tests/sgx/thread_interrupt/thread_interrupt.edl
@@ -8,7 +8,7 @@ enclave {
 
     trusted {
         public void enc_thread_interrupt_nonblocking();
-        public void enc_thread_interrupt_blocking();
+        public void enc_thread_interrupt_blocking([in, out] int* ret);
         public void enc_run_thread_nonblocking(int tid);
         public void enc_run_thread_blocking(int tid);
     };