only cancel thread if it's outside the enclave

Signed-off-by: Thomas Tendyck <[email protected]>

3rdparty/openenclave/ert.patch

@@ -600,7 +600,7 @@ index e3255a942..77351b2ac 100644
      if (!(fds = calloc(nfds, sizeof(struct oe_pollfd))))
      {
 diff --git a/host/sgx/calls.c b/host/sgx/calls.c
-index eed0c4dcf..fb791663e 100644
+index eed0c4dcf..7874c3585 100644
 --- a/host/sgx/calls.c
 +++ b/host/sgx/calls.c
 @@ -36,6 +36,8 @@
@@ -621,7 +621,19 @@ index eed0c4dcf..fb791663e 100644
      OE_CHECK(oe_safe_add_u64(
          args_ptr->input_buffer_size,
          args_ptr->output_buffer_size,
-@@ -365,22 +369,27 @@ static oe_result_t _handle_ocall(
+@@ -358,6 +362,11 @@ static oe_result_t _handle_ocall(
+         func == OE_OCALL_CALL_HOST_FUNCTION ? "EDL_OCALL" : "OE_OCALL",
+         oe_ocall_str(func));
+
++    // EDG: During an ocall, allow cancelation by the EnclaveThreadManager
++    // because this thread may block in a syscall.
++    void ert_set_cancelable(bool);
++    ert_set_cancelable(true);
++
+     switch ((oe_func_t)func)
+     {
+         case OE_OCALL_CALL_HOST_FUNCTION:
+@@ -365,22 +374,27 @@ static oe_result_t _handle_ocall(
              break;
 
          case OE_OCALL_MALLOC:
@@ -649,13 +661,14 @@ index eed0c4dcf..fb791663e 100644
              oe_handle_get_time(arg_in, arg_out);
              break;
 
-@@ -391,6 +400,10 @@ static oe_result_t _handle_ocall(
+@@ -391,6 +405,11 @@ static oe_result_t _handle_ocall(
          }
      }
 
-+    // EDG: Insert cancellation point after ocall (esp. after THREAD_WAIT) so
-+    // that EnclaveThreadManager is able to cancel enclave threads.
-+    pthread_testcancel();
++    // EDG: Disable cancelation before returning from ocall because the thread
++    // context will be swapped to the enclave values and canceling would cause a
++    // segfault then.
++    ert_set_cancelable(false);
 +
      result = OE_OK;
 

src/ert/host/enclave_thread_manager.cpp

@@ -14,6 +14,8 @@ extern "C" void HandleThreadWake(oe_enclave_t* enclave, uint64_t arg);
 
 namespace ert::host
 {
+thread_local EnclaveThreadManager::Thread* EnclaveThreadManager::self;
+
 EnclaveThreadManager::~EnclaveThreadManager()
 {
     // Process is exiting. If any enclave has not been terminated, detach its
@@ -58,7 +60,8 @@ void EnclaveThreadManager::create_thread(
         new_thread = &threads.emplace_back();
     }
 
-    new_thread->thread = thread([=, &finished = new_thread->finished] {
+    new_thread->thread = thread([=] {
+        self = new_thread;
         try
         {
             func(enclave);
@@ -67,10 +70,28 @@ void EnclaveThreadManager::create_thread(
         {
             OE_TRACE_ERROR("%s", e.what());
         }
-        finished = true;
+        new_thread->finished = true;
     });
 }
 
+void EnclaveThreadManager::set_cancelable(bool cancelable)
+{
+    // self may be null if thread is not created by
+    // EnclaveThreadManager::create_thread (e.g., the main thread)
+    if (!self)
+        return;
+
+    unsigned int cancel_state = 0;
+    if (cancelable)
+        cancel_state = self->cancel_state |= Cancelable;
+    else
+        cancel_state = self->cancel_state &= ~Cancelable;
+
+    // this function is always a cancelation point
+    if (cancel_state & Canceling)
+        pthread_exit(PTHREAD_CANCELED);
+}
+
 void EnclaveThreadManager::join_all_threads(const oe_enclave_t* enclave)
 {
     assert(enclave);
@@ -120,11 +141,15 @@ void EnclaveThreadManager::cancel_all_threads(oe_enclave_t* enclave)
     {
         if (!t.finished)
         {
-            const auto handle = t.thread.native_handle();
-            if (pthread_cancel(handle) == 0)
-                // Thread may be waiting in HandleThreadWait which is not a
-                // cancellation point, so wake the thread.
-                _wake_thread(*enclave, handle);
+            const auto cancel_state = t.cancel_state |= Canceling;
+            if (cancel_state & Cancelable)
+            {
+                const auto handle = t.thread.native_handle();
+                if (pthread_cancel(handle) == 0)
+                    // Thread may be waiting in HandleThreadWait which is not a
+                    // cancellation point, so wake the thread.
+                    _wake_thread(*enclave, handle);
+            }
         }
         t.thread.join();
     }

src/ert/host/enclave_thread_manager.h

@@ -25,6 +25,11 @@ class EnclaveThreadManager final
     // *func*.
     void create_thread(oe_enclave_t* enclave, void (*func)(oe_enclave_t*));
 
+    // Sets whether it's allowed to call pthread_cancel on the current thread.
+    // This is also a cancelation point (irrespective of the value of
+    // *cancelable*).
+    void set_cancelable(bool cancelable);
+
     // Joins all threads that have been created with create_thread() for
     // *enclave*.
     void join_all_threads(const oe_enclave_t* enclave);
@@ -36,12 +41,21 @@ class EnclaveThreadManager final
   private:
     EnclaveThreadManager() = default;
 
+    enum : unsigned int
+    {
+        Cancelable = 1 << 0,
+        Canceling = 1 << 1,
+    };
+
     struct Thread
     {
         std::thread thread;
         std::atomic<bool> finished;
+        std::atomic<unsigned int> cancel_state;
     };
 
+    static thread_local Thread* self;
+
     std::map<const oe_enclave_t*, std::list<Thread>> threads_;
     std::mutex mutex_;
 };

src/ert/host/thread.cpp

@@ -12,6 +12,18 @@
 using namespace std;
 using namespace ert;
 
+extern "C" void ert_set_cancelable(bool cancelable)
+{
+    try
+    {
+        host::EnclaveThreadManager::get_instance().set_cancelable(cancelable);
+    }
+    catch (const exception& e)
+    {
+        OE_TRACE_ERROR("%s", e.what());
+    }
+}
+
 extern "C" oe_result_t ert_join_threads_created_inside_enclave(
     oe_enclave_t* enclave)
 {