Fix signed-int overflow in SamplingState::Init to prevent heap-buffer-overflow

onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h

@@ -6,6 +6,7 @@
 #include <vector>
 #include "contrib_ops/cpu/transformers/generation_shared.h"
 #include "contrib_ops/cpu/transformers/generate_impl_base.h"
+#include "contrib_ops/cpu/transformers/sampling_buffer_element_count.h"
 
 namespace onnxruntime {
 namespace contrib {
@@ -22,25 +23,32 @@ struct SamplingState : public ISamplingState<T> {
             int seed,
             bool is_cuda,
             Stream* stream) {
-    int total_count = batch_size * vocab_size;
+    // Compute the product entirely in SafeInt's checked domain; an `int * int` multiply
+    // with model-controlled operands can silently positive-wrap before any SafeInt cast
+    // sees it, leading to under-allocated buffers (heap-buffer-overflow on the downstream
+    // memcpy in SamplingCpuHelper::Sample). `SafeMul<size_t>` also rejects negative
+    // inputs (e.g. an unvalidated `vocab_size` of -1) up front, which a
+    // `static_cast<size_t>` would silently turn into a huge value. `SamplingBufferElementCount`
+    // is the shared helper that regression tests exercise so a revert breaks them.
+    const size_t total_count = SamplingBufferElementCount(batch_size, vocab_size);
 
-    this->h_softmaxed_score = AllocateBuffer<float>(cpu_allocator, h_softmaxed_score_buffer_, SafeInt<size_t>(total_count), stream);
+    this->h_softmaxed_score = AllocateBuffer<float>(cpu_allocator, h_softmaxed_score_buffer_, total_count, stream);
 
     this->generator = std::default_random_engine{gsl::narrow_cast<uint32_t>(seed)};
 
     if (is_cuda) {
-      this->d_index_in = AllocateBuffer<int>(allocator, d_index_in_buffer_, SafeInt<size_t>(total_count), stream);
-      this->d_index_out = AllocateBuffer<int>(allocator, d_index_out_buffer_, SafeInt<size_t>(total_count), stream);
-      this->d_offset = AllocateBuffer<int>(allocator, d_offset_buffer_, SafeInt<size_t>(batch_size + 1), stream);
-      this->d_sorted_score = AllocateBuffer<T>(allocator, d_sorted_score_buffer_, SafeInt<size_t>(total_count), stream);
-      this->d_sorted_softmaxed_score = AllocateBuffer<float>(allocator, d_sorted_softmaxed_score_buffer_, SafeInt<size_t>(total_count), stream);
-      this->d_softmaxed_score = AllocateBuffer<float>(allocator, d_softmaxed_score_buffer_, SafeInt<size_t>(total_count), stream);
+      this->d_index_in = AllocateBuffer<int>(allocator, d_index_in_buffer_, total_count, stream);
+      this->d_index_out = AllocateBuffer<int>(allocator, d_index_out_buffer_, total_count, stream);
+      this->d_offset = AllocateBuffer<int>(allocator, d_offset_buffer_, SafeInt<size_t>(batch_size) + 1, stream);
+      this->d_sorted_score = AllocateBuffer<T>(allocator, d_sorted_score_buffer_, total_count, stream);
+      this->d_sorted_softmaxed_score = AllocateBuffer<float>(allocator, d_sorted_softmaxed_score_buffer_, total_count, stream);
+      this->d_softmaxed_score = AllocateBuffer<float>(allocator, d_softmaxed_score_buffer_, total_count, stream);
       this->d_sampled = AllocateBuffer<float>(allocator, d_sampled_buffer_, SafeInt<size_t>(batch_size), stream);
-      this->h_sampled_all = AllocateBuffer<float>(cpu_allocator, h_sampled_all_buffer_, SafeInt<size_t>(batch_size * max_iter), stream);
+      this->h_sampled_all = AllocateBuffer<float>(cpu_allocator, h_sampled_all_buffer_, SamplingBufferElementCount(batch_size, max_iter), stream);
       this->d_indices = AllocateBuffer<int32_t>(allocator, d_indices_buffer_, SafeInt<size_t>(batch_size), stream);
       this->temp_storage_bytes = 0;
       // TODO: Do not allocate this buffer if there's no presence_mask
-      this->d_presence_mask = AllocateBuffer<int>(allocator, d_presence_mask_buffer_, SafeInt<size_t>(total_count), stream);
+      this->d_presence_mask = AllocateBuffer<int>(allocator, d_presence_mask_buffer_, total_count, stream);
 
       std::uniform_real_distribution<float> distribution(0.0, 1.0);
       static_cast<void>(distribution(this->generator));
@@ -49,8 +57,8 @@ struct SamplingState : public ISamplingState<T> {
       }
     } else {
       // TODO: Some buffer can be reused for CPU
-      this->sorted_scores = AllocateBuffer<T>(cpu_allocator, sorted_scores_buffer_, SafeInt<size_t>(total_count), stream);
-      this->cumulative_probs = AllocateBuffer<T>(cpu_allocator, cumulative_probs_buffer_, SafeInt<size_t>(total_count), stream);
+      this->sorted_scores = AllocateBuffer<T>(cpu_allocator, sorted_scores_buffer_, total_count, stream);
+      this->cumulative_probs = AllocateBuffer<T>(cpu_allocator, cumulative_probs_buffer_, total_count, stream);
     }
   }
 

onnxruntime/contrib_ops/cpu/transformers/sampling_buffer_element_count.h

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/common/safeint.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace transformers {
+
+// Computes the element count for a `SamplingState` buffer using
+// `SafeMul<size_t>` so that overflow (from a plain `int * int` multiply) or
+// negative-to-unsigned conversion (e.g. an unvalidated `-1`) of the
+// model-controlled operands is detected up front (throws
+// `OnnxRuntimeException`) rather than silently producing an under-sized or
+// absurdly large allocation. Extracted into a self-contained header so that
+// both `SamplingState::Init` and its regression tests call the exact same
+// production code path; reverting this computation to `int * int` or
+// `static_cast<size_t>(...)` will fail the tests.
+inline size_t SamplingBufferElementCount(int batch_size, int per_batch_count) {
+  return SafeMul<size_t>(batch_size, per_batch_count);
+}
+
+}  // namespace transformers
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/test/contrib_ops/sampling_state_test.cc

@@ -0,0 +1,77 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+// Regression tests for the buffer-size arithmetic in
+// onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h
+// (SamplingState::Init). The bug being guarded was that a plain `int * int`
+// multiply of `batch_size * vocab_size`, or a `static_cast<size_t>(...)` of a
+// negative/unvalidated operand, could silently wrap and lead to under-allocated
+// buffers (heap-buffer-overflow on the downstream memcpy in
+// SamplingCpuHelper::Sample).
+//
+// The production computation lives in `sampling_buffer_element_count.h` as
+// `SamplingBufferElementCount`, which `SamplingState::Init` calls directly.
+// These tests call the same helper, so reverting the production code to
+// `int * int` or `static_cast<size_t>` on operands that may be negative will
+// fail these tests.
+
+#include "gtest/gtest.h"
+
+#include "core/common/common.h"
+#include "core/common/safeint.h"
+#include "contrib_ops/cpu/transformers/sampling_buffer_element_count.h"
+
+namespace onnxruntime {
+namespace test {
+
+using contrib::transformers::SamplingBufferElementCount;
+
+// Sanity check: well-formed inputs produce the expected element count.
+TEST(SamplingStateArithmeticTest, ProducesExpectedTotalCountForValidInputs) {
+  EXPECT_EQ(SamplingBufferElementCount(4, 32), static_cast<size_t>(4) * 32u);
+  EXPECT_EQ(SamplingBufferElementCount(1, 50257), static_cast<size_t>(50257));
+  // `h_sampled_all` uses the same helper with `max_iter` as the second operand.
+  EXPECT_EQ(SamplingBufferElementCount(8, 16), static_cast<size_t>(8) * 16u);
+}
+
+// A negative `vocab_size` (e.g. an unvalidated default of -1) used to be turned
+// into SIZE_MAX by `static_cast<size_t>(vocab_size)`, yielding a multiplication
+// result that either silently wrapped or requested an absurdly large buffer.
+// SafeInt<size_t> rejects the negative-to-unsigned conversion up front.
+TEST(SamplingStateArithmeticTest, ThrowsOnNegativeVocabSize) {
+  EXPECT_THROW(SamplingBufferElementCount(4, -1), OnnxRuntimeException);
+}
+
+// Symmetric check for a negative `batch_size`.
+TEST(SamplingStateArithmeticTest, ThrowsOnNegativeBatchSize) {
+  EXPECT_THROW(SamplingBufferElementCount(-1, 32), OnnxRuntimeException);
+}
+
+// `max_iter` flows through the same helper for the `h_sampled_all`
+// allocation, so a negative value in the second operand slot must also be
+// rejected. Uses a valid `batch_size` and a negative `max_iter` to exercise
+// the operand distinctly from `ThrowsOnNegativeVocabSize` above.
+TEST(SamplingStateArithmeticTest, ThrowsOnNegativeMaxIter) {
+  constexpr int batch_size = 4;
+  constexpr int max_iter = -1;
+  EXPECT_THROW(SamplingBufferElementCount(batch_size, max_iter), OnnxRuntimeException);
+}
+
+// The core bug this PR guards against: a plain `int * int` multiply of
+// `batch_size * vocab_size` silently wraps once the mathematical product
+// exceeds `INT_MAX`, producing an under-allocated buffer. `SafeMul<size_t>`
+// promotes both operands to `size_t` before multiplying, so the helper must
+// return the correct non-wrapped product. Reverting the production code to
+// `int * int` (or `static_cast<size_t>(int_product)`) fails this test.
+TEST(SamplingStateArithmeticTest, ReturnsCorrectProductWhenIntMultiplyWouldOverflow) {
+  // 50000 * 50000 = 2.5e9, which exceeds INT_MAX (~2.147e9) and would
+  // signed-overflow if computed in `int`, but still fits in a 32-bit
+  // `size_t` so the test is portable to 32-bit Windows builds where
+  // `size_t` is 32-bit (SIZE_MAX ~ 4.29e9).
+  constexpr int large_operand = 50000;
+  const size_t expected = static_cast<size_t>(large_operand) * static_cast<size_t>(large_operand);
+  EXPECT_EQ(SamplingBufferElementCount(large_operand, large_operand), expected);
+}
+
+}  // namespace test
+}  // namespace onnxruntime