CVE-2022-21725: Divide By Zero (CWE-369)
value: Tensorflow is an Open Source Machine Learning Framework. The estimator for the cost of some convolution operations can be made to execute a division by 0. The function fails to check that the stride argument is strictly positive. Hence, the fix is to add a check for the stride argument to ensure it is valid. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.
CWE description: The product divides a value by zero.
CVSS v3 base score: 6.5
Published: 2022-02-03T13:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/grappler/costs/BUILD, tensorflow/core/grappler/costs/op_level_cost_estimator.cc, tensorflow/core/grappler/costs/op_level_cost_estimator.h, tensorflow/core/grappler/costs/op_level_cost_estimator_test.cc
Commit date: 2021-11-23 14:37:54 -0800
Commit message: Internal change

PiperOrigin-RevId: 411896058
Change-Id: Ia031058247e3cf382957a6662d3f9e1cbb481ca2

Diff: diff --git a/tensorflow/core/grappler/costs/BUILD b/tensorflow/core/grappler/costs/BUILD
index e306aa1eaccb2..34bd9c7fc2120 100644
--- a/tensorflow/core/grappler/costs/BUILD
+++ b/tensorflow/core/grappler/costs/BUILD
@@ -353,10 +353,11 @@ tf_cc_test(
         ":op_level_cost_estimator",
         "//tensorflow/core:framework",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:status_matchers",
     ],
 )
 
 cc_library(
     name = "analytical_cost_estimator",
diff --git a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
index 6e516e2a68612..3dab8bccc543e 100644
--- a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
+++ b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
@@ -2151,11 +2151,11 @@ OpInfo::TensorProperties OpLevelCostEstimator::DescribeTensor(
 
   return ret;
 }
 
 /* static */
-OpLevelCostEstimator::ConvolutionDimensions
+StatusOr<OpLevelCostEstimator::ConvolutionDimensions>
 OpLevelCostEstimator::OpDimensionsFromInputs(
     const TensorShapeProto& original_image_shape, const OpInfo& op_info,
     bool* found_unknown_shapes) {
   VLOG(2) << "op features: " << op_info.DebugString();
   VLOG(2) << "Original image shape: " << original_image_shape.DebugString();
@@ -2188,10 +2188,15 @@ OpLevelCostEstimator::OpDimensionsFromInputs(
   int64_t kz = iz;
 
   std::vector<int64_t> strides = GetStrides(op_info);
   int64_t sx = strides[x_index];
   int64_t sy = strides[y_index];
+  if (sx == 0 || sy == 0) {
+    return errors::InvalidArgument(
+        "Stride must be > 0 for Height and Width, but got (", sy, ", ", sx,
+        ")");
+  }
   const auto padding = GetPadding(op_info);
 
   int64_t ox = GetOutputSize(ix, kx, sx, padding);
   int64_t oy = GetOutputSize(iy, ky, sy, padding);
   int64_t oz = iz;
@@ -2204,12 +2209,13 @@ OpLevelCostEstimator::OpDimensionsFromInputs(
 Status OpLevelCostEstimator::PredictMaxPool(const OpContext& op_context,
                                             NodeCosts* node_costs) const {
   bool found_unknown_shapes = false;
   const auto& op_info = op_context.op_info;
   // x: op_info.inputs(0)
-  ConvolutionDimensions dims = OpDimensionsFromInputs(
-      op_info.inputs(0).shape(), op_info, &found_unknown_shapes);
+  TF_ASSIGN_OR_RETURN(ConvolutionDimensions dims,
+                      OpDimensionsFromInputs(op_info.inputs(0).shape(), op_info,
+                                             &found_unknown_shapes));
   // kx * ky - 1 comparisons per output (kx * xy > 1)
   // or 1 copy per output (kx * k1 = 1).
   int per_output_ops = dims.kx * dims.ky == 1 ? 1 : dims.kx * dims.ky - 1;
   int64_t ops = dims.batch * dims.ox * dims.oy * dims.oz * per_output_ops;
   node_costs->num_compute_ops = ops;
@@ -2246,12 +2252,13 @@ Status OpLevelCostEstimator::PredictMaxPoolGrad(const OpContext& op_context,
   if (op_info.inputs_size() < 3) {
     return errors::InvalidArgument("MaxPoolGrad op has invalid inputs: ",
                                    op_info.ShortDebugString());
   }
 
-  ConvolutionDimensions dims = OpDimensionsFromInputs(
-      op_info.inputs(0).shape(), op_info, &found_unknown_shapes);
+  TF_ASSIGN_OR_RETURN(ConvolutionDimensions dims,
+                      OpDimensionsFromInputs(op_info.inputs(0).shape(), op_info,
+                                             &found_unknown_shapes));
 
   int64_t ops = 0;
   if (dims.kx == 1 && dims.ky == 1) {
     // 1x1 window. No need to know which input was max.
     ops = dims.batch * dims.ix * dims.iy * dims.iz;
@@ -2322,12 +2329,13 @@ Status OpLevelCostEstimator::PredictAssignVariableOps(
 Status OpLevelCostEstimator::PredictAvgPool(const OpContext& op_context,
                                             NodeCosts* node_costs) const {
   bool found_unknown_shapes = false;
   const auto& op_info = op_context.op_info;
   // x: op_info.inputs(0)
-  ConvolutionDimensions dims = OpDimensionsFromInputs(
-      op_info.inputs(0).shape(), op_info, &found_unknown_shapes);
+  TF_ASSIGN_OR_RETURN(ConvolutionDimensions dims,
+                      OpDimensionsFromInputs(op_info.inputs(0).shape(), op_info,
+                                             &found_unknown_shapes));
 
   // kx * ky - 1 additions and 1 multiplication per output.
   int64_t ops = dims.batch * dims.ox * dims.oy * dims.oz * dims.kx * dims.ky;
   node_costs->num_compute_ops = ops;
 
@@ -2380,12 +2388,13 @@ Status OpLevelCostEstimator::PredictAvgPoolGrad(const OpContext& op_context,
       x_shape.add_dim()->set_size(1);
     }
     found_unknown_shapes = true;
   }
 
-  ConvolutionDimensions dims =
-      OpDimensionsFromInputs(x_shape, op_info, &found_unknown_shapes);
+  TF_ASSIGN_OR_RETURN(
+      ConvolutionDimensions dims,
+      OpDimensionsFromInputs(x_shape, op_info, &found_unknown_shapes));
 
   int64_t ops = 0;
   if (dims.kx <= dims.sx && dims.ky <= dims.sy) {
     // Non-overlapping window.
     ops = dims.batch * dims.iz * (dims.ix * dims.iy + dims.ox * dims.oy);
@@ -2407,12 +2416,13 @@ Status OpLevelCostEstimator::PredictFusedBatchNorm(
   // x: op_info.inputs(0)
   // scale: op_info.inputs(1)
   // offset: op_info.inputs(2)
   // mean: op_info.inputs(3)  --> only for inference
   // variance: op_info.inputs(4) --> only for inference
-  ConvolutionDimensions dims = OpDimensionsFromInputs(
-      op_info.inputs(0).shape(), op_info, &found_unknown_shapes);
+  TF_ASSIGN_OR_RETURN(ConvolutionDimensions dims,
+                      OpDimensionsFromInputs(op_info.inputs(0).shape(), op_info,
+                                             &found_unknown_shapes));
   const bool is_training = IsTraining(op_info);
 
   int64_t ops = 0;
   const auto rsqrt_cost = Eigen::internal::functor_traits<
       Eigen::internal::scalar_rsqrt_op<float>>::Cost;
@@ -2457,12 +2467,13 @@ Status OpLevelCostEstimator::PredictFusedBatchNormGrad(
   // y_backprop: op_info.inputs(0)
   // x: op_info.inputs(1)
   // scale: op_info.inputs(2)
   // mean: op_info.inputs(3)
   // variance or inverse of variance: op_info.inputs(4)
-  ConvolutionDimensions dims = OpDimensionsFromInputs(
-      op_info.inputs(1).shape(), op_info, &found_unknown_shapes);
+  TF_ASSIGN_OR_RETURN(ConvolutionDimensions dims,
+                      OpDimensionsFromInputs(op_info.inputs(1).shape(), op_info,
+                                             &found_unknown_shapes));
 
   int64_t ops = 0;
   const auto rsqrt_cost = Eigen::internal::functor_traits<
       Eigen::internal::scalar_rsqrt_op<float>>::Cost;
   ops = dims.iz * (dims.batch * dims.ix * dims.iy * 11 + 5 + rsqrt_cost);
diff --git a/tensorflow/core/grappler/costs/op_level_cost_estimator.h b/tensorflow/core/grappler/costs/op_level_cost_estimator.h
index 6c5764f3fe3a8..f7de566548a78 100644
--- a/tensorflow/core/grappler/costs/op_level_cost_estimator.h
+++ b/tensorflow/core/grappler/costs/op_level_cost_estimator.h
@@ -288,11 +288,11 @@ class OpLevelCostEstimator {
       const TensorShapeProto& original_image_shape,
       const TensorShapeProto& original_filter_shape, const OpInfo& op_info,
       bool* found_unknown_shapes);
 
   // For Pooling, FusedBatchNorm, and their grad ops.
-  static ConvolutionDimensions OpDimensionsFromInputs(
+  static StatusOr<ConvolutionDimensions> OpDimensionsFromInputs(
       const TensorShapeProto& original_image_shape, const OpInfo& op_info,
       bool* found_unknown_shapes);
 
   // Helper to construct child operation contexts for the component operations
   // of fused ops.
diff --git a/tensorflow/core/grappler/costs/op_level_cost_estimator_test.cc b/tensorflow/core/grappler/costs/op_level_cost_estimator_test.cc
index 560845ee980eb..0cddce42df2f0 100644
--- a/tensorflow/core/grappler/costs/op_level_cost_estimator_test.cc
+++ b/tensorflow/core/grappler/costs/op_level_cost_estimator_test.cc
@@ -22,10 +22,11 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor.pb.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/protobuf/device_properties.pb.h"
 
 namespace tensorflow {
 namespace grappler {
@@ -556,13 +557,14 @@ class OpLevelCostEstimatorTest : public ::testing::Test {
       ho = op_context.op_info.outputs(0).shape().dim(2).size();
       wo = op_context.op_info.outputs(0).shape().dim(3).size();
     }
 
     bool found_unknown_shapes;
-    auto dims = OpLevelCostEstimator::OpDimensionsFromInputs(
-        op_context.op_info.inputs(0).shape(), op_context.op_info,
-        &found_unknown_shapes);
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto dims, OpLevelCostEstimator::OpDimensionsFromInputs(
+                       op_context.op_info.inputs(0).shape(), op_context.op_info,
+                       &found_unknown_shapes));
     Padding padding_enum;
     if (padding == "VALID") {
       padding_enum = Padding::VALID;
     } else {
       padding_enum = Padding::SAME;
@@ -579,10 +581,42 @@ class OpLevelCostEstimatorTest : public ::testing::Test {
     EXPECT_EQ(wo, dims.oy);
     EXPECT_EQ(c, dims.oz);
     EXPECT_EQ(padding_enum, dims.padding);
   }
 
+  StatusOr<OpLevelCostEstimator::ConvolutionDimensions>
+  CallOpDimensionsFromInputs(const int n, const int h, const int w, const int c,
+                             const int kx, const int ky, const int sx,
+                             const int sy, const string& data_format,
+                             const string& padding) {
+    OpContext op_context;
+
+    const std::vector<int> x = {n, h, w, c};
+    const std::vector<int> ksize = {1, kx, ky, 1};
+    std::vector<int> strides;
+    if (data_format == "NHWC") {
+      strides = {1, sy, sx, 1};
+    } else {
+      strides = {1, 1, sy, sx};
+    }
+
+    auto& op_info = op_context.op_info;
+    SetCpuDevice(&op_info);
+    op_info.set_op("MaxPool");
+
+    DescribeTensor4D(x[0], x[1], x[2], x[3], op_info.add_inputs());
+    auto* attr = op_info.mutable_attr();
+    SetAttrValue(data_format, &(*attr)["data_format"]);
+    SetAttrValue(padding, &(*attr)["padding"]);
+    SetAttrValue(strides, &(*attr)["strides"]);
+    SetAttrValue(ksize, &(*attr)["ksize"]);
+    bool found_unknown_shapes;
+    return OpLevelCostEstimator::OpDimensionsFromInputs(
+        op_context.op_info.inputs(0).shape(), op_context.op_info,
+        &found_unknown_shapes);
+  }
+
   OpLevelCostEstimator estimator_;
 };
 
 class OpLevelBatchMatMulCostEstimatorTest
     : public OpLevelCostEstimatorTest,
@@ -1381,10 +1415,30 @@ TEST_F(OpLevelCostEstimatorTest, OpDimensionsFromInputs) {
       ValidateOpDimensionsFromInputs(10, 14, 14, 3840, 3, 3, 2, 2, f, p);
     }
   }
 }
 
+TEST_F(OpLevelCostEstimatorTest, OpDimensionsFromInputsError) {
+  std::vector<string> paddings = {"VALID", "SAME"};
+  std::vector<string> formats = {"NHWC", "NCHW"};
+  for (const auto& p : paddings) {
+    for (const auto& f : formats) {
+      // n, h, w, c, kx, ky, sx, sy, data_format, padding.
+      ASSERT_THAT(
+          CallOpDimensionsFromInputs(10, 14, 14, 3840, 3, 3, 0, 2, f, p),
+          testing::StatusIs(
+              error::INVALID_ARGUMENT,
+              "Stride must be > 0 for Height and Width, but got (2, 0)"));
+      ASSERT_THAT(
+          CallOpDimensionsFromInputs(10, 14, 14, 3840, 3, 3, 2, 0, f, p),
+          testing::StatusIs(
+              error::INVALID_ARGUMENT,
+              "Stride must be > 0 for Height and Width, but got (0, 2)"));
+    }
+  }
+}
+
 TEST_F(OpLevelCostEstimatorTest, PredictMaxPool) {
   auto predict_max_pool = [this](const int n, const int in, const int c,
                                  const int k, const int s,
                                  const string& padding) -> Costs {
     OpContext op_context = DescribePoolingOp(

Vulnerable code: OpLevelCostEstimator::ConvolutionDimensions
  ConvolutionDimensions dims = OpDimensionsFromInputs(
      op_info.inputs(0).shape(), op_info, &found_unknown_shapes);
  ConvolutionDimensions dims = OpDimensionsFromInputs(
      op_info.inputs(0).shape(), op_info, &found_unknown_shapes);
  ConvolutionDimensions dims = OpDimensionsFromInputs(
      op_info.inputs(0).shape(), op_info, &found_unknown_shapes);
  ConvolutionDimensions dims =
      OpDimensionsFromInputs(x_shape, op_info, &found_unknown_shapes);
  ConvolutionDimensions dims = OpDimensionsFromInputs(
      op_info.inputs(0).shape(), op_info, &found_unknown_shapes);
  ConvolutionDimensions dims = OpDimensionsFromInputs(
      op_info.inputs(1).shape(), op_info, &found_unknown_shapes);
  static ConvolutionDimensions OpDimensionsFromInputs(
    auto dims = OpLevelCostEstimator::OpDimensionsFromInputs(
        op_context.op_info.inputs(0).shape(), op_context.op_info,
        &found_unknown_shapes);
Fixed code:         "//tensorflow/core/platform:status_matchers",
StatusOr<OpLevelCostEstimator::ConvolutionDimensions>
  if (sx == 0 || sy == 0) {
    return errors::InvalidArgument(
        "Stride must be > 0 for Height and Width, but got (", sy, ", ", sx,
        ")");
  }
  TF_ASSIGN_OR_RETURN(ConvolutionDimensions dims,
                      OpDimensionsFromInputs(op_info.inputs(0).shape(), op_info,
                                             &found_unknown_shapes));
  TF_ASSIGN_OR_RETURN(ConvolutionDimensions dims,
                      OpDimensionsFromInputs(op_info.inputs(0).shape(), op_info,
                                             &found_unknown_shapes));
  TF_ASSIGN_OR_RETURN(ConvolutionDimensions dims,
                      OpDimensionsFromInputs(op_info.inputs(0).shape(), op_info,
                                             &found_unknown_shapes));
  TF_ASSIGN_OR_RETURN(
      ConvolutionDimensions dims,
      OpDimensionsFromInputs(x_shape, op_info, &found_unknown_shapes));
  TF_ASSIGN_OR_RETURN(ConvolutionDimensions dims,
                      OpDimensionsFromInputs(op_info.inputs(0).shape(), op_info,
                                             &found_unknown_shapes));
  TF_ASSIGN_OR_RETURN(ConvolutionDimensions dims,
                      OpDimensionsFromInputs(op_info.inputs(1).shape(), op_info,
                                             &found_unknown_shapes));
  static StatusOr<ConvolutionDimensions> OpDimensionsFromInputs(
#include "tensorflow/core/platform/status_matchers.h"
    TF_ASSERT_OK_AND_ASSIGN(
        auto dims, OpLevelCostEstimator::OpDimensionsFromInputs(
                       op_context.op_info.inputs(0).shape(), op_context.op_info,
                       &found_unknown_shapes));
  StatusOr<OpLevelCostEstimator::ConvolutionDimensions>
  CallOpDimensionsFromInputs(const int n, const int h, const int w, const int c,
                             const int kx, const int ky, const int sx,
                             const int sy, const string& data_format,
                             const string& padding) {
    OpContext op_context;

    const std::vector<int> x = {n, h, w, c};
    const std::vector<int> ksize = {1, kx, ky, 1};
    std::vector<int> strides;
    if (data_format == "NHWC") {
      strides = {1, sy, sx, 1};
    } else {
      strides = {1, 1, sy, sx};
    }

    auto& op_info = op_context.op_info;
    SetCpuDevice(&op_info);
    op_info.set_op("MaxPool");

    DescribeTensor4D(x[0], x[1], x[2], x[3], op_info.add_inputs());
    auto* attr = op_info.mutable_attr();
    SetAttrValue(data_format, &(*attr)["data_format"]);
    SetAttrValue(padding, &(*attr)["padding"]);
    SetAttrValue(strides, &(*attr)["strides"]);
    SetAttrValue(ksize, &(*attr)["ksize"]);
    bool found_unknown_shapes;
    return OpLevelCostEstimator::OpDimensionsFromInputs(
        op_context.op_info.inputs(0).shape(), op_context.op_info,
        &found_unknown_shapes);
  }

TEST_F(OpLevelCostEstimatorTest, OpDimensionsFromInputsError) {
  std::vector<string> paddings = {"VALID", "SAME"};
  std::vector<string> formats = {"NHWC", "NCHW"};
  for (const auto& p : paddings) {
    for (const auto& f : formats) {
      // n, h, w, c, kx, ky, sx, sy, data_format, padding.
      ASSERT_THAT(
          CallOpDimensionsFromInputs(10, 14, 14, 3840, 3, 3, 0, 2, f, p),
          testing::StatusIs(
              error::INVALID_ARGUMENT,
              "Stride must be > 0 for Height and Width, but got (2, 0)"));
      ASSERT_THAT(
          CallOpDimensionsFromInputs(10, 14, 14, 3840, 3, 3, 2, 0, f, p),
          testing::StatusIs(
              error::INVALID_ARGUMENT,
              "Stride must be > 0 for Height and Width, but got (0, 2)"));
    }
  }
}