From 7c8cc4ec69cd348e44ad6a2699057ca88faad3e5 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Thu, 29 Apr 2021 19:43:09 -0700
Subject: [PATCH] Fix a dangerous integer overflow and a malloc of negative
 size.

PiperOrigin-RevId: 371254154
Change-Id: I250a98a3df26328770167025670235a963a72da0
---
 tensorflow/lite/c/common.c                         | 6 ++++--
 tensorflow/lite/kernels/embedding_lookup_sparse.cc | 1 +
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/c/common.c b/tensorflow/lite/c/common.c
index aaa98a98ebe69..00dd0260cbcc9 100644
--- a/tensorflow/lite/c/common.c
+++ b/tensorflow/lite/c/common.c
@@ -45,8 +45,10 @@ int TfLiteIntArrayEqualsArray(const TfLiteIntArray* a, int b_size,
 #ifndef TF_LITE_STATIC_MEMORY
 
 TfLiteIntArray* TfLiteIntArrayCreate(int size) {
-  TfLiteIntArray* ret =
-      (TfLiteIntArray*)malloc(TfLiteIntArrayGetSizeInBytes(size));
+  int alloc_size = TfLiteIntArrayGetSizeInBytes(size);
+  if (alloc_size <= 0) return NULL;
+  TfLiteIntArray* ret = (TfLiteIntArray*)malloc(alloc_size);
+  if (!ret) return ret;
   ret->size = size;
   return ret;
 }
diff --git a/tensorflow/lite/kernels/embedding_lookup_sparse.cc b/tensorflow/lite/kernels/embedding_lookup_sparse.cc
index e9ad7e50cf133..4ad1054340c9c 100644
--- a/tensorflow/lite/kernels/embedding_lookup_sparse.cc
+++ b/tensorflow/lite/kernels/embedding_lookup_sparse.cc
@@ -173,6 +173,7 @@ 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;