From 709127d631aeef901ecca617439766935abeecc6 Mon Sep 17 00:00:00 2001 From: Enrico Loparco Date: Tue, 5 Sep 2023 10:41:52 +0200 Subject: [PATCH] Add callback to handle memory.grow failures (#2522) When embedding WAMR, this PR allows to register a callback that is invoked when memory.grow fails. In case of memory allocation failures, some languages allow to handle the error (e.g. by checking the return code of malloc/calloc in C), some others (e.g. Rust) just panic. --- core/iwasm/aot/aot_runtime.c | 3 + core/iwasm/common/wasm_memory.c | 107 ++++++++++++++++++++++---- core/iwasm/common/wasm_memory.h | 4 + core/iwasm/include/wasm_export.h | 17 ++++ core/iwasm/interpreter/wasm_runtime.c | 3 + core/iwasm/interpreter/wasm_runtime.h | 2 + 6 files changed, 119 insertions(+), 17 deletions(-) diff --git a/core/iwasm/aot/aot_runtime.c b/core/iwasm/aot/aot_runtime.c index 2c647d345..55c399a8b 100644 --- a/core/iwasm/aot/aot_runtime.c +++ b/core/iwasm/aot/aot_runtime.c @@ -1443,6 +1443,9 @@ aot_call_function(WASMExecEnv *exec_env, AOTFunctionInstance *function, /* set thread handle and stack boundary */ wasm_exec_env_set_thread_info(exec_env); + /* set exec env so it can be later retrieved from instance */ + ((AOTModuleInstanceExtra *)module_inst->e)->common.cur_exec_env = exec_env; + if (ext_ret_count > 0) { uint32 cell_num = 0, i; uint8 *ext_ret_types = func_type->types + func_type->param_count + 1; diff --git a/core/iwasm/common/wasm_memory.c b/core/iwasm/common/wasm_memory.c index 7d5cb4353..a083a214e 100644 --- a/core/iwasm/common/wasm_memory.c +++ b/core/iwasm/common/wasm_memory.c @@ -8,6 +8,7 @@ #include "../aot/aot_runtime.h" #include "bh_platform.h" #include "mem_alloc.h" +#include "wasm_memory.h" #if WASM_ENABLE_SHARED_MEMORY != 0 #include "../common/wasm_shared_memory.h" @@ -24,6 +25,8 @@ static Memory_Mode memory_mode = MEMORY_MODE_UNKNOWN; static mem_allocator_t pool_allocator = NULL; +static enlarge_memory_error_callback_t enlarge_memory_error_cb; + #if WASM_MEM_ALLOC_WITH_USER_DATA != 0 static void *allocator_user_data = NULL; static void *(*malloc_func)(void *user_data, unsigned int size) = NULL; @@ -570,13 +573,16 @@ wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count) { WASMMemoryInstance *memory = wasm_get_default_memory(module); uint8 *memory_data_old, *memory_data_new, *heap_data_old; - uint32 num_bytes_per_page, heap_size, total_size_old; + uint32 num_bytes_per_page, heap_size, total_size_old = 0; uint32 cur_page_count, max_page_count, total_page_count; uint64 total_size_new; bool ret = true; + enlarge_memory_error_reason_t failure_reason = INTERNAL_ERROR; - if (!memory) - return false; + if (!memory) { + ret = false; + goto return_func; + } heap_data_old = memory->heap_data; heap_size = (uint32)(memory->heap_data_end - memory->heap_data); @@ -594,9 +600,15 @@ wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count) /* No need to enlarge memory */ return true; - if (total_page_count < cur_page_count /* integer overflow */ - || total_page_count > max_page_count) { - return false; + if (total_page_count < cur_page_count) { /* integer overflow */ + ret = false; + goto return_func; + } + + if (total_page_count > max_page_count) { + failure_reason = MAX_SIZE_REACHED; + ret = false; + goto return_func; } bh_assert(total_size_new <= 4 * (uint64)BH_GB); @@ -622,14 +634,16 @@ wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count) if (heap_size > 0) { if (mem_allocator_is_heap_corrupted(memory->heap_handle)) { wasm_runtime_show_app_heap_corrupted_prompt(); - return false; + ret = false; + goto return_func; } } if (!(memory_data_new = wasm_runtime_realloc(memory_data_old, (uint32)total_size_new))) { if (!(memory_data_new = wasm_runtime_malloc((uint32)total_size_new))) { - return false; + ret = false; + goto return_func; } if (memory_data_old) { bh_memcpy_s(memory_data_new, (uint32)total_size_new, @@ -685,6 +699,26 @@ wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count) os_writegsbase(memory_data_new); #endif +return_func: + if (!ret && enlarge_memory_error_cb) { + WASMExecEnv *exec_env = NULL; + +#if WASM_ENABLE_INTERP != 0 + if (module->module_type == Wasm_Module_Bytecode) + exec_env = + ((WASMModuleInstanceExtra *)module->e)->common.cur_exec_env; +#endif +#if WASM_ENABLE_AOT != 0 + if (module->module_type == Wasm_Module_AoT) + exec_env = + ((AOTModuleInstanceExtra *)module->e)->common.cur_exec_env; +#endif + + enlarge_memory_error_cb(inc_page_count, total_size_old, 0, + failure_reason, + (WASMModuleInstanceCommon *)module, exec_env); + } + return ret; } #else @@ -692,12 +726,16 @@ bool wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count) { WASMMemoryInstance *memory = wasm_get_default_memory(module); - uint32 num_bytes_per_page, total_size_old; + uint32 num_bytes_per_page, total_size_old = 0; uint32 cur_page_count, max_page_count, total_page_count; uint64 total_size_new; + bool ret = true; + enlarge_memory_error_reason_t failure_reason = INTERNAL_ERROR; - if (!memory) - return false; + if (!memory) { + ret = false; + goto return_func; + } num_bytes_per_page = memory->num_bytes_per_page; cur_page_count = memory->cur_page_count; @@ -710,9 +748,15 @@ wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count) /* No need to enlarge memory */ return true; - if (total_page_count < cur_page_count /* integer overflow */ - || total_page_count > max_page_count) { - return false; + if (total_page_count < cur_page_count) { /* integer overflow */ + ret = false; + goto return_func; + } + + if (total_page_count > max_page_count) { + failure_reason = MAX_SIZE_REACHED; + ret = false; + goto return_func; } bh_assert(total_size_new <= 4 * (uint64)BH_GB); @@ -727,7 +771,8 @@ wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count) if (!os_mem_commit(memory->memory_data_end, (uint32)total_size_new - total_size_old, MMAP_PROT_READ | MMAP_PROT_WRITE)) { - return false; + ret = false; + goto return_func; } #endif @@ -739,7 +784,8 @@ wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count) os_mem_decommit(memory->memory_data_end, (uint32)total_size_new - total_size_old); #endif - return false; + ret = false; + goto return_func; } /* The increased pages are filled with zero by the OS when os_mmap, @@ -759,10 +805,37 @@ wasm_enlarge_memory_internal(WASMModuleInstance *module, uint32 inc_page_count) memory->mem_bound_check_16bytes.u64 = total_size_new - 16; #endif - return true; +return_func: + if (!ret && enlarge_memory_error_cb) { + WASMExecEnv *exec_env = NULL; + +#if WASM_ENABLE_INTERP != 0 + if (module->module_type == Wasm_Module_Bytecode) + exec_env = + ((WASMModuleInstanceExtra *)module->e)->common.cur_exec_env; +#endif +#if WASM_ENABLE_AOT != 0 + if (module->module_type == Wasm_Module_AoT) + exec_env = + ((AOTModuleInstanceExtra *)module->e)->common.cur_exec_env; +#endif + + enlarge_memory_error_cb(inc_page_count, total_size_old, 0, + failure_reason, + (WASMModuleInstanceCommon *)module, exec_env); + } + + return ret; } #endif /* end of OS_ENABLE_HW_BOUND_CHECK */ +void +wasm_runtime_set_enlarge_mem_error_callback( + const enlarge_memory_error_callback_t callback) +{ + enlarge_memory_error_cb = callback; +} + bool wasm_enlarge_memory(WASMModuleInstance *module, uint32 inc_page_count) { diff --git a/core/iwasm/common/wasm_memory.h b/core/iwasm/common/wasm_memory.h index 516ee8497..cf7d7ffcd 100644 --- a/core/iwasm/common/wasm_memory.h +++ b/core/iwasm/common/wasm_memory.h @@ -24,6 +24,10 @@ wasm_runtime_memory_destroy(); unsigned wasm_runtime_memory_pool_size(); +void +wasm_runtime_set_enlarge_mem_error_callback( + const enlarge_memory_error_callback_t callback); + #ifdef __cplusplus } #endif diff --git a/core/iwasm/include/wasm_export.h b/core/iwasm/include/wasm_export.h index 18e15782f..79edc68ed 100644 --- a/core/iwasm/include/wasm_export.h +++ b/core/iwasm/include/wasm_export.h @@ -1439,6 +1439,23 @@ WASM_RUNTIME_API_EXTERN bool wasm_runtime_is_import_global_linked(const char *module_name, const char *global_name); +typedef enum { + INTERNAL_ERROR, + MAX_SIZE_REACHED, +} enlarge_memory_error_reason_t; + +typedef void (*enlarge_memory_error_callback_t)( + uint32_t inc_page_count, uint64_t current_memory_size, + uint32_t memory_index, enlarge_memory_error_reason_t failure_reason, + wasm_module_inst_t instance, wasm_exec_env_t exec_env); + +/** + * Setup callback invoked when memory.grow fails + */ +WASM_RUNTIME_API_EXTERN void +wasm_runtime_set_enlarge_mem_error_callback( + const enlarge_memory_error_callback_t callback); + /* clang-format on */ #ifdef __cplusplus diff --git a/core/iwasm/interpreter/wasm_runtime.c b/core/iwasm/interpreter/wasm_runtime.c index b71b608a2..f395b1eaa 100644 --- a/core/iwasm/interpreter/wasm_runtime.c +++ b/core/iwasm/interpreter/wasm_runtime.c @@ -2400,6 +2400,9 @@ wasm_call_function(WASMExecEnv *exec_env, WASMFunctionInstance *function, /* set thread handle and stack boundary */ wasm_exec_env_set_thread_info(exec_env); + /* set exec env so it can be later retrieved from instance */ + module_inst->e->common.cur_exec_env = exec_env; + interp_call_wasm(module_inst, exec_env, function, argc, argv); return !wasm_copy_exception(module_inst, NULL); } diff --git a/core/iwasm/interpreter/wasm_runtime.h b/core/iwasm/interpreter/wasm_runtime.h index 784febd44..aa33dc18c 100644 --- a/core/iwasm/interpreter/wasm_runtime.h +++ b/core/iwasm/interpreter/wasm_runtime.h @@ -213,6 +213,8 @@ typedef struct CApiFuncImport { /* The common part of WASMModuleInstanceExtra and AOTModuleInstanceExtra */ typedef struct WASMModuleInstanceExtraCommon { CApiFuncImport *c_api_func_imports; + /* pointer to the exec env currently used */ + WASMExecEnv *cur_exec_env; #if WASM_CONFIGUABLE_BOUNDS_CHECKS != 0 /* Disable bounds checks or not */ bool disable_bounds_checks;