36 lines
1.5 KiB
Diff
36 lines
1.5 KiB
Diff
From bb6a0383ed553c286f87ca88c207f6774d5c4a8f Mon Sep 17 00:00:00 2001
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
|
Date: Tue, 27 Jul 2021 15:20:26 -0700
|
|
Subject: [PATCH] Prevent heap OOB read in TFLite's `gather_nd.cc`.
|
|
|
|
Passing negative indices is illegal but there was a missing check so that resulted in OOB accesses.
|
|
|
|
PiperOrigin-RevId: 387208551
|
|
Change-Id: I6b7a8a62d3e7c13a16d81619e5bc23ae2cdbc7fd
|
|
---
|
|
tensorflow/lite/kernels/gather_nd.cc | 11 +++++++++++
|
|
1 file changed, 11 insertions(+)
|
|
|
|
diff --git a/tensorflow/lite/kernels/gather_nd.cc b/tensorflow/lite/kernels/gather_nd.cc
|
|
index 3ded771382569..c39917b478505 100644
|
|
--- a/tensorflow/lite/kernels/gather_nd.cc
|
|
+++ b/tensorflow/lite/kernels/gather_nd.cc
|
|
@@ -123,6 +123,17 @@ TfLiteStatus GatherNdString(const TfLiteTensor* params,
|
|
template <typename IndicesT>
|
|
TfLiteStatus EvalGatherNd(TfLiteContext* context, const TfLiteTensor* params,
|
|
const TfLiteTensor* indices, TfLiteTensor* output) {
|
|
+ bool indices_has_only_positive_elements = true;
|
|
+ const auto* indices_values = GetTensorData<IndicesT>(indices);
|
|
+ const size_t num_indices = indices->bytes / sizeof(IndicesT);
|
|
+ for (size_t i = 0; i < num_indices; i++) {
|
|
+ if (indices_values[i] < 0) {
|
|
+ indices_has_only_positive_elements = false;
|
|
+ break;
|
|
+ }
|
|
+ }
|
|
+ TF_LITE_ENSURE(context, indices_has_only_positive_elements);
|
|
+
|
|
switch (params->type) {
|
|
case kTfLiteFloat32:
|
|
return GatherNd<float, IndicesT>(params, indices, output);
|