CVE-2021-29586: Divide By Zero (CWE-369)
value: TensorFlow is an end-to-end open source platform for machine learning. Optimized pooling implementations in TFLite fail to check that the stride arguments are not 0 before calling `ComputePaddingHeightWidth`(https://github.com/tensorflow/tensorflow/blob/3f24ccd932546416ec906a02ddd183b48a1d2c83/tensorflow/lite/kernels/pooling.cc#L90). Since users can craft special models which will have `params->stride_{height,width}` be zero, this will result in a division by zero. 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/pooling.cc, tensorflow/lite/kernels/pooling_test.cc
Commit date: 2021-04-27 17:54:29 -0700
Commit message: Prevent another div by 0 in optimized pooling implementations TFLite

PiperOrigin-RevId: 370800091
Change-Id: I2119352f57fb5ca4f2051e0e2d749403304a979b

Diff: diff --git a/tensorflow/lite/kernels/pooling.cc b/tensorflow/lite/kernels/pooling.cc
index 1ae3d207b135e..474bd3825f4ff 100644
--- a/tensorflow/lite/kernels/pooling.cc
+++ b/tensorflow/lite/kernels/pooling.cc
@@ -85,10 +85,14 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) {
 
   // Matching GetWindowedOutputSize in TensorFlow.
   auto padding = params->padding;
   int out_width, out_height;
 
+  // Prevent division by 0 in optimized pooling implementations
+  TF_LITE_ENSURE(context, params->stride_height > 0);
+  TF_LITE_ENSURE(context, params->stride_width > 0);
+
   data->padding = ComputePaddingHeightWidth(
       params->stride_height, params->stride_width, 1, 1, height, width,
       params->filter_height, params->filter_width, padding, &out_height,
       &out_width);
 
diff --git a/tensorflow/lite/kernels/pooling_test.cc b/tensorflow/lite/kernels/pooling_test.cc
index e614fedccfd50..108195388141d 100644
--- a/tensorflow/lite/kernels/pooling_test.cc
+++ b/tensorflow/lite/kernels/pooling_test.cc
@@ -1149,7 +1149,20 @@ TEST(FloatPoolingOpTest, L2PoolPaddingValidSlide1) {
   });
   m.Invoke();
   EXPECT_THAT(m.GetOutput(), ElementsAreArray({3.5, 6.0, 6.5}));
 }
 
+#ifdef GTEST_HAS_DEATH_TEST
+TEST(FloatPoolingOpTest, MaxPoolWithZeroStride) {
+  EXPECT_DEATH(
+      FloatPoolingOpModel m(BuiltinOperator_MAX_POOL_2D,
+                            /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}},
+                            /*filter_width=*/2, /*filter_height=*/2,
+                            /*output=*/{TensorType_FLOAT32, {}},
+                            /*padding=*/Padding_VALID,
+                            /*stride_w=*/0, /*stride_h=*/0),
+      "Cannot allocate tensors");
+}
+#endif
+
 }  // namespace
 }  // namespace tflite

Fixed code:   // Prevent division by 0 in optimized pooling implementations
  TF_LITE_ENSURE(context, params->stride_height > 0);
  TF_LITE_ENSURE(context, params->stride_width > 0);

#ifdef GTEST_HAS_DEATH_TEST
TEST(FloatPoolingOpTest, MaxPoolWithZeroStride) {
  EXPECT_DEATH(
      FloatPoolingOpModel m(BuiltinOperator_MAX_POOL_2D,
                            /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}},
                            /*filter_width=*/2, /*filter_height=*/2,
                            /*output=*/{TensorType_FLOAT32, {}},
                            /*padding=*/Padding_VALID,
                            /*stride_w=*/0, /*stride_h=*/0),
      "Cannot allocate tensors");
}
#endif