Fix signed-int overflow in SamplingState::Init to prevent heap-buffer-overflow
Merged
AP
4mo ago+124+124+012345678901234567890123456789-12-12-01234567890123456789
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_sizeandvocab_sizenow useSafeInt<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 inSamplingState. [1] [2] - The calculation for the buffer size of
h_sampled_allnow also safely castsmax_itertosize_tbefore multiplication, further protecting against overflow.
These changes make the code more robust and secure, especially when handling large or model-controlled input sizes.