add a validator for aot module (#3995)

- Add AOT module validation to ensure memory constraints are met - Enable AOT validator in build configuration and update related source files
bytecodealliance · Feb 5, 2025 · c6712b4 · c6712b4
1 parent b2c7cb2
commit c6712b4
Show file tree

Hide file tree

Showing 8 changed files with 92 additions and 6 deletions.
diff --git a/build-scripts/config_common.cmake b/build-scripts/config_common.cmake
@@ -610,4 +610,8 @@ if (WAMR_BUILD_SHRUNK_MEMORY EQUAL 1)
 else ()
   add_definitions (-DWASM_ENABLE_SHRUNK_MEMORY=0)
   message ("     Shrunk memory disabled")
+endif()
+if (WAMR_BUILD_AOT_VALIDATOR EQUAL 1)
+  message ("     AOT validator enabled")
+  add_definitions (-DWASM_ENABLE_AOT_VALIDATOR=1)
 endif ()
diff --git a/core/config.h b/core/config.h
@@ -702,4 +702,8 @@
 #define WASM_ENABLE_SHRUNK_MEMORY 1
 #endif
 
+#ifndef WASM_ENABLE_AOT_VALIDATOR
+#define WASM_ENABLE_AOT_VALIDATOR 0
+#endif
+
 #endif /* end of _CONFIG_H_ */
diff --git a/core/iwasm/aot/aot_loader.c b/core/iwasm/aot/aot_loader.c
@@ -10,6 +10,9 @@
 #include "../common/wasm_native.h"
 #include "../common/wasm_loader_common.h"
 #include "../compilation/aot.h"
+#if WASM_ENABLE_AOT_VALIDATOR != 0
+#include "aot_validator.h"
+#endif
 
 #if WASM_ENABLE_DEBUG_AOT != 0
 #include "debug/elf_parser.h"
@@ -1106,9 +1109,6 @@ load_memory_info(const uint8 **p_buf, const uint8 *buf_end, AOTModule *module,
     const uint8 *buf = *p_buf;
 
     read_uint32(buf, buf_end, module->import_memory_count);
-    /* We don't support import_memory_count > 0 currently */
-    if (module->import_memory_count > 0)
-        return false;
 
     read_uint32(buf, buf_end, module->memory_count);
     total_size = sizeof(AOTMemory) * (uint64)module->memory_count;
@@ -4403,6 +4403,13 @@ aot_load_from_aot_file(const uint8 *buf, uint32 size, const LoadArgs *args,
     os_thread_jit_write_protect_np(true); /* Make memory executable */
     os_icache_flush(module->code, module->code_size);
 
+#if WASM_ENABLE_AOT_VALIDATOR != 0
+    if (!aot_module_validate(module, error_buf, error_buf_size)) {
+        aot_unload(module);
+        return NULL;
+    }
+#endif /* WASM_ENABLE_AOT_VALIDATOR != 0 */
+
     LOG_VERBOSE("Load module success.\n");
     return module;
 }

diff --git a/core/iwasm/aot/aot_perf_map.c b/core/iwasm/aot/aot_perf_map.c
@@ -7,7 +7,6 @@
 #include "bh_log.h"
 #include "bh_platform.h"
 
-#if WASM_ENABLE_LINUX_PERF != 0
 struct func_info {
     uint32 idx;
     void *ptr;
@@ -117,4 +116,3 @@ aot_create_perf_map(const AOTModule *module, char *error_buf,
 
     return ret;
 }
-#endif /* WASM_ENABLE_LINUX_PERF != 0 */
diff --git a/core/iwasm/aot/aot_validator.c b/core/iwasm/aot/aot_validator.c
@@ -0,0 +1,45 @@
+/*
+ * Copyright (C) 2025 Intel Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ */
+
+#include "aot_validator.h"
+
+static void
+set_error_buf(char *error_buf, uint32 error_buf_size, const char *string)
+{
+    if (error_buf != NULL) {
+        snprintf(error_buf, error_buf_size,
+                 "AOT module load failed: from validator. %s", string);
+    }
+}
+
+static bool
+aot_memory_info_validate(const AOTModule *module, char *error_buf,
+                         uint32 error_buf_size)
+{
+    if (module->import_memory_count > 0) {
+        set_error_buf(error_buf, error_buf_size,
+                      "import memory is not supported");
+        return false;
+    }
+
+    if (module->memory_count < 1) {
+        set_error_buf(error_buf, error_buf_size,
+                      "there should be >=1 memory in one aot module");
+        return false;
+    }
+
+    return true;
+}
+
+bool
+aot_module_validate(const AOTModule *module, char *error_buf,
+                    uint32 error_buf_size)
+{
+    if (!aot_memory_info_validate(module, error_buf, error_buf_size)) {
+        return false;
+    }
+
+    return true;
+}
diff --git a/core/iwasm/aot/aot_validator.h b/core/iwasm/aot/aot_validator.h
@@ -0,0 +1,15 @@
+/*
+ * Copyright (C) 2025 Intel Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ */
+
+#ifndef _AOT_VALIDATOR_H_
+#define _AOT_VALIDATOR_H_
+
+#include "aot_runtime.h"
+
+bool
+aot_module_validate(const AOTModule *module, char *error_buf,
+                    uint32 error_buf_size);
+
+#endif /* _AOT_VALIDATOR_H_ */
diff --git a/core/iwasm/aot/iwasm_aot.cmake b/core/iwasm/aot/iwasm_aot.cmake
@@ -7,7 +7,19 @@ add_definitions (-DWASM_ENABLE_AOT=1)
 
 include_directories (${IWASM_AOT_DIR})
 
-file (GLOB c_source_all ${IWASM_AOT_DIR}/*.c)
+list (APPEND c_source_all
+  ${IWASM_AOT_DIR}/aot_intrinsic.c
+  ${IWASM_AOT_DIR}/aot_loader.c
+  ${IWASM_AOT_DIR}/aot_runtime.c
+)
+
+if (WAMR_BUILD_LINUX_PERF EQUAL 1)
+  list (APPEND c_source_all ${IWASM_AOT_DIR}/aot_perf_map.c)
+endif ()
+
+if (WAMR_BUILD_AOT_VALIDATOR EQUAL 1)
+  list (APPEND c_source_all ${IWASM_AOT_DIR}/aot_validator.c)
+endif ()
 
 if (WAMR_BUILD_TARGET STREQUAL "X86_64" OR WAMR_BUILD_TARGET STREQUAL "AMD_64")
   set (arch_source ${IWASM_AOT_DIR}/arch/aot_reloc_x86_64.c)

diff --git a/wamr-compiler/CMakeLists.txt b/wamr-compiler/CMakeLists.txt
@@ -58,6 +58,7 @@ if (WAMR_BUILD_LLVM_LEGACY_PM EQUAL 1)
 endif ()
 
 if (LINUX)
+  set(WAMR_BUILD_LINUX_PERF 1)
   add_definitions(-DWASM_ENABLE_LINUX_PERF=1)
 endif ()