Skip to content

Commit

Permalink
Add AOT module validation to ensure memory constraints are met
Browse files Browse the repository at this point in the history
  • Loading branch information
lum1n0us committed Dec 31, 2024
1 parent 7f3e0df commit 716fcb7
Show file tree
Hide file tree
Showing 3 changed files with 67 additions and 3 deletions.
10 changes: 7 additions & 3 deletions core/iwasm/aot/aot_loader.c
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include "../common/wasm_native.h"
#include "../common/wasm_loader_common.h"
#include "../compilation/aot.h"
#include "aot_validator.h"

#if WASM_ENABLE_DEBUG_AOT != 0
#include "debug/elf_parser.h"
Expand Down Expand Up @@ -1106,9 +1107,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;
Expand Down Expand Up @@ -4403,6 +4401,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;
}
Expand Down
45 changes: 45 additions & 0 deletions core/iwasm/aot/aot_validator.c
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
/*
* Copyright (C) 2019 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;
}
15 changes: 15 additions & 0 deletions core/iwasm/aot/aot_validator.h
Original file line number Diff line number Diff line change
@@ -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_ */

0 comments on commit 716fcb7

Please sign in to comment.