Clone data segments when specified with load args (#3463)

Follow-up on https://github.com/bytecodealliance/wasm-micro-runtime/pull/3389, specifically: https://github.com/bytecodealliance/wasm-micro-runtime/pull/3389#discussion_r1600872451

If we want to free the wasm binary buffer early, we need to clone the data segments into the module.
That's because, in case of [passive data segments](https://webassembly.github.io/threads/core/syntax/modules.html#syntax-data),
they can be referred during wasm execution.

core/iwasm/common/wasm_c_api.c

@@ -3007,13 +3007,12 @@ wasm_module_get_name(wasm_module_t *module)
 }
 
 bool
-wasm_module_is_underlying_binary_freeable(
-    const wasm_module_t *module, const struct wasm_instance_t *instance)
+wasm_module_is_underlying_binary_freeable(const wasm_module_t *module)
 {
     if (((wasm_module_ex_t *)module)->is_binary_cloned)
         return true;
 
-    return wasm_runtime_is_underlying_binary_freeable(instance->inst_comm_rt);
+    return wasm_runtime_is_underlying_binary_freeable(*module);
 }
 
 static wasm_func_t *

core/iwasm/common/wasm_runtime_common.c

@@ -7322,12 +7322,10 @@ wasm_runtime_detect_native_stack_overflow_size(WASMExecEnv *exec_env,
 }
 
 WASM_RUNTIME_API_EXTERN bool
-wasm_runtime_is_underlying_binary_freeable(const wasm_module_inst_t module_inst)
+wasm_runtime_is_underlying_binary_freeable(WASMModuleCommon *const module)
 {
-    uint32 i;
-
 #if WASM_ENABLE_INTERP != 0
-    if (module_inst->module_type == Wasm_Module_Bytecode) {
+    if (module->module_type == Wasm_Module_Bytecode) {
 #if (WASM_ENABLE_JIT != 0 || WASM_ENABLE_FAST_JIT != 0) \
     && (WASM_ENABLE_LAZY_JIT != 0)
         return false;
@@ -7335,36 +7333,25 @@ wasm_runtime_is_underlying_binary_freeable(const wasm_module_inst_t module_inst)
         return false;
 #else
         /* Fast interpreter mode */
-        WASMModule *module = ((WASMModuleInstance *)module_inst)->module;
-        if (!module->is_binary_freeable)
+        if (!((WASMModule *)module)->is_binary_freeable)
             return false;
 #if WASM_ENABLE_GC != 0 && WASM_ENABLE_STRINGREF != 0
-        if (module->string_literal_ptrs)
+        if (((WASMModule *)module)->string_literal_ptrs)
             return false;
 #endif
-#if WASM_ENABLE_BULK_MEMORY != 0
-        for (i = 0; i < module->data_seg_count; i++)
-            if (!bh_bitmap_get_bit(
-                    ((WASMModuleInstance *)module_inst)->e->common.data_dropped,
-                    i))
-                return false;
-#endif
 #endif
     }
 #endif /* WASM_ENABLE_INTERP != 0 */
 #if WASM_ENABLE_AOT != 0
-    if (module_inst->module_type == Wasm_Module_AoT) {
-        AOTModule *module =
-            (AOTModule *)((AOTModuleInstance *)module_inst)->module;
-        if (!module->is_binary_freeable)
+    if (module->module_type == Wasm_Module_AoT) {
+        if (!((AOTModule *)module)->is_binary_freeable)
             return false;
 #if WASM_ENABLE_GC != 0 && WASM_ENABLE_STRINGREF != 0
-        if (module->string_literal_ptrs)
+        if (((AOTModule *)module)->string_literal_ptrs)
             return false;
 #endif
     }
 #endif /* WASM_ENABLE_AOT != 0 */
 
-    (void)i;
     return true;
 }

core/iwasm/common/wasm_runtime_common.h

@@ -1200,8 +1200,7 @@ wasm_runtime_detect_native_stack_overflow_size(WASMExecEnv *exec_env,
                                                uint32 requested_size);
 
 WASM_RUNTIME_API_EXTERN bool
-wasm_runtime_is_underlying_binary_freeable(
-    const wasm_module_inst_t module_inst);
+wasm_runtime_is_underlying_binary_freeable(WASMModuleCommon *const module);
 
 #if WASM_ENABLE_LINUX_PERF != 0
 bool

core/iwasm/include/wasm_c_api.h

@@ -563,7 +563,7 @@ WASM_API_EXTERN void wasm_shared_module_delete(own wasm_shared_module_t*);
 WASM_API_EXTERN bool wasm_module_set_name(wasm_module_t*, const char* name);
 WASM_API_EXTERN const char *wasm_module_get_name(wasm_module_t*);
 
-WASM_API_EXTERN bool wasm_module_is_underlying_binary_freeable(const wasm_module_t *module, const struct wasm_instance_t* instance);
+WASM_API_EXTERN bool wasm_module_is_underlying_binary_freeable(const wasm_module_t *module);
 
 
 // Function Instances

core/iwasm/include/wasm_export.h

@@ -1918,12 +1918,11 @@ wasm_runtime_detect_native_stack_overflow_size(wasm_exec_env_t exec_env,
 /**
  * Query whether the wasm binary buffer used to create the module can be freed
  *
- * @param module_inst the target module instance
+ * @param module the target module
  * @return true if the wasm binary buffer can be freed
  */
 WASM_RUNTIME_API_EXTERN bool
-wasm_runtime_is_underlying_binary_freeable(
-    const wasm_module_inst_t module_inst);
+wasm_runtime_is_underlying_binary_freeable(const wasm_module_t module);
 
 #ifdef __cplusplus
 }

core/iwasm/interpreter/wasm.h

@@ -752,6 +752,7 @@ typedef struct WASMDataSeg {
     bool is_passive;
 #endif
     uint8 *data;
+    bool is_data_cloned;
 } WASMDataSeg;
 
 typedef struct BlockAddr {

core/iwasm/interpreter/wasm_loader.c

@@ -4705,8 +4705,8 @@ fail:
 
 static bool
 load_data_segment_section(const uint8 *buf, const uint8 *buf_end,
-                          WASMModule *module, char *error_buf,
-                          uint32 error_buf_size)
+                          WASMModule *module, bool clone_data_seg,
+                          char *error_buf, uint32 error_buf_size)
 {
     const uint8 *p = buf, *p_end = buf_end;
     uint32 data_seg_count, i, mem_index, data_seg_len;
@@ -4836,7 +4836,19 @@ load_data_segment_section(const uint8 *buf, const uint8 *buf_end,
 
             dataseg->data_length = data_seg_len;
             CHECK_BUF(p, p_end, data_seg_len);
-            dataseg->data = (uint8 *)p;
+            if (clone_data_seg) {
+                if (!(dataseg->data = loader_malloc(
+                          dataseg->data_length, error_buf, error_buf_size))) {
+                    return false;
+                }
+
+                bh_memcpy_s(dataseg->data, dataseg->data_length, p,
+                            data_seg_len);
+            }
+            else {
+                dataseg->data = (uint8 *)p;
+            }
+            dataseg->is_data_cloned = clone_data_seg;
             p += data_seg_len;
         }
     }
@@ -5748,8 +5760,8 @@ static void **handle_table;
 
 static bool
 load_from_sections(WASMModule *module, WASMSection *sections,
-                   bool is_load_from_file_buf, char *error_buf,
-                   uint32 error_buf_size)
+                   bool is_load_from_file_buf, bool wasm_binary_freeable,
+                   char *error_buf, uint32 error_buf_size)
 {
     WASMExport *export;
     WASMSection *section = sections;
@@ -5764,6 +5776,8 @@ load_from_sections(WASMModule *module, WASMSection *sections,
     uint32 aux_heap_base_global_index = (uint32)-1;
     WASMFuncType *func_type;
     uint8 malloc_free_io_type = VALUE_TYPE_I32;
+    bool reuse_const_strings = is_load_from_file_buf && !wasm_binary_freeable;
+    bool clone_data_seg = is_load_from_file_buf && wasm_binary_freeable;
 
     /* Find code and function sections if have */
     while (section) {
@@ -5790,7 +5804,7 @@ load_from_sections(WASMModule *module, WASMSection *sections,
             case SECTION_TYPE_USER:
                 /* unsupported user section, ignore it. */
                 if (!load_user_section(buf, buf_end, module,
-                                       is_load_from_file_buf, error_buf,
+                                       reuse_const_strings, error_buf,
                                        error_buf_size))
                     return false;
                 break;
@@ -5801,7 +5815,7 @@ load_from_sections(WASMModule *module, WASMSection *sections,
                 break;
             case SECTION_TYPE_IMPORT:
                 if (!load_import_section(buf, buf_end, module,
-                                         is_load_from_file_buf, error_buf,
+                                         reuse_const_strings, error_buf,
                                          error_buf_size))
                     return false;
                 break;
@@ -5835,7 +5849,7 @@ load_from_sections(WASMModule *module, WASMSection *sections,
                 break;
             case SECTION_TYPE_EXPORT:
                 if (!load_export_section(buf, buf_end, module,
-                                         is_load_from_file_buf, error_buf,
+                                         reuse_const_strings, error_buf,
                                          error_buf_size))
                     return false;
                 break;
@@ -5855,7 +5869,8 @@ load_from_sections(WASMModule *module, WASMSection *sections,
                     return false;
                 break;
             case SECTION_TYPE_DATA:
-                if (!load_data_segment_section(buf, buf_end, module, error_buf,
+                if (!load_data_segment_section(buf, buf_end, module,
+                                               clone_data_seg, error_buf,
                                                error_buf_size))
                     return false;
                 break;
@@ -5869,7 +5884,7 @@ load_from_sections(WASMModule *module, WASMSection *sections,
 #if WASM_ENABLE_STRINGREF != 0
             case SECTION_TYPE_STRINGREF:
                 if (!load_stringref_section(buf, buf_end, module,
-                                            is_load_from_file_buf, error_buf,
+                                            reuse_const_strings, error_buf,
                                             error_buf_size))
                     return false;
                 break;
@@ -6321,7 +6336,7 @@ wasm_loader_load_from_sections(WASMSection *section_list, char *error_buf,
     if (!module)
         return NULL;
 
-    if (!load_from_sections(module, section_list, false, error_buf,
+    if (!load_from_sections(module, section_list, false, true, error_buf,
                             error_buf_size)) {
         wasm_loader_unload(module);
         return NULL;
@@ -6494,7 +6509,7 @@ load(const uint8 *buf, uint32 size, WASMModule *module,
     }
 
     if (!create_sections(buf, size, &section_list, error_buf, error_buf_size)
-        || !load_from_sections(module, section_list, !wasm_binary_freeable,
+        || !load_from_sections(module, section_list, true, wasm_binary_freeable,
                                error_buf, error_buf_size)) {
         destroy_sections(section_list);
         return false;
@@ -6823,8 +6838,11 @@ wasm_loader_unload(WASMModule *module)
 
     if (module->data_segments) {
         for (i = 0; i < module->data_seg_count; i++) {
-            if (module->data_segments[i])
+            if (module->data_segments[i]) {
+                if (module->data_segments[i]->is_data_cloned)
+                    wasm_runtime_free(module->data_segments[i]->data);
                 wasm_runtime_free(module->data_segments[i]);
+            }
         }
         wasm_runtime_free(module->data_segments);
     }

core/iwasm/interpreter/wasm_mini_loader.c

@@ -1740,8 +1740,8 @@ load_table_segment_section(const uint8 *buf, const uint8 *buf_end,
 
 static bool
 load_data_segment_section(const uint8 *buf, const uint8 *buf_end,
-                          WASMModule *module, char *error_buf,
-                          uint32 error_buf_size)
+                          WASMModule *module, bool clone_data_seg,
+                          char *error_buf, uint32 error_buf_size)
 {
     const uint8 *p = buf, *p_end = buf_end;
     uint32 data_seg_count, i, mem_index, data_seg_len;
@@ -1851,7 +1851,19 @@ load_data_segment_section(const uint8 *buf, const uint8 *buf_end,
 
             dataseg->data_length = data_seg_len;
             CHECK_BUF(p, p_end, data_seg_len);
-            dataseg->data = (uint8 *)p;
+            if (clone_data_seg) {
+                if (!(dataseg->data = loader_malloc(
+                          dataseg->data_length, error_buf, error_buf_size))) {
+                    return false;
+                }
+
+                bh_memcpy_s(dataseg->data, dataseg->data_length, p,
+                            data_seg_len);
+            }
+            else {
+                dataseg->data = (uint8 *)p;
+            }
+            dataseg->is_data_cloned = clone_data_seg;
             p += data_seg_len;
         }
     }
@@ -2549,8 +2561,8 @@ static void **handle_table;
 
 static bool
 load_from_sections(WASMModule *module, WASMSection *sections,
-                   bool is_load_from_file_buf, char *error_buf,
-                   uint32 error_buf_size)
+                   bool is_load_from_file_buf, bool wasm_binary_freeable,
+                   char *error_buf, uint32 error_buf_size)
 {
     WASMExport *export;
     WASMSection *section = sections;
@@ -2565,6 +2577,8 @@ load_from_sections(WASMModule *module, WASMSection *sections,
     uint32 aux_heap_base_global_index = (uint32)-1;
     WASMFuncType *func_type;
     uint8 malloc_free_io_type = VALUE_TYPE_I32;
+    bool reuse_const_strings = is_load_from_file_buf && !wasm_binary_freeable;
+    bool clone_data_seg = is_load_from_file_buf && wasm_binary_freeable;
 
     /* Find code and function sections if have */
     while (section) {
@@ -2588,7 +2602,7 @@ load_from_sections(WASMModule *module, WASMSection *sections,
             case SECTION_TYPE_USER:
                 /* unsupported user section, ignore it. */
                 if (!load_user_section(buf, buf_end, module,
-                                       is_load_from_file_buf, error_buf,
+                                       reuse_const_strings, error_buf,
                                        error_buf_size))
                     return false;
                 break;
@@ -2599,7 +2613,7 @@ load_from_sections(WASMModule *module, WASMSection *sections,
                 break;
             case SECTION_TYPE_IMPORT:
                 if (!load_import_section(buf, buf_end, module,
-                                         is_load_from_file_buf, error_buf,
+                                         reuse_const_strings, error_buf,
                                          error_buf_size))
                     return false;
                 break;
@@ -2625,7 +2639,7 @@ load_from_sections(WASMModule *module, WASMSection *sections,
                 break;
             case SECTION_TYPE_EXPORT:
                 if (!load_export_section(buf, buf_end, module,
-                                         is_load_from_file_buf, error_buf,
+                                         reuse_const_strings, error_buf,
                                          error_buf_size))
                     return false;
                 break;
@@ -2645,7 +2659,8 @@ load_from_sections(WASMModule *module, WASMSection *sections,
                     return false;
                 break;
             case SECTION_TYPE_DATA:
-                if (!load_data_segment_section(buf, buf_end, module, error_buf,
+                if (!load_data_segment_section(buf, buf_end, module,
+                                               clone_data_seg, error_buf,
                                                error_buf_size))
                     return false;
                 break;
@@ -3038,7 +3053,7 @@ wasm_loader_load_from_sections(WASMSection *section_list, char *error_buf,
     if (!module)
         return NULL;
 
-    if (!load_from_sections(module, section_list, false, error_buf,
+    if (!load_from_sections(module, section_list, false, true, error_buf,
                             error_buf_size)) {
         wasm_loader_unload(module);
         return NULL;
@@ -3193,7 +3208,7 @@ load(const uint8 *buf, uint32 size, WASMModule *module,
     }
 
     if (!create_sections(buf, size, &section_list, error_buf, error_buf_size)
-        || !load_from_sections(module, section_list, !wasm_binary_freeable,
+        || !load_from_sections(module, section_list, true, wasm_binary_freeable,
                                error_buf, error_buf_size)) {
         destroy_sections(section_list);
         return false;
@@ -3332,8 +3347,11 @@ wasm_loader_unload(WASMModule *module)
 
     if (module->data_segments) {
         for (i = 0; i < module->data_seg_count; i++) {
-            if (module->data_segments[i])
+            if (module->data_segments[i]) {
+                if (module->data_segments[i]->is_data_cloned)
+                    wasm_runtime_free(module->data_segments[i]->data);
                 wasm_runtime_free(module->data_segments[i]);
+            }
         }
         wasm_runtime_free(module->data_segments);
     }

samples/basic/src/free_buffer_early.c

@@ -89,6 +89,12 @@ main(int argc, char *argv_main[])
         goto fail;
     }
 
+    if (wasm_runtime_is_underlying_binary_freeable(module)) {
+        printf("Able to free wasm binary buffer.\n");
+        wasm_runtime_free(buffer);
+        buffer = NULL;
+    }
+
     module_inst = wasm_runtime_instantiate(module, stack_size, heap_size,
                                            error_buf, sizeof(error_buf));
     if (!module_inst) {
@@ -96,12 +102,6 @@ main(int argc, char *argv_main[])
         goto fail;
     }
 
-    if (wasm_runtime_is_underlying_binary_freeable(module_inst)) {
-        printf("Able to free wasm binary buffer.\n");
-        wasm_runtime_free(buffer);
-        buffer = NULL;
-    }
-
     char *args[1] = { "3" };
     success = wasm_application_execute_func(module_inst, "mul7", 1, args);
     if (!success) {