From 35972493d3016f46894ee29f16bde115d675c148 Mon Sep 17 00:00:00 2001
From: "liang.he@intel.com" <liang.he@intel.com>
Date: Tue, 31 Dec 2024 00:37:16 +0000
Subject: [PATCH] Add AOT module validation to ensure memory constraints are
 met

---
 core/iwasm/aot/aot_loader.c    | 15 +++++++-------
 core/iwasm/aot/aot_validator.c | 37 ++++++++++++++++++++++++++++++++++
 core/iwasm/aot/aot_validator.h | 15 ++++++++++++++
 3 files changed, 59 insertions(+), 8 deletions(-)
 create mode 100644 core/iwasm/aot/aot_validator.c
 create mode 100644 core/iwasm/aot/aot_validator.h

diff --git a/core/iwasm/aot/aot_loader.c b/core/iwasm/aot/aot_loader.c
index 353f787540..d783f2f565 100644
--- a/core/iwasm/aot/aot_loader.c
+++ b/core/iwasm/aot/aot_loader.c
@@ -11,6 +11,7 @@
 #include "../common/wasm_loader_common.h"
 #include "../compilation/aot.h"
 #include "aot_utility.h"
+#include "aot_validator.h"
 
 #if WASM_ENABLE_DEBUG_AOT != 0
 #include "debug/elf_parser.h"
@@ -1084,16 +1085,8 @@ 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);
-    if (module->import_memory_count > 0) {
-        LOG_WARNING("Don't support >= 0 import memory");
-        return false;
-    }
 
     read_uint32(buf, buf_end, module->memory_count);
-    if (module->memory_count < 1) {
-        LOG_WARNING("Always assume >= 1 memory");
-        return false;
-    }
 
     total_size = sizeof(AOTMemory) * (uint64)module->memory_count;
     if (!(module->memories =
@@ -4387,6 +4380,12 @@ 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);
 
+    /*TODO: use a CLI option to control? */
+    if (!aot_module_validate(module, error_buf, error_buf_size)) {
+        aot_unload(module);
+        return NULL;
+    }
+
     LOG_VERBOSE("Load module success.\n");
     return module;
 }
diff --git a/core/iwasm/aot/aot_validator.c b/core/iwasm/aot/aot_validator.c
new file mode 100644
index 0000000000..500804bcd2
--- /dev/null
+++ b/core/iwasm/aot/aot_validator.c
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2019 Intel Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ */
+
+#include "aot_validator.h"
+#include "aot_utility.h"
+
+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;
+}
\ No newline at end of file
diff --git a/core/iwasm/aot/aot_validator.h b/core/iwasm/aot/aot_validator.h
new file mode 100644
index 0000000000..baec27640a
--- /dev/null
+++ b/core/iwasm/aot/aot_validator.h
@@ -0,0 +1,15 @@
+/*
+ * Copyright (C) 2019 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_ */