Skip to content

Commit ee056d8

Browse files
authored
wasi_nn_llamacpp.c: validate input tensor type/dimensions (#4442)
1 parent 68d5ae1 commit ee056d8

File tree

1 file changed

+12
-0
lines changed

1 file changed

+12
-0
lines changed

core/iwasm/libraries/wasi-nn/src/wasi_nn_llamacpp.c

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -411,6 +411,18 @@ set_input(void *ctx, graph_execution_context exec_ctx, uint32_t index,
411411
char *prompt_text = (char *)wasi_nn_tensor->data.buf;
412412
uint32_t prompt_text_len = wasi_nn_tensor->data.size;
413413

414+
// note: buf[0] == 1 is a workaround for
415+
// https://github.com/second-state/WasmEdge-WASINN-examples/issues/196.
416+
// we may remove it in future.
417+
if (wasi_nn_tensor->type != u8 || wasi_nn_tensor->dimensions->size != 1
418+
|| !(wasi_nn_tensor->dimensions->buf[0] == 1
419+
|| wasi_nn_tensor->dimensions->buf[0] == prompt_text_len)) {
420+
return invalid_argument;
421+
}
422+
if (wasi_nn_tensor->dimensions->buf[0] == 1 && prompt_text_len != 1) {
423+
NN_WARN_PRINTF("Ignoring seemingly wrong input tensor dimensions.");
424+
}
425+
414426
#ifndef NDEBUG
415427
NN_DBG_PRINTF("--------------------------------------------------");
416428
NN_DBG_PRINTF("prompt_text: %.*s", (int)prompt_text_len, prompt_text);

0 commit comments

Comments
 (0)