Skip to content

Commit 3d707f9

Browse files
authored
Add bounds checking for output tensor buffer in wasi-nn llama.cpp (#4847)
* Add bounds checking for output tensor buffer in wasi-nn llama.cpp 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. Instead of silently truncating output when the buffer is too small, return the too_large error with a diagnostic message. This makes the behavior consistent with the OpenVINO backend's get_output and allows callers to distinguish between successful completion and insufficient buffer size.
1 parent aea4da0 commit 3d707f9

File tree

1 file changed

+18
-4
lines changed

1 file changed

+18
-4
lines changed

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

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -623,8 +623,15 @@ get_output(void *ctx, graph_execution_context exec_ctx, uint32_t index,
623623
printf("%s\n", output_metadata);
624624
}
625625

626-
memcpy(output_tensor->buf, output_metadata, strlen(output_metadata));
627-
*output_tensor_size = strlen(output_metadata);
626+
size_t metadata_len = strlen(output_metadata);
627+
if (metadata_len > output_tensor->size) {
628+
NN_ERR_PRINTF("Output buffer too small for metadata: "
629+
"need %zu, got %zu",
630+
metadata_len, output_tensor->size);
631+
return too_large;
632+
}
633+
memcpy(output_tensor->buf, output_metadata, metadata_len);
634+
*output_tensor_size = metadata_len;
628635
return success;
629636
}
630637

@@ -643,8 +650,15 @@ get_output(void *ctx, graph_execution_context exec_ctx, uint32_t index,
643650
printf("%s", buf);
644651
}
645652

646-
memcpy(output_tensor->buf + end_pos, buf, strlen(buf));
647-
end_pos += strlen(buf);
653+
size_t piece_len = strlen(buf);
654+
if (end_pos + piece_len > output_tensor->size) {
655+
NN_ERR_PRINTF("Output buffer too small: need at least %zu,"
656+
" got %zu",
657+
end_pos + piece_len, output_tensor->size);
658+
return too_large;
659+
}
660+
memcpy(output_tensor->buf + end_pos, buf, piece_len);
661+
end_pos += piece_len;
648662
}
649663

650664
if (backend_ctx->config.stream_stdout) {

0 commit comments

Comments
 (0)