Fix OOB reads in tokenizer decoders, image transforms, and speech features

[Sayan Shaw (sayanshaw24)]

Fix OOB reads in tokenizer decoders, image transforms, and speech features

Summary

This PR adds missing bounds checks and input validation across multiple custom operators to prevent heap out-of-bounds reads from malformed inputs or crafted models.

Changes

operators/tokenizer/bpe_decoder.hpp

• Fix wrong loop bound in KernelBpeDecoder::Compute: the inner loop iterated over the entire tensor element count instead of the per-row sequence length, causing OOB reads when processing batched inputs (batch >= 2).

operators/tokenizer/wordpiece_tokenizer.cc

• Fix misplaced bounds check in KernelWordpieceTokenizer_Tokenizer: the row_index >= n_existing_rows guard ran one line after the dereference of existing_rows[row_index]. Moved the check into the condition before the access.

operators/tokenizer/bert_tokenizer_decoder.cc

• Fix unbounded slice in KernelBertTokenizerDecoder::Compute: start/end indices from the attacker-controlled positions input were used to slice the ids buffer without any bounds validation. Added 0 <= start <= end <= ids.NumberOfElement() check.

operators/tokenizer/trie_tokenizer.hpp

• Fix shape vector OOB read in TrieDetokenizer::Compute: the code assumed rank-2 input and accessed ids_dim[1] without verifying ids_dim.size() >= 2. Rank-1 inputs now treated as [1, N].

shared/api/image_transforms.hpp

• Fix OOB read in Resize::Compute: added c != 3 channel validation. The copy loop uses a hardcoded stride of 3, so non-3-channel inputs read past the buffer.
• Fix OOB read in Normalize::Compute: added C == mean_.size() validation. The channel loop indexed into fixed 3-element mean_/std_ vectors using the input's channel dimension.
• Fix OOB read in Permute3D::Init: added validation that each dims_[i] is in [0, 2]. Out-of-range values would index past the 3-element shape vector.

shared/api/speech_features.hpp

• Fix OOB read in SpeechLibSTFTNorm: added validation that fft_win_.size() >= frame_length_ after initialization. A caller-supplied hann_win shorter than frame_length_ caused the STFT loop to read past the window buffer.

test/pp_api_test/test_tokenizer_impl.cc

• Add BpeDecoderBatchedDetokenize regression test: verifies batched detokenization (batch=2) produces correct output without OOB access.

Testing

The BpeDecoderBatchedDetokenize test exercises the BPE decoder fix through the TokenizerImpl::Detokenize API. The remaining fixes are in custom op kernels or config-driven pipelines that require either a crafted ONNX model with an ORT session or attacker-controlled preprocessor config bundles, not worth adding into the repository.

[Copilot]

Pull request overview

This PR hardens several tokenizer, image pre-processing, and speech feature kernels against malformed inputs by adding bounds checks and validating assumptions that previously could lead to out-of-bounds reads.

Changes:

• Fix batched BPE decoder loop bound to use per-row seq_len instead of total element count.
• Add/adjust bounds checks for tokenizer decoders (WordPiece row guard, BERT positions slicing, Trie detokenizer rank handling).
• Add input validation in image transforms (Resize channel count, Normalize channel-vs-mean/std alignment, Permute dims range) and speech features (window length vs frame length), plus a regression test for batched detokenization.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
operators/tokenizer/bpe_decoder.hpp	Fix per-row iteration bound for batched decoding to prevent OOB reads.
operators/tokenizer/wordpiece_tokenizer.cc	Move bounds check ahead of existing_rows[row_index] access.
operators/tokenizer/bert_tokenizer_decoder.cc	Validate attacker-controlled positions slices before slicing ids.
operators/tokenizer/trie_tokenizer.hpp	Handle rank-1 inputs without indexing ids_dim[1]; adjust indexing by seq_len.
shared/api/image_transforms.hpp	Add validation for Resize/Normalize/Permute to prevent channel/dim-driven OOB reads.
shared/api/speech_features.hpp	Validate FFT window length is at least frame_length_ to prevent OOB reads.
test/pp_api_test/test_tokenizer_impl.cc	Add regression test for batched BPE detokenization (batch=2).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.