CVE-2022-35937: 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 is triggered. This issue has been patched in GitHub commit 595a65a3e224a0362d7e68c2213acfc2b499a196. 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/tensorflow
Language: C++
Files changed: tensorflow/lite/kernels/gather_nd.cc, tensorflow/lite/kernels/gather_nd_test.cc, tensorflow/lite/kernels/internal/reference/reference_ops.h
Commit date: 2022-07-25 04:40:43 -0700
Commit message: Return a TFLite error if gather_nd will result in reading invalid memory

PiperOrigin-RevId: 463054033

Diff: diff --git a/tensorflow/lite/kernels/gather_nd.cc b/tensorflow/lite/kernels/gather_nd.cc
index 9796028f85964..dbe1ddb9ed396 100644
--- a/tensorflow/lite/kernels/gather_nd.cc
+++ b/tensorflow/lite/kernels/gather_nd.cc
@@ -12,10 +12,11 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 #include <stdint.h>
 
+#include "tensorflow/lite/c/c_api_types.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
@@ -100,17 +101,20 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   }
   return context->ResizeTensor(context, output, output_shape);
 }
 
 template <typename ParamsT, typename IndicesT>
-TfLiteStatus GatherNd(const TfLiteTensor* params, const TfLiteTensor* indices,
-                      TfLiteTensor* output) {
-  reference_ops::GatherNd(
+TfLiteStatus GatherNd(TfLiteContext* context, const TfLiteTensor* params,
+                      const TfLiteTensor* indices, TfLiteTensor* output) {
+  const TfLiteStatus status = reference_ops::GatherNd(
       GetTensorShape(params), GetTensorData<ParamsT>(params),
       GetTensorShape(indices), GetTensorData<IndicesT>(indices),
       GetTensorShape(output), GetTensorData<ParamsT>(output));
-  return kTfLiteOk;
+  if (status != kTfLiteOk) {
+    TF_LITE_KERNEL_LOG(context, "gather_nd index out of bounds");
+  }
+  return status;
 }
 
 template <typename IndicesT>
 TfLiteStatus GatherNdString(const TfLiteTensor* params,
                             const TfLiteTensor* indices, TfLiteTensor* output) {
@@ -134,21 +138,21 @@ TfLiteStatus EvalGatherNd(TfLiteContext* context, const TfLiteTensor* params,
   }
   TF_LITE_ENSURE(context, indices_has_only_positive_elements);
 
   switch (params->type) {
     case kTfLiteFloat32:
-      return GatherNd<float, IndicesT>(params, indices, output);
+      return GatherNd<float, IndicesT>(context, params, indices, output);
     case kTfLiteUInt8:
-      return GatherNd<uint8_t, IndicesT>(params, indices, output);
+      return GatherNd<uint8_t, IndicesT>(context, params, indices, output);
     case kTfLiteInt8:
-      return GatherNd<int8_t, IndicesT>(params, indices, output);
+      return GatherNd<int8_t, IndicesT>(context, params, indices, output);
     case kTfLiteInt16:
-      return GatherNd<int16_t, IndicesT>(params, indices, output);
+      return GatherNd<int16_t, IndicesT>(context, params, indices, output);
     case kTfLiteInt32:
-      return GatherNd<int32_t, IndicesT>(params, indices, output);
+      return GatherNd<int32_t, IndicesT>(context, params, indices, output);
     case kTfLiteInt64:
-      return GatherNd<int64_t, IndicesT>(params, indices, output);
+      return GatherNd<int64_t, IndicesT>(context, params, indices, output);
     case kTfLiteString:
       return GatherNdString<IndicesT>(params, indices, output);
     default:
       TF_LITE_KERNEL_LOG(context,
                          "Params type '%s' are not supported by gather_nd.",
diff --git a/tensorflow/lite/kernels/gather_nd_test.cc b/tensorflow/lite/kernels/gather_nd_test.cc
index d20e46a9f3495..5473ae0b8ccc2 100644
--- a/tensorflow/lite/kernels/gather_nd_test.cc
+++ b/tensorflow/lite/kernels/gather_nd_test.cc
@@ -71,10 +71,26 @@ TEST(GatherNdOpTest, ElementIndexingIntoMatrix) {
   ASSERT_EQ(m.Invoke(), kTfLiteOk);
 
   EXPECT_THAT(m.GetOutput<float>(), ElementsAreArray({1.1, 2.2}));
 }
 
+TEST(GatherNdOpTest, ErrorOnOutOfBoundsTooLarge) {
+  GatherNdOpModel m({TensorType_FLOAT32, {2, 2}}, {TensorType_INT32, {2, 2}});
+  m.SetInput<float>({1.1, 1.2, 2.1, 2.2});
+  m.SetPositions<int32_t>({0, 0, 2, 0});
+  EXPECT_EQ(m.Invoke(), kTfLiteError);
+  m.SetPositions<int32_t>({0, 0, 1, 2});
+  EXPECT_EQ(m.Invoke(), kTfLiteError);
+}
+
+TEST(GatherNdOpTest, ErrorOnOutOfBoundsNegative) {
+  GatherNdOpModel m({TensorType_FLOAT32, {2, 2}}, {TensorType_INT32, {2, 2}});
+  m.SetInput<float>({1.1, 1.2, 2.1, 2.2});
+  m.SetPositions<int32_t>({1, -1, 1, 1});
+  EXPECT_EQ(m.Invoke(), kTfLiteError);
+}
+
 TEST(GatherNdOpTest, SliceIndexingIntoMatrix) {
   GatherNdOpModel m({TensorType_FLOAT32, {2, 2}}, {TensorType_INT32, {2, 1}});
   m.SetInput<float>({1.1, 1.2, 2.1, 2.2});
   m.SetPositions<int32_t>({1, 0});
   ASSERT_EQ(m.Invoke(), kTfLiteOk);
diff --git a/tensorflow/lite/kernels/internal/reference/reference_ops.h b/tensorflow/lite/kernels/internal/reference/reference_ops.h
index baf74497b6b99..2879212d2fc5a 100644
--- a/tensorflow/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/lite/kernels/internal/reference/reference_ops.h
@@ -27,10 +27,11 @@ limitations under the License.
 #include <type_traits>
 
 #include "third_party/eigen3/Eigen/Core"
 #include "fixedpoint/fixedpoint.h"
 #include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/c/c_api_types.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/common.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/reference/add.h"
 #include "tensorflow/lite/kernels/internal/reference/add_n.h"
@@ -593,27 +594,35 @@ inline GatherNdHelperResult GatherNdHelper(const RuntimeShape& params_shape,
   }
 
   return ret;
 }
 
+// Implements GatherNd.
+// Returns an error if any of the indices_data would cause an out of bounds
+// memory read.
 template <typename ParamsT, typename IndicesT = int32>
-inline void GatherNd(const RuntimeShape& params_shape,
-                     const ParamsT* params_data,
-                     const RuntimeShape& indices_shape,
-                     const IndicesT* indices_data,
-                     const RuntimeShape& output_shape, ParamsT* output_data) {
+inline TfLiteStatus GatherNd(const RuntimeShape& params_shape,
+                             const ParamsT* params_data,
+                             const RuntimeShape& indices_shape,
+                             const IndicesT* indices_data,
+                             const RuntimeShape& output_shape,
+                             ParamsT* output_data) {
   ruy::profiler::ScopeLabel label("GatherNd");
 
   const GatherNdHelperResult res = GatherNdHelper(params_shape, indices_shape);
   for (int i = 0; i < res.n_slices; ++i) {
-    int from_pos = 0;
+    int64_t from_pos = 0;
     for (int j = 0; j < res.indices_nd; ++j) {
       from_pos += indices_data[i * res.indices_nd + j] * res.dims_to_count[j];
     }
+    if (from_pos < 0 || from_pos + res.slice_size > params_shape.FlatSize()) {
+      return kTfLiteError;
+    }
     std::memcpy(output_data + i * res.slice_size, params_data + from_pos,
                 sizeof(ParamsT) * res.slice_size);
   }
+  return kTfLiteOk;
 }
 
 #ifndef TF_LITE_STATIC_MEMORY
 template <typename IndicesT = int32>
 inline void GatherNdString(const RuntimeShape& params_shape,

Vulnerable code: TfLiteStatus GatherNd(const TfLiteTensor* params, const TfLiteTensor* indices,
                      TfLiteTensor* output) {
  reference_ops::GatherNd(
  return kTfLiteOk;
      return GatherNd<float, IndicesT>(params, indices, output);
      return GatherNd<uint8_t, IndicesT>(params, indices, output);
      return GatherNd<int8_t, IndicesT>(params, indices, output);
      return GatherNd<int16_t, IndicesT>(params, indices, output);
      return GatherNd<int32_t, IndicesT>(params, indices, output);
      return GatherNd<int64_t, IndicesT>(params, indices, output);
inline void GatherNd(const RuntimeShape& params_shape,
                     const ParamsT* params_data,
                     const RuntimeShape& indices_shape,
                     const IndicesT* indices_data,
                     const RuntimeShape& output_shape, ParamsT* output_data) {
    int from_pos = 0;
Fixed code: #include "tensorflow/lite/c/c_api_types.h"
TfLiteStatus GatherNd(TfLiteContext* context, const TfLiteTensor* params,
                      const TfLiteTensor* indices, TfLiteTensor* output) {
  const TfLiteStatus status = reference_ops::GatherNd(
  if (status != kTfLiteOk) {
    TF_LITE_KERNEL_LOG(context, "gather_nd index out of bounds");
  }
  return status;
      return GatherNd<float, IndicesT>(context, params, indices, output);
      return GatherNd<uint8_t, IndicesT>(context, params, indices, output);
      return GatherNd<int8_t, IndicesT>(context, params, indices, output);
      return GatherNd<int16_t, IndicesT>(context, params, indices, output);
      return GatherNd<int32_t, IndicesT>(context, params, indices, output);
      return GatherNd<int64_t, IndicesT>(context, params, indices, output);
TEST(GatherNdOpTest, ErrorOnOutOfBoundsTooLarge) {
  GatherNdOpModel m({TensorType_FLOAT32, {2, 2}}, {TensorType_INT32, {2, 2}});
  m.SetInput<float>({1.1, 1.2, 2.1, 2.2});
  m.SetPositions<int32_t>({0, 0, 2, 0});
  EXPECT_EQ(m.Invoke(), kTfLiteError);
  m.SetPositions<int32_t>({0, 0, 1, 2});
  EXPECT_EQ(m.Invoke(), kTfLiteError);
}

TEST(GatherNdOpTest, ErrorOnOutOfBoundsNegative) {
  GatherNdOpModel m({TensorType_FLOAT32, {2, 2}}, {TensorType_INT32, {2, 2}});
  m.SetInput<float>({1.1, 1.2, 2.1, 2.2});
  m.SetPositions<int32_t>({1, -1, 1, 1});
  EXPECT_EQ(m.Invoke(), kTfLiteError);
}

#include "tensorflow/lite/c/c_api_types.h"
// Implements GatherNd.
// Returns an error if any of the indices_data would cause an out of bounds
// memory read.
inline TfLiteStatus GatherNd(const RuntimeShape& params_shape,
                             const ParamsT* params_data,
                             const RuntimeShape& indices_shape,
                             const IndicesT* indices_data,
                             const RuntimeShape& output_shape,
                             ParamsT* output_data) {
    int64_t from_pos = 0;
    if (from_pos < 0 || from_pos + res.slice_size > params_shape.FlatSize()) {
      return kTfLiteError;
    }
  return kTfLiteOk;