Skip to content

Fix integer overflow in scatter_nd / sparse_to_dense output shape#115460

Open
mohammadmseet-hue wants to merge 3 commits intotensorflow:masterfrom
mohammadmseet-hue:fix-scatter-nd-sparse-to-dense-overflow
Open

Fix integer overflow in scatter_nd / sparse_to_dense output shape#115460
mohammadmseet-hue wants to merge 3 commits intotensorflow:masterfrom
mohammadmseet-hue:fix-scatter-nd-sparse-to-dense-overflow

Conversation

@mohammadmseet-hue
Copy link
Copy Markdown

Summary

scatter_nd::ResizeOutputTensor and sparse_to_dense::Resize both copy the contents of an attacker-controlled shape / output_shape input tensor directly into TfLiteIntArray::data[i] (int) without validating that each value is non-negative or fits in the int32 range used by the TfLiteIntArray storage.

scatter_nd

template <typename IndicesT>
TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
                                const TfLiteTensor* shape,
                                TfLiteTensor* output) {
  const int shape_rank = SizeOfDimension(shape, 0);
  TfLiteIntArray* output_shape = TfLiteIntArrayCreate(shape_rank);
  const auto* shape_data = GetTensorData<IndicesT>(shape);
  for (int i = 0; i < shape_rank; i++) {
    output_shape->data[i] = shape_data[i];                    // unchecked narrowing
  }
  return context->ResizeTensor(context, output, output_shape);
}

shape_data is IndicesT* where IndicesT can be int32_t or int64_t. For int64 a positive value > INT32_MAX silently narrows; negatives pass unchecked. The wrapped per-dim value flows into ResizeTensor and the kernel Eval (reference_ops::ScatterNd) later writes through indices computed from the indices tensor that are not bounds-checked against the wrapped output dims — heap-buffer-overflow write whose length and content are controlled by the model.

sparse_to_dense

template <typename T>
TfLiteStatus Resize(TfLiteContext* context, const TfLiteTensor* output_shape,
                    TfLiteTensor* output) {
  const int output_dimensions = NumElements(output_shape);
  TfLiteIntArray* output_shape_array = TfLiteIntArrayCreate(output_dimensions);
  for (int i = 0; i < output_dimensions; ++i) {
    output_shape_array->data[i] = GetTensorData<T>(output_shape)[i];   // unchecked narrowing
  }
  return context->ResizeTensor(context, output, output_shape_array);
}

Same shape, same template-T narrowing. T = int64 silently truncates; negatives pass through. reference_ops::SparseToDense then writes at indices computed from sparse indices (also unchecked) that are nominally within the un-truncated output region — heap OOB write.

Fix

In both files, accumulate each dim into an int64_t temporary and reject values that are negative or exceed std::numeric_limits<int32_t>::max() before assigning into TfLiteIntArray::data[i]. On any failure, free the partially-built TfLiteIntArray, log a kernel error, and return kTfLiteError before ResizeTensor.

   for (int i = 0; i < shape_rank; i++) {
-    output_shape->data[i] = shape_data[i];
+    const int64_t dim = static_cast<int64_t>(shape_data[i]);
+    if (dim < 0 || dim > std::numeric_limits<int32_t>::max()) {
+      TfLiteIntArrayFree(output_shape);
+      TF_LITE_KERNEL_LOG(context, ...);
+      return kTfLiteError;
+    }
+    output_shape->data[i] = static_cast<int>(dim);
   }

Drop <stdint.h> in favor of <cstdint> for both files (C++ only TUs) per the style review on PR #115031.

Relationship to other PRs in this series

This is the same family of fix as PRs #115031, #115452, #115453, #115454, #115455, #115456, #115457, #115458, #115459. Eleven tflite kernels in the CheckedInt incomplete-pattern hunt now share the same bug class, the same downstream narrowing into TfLiteIntArray::data[], the same heap-OOB-write outcome, and the same fix template (validate + early-return before ResizeTensor).

Files changed

File Lines
tensorflow/lite/kernels/scatter_nd.cc +25 / -2
tensorflow/lite/kernels/sparse_to_dense.cc +20 / -3

Test plan

  • No public API change.
  • No new dependencies.
  • Existing scatter_nd_test and sparse_to_dense_test tests pass against the patched kernels.
  • Happy to add regression tests for the int64 narrowing and negative-shape cases on request.

gather_nd::Prepare reads `indices_nd` from `indices->dims->data[
indices_rank - 1]`, which is an attacker-controlled int32 value coming
from the .tflite model's `indices` tensor shape. The existing bound
check only validated the upper bound:

  if (indices_nd > params_rank) { error }

A *negative* value passes this check (negative is not > positive) and
propagates into the rest of Prepare:

  // output_rank wraps to a huge positive int because subtracting a
  // negative is the same as adding a positive
  const int output_rank = indices_rank + params_rank - indices_nd - 1;
  TfLiteIntArray* output_shape = TfLiteIntArrayCreate(output_rank);
  ...
  for (int i = indices_nd; i < params_rank; ++i) {
    output_shape->data[output_index++] = params->dims->data[i];   // OOB read
  }

Two distinct memory-safety primitives result:

1. The loop iterates `for (int i = indices_nd; i < params_rank; ++i)`
   with i starting at a negative int32 and ending at a small positive
   value. Each iteration reads `params->dims->data[i]` for negative
   i — an OOB read of params->dims->data on the order of ~2^31 entries.

2. `output_index` runs past `output_rank` and writes into
   `output_shape->data[]` past the just-allocated TfLiteIntArray
   storage, corrupting heap memory adjacent to the allocation.

3. The wrapped `output_rank` also produces a garbage TfLiteIntArray
   shape, which `ResizeTensor` then uses to size the output tensor;
   the kernel Eval path later walks `reference_ops::GatherNd` with the
   un-wrapped intended sizes, producing a heap-buffer-overflow write
   whose length and content are controlled by the model.

A malicious .tflite that contains a GatherNd op whose `indices` tensor
has its innermost dimension set to a negative int32 (e.g. via a crafted
flatbuffer) is enough to trigger the chain.

Fix
---

Validate `indices_nd >= 0` in addition to the existing
`indices_nd > params_rank` check, and add a defensive
`TF_LITE_ENSURE(context, output_rank >= 0)` after the subtraction so
any future regression that introduces a wrap is caught before the
TfLiteIntArrayCreate / loop combination.

Also drop <stdint.h> in favor of <cstdint> per the style review on
PR tensorflow#115031.

This is the same family of fix as PRs tensorflow#115031 / tensorflow#115452 / tensorflow#115453 /
tensorflow#115454 / tensorflow#115455 / tensorflow#115456 / tensorflow#115457 / tensorflow#115458, with the addition
that gather_nd carries a *second* OOB primitive (the loop reading
params->dims->data[negative]).
scatter_nd::ResizeOutputTensor and sparse_to_dense::Resize both copy
the contents of an attacker-controlled `shape` / `output_shape` input
tensor directly into TfLiteIntArray::data[i] (`int`) without
validating that each value is non-negative or fits in the int32 range
used by the TfLiteIntArray storage.

scatter_nd
----------

  for (int i = 0; i < shape_rank; i++) {
    output_shape->data[i] = shape_data[i];
  }

`shape_data` is `IndicesT*` where IndicesT can be int32 or int64. For
int64 a positive value > INT32_MAX silently narrows; negatives pass
unchecked. The wrapped per-dim value flows into ResizeTensor and the
kernel Eval (reference_ops::ScatterNd) later writes through indices
that are not bounds-checked against the wrapped output dims —
heap-buffer-overflow write whose length and content are controlled by
the model.

sparse_to_dense
---------------

  for (int i = 0; i < output_dimensions; ++i) {
    output_shape_array->data[i] = GetTensorData<T>(output_shape)[i];
  }

Same shape, same template-T narrowing. T = int64 silently truncates;
negatives pass through. reference_ops::SparseToDense then writes at
indices computed from sparse `indices` (also unchecked) that are
nominally within the un-truncated output region — heap OOB write.

Fix
---

In both files, accumulate each dim into an int64 temporary and reject
values that are negative or exceed std::numeric_limits<int32_t>::max()
before assigning into TfLiteIntArray::data[i]. On any failure, free
the partially-built TfLiteIntArray, log a kernel error, and return
kTfLiteError before ResizeTensor.

Drop <stdint.h> in favor of <cstdint> for both files (C++ only TUs)
per the style review on PR tensorflow#115031.

This is the same family of fix as PRs tensorflow#115031 / tensorflow#115452 / tensorflow#115453 /
tensorflow#115454 / tensorflow#115455 / tensorflow#115456 / tensorflow#115457 / tensorflow#115458 / tensorflow#115459.
@google-ml-butler google-ml-butler Bot added the size:M CL Change Size: Medium label Apr 8, 2026
mohammadmseet-hue added a commit to mohammadmseet-hue/tensorflow that referenced this pull request Apr 8, 2026
…ites

slice::CalculateOutputShapeVector<T> reads `begin` and `size` values
from attacker-controlled int32 / int64 input tensors. The existing
validation has three independent gaps:

  1. begin[idx] is never validated to be >= 0 or <= input_dim. With a
     negative begin and size_value == -1, the code computes
     `size_value = input_dim - begin` which is larger than input_dim,
     producing an output shape that asks the kernel Eval path to read
     past the end of input. With a begin > input_dim the same
     computation underflows or overshoots.

  2. The `else` branch checks `input_dim < begin + size` in template-T
     arithmetic, where T can be int64. With begin = INT64_MAX and
     size = 1, the addition signed-overflows to INT64_MIN and the
     check `input_dim < INT64_MIN` is false → bypass. The unchecked
     begin then propagates to GetBeginAndSizeVectors and into
     reference_ops::Slice as `op_params.begin[i]`, where it is used
     to compute `input_offset + begin * stride` for the read pointer.
     Result: OOB read on the input buffer.

  3. The final `static_cast<int>(size_value)` silently truncates int64
     to int. A large positive int64 size value becomes a small or
     negative int written into the output_shape_vector and on into
     ResizeTensor, producing an undersized buffer that the kernel
     later overruns.

A malicious .tflite with crafted begin / size constant tensors can
therefore drive any of these into a heap-buffer-overflow read or
write, depending on which branch is taken.

Fix
---

  * Validate begin in [0, input_dim] before either branch.
  * Compute `begin + size` in int64 in the else branch so the
    comparison cannot wrap.
  * Bounds-check the final size_value against the int range used by
    output_shape_vector before the static_cast.

Drop <stdint.h> in favor of <cstdint> per the style review on
PR tensorflow#115031.

This is the same family of fix as PRs tensorflow#115031, tensorflow#115452, tensorflow#115453,
tensorflow#115455, tensorflow#115456, tensorflow#115457, tensorflow#115458, tensorflow#115459, tensorflow#115460. Twelve tflite
kernels in the CheckedInt incomplete-pattern hunt now share the same
bug class.
@google-ml-butler google-ml-butler Bot added the awaiting review Pull request awaiting review label Apr 9, 2026
@github-project-automation github-project-automation Bot moved this to Assigned Reviewer in PR Queue Apr 9, 2026
@keerthanakadiri keerthanakadiri added the comp:lite TF Lite related issues label Apr 9, 2026
@mohammadmseet-hue mohammadmseet-hue force-pushed the fix-scatter-nd-sparse-to-dense-overflow branch from 71b91df to c4fd31d Compare April 9, 2026 14:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

awaiting review Pull request awaiting review comp:lite TF Lite related issues size:M CL Change Size: Medium

Projects

Status: Assigned Reviewer

Development

Successfully merging this pull request may close these issues.

3 participants