Skip to content

Commit acd36fc

Browse files
committed
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.
1 parent c46b10d commit acd36fc

1 file changed

Lines changed: 10 additions & 4 deletions

File tree

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

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -623,8 +623,11 @@ 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+
metadata_len = output_tensor->size;
629+
memcpy(output_tensor->buf, output_metadata, metadata_len);
630+
*output_tensor_size = metadata_len;
628631
return success;
629632
}
630633

@@ -643,8 +646,11 @@ get_output(void *ctx, graph_execution_context exec_ctx, uint32_t index,
643646
printf("%s", buf);
644647
}
645648

646-
memcpy(output_tensor->buf + end_pos, buf, strlen(buf));
647-
end_pos += strlen(buf);
649+
size_t piece_len = strlen(buf);
650+
if (end_pos + piece_len > output_tensor->size)
651+
break;
652+
memcpy(output_tensor->buf + end_pos, buf, piece_len);
653+
end_pos += piece_len;
648654
}
649655

650656
if (backend_ctx->config.stream_stdout) {

0 commit comments

Comments
 (0)