gnh1201/wasm-micro-runtime
1
0
Fork 0
mirror of https://github.com/bytecodealliance/wasm-micro-runtime.git synced 2026-03-03 01:51:32 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

Add bounds checking for output tensor buffer in wasi-nn llama.cpp

Browse Source 
The get_output function copies LLM output into output_tensor->buf
without checking against output_tensor->size, allowing writes
past the buffer when the model generates output longer than the
caller-provided buffer. Add size checks for both the metadata
path and the token output loop.
This commit is contained in:
Yi LIU 2026-02-25 23:05:35 +08:00
parent c46b10dcbc
commit acd36fc0ab
No known key found for this signature in database
GPG Key ID: 138FFD9517675924
1 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -623,8 +623,11 @@ get_output(void *ctx, graph_execution_context exec_ctx, uint32_t index,
printf("%s\n", output_metadata);
}
memcpy(output_tensor->buf, output_metadata, strlen(output_metadata));
*output_tensor_size = strlen(output_metadata);
size_t metadata_len = strlen(output_metadata);
if (metadata_len > output_tensor->size)
metadata_len = output_tensor->size;
memcpy(output_tensor->buf, output_metadata, metadata_len);
*output_tensor_size = metadata_len;
return success;
}
@ -643,8 +646,11 @@ get_output(void *ctx, graph_execution_context exec_ctx, uint32_t index,
printf("%s", buf);
}
memcpy(output_tensor->buf + end_pos, buf, strlen(buf));
end_pos += strlen(buf);
size_t piece_len = strlen(buf);
if (end_pos + piece_len > output_tensor->size)
break;
memcpy(output_tensor->buf + end_pos, buf, piece_len);
end_pos += piece_len;
}
if (backend_ctx->config.stream_stdout) {