Create more robust shutdown sequence for cooperative scheduling

src/enclave/enclave_oe.c

@@ -216,7 +216,7 @@ static void _sgxlkl_enclave_show_attribute(const void* sgxlkl_enclave_base)
 }
 #endif
 
-void sgxlkl_ethread_init(void)
+int sgxlkl_ethread_init(void)
 {
     void* tls_page;
     __asm__ __volatile__("mov %%fs:0,%0" : "=r"(tls_page));
@@ -238,7 +238,7 @@ void sgxlkl_ethread_init(void)
     _lthread_sched_init(sgxlkl_enclave_state.config->stacksize);
     lthread_run();
 
-    return;
+    return sgxlkl_enclave_state.exit_status;
 }
 
 static void _read_eeid_config()

src/lkl/setup.c

@@ -901,7 +901,7 @@ void lkl_mount_disks(
 
     if (cwd)
     {
-        SGXLKL_VERBOSE("Set working directory: %s\n", cwd);
+        SGXLKL_VERBOSE("Set working directory: \'%s\'\n", cwd);
         int ret = lkl_sys_chdir(cwd);
         if (ret != 0)
         {
@@ -1284,16 +1284,20 @@ static void* lkl_termination_thread(void* args)
     SGXLKL_VERBOSE("calling lkl_virtio_netdev_remove()\n");
     lkl_virtio_netdev_remove();
 
-    SGXLKL_VERBOSE("calling lkl_sys_halt()\n");
-    res = lkl_sys_halt();
-    if (res < 0)
-    {
-        sgxlkl_fail("LKL halt, %s\n", lkl_strerror(res));
-    }
+    /**
+     * If kernel threads are stuck, this may block indefinitely under
+     * cooperative scheduling.
+     */
+    // SGXLKL_VERBOSE("calling lkl_sys_halt()\n");
+    // res = lkl_sys_halt();
+    // if (res < 0)
+    //     sgxlkl_fail("LKL halt, %s\n", lkl_strerror(res));
 
+    SGXLKL_VERBOSE("calling sgxlkl_host_shutdown_notification()\n");
     /* Notify host about the guest shutdown */
     sgxlkl_host_shutdown_notification();
 
+    SGXLKL_VERBOSE("calling lthread_notify_completion()\n");
     /* Set termination flag to notify lthread scheduler to bail out. */
     lthread_notify_completion();
 

src/main-oe/sgxlkl_run_oe.c

@@ -88,6 +88,10 @@ typedef struct ethread_args
     oe_enclave_t* oe_enclave;
 } ethread_args_t;
 
+static pthread_cond_t first_ethread_exited_cv;
+static pthread_mutex_t first_ethread_exited_mtx;
+static long enclave_return_status = 0;
+
 /**************************************************************************************************************************/
 
 #ifdef DEBUG
@@ -1304,18 +1308,18 @@ void parse_cpu_affinity_params(char* config, int** cores, size_t* cores_len)
     }
 }
 
-static _Atomic(bool) first_thread = true;
-
 void* ethread_init(ethread_args_t* args)
 {
-    oe_result_t result = sgxlkl_ethread_init(args->oe_enclave);
-    bool current_value = true;
-    if (sgxlkl_config_bool(SGXLKL_VERBOSE) &&
-        atomic_compare_exchange_strong(&first_thread, &current_value, false))
-    {
-        sgxlkl_host_verbose("");
-    }
-    sgxlkl_host_verbose_raw("ethread (%i: %u) ", args->ethread_id, result);
+    int exit_status = 0;
+    oe_result_t result = sgxlkl_ethread_init(args->oe_enclave, &exit_status);
+
+    pthread_mutex_lock(&first_ethread_exited_mtx);
+    int ret = pthread_cond_signal(&first_ethread_exited_cv);
+    if (ret != 0)
+        sgxlkl_host_fail("Failed to signal enclave termination: ret=%i\n", ret);
+    enclave_return_status = exit_status;
+    pthread_mutex_unlock(&first_ethread_exited_mtx);
+
     if (result != OE_OK)
     {
         sgxlkl_host_fail(
@@ -1331,17 +1335,17 @@ void* enclave_init(ethread_args_t* args)
 {
     sgxlkl_host_verbose(
         "sgxlkl_enclave_init(ethread_id=%i)\n", args->ethread_id);
-    int exit_status;
+
+    int exit_status = 0;
     oe_result_t result =
         sgxlkl_enclave_init(args->oe_enclave, &exit_status, args->shm);
-    bool current_value = true;
-    if (sgxlkl_config_bool(SGXLKL_VERBOSE) &&
-        atomic_compare_exchange_strong(&first_thread, &current_value, false))
-    {
-        sgxlkl_host_verbose("");
-    }
-    sgxlkl_host_verbose_raw(
-        "init (%i: %u exit=%i) ", args->ethread_id, result, exit_status);
+
+    pthread_mutex_lock(&first_ethread_exited_mtx);
+    int ret = pthread_cond_signal(&first_ethread_exited_cv);
+    if (ret != 0)
+        sgxlkl_host_fail("Failed to signal enclave termination: ret=%i\n", ret);
+    enclave_return_status = exit_status;
+    pthread_mutex_unlock(&first_ethread_exited_mtx);
 
     if (result != OE_OK)
     {
@@ -1351,8 +1355,7 @@ void* enclave_init(ethread_args_t* args)
             result,
             oe_result_str(result));
     }
-
-    return (void*)(long)exit_status;
+    return NULL;
 }
 
 /* Creates an SGX-LKL enclave with enclave configuration in the EEID. */
@@ -1694,7 +1697,6 @@ int main(int argc, char* argv[], char* envp[])
     size_t ethreads_cores_len;
     pthread_attr_t eattr;
     cpu_set_t set;
-    void* return_value;
     bool enclave_image_provided = false;
 
     oe_enclave_t* oe_enclave = NULL;
@@ -1994,6 +1996,9 @@ int main(int argc, char* argv[], char* envp[])
     __gdb_hook_starter_ready(base_addr, econf->mode, libsgxlkl);
 #endif
 
+    pthread_mutex_init(&first_ethread_exited_mtx, NULL);
+    pthread_cond_init(&first_ethread_exited_cv, NULL);
+
     ethread_args_t ethreads_args[econf->ethreads];
 
     for (int i = 0; i < econf->ethreads; i++)
@@ -2032,24 +2037,26 @@ int main(int argc, char* argv[], char* envp[])
         pthread_setname_np(sgxlkl_threads[i], "ENCLAVE");
     }
 
-    long exit_status = 0;
-
-    for (int i = 0; i < econf->ethreads; i++)
-    {
-        pthread_join(sgxlkl_threads[i], &return_value);
-        if (i == 0)
-        {
-            exit_status = (long)return_value;
-        }
-    }
-    sgxlkl_host_verbose_raw("\n");
-
-    if (oe_enclave)
-    {
-        sgxlkl_host_verbose("oe_terminate_enclave... ");
-        oe_terminate_enclave(oe_enclave);
-        sgxlkl_host_verbose_raw("done\n");
-    }
+    pthread_mutex_lock(&first_ethread_exited_mtx);
+    ret =
+        pthread_cond_wait(&first_ethread_exited_cv, &first_ethread_exited_mtx);
+    if (ret != 0)
+        sgxlkl_host_fail("Failed to wait for enclave to finish: ret=%i\n", ret);
+    long exit_status = (long)enclave_return_status;
+    pthread_mutex_unlock(&first_ethread_exited_mtx);
+
+    /**
+     * This fails with an error when we try to terminate the enclave while
+     * ethreads are still executing (?). We need to figure out a way to shutdown
+     * the enclave under coperative scheduling, i.e., when an ethread may be
+     * stuck executing an application busy-loop.
+     */
+    // if (oe_enclave)
+    // {
+    //     sgxlkl_host_verbose("oe_terminate_enclave... ");
+    //     oe_terminate_enclave(oe_enclave);
+    //     sgxlkl_host_verbose_raw("done\n");
+    // }
 
     sgxlkl_host_verbose("SGX-LKL-OE exit: exit_status=%i\n", exit_status);
     return exit_status;

src/sched/lthread.c

@@ -235,6 +235,17 @@ void lthread_run(void)
                     "[%4d] lthread_run(): lthread_resume (dequeue)\n",
                     lt ? lt->tid : -1);
                 _lthread_resume(lt);
+
+                // Break out of scheduler loop when previous thread triggered
+                // termination
+                if (_lthread_should_stop)
+                {
+                    SGXLKL_TRACE_THREAD(
+                        "[%4d] lthread_run(): quitting\n", lt ? lt->tid : -1);
+                    // We only need this ethread to leave the enclave
+                    _lthread_should_stop = false;
+                    return;
+                }
             }
 
             if (vio_enclave_wakeup_event_channel())
@@ -246,12 +257,6 @@ void lthread_run(void)
             spins--;
             if (spins <= 0)
             {
-                /* Do not handle futexes when enclave is terminating */
-                if (_lthread_should_stop)
-                {
-                    break;
-                }
-
                 futex_tick();
                 spins = futex_wake_spins;
             }
@@ -265,14 +270,6 @@ void lthread_run(void)
             /* sleep outside the enclave */
             sgxlkl_host_idle_ethread(sleeptime_ns);
         }
-
-        /* Break out of scheduler loop when enclave is terminating */
-        if (_lthread_should_stop)
-        {
-            SGXLKL_TRACE_THREAD(
-                "[%4d] lthread_run(): quitting\n", lt ? lt->tid : -1);
-            break;
-        }
     }
 }
 

src/sgxlkl.edl

@@ -29,7 +29,7 @@ enclave {
             [in] const sgxlkl_shared_memory_t* shared_memory);
 
         // Enclave call for initializing ethreads to enter enclave
-        public void sgxlkl_ethread_init(void);
+        public int sgxlkl_ethread_init(void);
 
         // Enclave call to dump stack traces for all lthreads (DEBUG only)
         // TODO: This should only be included for a DEBUG build, but EDL doesn't seem to