Skip to content

Commit 6b51c61

Browse files
authored
wasi-nn: fix set_input memory address validation for the legacy abi (#4534)
cf. #3223
1 parent 5b63f35 commit 6b51c61

File tree

1 file changed

+10
-1
lines changed

1 file changed

+10
-1
lines changed

core/iwasm/libraries/wasi-nn/src/utils/wasi_nn_app_native.c

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,9 +108,18 @@ tensor_data_app_native(wasm_module_inst_t instance, uint32_t total_elements,
108108
#define data_size total_elements
109109
#endif
110110

111+
uint64 data_size_in_bytes = data_size;
112+
#if WASM_ENABLE_WASI_EPHEMERAL_NN == 0
113+
data_size_in_bytes *= sizeof(float);
114+
if (data_size_in_bytes / sizeof(float) != data_size) {
115+
/* overflow */
116+
return invalid_argument;
117+
}
118+
#endif
119+
111120
if (!wasm_runtime_validate_app_addr(instance,
112121
(uint64)input_tensor_wasm->data_offset,
113-
(uint64)data_size)) {
122+
data_size_in_bytes)) {
114123
NN_ERR_PRINTF("input_tensor_wasm->data_offset is invalid");
115124
return invalid_argument;
116125
}

0 commit comments

Comments
 (0)