Fix out of bounds read issue in cross.entropy.cc

Merged
VR
7mo ago
+54-0

Description

Add bounds checking for label tensor values in SparseSoftmaxCrossEntropy::Compute to prevent out-of-bounds memory reads.

The SparseSoftmaxCrossEntropy operator uses label_data[i] (int64_t) directly as an array index into the log-probability buffer without validating that the value falls within [0, D) where D is the number of classes. A malicious ONNX model can embed arbitrary label values in a model initializer, causing the operator to read heap memory beyond the log-probability buffer.

Affected expressions in cross_entropy.cc:

loss_sample[i] = -log_prob_data[i * d + label_data[i]] * weight_data[i];  // weighted path
loss_sample[i] = -log_prob_data[i * d + label_data[i]];                   // unweighted path

Existing shape validation confirms label and logit dimensions are compatible, but never validates label values against the class dimension.

Fix

Added a validation loop before the loss computation that returns an error status if any label value is outside [0, D):

for (ptrdiff_t i = 0; i < n; i++) {
  ORT_RETURN_IF(label_data[i] < 0 || label_data[i] >= d,
                "SparseSoftmaxCrossEntropy: label value ", label_data[i],
                " at index ", i, " is out of range [0, ", d, ")");
}