onnxruntime
d762130a - Fix heap OOB write in EmbedLayerNormalizationShapeInference (#28176)

Commit
20 days ago
Fix heap OOB write in EmbedLayerNormalizationShapeInference (#28176) ### Description `EmbedLayerNormalizationShapeInference` unconditionally wrote to output index 2 when `getNumOutputs() == 2 && mask_index_type == 0`, causing a heap out-of-bounds write during model loading — no `session.run()` required. **Fix:** Replace the flawed condition with a simple bounds check: ```cpp // Before (vulnerable): if (ctx.getNumOutputs() == 3 || (ctx.getNumOutputs() == 2 && mask_index_type == 0)) { updateOutputShape(ctx, 2, output_shape); // OOB when numOutputs == 2 propagateElemTypeFromInputToOutput(ctx, 0, 2); } // After (fixed): if (ctx.getNumOutputs() > 2) { updateOutputShape(ctx, 2, output_shape); propagateElemTypeFromInputToOutput(ctx, 0, 2); } ``` A regression test `EmbedLayerNormBatch1_NoMaskIndex_NoSumOutput` has been added to `embed_layer_norm_op_test.cc` to cover the previously vulnerable path: `mask_index_type=0` with exactly 2 outputs (no `embedding_sum`). ### Motivation and Context A crafted ONNX model with an `EmbedLayerNormalization` node declaring 2 outputs and `mask_index_type=0` triggers the vulnerable path. `getOutputType(2)` returns a pointer one past the end of the internal `node_output_types_` vector; subsequent writes through that pointer corrupt adjacent heap memory. In release builds this is silent — no assertion, no crash, exploitable via heap shaping. The `embedding_sum` output is always at index 2 by definition (confirmed in the CPU kernel). The old special-casing for `mask_index_type == 0` was both incorrect and unnecessary. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com> Co-authored-by: Xavier Dupré <xadupre@microsoft.com>
Author
Parents
Loading