onnxruntime-sampling-custom-vocab1-oob-read-poc

PoC: ONNX Runtime com.microsoft.Sampling custom-sampling heap OOB read (vocab_size = 1)

Reproduces an AddressSanitizer-confirmed heap-buffer-overflow READ at:

onnxruntime/contrib_ops/cpu/transformers/sampling_cpu_helper.h:15

A com.microsoft.Sampling node with custom = 1 and vocab_size = 1 makes cumulate_and_filter_custom call filter_scores(..., offset, 1), which reads sorted_indice[offset + 1] — one size_t (8 bytes) past a one-element sorted_indices row — on a normal InferenceSession.Run.

Reproduce

Build ONNX Runtime with Debug + AddressSanitizer, compile run_sampling_model_full.cpp against that build, then:

# Triggers the OOB read
ASAN_OPTIONS=halt_on_error=0:detect_leaks=0 ./run_sampling_model_full sampling_custom_vocab1_bad.onnx

Expected:

ERROR: AddressSanitizer: heap-buffer-overflow
READ of size 8
  #0 ... SamplingCpuHelper::filter_scores<float>(std::vector<unsigned long>&, gsl::span<float>&, ...) ... sampling_cpu_helper.h:15
  #1 ... cumulate_and_filter_custom<float> ... sampling_cpu_helper.h:27
  ... Sampling::Compute ... sampling.cc:148
SUMMARY: AddressSanitizer: heap-buffer-overflow ... sampling_cpu_helper.h:15

Safe control (identical model, vocab_size = 2):

ASAN_OPTIONS=halt_on_error=0:detect_leaks=0 ./run_sampling_model_full sampling_custom_vocab2_safe.onnx
# run completed, outputs=1

Files

• sampling_custom_vocab1_bad.onnx — crafted model that triggers the OOB read (custom=1, vocab_size=1).
• sampling_custom_vocab2_safe.onnx — safe control (vocab_size=2).
• run_sampling_model_full.cpp — C++ InferenceSession.Run harness used for validation.
• recover_custom_vocab1_bad.log — ASan trace for the crafted model.
• recover_custom_vocab2_safe.log — safe-control output.

Scope

The confirmed primitive is an out-of-bounds read at line 15 (on the raw std::vector<size_t> sorted_indice). The over-read value is used as the index for the line 16 write into next_token_score, which is a bounds-checked gsl::span; an out-of-range index there triggers the span contract and terminates the process rather than corrupting memory. This PoC therefore demonstrates CWE-125 (out-of-bounds read), not a heap write.