[Security] Validate CPU cache_indirection beam indices
Merged
TI
5mo ago+124+124+012345678901234567890123456789-0-0-0123456789
Summary
- reject out-of-range
cache_indirectionbeam indices in the CPU beam-attention path before they are converted into past KV offsets - keep
DecoderMaskedMultiHeadAttentionbeam-width handling consistent with thecache_indirectionshape - add CPU regression tests for
MultiHeadAttentionandDecoderMaskedMultiHeadAttention
Motivation
MultiHeadAttention and DecoderMaskedMultiHeadAttention on the CPU provider could consume attacker-controlled cache_indirection values as beam indices without validating that each element stayed within [0, beam_width). That let malformed models compute offsets past the past key/value buffers. This change rejects invalid indices up front and adds focused tests for the failure path.
Key Changes
- add shared CPU validation in
AttentionCPUBase::ApplyAttentionWithBeamsso the beam path fails before any past-key or past-value reads occur - report an
INVALID_ARGUMENTerror that identifies the offending beam index and its position - validate that an explicit decoder
beam_widthinput matchescache_indirectiondimension 1 when both are present - add contrib-op tests that exercise invalid cache indirection values on the CPU execution provider
Testing
lintrunner -acd build/Linux/Debug && make -j4 CMakeFiles/onnxruntime_providers.dir/home/tlwu/onnxruntime/onnxruntime/contrib_ops/cpu/bert/multihead_attention.cc.o CMakeFiles/onnxruntime_providers.dir/home/tlwu/onnxruntime/onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.cc.o CMakeFiles/onnxruntime_provider_test.dir/home/tlwu/onnxruntime/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc.o CMakeFiles/onnxruntime_provider_test.dir/home/tlwu/onnxruntime/onnxruntime/test/contrib_ops/decoder_masked_multihead_attention_op_test.cc.o- full
onnxruntime_provider_testrelink/run was not completed locally