onnxruntime
fee68589 - Fix signed-int overflow in SamplingState::Init to prevent heap-buffer-overflow (#29443)

Commit
100 days ago
Fix signed-int overflow in SamplingState::Init to prevent heap-buffer-overflow (#29443) This pull request improves the safety of buffer size calculations in the `SamplingState` initialization logic by ensuring that all multiplications involving `batch_size` and `vocab_size` are safely performed using `SafeInt<size_t>`. This prevents potential integer overflow bugs that could lead to under-allocated buffers and memory errors. **Buffer allocation safety improvements:** * All buffer size calculations that multiply `batch_size` and `vocab_size` now use `SafeInt<size_t>` to ensure checked arithmetic, preventing silent integer overflows that could cause heap-buffer-overflow issues. This includes allocations for both CPU and CUDA buffers in `SamplingState`. [[1]](diffhunk://#diff-ad3815054e84321b726b1e4c36d32cf2ab224301f699094f33b6ffd81b91eb64L25-R48) [[2]](diffhunk://#diff-ad3815054e84321b726b1e4c36d32cf2ab224301f699094f33b6ffd81b91eb64L52-R58) * The calculation for the buffer size of `h_sampled_all` now also safely casts `max_iter` to `size_t` before multiplication, further protecting against overflow. These changes make the code more robust and secure, especially when handling large or model-controlled input sizes.
Author
Parents
Loading