CVE-2021-29595: Divide By Zero (CWE-369)
value: TensorFlow is an end-to-end open source platform for machine learning. The implementation of the `DepthToSpace` TFLite operator is vulnerable to a division by zero error(https://github.com/tensorflow/tensorflow/blob/0d45ea1ca641b21b73bcf9c00e0179cda284e7e7/tensorflow/lite/kernels/depth_to_space.cc#L63-L69). An attacker can craft a model such that `params->block_size` is 0. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range.
CWE description: The product divides a value by zero.
CVSS v3 base score: 7.8
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/lite/kernels/depth_to_space.cc, tensorflow/lite/kernels/depth_to_space_test.cc, tensorflow/lite/micro/kernels/depth_to_space.cc
Commit date: 2021-04-27 18:25:13 -0700
Commit message: Prevent division by 0 in TFLite

PiperOrigin-RevId: 370800311
Change-Id: I21ccdbd31c30118acc67df8751807ee2e0b12f91

Diff: diff --git a/tensorflow/lite/kernels/depth_to_space.cc b/tensorflow/lite/kernels/depth_to_space.cc
index d61049f85a9ad..9b552541c74b7 100644
--- a/tensorflow/lite/kernels/depth_to_space.cc
+++ b/tensorflow/lite/kernels/depth_to_space.cc
@@ -59,10 +59,11 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                      data_type == kTfLiteInt8 || data_type == kTfLiteInt32 ||
                      data_type == kTfLiteInt64);
   TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   const int block_size = params->block_size;
+  TF_LITE_ENSURE(context, block_size > 0);
   const int input_height = input->dims->data[1];
   const int input_width = input->dims->data[2];
   const int input_channels = input->dims->data[3];
   int output_height = input_height * block_size;
   int output_width = input_width * block_size;
diff --git a/tensorflow/lite/kernels/depth_to_space_test.cc b/tensorflow/lite/kernels/depth_to_space_test.cc
index 4429faf990917..c03512dd710ad 100644
--- a/tensorflow/lite/kernels/depth_to_space_test.cc
+++ b/tensorflow/lite/kernels/depth_to_space_test.cc
@@ -58,10 +58,15 @@ class DepthToSpaceOpModel : public SingleOpModel {
 #ifdef GTEST_HAS_DEATH_TEST
 TEST(DepthToSpaceOpModel, BadBlockSize) {
   EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 4),
                "Cannot allocate tensors");
 }
+
+TEST(DepthToSpaceOpModel, NoBlockSize) {
+  EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 0),
+               "Cannot allocate tensors");
+}
 #endif
 
 TEST(DepthToSpaceOpModel, Float32) {
   DepthToSpaceOpModel m({TensorType_FLOAT32, {1, 1, 1, 4}}, 2);
   m.SetInput<float>({1.4, 2.3, 3.2, 4.1});
diff --git a/tensorflow/lite/micro/kernels/depth_to_space.cc b/tensorflow/lite/micro/kernels/depth_to_space.cc
index b3960f1f0c6b1..ae42ee1b331db 100644
--- a/tensorflow/lite/micro/kernels/depth_to_space.cc
+++ b/tensorflow/lite/micro/kernels/depth_to_space.cc
@@ -52,10 +52,11 @@ TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE(context,
                  data_type == kTfLiteFloat32 || data_type == kTfLiteInt8);
   TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   const int block_size = params->block_size;
+  TF_LITE_ENSURE(context, block_size > 0);
   const int input_height = input->dims->data[kHeightRank];
   const int input_width = input->dims->data[kWidthRank];
   const int input_channels = input->dims->data[kDepthRank];
   int output_height = input_height * block_size;
   int output_width = input_width * block_size;

Fixed code:   TF_LITE_ENSURE(context, block_size > 0);

TEST(DepthToSpaceOpModel, NoBlockSize) {
  EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 0),
               "Cannot allocate tensors");
}
  TF_LITE_ENSURE(context, block_size > 0);