CVE-2022-23559: Integer Overflow or Wraparound (CWE-190)
value: 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.
CWE 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.
CVSS v3 base score: 8.8
Published: 2022-02-04T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/lite/kernels/embedding_lookup_sparse.cc
Commit date: 2021-12-21 08:52:29 -0800
Commit message: [lite] Check for overflow when creating required bytes.

PiperOrigin-RevId: 417629001
Change-Id: Ia7feb3ea8e988f4fd4b3c98c1a1fed4557d99fd7

Diff: diff --git a/tensorflow/lite/kernels/embedding_lookup_sparse.cc b/tensorflow/lite/kernels/embedding_lookup_sparse.cc
index 4ad1054340c9c..a0b9586203a93 100644
--- a/tensorflow/lite/kernels/embedding_lookup_sparse.cc
+++ b/tensorflow/lite/kernels/embedding_lookup_sparse.cc
@@ -70,10 +70,11 @@ limitations under the License.
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #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 {
 namespace builtin {
 
@@ -173,29 +174,37 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
 
   // Resize output tensor.
   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);
 
   // Keep track of the current bucket for aggregation/combination.
   int current_output_offset = 0;

Vulnerable code:   int embedding_size = 1;
  int lookup_size = 1;
    const int dim = dense_shape->data.i32[i];
    lookup_size *= dim;
    const int dim = SizeOfDimension(value, i);
    embedding_size *= dim;
  const int output_size = lookup_size * embedding_size;
Fixed code: #include "tensorflow/lite/util.h"
  size_t embedding_size = 1;
  size_t lookup_size = 1;
    const size_t dim = dense_shape->data.i32[i];
    TF_LITE_ENSURE_MSG(
        context,
        MultiplyAndCheckOverflow(lookup_size, dim, &lookup_size) == kTfLiteOk,
        "Lookup size overflowed.");
    const size_t dim = SizeOfDimension(value, i);
    TF_LITE_ENSURE_MSG(context,
                       MultiplyAndCheckOverflow(embedding_size, dim,
                                                &embedding_size) == kTfLiteOk,
                       "Embedding size overflowed.");
  const size_t output_size = lookup_size * embedding_size;
  // Makes sure reallocation was successful.
  TF_LITE_ENSURE(context, output_ptr != nullptr);