refactor(enclave): add module and instance handle release functions for better resource management

Author: lum1n0us

.gitignore

@@ -21,9 +21,9 @@ core/iwasm/libraries/lib-wasi-threads/test/*.wasm
 core/iwasm/libraries/lib-socket/test/*.wasm
 
 product-mini/app-samples/hello-world/test.wasm
-product-mini/platforms/linux-sgx/enclave-sample/App/
-product-mini/platforms/linux-sgx/enclave-sample/Enclave/
-product-mini/platforms/linux-sgx/enclave-sample/iwasm
+product-mini/platforms/linux-sgx/enclave-sample/
+!product-mini/platforms/linux-sgx/enclave-sample/App/App.*
+!product-mini/platforms/linux-sgx/enclave-sample/Enclave/Enclave.*
 
 build_out
 tests/wamr-test-suites/workspace

product-mini/platforms/linux-sgx/enclave-sample/Enclave/Enclave.cpp

@@ -148,6 +148,23 @@ lookup_module_by_handle(uint32 handle_id)
     return module;
 }
 
+static void
+release_module_handle(uint32 handle_id)
+{
+    os_mutex_lock(&module_table_lock);
+
+    for (uint32 i = 0; i < MAX_MODULES; i++) {
+        if (module_table[i].in_use && module_table[i].id == handle_id) {
+            module_table[i].id = 0;
+            module_table[i].module_ref = NULL;
+            module_table[i].in_use = false;
+            break;
+        }
+    }
+
+    os_mutex_unlock(&module_table_lock);
+}
+
 /* SECURITY FIX: Handle table for secure wasm_module_inst_t reference management
  */
 #define MAX_INSTANCES 128
@@ -215,6 +232,23 @@ lookup_instance_by_handle(uint32 handle_id)
     return inst;
 }
 
+static void
+release_instance_handle(uint32 handle_id)
+{
+    os_mutex_lock(&instance_table_lock);
+
+    for (uint32 i = 0; i < MAX_INSTANCES; i++) {
+        if (instance_table[i].in_use && instance_table[i].id == handle_id) {
+            instance_table[i].id = 0;
+            instance_table[i].inst_ref = NULL;
+            instance_table[i].in_use = false;
+            break;
+        }
+    }
+
+    os_mutex_unlock(&instance_table_lock);
+}
+
 #if WASM_ENABLE_GLOBAL_HEAP_POOL != 0
 static char global_heap_buf[WASM_GLOBAL_HEAP_SIZE] = { 0 };
 #endif
@@ -444,7 +478,8 @@ handle_cmd_load_module(uint64 *args, uint32 argc)
 static void
 handle_cmd_unload_module(uint64 *args, uint32 argc)
 {
-    EnclaveModule *enclave_module = *(EnclaveModule **)args++;
+    uint32 module_handle_id = *(uint32 *)args++;
+    EnclaveModule *enclave_module = lookup_module_by_handle(module_handle_id);
 
     bh_assert(argc == 1);
 
@@ -473,6 +508,9 @@ handle_cmd_unload_module(uint64 *args, uint32 argc)
     os_mutex_unlock(&enclave_module_list_lock);
 #endif
 
+    /* Release module handle */
+    release_module_handle(module_handle_id);
+
     /* Destroy enclave module resources */
     if (enclave_module->wasi_arg_buf)
         wasm_runtime_free(enclave_module->wasi_arg_buf);
@@ -552,7 +590,8 @@ handle_cmd_instantiate_module(uint64 *args, uint32 argc)
 static void
 handle_cmd_deinstantiate_module(uint64 *args, uint32 argc)
 {
-    wasm_module_inst_t module_inst = *(wasm_module_inst_t *)args++;
+    uint32 instance_handle_id = *(uint32 *)args++;
+    wasm_module_inst_t module_inst = lookup_instance_by_handle(instance_handle_id);
 
     bh_assert(argc == 1);
 
@@ -562,14 +601,17 @@ handle_cmd_deinstantiate_module(uint64 *args, uint32 argc)
 
     wasm_runtime_deinstantiate(module_inst);
 
+    release_instance_handle(instance_handle_id);
+
     LOG_VERBOSE("Deinstantiate module success.\n");
 }
 
 static void
 handle_cmd_get_exception(uint64 *args, uint32 argc)
 {
     uint64 *args_org = args;
-    wasm_module_inst_t module_inst = *(wasm_module_inst_t *)args++;
+    uint32 instance_handle_id = *(uint32 *)args++;
+    wasm_module_inst_t module_inst = lookup_instance_by_handle(instance_handle_id);
     char *exception = *(char **)args++;
     uint32 exception_size = *(uint32 *)args++;
     const char *exception1;
@@ -593,7 +635,8 @@ handle_cmd_get_exception(uint64 *args, uint32 argc)
 static void
 handle_cmd_exec_app_main(uint64 *args, int32 argc)
 {
-    wasm_module_inst_t module_inst = *(wasm_module_inst_t *)args++;
+    uint32 instance_handle_id = *(uint32*)args++;
+    wasm_module_inst_t module_inst = lookup_instance_by_handle(instance_handle_id);
     uint32 app_argc = *(uint32 *)args++;
     char **app_argv = NULL;
     uint64 total_size;
@@ -626,7 +669,8 @@ handle_cmd_exec_app_main(uint64 *args, int32 argc)
 static void
 handle_cmd_exec_app_func(uint64 *args, int32 argc)
 {
-    wasm_module_inst_t module_inst = *(wasm_module_inst_t *)args++;
+    uint32 instance_handle_id = *(uint32*)args++;
+    wasm_module_inst_t module_inst = lookup_instance_by_handle(instance_handle_id);
     char *func_name = *(char **)args++;
     uint32 app_argc = *(uint32 *)args++;
     char **app_argv = NULL;
@@ -894,17 +938,19 @@ ecall_handle_command(unsigned cmd, unsigned char *cmd_buf,
                      unsigned cmd_buf_size)
 {
     /* SECURITY FIX: Validate buffer before processing */
-    if (!cmd_buf || cmd_buf_size < sizeof(uint64)) {
+    // cmd_buf can be NULL if cmd_buf_size is 0, but if cmd_buf_size is non-zero, cmd_buf must be valid
+    if ((!cmd_buf && cmd_buf_size > 0) || (cmd_buf && cmd_buf_size == 0)) {
         int bytes_written = 0;
         ocall_print(&bytes_written,
                     "SECURITY ERROR: Invalid buffer parameters\n");
         return;
     }
 
-    if (!sgx_is_outside_enclave(cmd_buf, cmd_buf_size)) {
+    // Because of [in, out] cmd_buf in edl, it is allocated inside enclave.
+    if (cmd_buf && sgx_is_outside_enclave(cmd_buf, cmd_buf_size)) {
         int bytes_written = 0;
         ocall_print(&bytes_written,
-                    "SECURITY ERROR: Buffer not outside enclave\n");
+                    "SECURITY ERROR: Buffer should be inside enclave\n");
         return;
     }
 
@@ -918,6 +964,8 @@ ecall_handle_command(unsigned cmd, unsigned char *cmd_buf,
     uint64 *args = (uint64 *)cmd_buf;
     uint32 argc = cmd_buf_size / sizeof(uint64);
 
+    LOG_VERBOSE("Received command %d with %u arguments.\n", cmd, argc);
+
     switch (cmd) {
         case CMD_INIT_RUNTIME:
             handle_cmd_init_runtime(args, argc);

product-mini/platforms/linux-sgx/enclave-sample/Enclave/Enclave.edl

@@ -18,6 +18,7 @@ enclave {
     from "sgx_tprotected_fs.edl" import *;
 #endif
 
+    //TODO: replace void with an int as error code
     trusted {
         /* define ECALLs here. */
         public void ecall_handle_command(unsigned cmd,