gemini-code-assist[bot] commented on code in PR #19997:
URL: https://github.com/apache/tvm/pull/19997#discussion_r3572902702
##########
src/runtime/vm/tensor_cache_support.cc:
##########
@@ -153,6 +153,16 @@ void CopyTensorFromBytes(Tensor param, const void* data,
size_t nbytes,
Tensor TensorCacheMetadata::FileRecord::ParamRecord::Load(
Device device, const std::string* raw_data, ffi::Optional<Tensor>*
staging_buffer) const {
+ // `byte_offset` and `nbytes` come from tensor-cache.json; FileRecord::Load
only validates the shard
+ // total, not each parameter's range. Bounds-check this parameter against
the shard buffer before the
+ // copy so a malformed/inconsistent cache fails with a clear error instead
of reading out of bounds.
+ // (Written to avoid signed overflow on `byte_offset + nbytes`.)
+ int64_t shard_len = static_cast<int64_t>(raw_data->length());
+ TVM_FFI_CHECK(byte_offset >= 0 && nbytes >= 0 && byte_offset <= shard_len &&
+ nbytes <= shard_len - byte_offset,
+ ValueError)
+ << "Parameter byte range [" << byte_offset << ", " << (byte_offset +
nbytes)
+ << ") is out of bounds of the " << shard_len << "-byte shard buffer.";
Review Comment:
There are two potential issues here:
1. **Null Pointer Dereference**: `raw_data` is a pointer and is not checked
for `nullptr` before calling `raw_data->length()`. If a caller passes
`nullptr`, this will result in a crash.
2. **Buffer Overflow**: In the `f32-to-bf16` format path below,
`std::vector<uint16_t> buffer(nbytes / 2)` is allocated, and
`std::memcpy(buffer.data(), raw_data->data() + byte_offset, nbytes)` is called.
If `nbytes` is odd, `nbytes / 2` will allocate a buffer of size `nbytes - 1`
bytes, leading to an out-of-bounds write (buffer overflow) during `std::memcpy`.
We should add a null check for `raw_data` and ensure `nbytes` is even when
using the `f32-to-bf16` format.
```c
TVM_FFI_CHECK(raw_data != nullptr, ValueError) << "raw_data pointer cannot
be nullptr.";
int64_t shard_len = static_cast<int64_t>(raw_data->length());
TVM_FFI_CHECK(byte_offset >= 0 && nbytes >= 0 && byte_offset <= shard_len
&&
nbytes <= shard_len - byte_offset,
ValueError)
<< "Parameter byte range [" << byte_offset << ", " << (byte_offset +
nbytes)
<< ") is out of bounds of the " << shard_len << "-byte shard buffer.";
if (dtype == DLDataType{kDLFloat, 32, 1} && format == "f32-to-bf16") {
TVM_FFI_CHECK(nbytes % 2 == 0, ValueError)
<< "Parameter nbytes must be even for f32-to-bf16 format, but got "
<< nbytes;
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
