43 lines
1.7 KiB
Diff
43 lines
1.7 KiB
Diff
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;
|