CVE-2022-23559 (published 2022-02-04T23:15Z)

Tensorflow is an Open Source Machine Learning Framework. An attacker can craft a TFLite model that would cause an integer overflow in embedding lookup operations. Both `embedding_size` and `lookup_size` are products of values provided by the user. Hence, a malicious user could trigger overflows in the multiplication. In certain scenarios, this can then result in heap OOB read/write. Users are advised to upgrade to a patched version.

Severity: MEDIUM
CVSS v2 base score: 6.5 (AV:N/AC:L/Au:S/C:P/I:P/A:P)
CVSS v3 base score: 8.8, HIGH (CVSS:3.1/AV:N/AC:L/PR:L/UI:N/S:U/C:H/I:H/A:H)

Weakness information:
cwe_id: CWE-190
cwe_name: Integer Overflow or Wraparound
description: The product performs a calculation that can
         produce an integer overflow or wraparound when the logic
         assumes that the resulting value will always be larger than
         the original value. This occurs when an integer value is
         incremented to a value that is too large to store in the
         associated representation. When this occurs, the value may
         become a very small or negative number.
extended_description: 
            
            
            
            
         
         

Fix information:
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Mihai Maruseac
author_date: 2021-12-21 13:06:51-08:00
msg: Prevent segfault in `embedding_lookup_sparse.cc`

Previous fixes missed one additional case.

PiperOrigin-RevId: 417676944
Change-Id: I8ab412155cf9b1e897448a6611d209eaa7ca9e66
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: embedding_lookup_sparse.cc
change_type: ModificationType.MODIFY
diff: @@ -159,6 +159,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 3, &weights));
   const TfLiteTensor* value;
   TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 4, &value));
+  const size_t values_size = NumElements(value);
 
   const int lookup_rank = SizeOfDimension(indices, 1);
   const int embedding_rank = NumDimensions(value);
@@ -253,6 +254,11 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
     current_squares_weight += w * w;
     current_total_weight += w;
     for (int k = 0; k < embedding_size; k++) {
+      // only index if indices are valid
+      if (current_output_offset + k < 0) continue;
+      if (current_output_offset + k >= output_size) continue;
+      if (example_embedding_offset + k < 0) continue;
+      if (example_embedding_offset + k >= values_size) continue;
       output_ptr[current_output_offset + k] +=
           value_ptr[example_embedding_offset + k] * w;
     }

programming_language: C++
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Karim Nosir
author_date: 2021-12-21 08:48:11-08:00
msg: [lite] Check for overflow when creating required bytes.

PiperOrigin-RevId: 417629001
Change-Id: Ia7feb3ea8e988f4fd4b3c98c1a1fed4557d99fd7
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: embedding_lookup_sparse.cc
change_type: ModificationType.MODIFY
diff: @@ -72,6 +72,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/tensor_utils.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/util.h"
 
 namespace tflite {
 namespace ops {
@@ -175,25 +176,33 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   TfLiteIntArray* output_shape = TfLiteIntArrayCreate(output_rank);
   TF_LITE_ENSURE(context, output_shape != nullptr);
   int k = 0;
-  int embedding_size = 1;
-  int lookup_size = 1;
+  size_t embedding_size = 1;
+  size_t lookup_size = 1;
   for (int i = 0; i < lookup_rank - 1; i++, k++) {
-    const int dim = dense_shape->data.i32[i];
-    lookup_size *= dim;
+    const size_t dim = dense_shape->data.i32[i];
+    TF_LITE_ENSURE_MSG(
+        context,
+        MultiplyAndCheckOverflow(lookup_size, dim, &lookup_size) == kTfLiteOk,
+        "Lookup size overflowed.");
     output_shape->data[k] = dim;
   }
   for (int i = 1; i < embedding_rank; i++, k++) {
-    const int dim = SizeOfDimension(value, i);
-    embedding_size *= dim;
+    const size_t dim = SizeOfDimension(value, i);
+    TF_LITE_ENSURE_MSG(context,
+                       MultiplyAndCheckOverflow(embedding_size, dim,
+                                                &embedding_size) == kTfLiteOk,
+                       "Embedding size overflowed.");
     output_shape->data[k] = dim;
   }
   TF_LITE_ENSURE_STATUS(context->ResizeTensor(context, output, output_shape));
-  const int output_size = lookup_size * embedding_size;
+  const size_t output_size = lookup_size * embedding_size;
   TfLiteTensorRealloc(output_size * sizeof(float), output);
 
   float* output_ptr = GetTensorData<float>(output);
   const float* weights_ptr = GetTensorData<float>(weights);
   const float* value_ptr = GetTensorData<float>(value);
+  // Makes sure reallocation was successful.
+  TF_LITE_ENSURE(context, output_ptr != nullptr);
 
   std::fill_n(output_ptr, output_size, 0.0f);
 

programming_language: C++
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Karim Nosir
author_date: 2021-12-16 14:32:07-08:00
msg: [lite] Move MultiplyAndCheckOverflow to util to be able to share it.

PiperOrigin-RevId: 416897229
Change-Id: I5feb44881bdcbb6ed911da4f17c55bb978754059
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: BUILD
change_type: ModificationType.MODIFY
diff: @@ -1030,6 +1030,7 @@ cc_library(
     copts = tflite_copts_warnings() + tflite_copts(),
     deps = [
         ":kernel_api",
+        ":macros",
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/schema:schema_fbs",
     ],
@@ -1083,6 +1084,7 @@ cc_test(
     features = ["-dynamic_link_test_srcs"],  # see go/dynamic_link_test_srcs
     deps = [
         ":util",
+        "//tensorflow/lite/c:c_api_types",
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest_main",

programming_language: Python
filename: subgraph.cc
change_type: ModificationType.MODIFY
diff: @@ -690,27 +690,6 @@ TfLiteStatus Subgraph::CheckInputAndOutputForOverlap(const int* input_indices,
   return kTfLiteOk;
 }
 
-namespace {
-// Multiply two sizes and return true if overflow occurred;
-// This is based off tensorflow/overflow.h but is simpler as we already
-// have unsigned numbers. It is also generalized to work where sizeof(size_t)
-// is not 8.
-TfLiteStatus MultiplyAndCheckOverflow(size_t a, size_t b, size_t* product) {
-  // Multiplying a * b where a and b are size_t cannot result in overflow in a
-  // size_t accumulator if both numbers have no non-zero bits in their upper
-  // half.
-  constexpr size_t size_t_bits = 8 * sizeof(size_t);
-  constexpr size_t overflow_upper_half_bit_position = size_t_bits / 2;
-  *product = a * b;
-  // If neither integers have non-zero bits past 32 bits can't overflow.
-  // Otherwise check using slow devision.
-  if (TFLITE_EXPECT_FALSE((a | b) >> overflow_upper_half_bit_position != 0)) {
-    if (a != 0 && *product / a != b) return kTfLiteError;
-  }
-  return kTfLiteOk;
-}
-}  // namespace
-
 TfLiteStatus Subgraph::BytesRequired(TfLiteType type, const int* dims,
                                      size_t dims_size, size_t* bytes) {
   TF_LITE_ENSURE(&context_, bytes != nullptr);

programming_language: C++
filename: util.cc
change_type: ModificationType.MODIFY
diff: @@ -27,6 +27,7 @@ limitations under the License.
 
 #include "tensorflow/lite/builtin_ops.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/core/macros.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
@@ -176,4 +177,19 @@ bool IsValidationSubgraph(const char* name) {
   // NOLINTNEXTLINE: can't use absl::StartsWith as absl is not allowed.
   return name && std::string(name).find(kValidationSubgraphNamePrefix) == 0;
 }
+
+TfLiteStatus MultiplyAndCheckOverflow(size_t a, size_t b, size_t* product) {
+  // Multiplying a * b where a and b are size_t cannot result in overflow in a
+  // size_t accumulator if both numbers have no non-zero bits in their upper
+  // half.
+  constexpr size_t size_t_bits = 8 * sizeof(size_t);
+  constexpr size_t overflow_upper_half_bit_position = size_t_bits / 2;
+  *product = a * b;
+  // If neither integers have non-zero bits past 32 bits can't overflow.
+  // Otherwise check using slow devision.
+  if (TFLITE_EXPECT_FALSE((a | b) >> overflow_upper_half_bit_position != 0)) {
+    if (a != 0 && *product / a != b) return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
 }  // namespace tflite

programming_language: C++
filename: util.h
change_type: ModificationType.MODIFY
diff: @@ -99,6 +99,12 @@ constexpr char kValidationSubgraphNamePrefix[] = "VALIDATION:";
 // Checks whether the prefix of the subgraph name indicates the subgraph is a
 // validation subgraph.
 bool IsValidationSubgraph(const char* name);
+
+// Multiply two sizes and return true if overflow occurred;
+// This is based off tensorflow/overflow.h but is simpler as we already
+// have unsigned numbers. It is also generalized to work where sizeof(size_t)
+// is not 8.
+TfLiteStatus MultiplyAndCheckOverflow(size_t a, size_t b, size_t* product);
 }  // namespace tflite
 
 #endif  // TENSORFLOW_LITE_UTIL_H_

programming_language: C++
filename: util_test.cc
change_type: ModificationType.MODIFY
diff: @@ -22,6 +22,7 @@ limitations under the License.
 #include <vector>
 
 #include <gtest/gtest.h>
+#include "tensorflow/lite/c/c_api_types.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 
@@ -130,5 +131,12 @@ TEST(ValidationSubgraph, NameIsDetected) {
   EXPECT_TRUE(IsValidationSubgraph("VALIDATION:main"));
 }
 
+TEST(MultiplyAndCheckOverflow, Validate) {
+  size_t res = 0;
+  EXPECT_TRUE(MultiplyAndCheckOverflow(1, 2, &res) == kTfLiteOk);
+  EXPECT_FALSE(MultiplyAndCheckOverflow(static_cast<size_t>(123456789023),
+                                        1223423425, &res) == kTfLiteOk);
+}
+
 }  // namespace
 }  // namespace tflite

programming_language: C++