CVE-2022-35938: Out-of-bounds Read (CWE-125)
value: TensorFlow is an open source platform for machine learning. The `GatherNd` function takes arguments that determine the sizes of inputs and outputs. If the inputs given are greater than or equal to the sizes of the outputs, an out-of-bounds memory read or a crash is triggered. This issue has been patched in GitHub commit 4142e47e9e31db481781b955ed3ff807a781b494. 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 reads data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 9.1
Published: 2022-09-16T20:15Z
Repository: https://github.com/tensorflow/tflite-micro
Language: C++
Files changed: tensorflow/lite/micro/kernels/gather_nd.cc, tensorflow/lite/micro/kernels/gather_nd_test.cc
Commit date: 2022-07-26 19:54:47 +0000
Commit message: GatherNd verifies that an index is valid before reading. (#1286)


Diff: diff --git a/tensorflow/lite/micro/kernels/gather_nd.cc b/tensorflow/lite/micro/kernels/gather_nd.cc
index eaa1abc..4327177 100644
--- a/tensorflow/lite/micro/kernels/gather_nd.cc
+++ b/tensorflow/lite/micro/kernels/gather_nd.cc
@@ -129,11 +129,12 @@ TfLiteStatus GatherNd(const TfLiteEvalTensor* params,
   int slice_size = 1;
   for (int i = indices_nd; i < params_dims; ++i) {
     slice_size *= params->dims->data[i];
   }
 
-  int remain_flat_size = ElementCount(*params->dims);
+  int params_flat_size = ElementCount(*params->dims);
+  int remain_flat_size = params_flat_size;
 
   // Number of elements per dimension
   int dims_to_count[MAX_INDICES_ND];
   for (int i = 0; i < indices_nd; ++i) {
     dims_to_count[i] = remain_flat_size / params->dims->data[i];
@@ -145,10 +146,13 @@ TfLiteStatus GatherNd(const TfLiteEvalTensor* params,
     for (int j = 0; j < indices_nd; ++j) {
       int offset = i * indices_nd + j;
       IndicesT index = index_data[offset];
       from_pos += index * dims_to_count[j];
     }
+    if (from_pos < 0 || from_pos + slice_size > params_flat_size) {
+      return kTfLiteError;
+    }
     std::memcpy(output_data + i * slice_size, param_data + from_pos,
                 sizeof(ParamsT) * slice_size);
   }
   return kTfLiteOk;
 }
@@ -156,23 +160,28 @@ TfLiteStatus GatherNd(const TfLiteEvalTensor* params,
 template <typename IndicesT>
 TfLiteStatus EvalGatherNd(TfLiteContext* context,
                           const TfLiteEvalTensor* params,
                           const TfLiteEvalTensor* indices,
                           TfLiteEvalTensor* output) {
+  TfLiteStatus status = kTfLiteError;
   switch (params->type) {
     case kTfLiteFloat32:
-      return GatherNd<float, IndicesT>(params, indices, output);
+      status = GatherNd<float, IndicesT>(params, indices, output);
       break;
     case kTfLiteInt8:
-      return GatherNd<int8_t, IndicesT>(params, indices, output);
+      status = GatherNd<int8_t, IndicesT>(params, indices, output);
       break;
     default:
       TF_LITE_KERNEL_LOG(context,
                          "Params type '%s' are not supported by gather_nd.",
                          TfLiteTypeGetName(params->type));
       return kTfLiteError;
   }
+  if (status != kTfLiteOk) {
+    TF_LITE_KERNEL_LOG(context, "gather_nd index out of bounds");
+  }
+  return status;
 }
 
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteEvalTensor* params =
       tflite::micro::GetEvalInput(context, node, kParams);
diff --git a/tensorflow/lite/micro/kernels/gather_nd_test.cc b/tensorflow/lite/micro/kernels/gather_nd_test.cc
index 1d3db2f..3dd42c7 100644
--- a/tensorflow/lite/micro/kernels/gather_nd_test.cc
+++ b/tensorflow/lite/micro/kernels/gather_nd_test.cc
@@ -24,12 +24,12 @@ namespace testing {
 namespace {
 
 template <typename ParamType, typename IndexType>
 void TestGatherNd(int* param_dims, const ParamType* param_data, int* index_dims,
                   const IndexType* index_data, int* output_dims,
-                  ParamType* output_data,
-                  const ParamType* expected_output_data) {
+                  ParamType* output_data, const ParamType* expected_output_data,
+                  const TfLiteStatus expected_status = kTfLiteOk) {
   TfLiteIntArray* pdims = IntArrayFromInts(param_dims);
   TfLiteIntArray* idims = IntArrayFromInts(index_dims);
   TfLiteIntArray* odims = IntArrayFromInts(output_dims);
 
   constexpr int inputs_size = 2;
@@ -47,18 +47,20 @@ void TestGatherNd(int* param_dims, const ParamType* param_data, int* index_dims,
 
   const TfLiteRegistration registration = Register_GATHER_ND();
   micro::KernelRunner runner(registration, tensors, tensors_size, inputs_array,
                              outputs_array, /*builtin_data=*/nullptr);
   TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, runner.InitAndPrepare());
-  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, runner.Invoke());
+  TF_LITE_MICRO_EXPECT_EQ(expected_status, runner.Invoke());
 
-  // The output tensor's data and shape have been updated by the kernel.
-  TfLiteTensor* actual_output_tensor = &tensors[2];
-  TfLiteIntArray* actual_output_dims = actual_output_tensor->dims;
-  const int output_size = ElementCount(*actual_output_dims);
-  for (int i = 0; i < output_size; ++i) {
-    TF_LITE_MICRO_EXPECT_EQ(expected_output_data[i], output_data[i]);
+  if (expected_status == kTfLiteOk) {
+    // The output tensor's data and shape have been updated by the kernel.
+    TfLiteTensor* actual_output_tensor = &tensors[2];
+    TfLiteIntArray* actual_output_dims = actual_output_tensor->dims;
+    const int output_size = ElementCount(*actual_output_dims);
+    for (int i = 0; i < output_size; ++i) {
+      TF_LITE_MICRO_EXPECT_EQ(expected_output_data[i], output_data[i]);
+    }
   }
 }
 
 }  // namespace
 }  // namespace testing
@@ -296,6 +298,34 @@ TF_LITE_MICRO_TEST(GatherNd_Int8Int32) {
   tflite::testing::TestGatherNd<int8_t, int32_t>(
       input_dims, input_data, index_dims, index_data, output_dims, output_data,
       golden_data);
 }
 
+TF_LITE_MICRO_TEST(GatherNd_ReadOOB) {
+  // For input_dims[], index_dims[], or output_dims[], element 0 is the
+  // number of dimensions in that array, not the actual dimension data.
+  int input_dims[] = {2, 2, 2};
+  int index_dims[] = {2, 2, 2};
+  const int32_t index_data[] = {0, 1, 2, 0};
+  const int8_t input_data[] = {1, -1, 1, -2};
+  int8_t output_data;
+  int output_dims[] = {1, 0, 0};
+  tflite::testing::TestGatherNd<int8_t, int32_t>(
+      input_dims, input_data, index_dims, index_data, output_dims, &output_data,
+      nullptr, kTfLiteError);
+}
+
+TF_LITE_MICRO_TEST(GatherNd_ReadOOBNegative) {
+  // For input_dims[], index_dims[], or output_dims[], element 0 is the
+  // number of dimensions in that array, not the actual dimension data.
+  int input_dims[] = {2, 2, 2};
+  int index_dims[] = {2, 2, 2};
+  const int32_t index_data[] = {0, -1, 1, 0};
+  const int8_t input_data[] = {1, -1, 1, -2};
+  int8_t output_data;
+  int output_dims[] = {1, 0, 0};
+  tflite::testing::TestGatherNd<int8_t, int32_t>(
+      input_dims, input_data, index_dims, index_data, output_dims, &output_data,
+      nullptr, kTfLiteError);
+}
+
 TF_LITE_MICRO_TESTS_END

Vulnerable code:   int remain_flat_size = ElementCount(*params->dims);
      return GatherNd<float, IndicesT>(params, indices, output);
      return GatherNd<int8_t, IndicesT>(params, indices, output);
                  ParamType* output_data,
                  const ParamType* expected_output_data) {
  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, runner.Invoke());
  // The output tensor's data and shape have been updated by the kernel.
  TfLiteTensor* actual_output_tensor = &tensors[2];
  TfLiteIntArray* actual_output_dims = actual_output_tensor->dims;
  const int output_size = ElementCount(*actual_output_dims);
  for (int i = 0; i < output_size; ++i) {
    TF_LITE_MICRO_EXPECT_EQ(expected_output_data[i], output_data[i]);
Fixed code:   int params_flat_size = ElementCount(*params->dims);
  int remain_flat_size = params_flat_size;
    if (from_pos < 0 || from_pos + slice_size > params_flat_size) {
      return kTfLiteError;
    }
  TfLiteStatus status = kTfLiteError;
      status = GatherNd<float, IndicesT>(params, indices, output);
      status = GatherNd<int8_t, IndicesT>(params, indices, output);
  if (status != kTfLiteOk) {
    TF_LITE_KERNEL_LOG(context, "gather_nd index out of bounds");
  }
  return status;
                  ParamType* output_data, const ParamType* expected_output_data,
                  const TfLiteStatus expected_status = kTfLiteOk) {
  TF_LITE_MICRO_EXPECT_EQ(expected_status, runner.Invoke());
  if (expected_status == kTfLiteOk) {
    // The output tensor's data and shape have been updated by the kernel.
    TfLiteTensor* actual_output_tensor = &tensors[2];
    TfLiteIntArray* actual_output_dims = actual_output_tensor->dims;
    const int output_size = ElementCount(*actual_output_dims);
    for (int i = 0; i < output_size; ++i) {
      TF_LITE_MICRO_EXPECT_EQ(expected_output_data[i], output_data[i]);
    }
TF_LITE_MICRO_TEST(GatherNd_ReadOOB) {
  // For input_dims[], index_dims[], or output_dims[], element 0 is the
  // number of dimensions in that array, not the actual dimension data.
  int input_dims[] = {2, 2, 2};
  int index_dims[] = {2, 2, 2};
  const int32_t index_data[] = {0, 1, 2, 0};
  const int8_t input_data[] = {1, -1, 1, -2};
  int8_t output_data;
  int output_dims[] = {1, 0, 0};
  tflite::testing::TestGatherNd<int8_t, int32_t>(
      input_dims, input_data, index_dims, index_data, output_dims, &output_data,
      nullptr, kTfLiteError);
}

TF_LITE_MICRO_TEST(GatherNd_ReadOOBNegative) {
  // For input_dims[], index_dims[], or output_dims[], element 0 is the
  // number of dimensions in that array, not the actual dimension data.
  int input_dims[] = {2, 2, 2};
  int index_dims[] = {2, 2, 2};
  const int32_t index_data[] = {0, -1, 1, 0};
  const int8_t input_data[] = {1, -1, 1, -2};
  int8_t output_data;
  int output_dims[] = {1, 0, 0};
  tflite::testing::TestGatherNd<int8_t, int32_t>(
      input_dims, input_data, index_dims, index_data, output_dims, &output_data,
      nullptr, kTfLiteError);
}