CVE-2022-35939: Out-of-bounds Write (CWE-787)
value: TensorFlow is an open source platform for machine learning. The `ScatterNd` function takes an input argument that determines the indices of of the output tensor. An input index greater than the output tensor or less than zero will either write content at the wrong index or trigger a crash. We have patched the issue in GitHub commit b4d4b4cb019bd7240a52daa4ba61e3cc814f0384. The fix will be included in TensorFlow 2.10.0. We will also cherrypick this commit on TensorFlow 2.9.1, TensorFlow 2.8.1, and TensorFlow 2.7.2, as these are also affected and still in supported range. There are no known workarounds for this issue.
CWE description: The product writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 9.8
Published: 2022-09-16T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/lite/kernels/internal/reference/reference_ops.h, tensorflow/lite/kernels/scatter_nd.cc, tensorflow/lite/kernels/scatter_nd_test.cc
Commit date: 2022-07-26 10:09:10 -0700
Commit message: Check bounds for reads and writes in scatter_nd

PiperOrigin-RevId: 463365479

Diff: diff --git a/tensorflow/lite/kernels/internal/reference/reference_ops.h b/tensorflow/lite/kernels/internal/reference/reference_ops.h
index 52d3814a05467..0a7b202f28a7b 100644
--- a/tensorflow/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/lite/kernels/internal/reference/reference_ops.h
@@ -654,15 +654,16 @@ inline TfLiteStatus GatherNdString(const RuntimeShape& params_shape,
   return kTfLiteOk;
 }
 #endif
 
 template <typename IndicesT, typename UpdatesT>
-inline void ScatterNd(const RuntimeShape& indices_shape,
-                      const IndicesT* indices_data,
-                      const RuntimeShape& updates_shape,
-                      const UpdatesT* updates_data,
-                      const RuntimeShape& output_shape, UpdatesT* output_data) {
+inline TfLiteStatus ScatterNd(const RuntimeShape& indices_shape,
+                              const IndicesT* indices_data,
+                              const RuntimeShape& updates_shape,
+                              const UpdatesT* updates_data,
+                              const RuntimeShape& output_shape,
+                              UpdatesT* output_data) {
   ruy::profiler::ScopeLabel label("ScatterNd");
 
   int n_slices = 1;
   int slice_size = 1;
   const int outer_dims = indices_shape.DimensionsCount() - 1;
@@ -681,22 +682,28 @@ inline void ScatterNd(const RuntimeShape& indices_shape,
   for (int i = 0; i < indices_nd; ++i) {
     dims_to_count[i] = remain_flat_size / output_shape.Dims(i);
     remain_flat_size = dims_to_count[i];
   }
 
+  if (n_slices * slice_size > updates_shape.FlatSize()) {
+    return kTfLiteError;
+  }
   memset(output_data, 0, sizeof(UpdatesT) * output_flat_size);
   for (int i = 0; i < n_slices; ++i) {
     int to_pos = 0;
     for (int j = 0; j < indices_nd; ++j) {
       IndicesT idx = indices_data[i * indices_nd + j];
-      TFLITE_DCHECK(0 <= idx && idx < output_shape.Dims(j));
       to_pos += idx * dims_to_count[j];
     }
+    if (to_pos < 0 || to_pos + slice_size > output_flat_size) {
+      return kTfLiteError;
+    }
     for (int j = 0; j < slice_size; j++) {
       output_data[to_pos + j] += updates_data[i * slice_size + j];
     }
   }
+  return kTfLiteOk;
 }
 
 template <typename T>
 void Minimum(const RuntimeShape& input1_shape, const T* input1_data,
              const T* input2_data, const RuntimeShape& output_shape,
diff --git a/tensorflow/lite/kernels/scatter_nd.cc b/tensorflow/lite/kernels/scatter_nd.cc
index ab02a206fd157..c5aab584093df 100644
--- a/tensorflow/lite/kernels/scatter_nd.cc
+++ b/tensorflow/lite/kernels/scatter_nd.cc
@@ -127,15 +127,14 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 }
 
 template <typename IndicesT, typename UpdatesT>
 TfLiteStatus ScatterNd(const TfLiteTensor* indices, const TfLiteTensor* updates,
                        TfLiteTensor* output) {
-  reference_ops::ScatterNd(
+  return reference_ops::ScatterNd(
       GetTensorShape(indices), GetTensorData<IndicesT>(indices),
       GetTensorShape(updates), GetTensorData<UpdatesT>(updates),
       GetTensorShape(output), GetTensorData<UpdatesT>(output));
-  return kTfLiteOk;
 }
 
 template <typename IndicesT>
 TfLiteStatus EvalScatterNd(TfLiteContext* context, const TfLiteTensor* indices,
                            const TfLiteTensor* updates,
@@ -147,29 +146,40 @@ TfLiteStatus EvalScatterNd(TfLiteContext* context, const TfLiteTensor* indices,
                      GetTensorShape(shape), GetTensorData<IndicesT>(shape)));
     TF_LITE_ENSURE_OK(context,
                       ResizeOutputTensor<IndicesT>(context, shape, output));
   }
 
+  TfLiteStatus status = kTfLiteError;
   switch (updates->type) {
     case kTfLiteFloat32:
-      return ScatterNd<IndicesT, float>(indices, updates, output);
+      status = ScatterNd<IndicesT, float>(indices, updates, output);
+      break;
     case kTfLiteUInt8:
-      return ScatterNd<IndicesT, uint8_t>(indices, updates, output);
+      status = ScatterNd<IndicesT, uint8_t>(indices, updates, output);
+      break;
     case kTfLiteBool:
-      return ScatterNd<IndicesT, bool>(indices, updates, output);
+      status = ScatterNd<IndicesT, bool>(indices, updates, output);
+      break;
     case kTfLiteInt8:
-      return ScatterNd<IndicesT, int8_t>(indices, updates, output);
+      status = ScatterNd<IndicesT, int8_t>(indices, updates, output);
+      break;
     case kTfLiteInt32:
-      return ScatterNd<IndicesT, int32_t>(indices, updates, output);
+      status = ScatterNd<IndicesT, int32_t>(indices, updates, output);
+      break;
     case kTfLiteInt64:
-      return ScatterNd<IndicesT, int64_t>(indices, updates, output);
+      status = ScatterNd<IndicesT, int64_t>(indices, updates, output);
+      break;
     default:
       TF_LITE_KERNEL_LOG(
           context, "Updates of type '%s' are not supported by scatter_nd.",
           TfLiteTypeGetName(updates->type));
       return kTfLiteError;
   }
+  if (status != kTfLiteOk) {
+    TF_LITE_KERNEL_LOG(context, "scatter_nd index out of bounds");
+  }
+  return status;
 }
 
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* indices;
   TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kIndices, &indices));
diff --git a/tensorflow/lite/kernels/scatter_nd_test.cc b/tensorflow/lite/kernels/scatter_nd_test.cc
index ed62c67f736d9..b0951efcbb2e7 100644
--- a/tensorflow/lite/kernels/scatter_nd_test.cc
+++ b/tensorflow/lite/kernels/scatter_nd_test.cc
@@ -359,7 +359,36 @@ TEST(ScatterNdOpTest, DynamicShape) {
                                 /*2, 1*/ 0,  0,  0,  0,  0,
                                 /*2, 2*/ 0,  0,  0,  0,  0,
                                 /*2, 3*/ 1,  2,  3,  4,  5}));
 }
 
+TEST(ScatterNdOpTest, ReadAndWriteArrayLimits) {
+  ScatterNdOpModel m({TensorType_INT32, {5, 1}}, {TensorType_INT32, {5}},
+                     {TensorType_INT32, {1}});
+  m.SetIndices<int32_t>({4, 3, 1, 0, 2});
+  m.SetUpdates<int32_t>({1, 2, 3, 7, 9});
+  m.SetShape<int32_t>({5});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({5}));
+  EXPECT_THAT(m.GetOutput<int32_t>(), ElementsAreArray({7, 3, 9, 2, 1}));
+}
+
+TEST(ScatterNdOpTest, OOBRead) {
+  ScatterNdOpModel m({TensorType_INT32, {1, 1}}, {TensorType_INT32, {1}},
+                     {TensorType_INT32, {1}});
+  m.SetIndices<int32_t>({4});
+  m.SetUpdates<int32_t>({1});
+  m.SetShape<int32_t>({1});
+  ASSERT_EQ(m.Invoke(), kTfLiteError);
+}
+
+TEST(ScatterNdOpTest, OOBWrites) {
+  ScatterNdOpModel m({TensorType_INT32, {5, 1}}, {TensorType_INT32, {5}},
+                     {TensorType_INT32, {1}});
+  m.SetIndices<int32_t>({4, 3, 1, -0x38, 0x38});
+  m.SetUpdates<int32_t>({1, 2, 3, 0x44444444, 0x55555555});
+  m.SetShape<int32_t>({1});
+  ASSERT_EQ(m.Invoke(), kTfLiteError);
+}
+
 }  // namespace
 }  // namespace tflite

Vulnerable code: inline void ScatterNd(const RuntimeShape& indices_shape,
                      const IndicesT* indices_data,
                      const RuntimeShape& updates_shape,
                      const UpdatesT* updates_data,
                      const RuntimeShape& output_shape, UpdatesT* output_data) {
      TFLITE_DCHECK(0 <= idx && idx < output_shape.Dims(j));
  reference_ops::ScatterNd(
  return kTfLiteOk;
      return ScatterNd<IndicesT, float>(indices, updates, output);
      return ScatterNd<IndicesT, uint8_t>(indices, updates, output);
      return ScatterNd<IndicesT, bool>(indices, updates, output);
      return ScatterNd<IndicesT, int8_t>(indices, updates, output);
      return ScatterNd<IndicesT, int32_t>(indices, updates, output);
      return ScatterNd<IndicesT, int64_t>(indices, updates, output);
Fixed code: inline TfLiteStatus ScatterNd(const RuntimeShape& indices_shape,
                              const IndicesT* indices_data,
                              const RuntimeShape& updates_shape,
                              const UpdatesT* updates_data,
                              const RuntimeShape& output_shape,
                              UpdatesT* output_data) {
  if (n_slices * slice_size > updates_shape.FlatSize()) {
    return kTfLiteError;
  }
    if (to_pos < 0 || to_pos + slice_size > output_flat_size) {
      return kTfLiteError;
    }
  return kTfLiteOk;
  return reference_ops::ScatterNd(
  TfLiteStatus status = kTfLiteError;
      status = ScatterNd<IndicesT, float>(indices, updates, output);
      break;
      status = ScatterNd<IndicesT, uint8_t>(indices, updates, output);
      break;
      status = ScatterNd<IndicesT, bool>(indices, updates, output);
      break;
      status = ScatterNd<IndicesT, int8_t>(indices, updates, output);
      break;
      status = ScatterNd<IndicesT, int32_t>(indices, updates, output);
      break;
      status = ScatterNd<IndicesT, int64_t>(indices, updates, output);
      break;
  if (status != kTfLiteOk) {
    TF_LITE_KERNEL_LOG(context, "scatter_nd index out of bounds");
  }
  return status;
TEST(ScatterNdOpTest, ReadAndWriteArrayLimits) {
  ScatterNdOpModel m({TensorType_INT32, {5, 1}}, {TensorType_INT32, {5}},
                     {TensorType_INT32, {1}});
  m.SetIndices<int32_t>({4, 3, 1, 0, 2});
  m.SetUpdates<int32_t>({1, 2, 3, 7, 9});
  m.SetShape<int32_t>({5});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({5}));
  EXPECT_THAT(m.GetOutput<int32_t>(), ElementsAreArray({7, 3, 9, 2, 1}));
}

TEST(ScatterNdOpTest, OOBRead) {
  ScatterNdOpModel m({TensorType_INT32, {1, 1}}, {TensorType_INT32, {1}},
                     {TensorType_INT32, {1}});
  m.SetIndices<int32_t>({4});
  m.SetUpdates<int32_t>({1});
  m.SetShape<int32_t>({1});
  ASSERT_EQ(m.Invoke(), kTfLiteError);
}

TEST(ScatterNdOpTest, OOBWrites) {
  ScatterNdOpModel m({TensorType_INT32, {5, 1}}, {TensorType_INT32, {5}},
                     {TensorType_INT32, {1}});
  m.SetIndices<int32_t>({4, 3, 1, -0x38, 0x38});
  m.SetUpdates<int32_t>({1, 2, 3, 0x44444444, 0x55555555});
  m.SetShape<int32_t>({1});
  ASSERT_EQ(m.Invoke(), kTfLiteError);
}