CVE-2022-21725 (C++, published 2022-02-03T13:15Z)
Description: 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.
Severity: MEDIUM (CVSS v3 base score: 6.5)
Weakness: CWE-369 \u2014 Divide By Zero
The product divides a value by zero.
Repository: https://github.com/tensorflow/tensorflow
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 message: Internal change

PiperOrigin-RevId: 411896058
Change-Id: Ia031058247e3cf382957a6662d3f9e1cbb481ca2


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)"));
    }
  }
}