gnh1201/wasm-micro-runtime
1
0
Fork 0
mirror of https://github.com/bytecodealliance/wasm-micro-runtime.git synced 2025-09-06 18:01:08 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

wasi_nn_llamacpp.c: validate input tensor type/dimensions (#4442)

Browse Source
This commit is contained in:
YAMAMOTO Takashi 2025-07-03 11:17:19 +09:00 committed by GitHub
parent 68d5ae10d4
commit ee056d8076
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 12 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
core/iwasm/libraries/wasi-nn/src/wasi_nn_llamacpp.c
View File

@ -411,6 +411,18 @@ set_input(void *ctx, graph_execution_context exec_ctx, uint32_t index,
char *prompt_text = (char *)wasi_nn_tensor->data.buf;
uint32_t prompt_text_len = wasi_nn_tensor->data.size;
// note: buf[0] == 1 is a workaround for
// https://github.com/second-state/WasmEdge-WASINN-examples/issues/196.
// we may remove it in future.
if (wasi_nn_tensor->type != u8 || wasi_nn_tensor->dimensions->size != 1
|| !(wasi_nn_tensor->dimensions->buf[0] == 1
|| wasi_nn_tensor->dimensions->buf[0] == prompt_text_len)) {
return invalid_argument;
}
if (wasi_nn_tensor->dimensions->buf[0] == 1 && prompt_text_len != 1) {
NN_WARN_PRINTF("Ignoring seemingly wrong input tensor dimensions.");
}
#ifndef NDEBUG
NN_DBG_PRINTF("--------------------------------------------------");
NN_DBG_PRINTF("prompt_text: %.*s", (int)prompt_text_len, prompt_text);