CVE-2022-35968: Reachable Assertion (CWE-617)
value: TensorFlow is an open source platform for machine learning. The implementation of `AvgPoolGrad` does not fully validate the input `orig_input_shape`. This results in a `CHECK` failure which can be used to trigger a denial of service attack. We have patched the issue in GitHub commit 3a6ac52664c6c095aa2b114e742b0aa17fdce78f. The fix will be included in TensorFlow 2.10.0. We will also cherrypick this commit on TensorFlow 2.9.1, TensorFlow 2.8.1, and TensorFlow 2.7.2, as these are also affected and still in supported range. There are no known workarounds for this issue.
CWE description: The product contains an assert() or similar statement that can be triggered by an attacker, which leads to an application exit or other behavior that is more severe than necessary.
CVSS v3 base score: 7.5
Published: 2022-09-16T21:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/avgpooling_op.cc, tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
Commit date: 2022-07-18 15:08:51 -0700
Commit message: Fix security vulnerability with AvgPoolGrad

PiperOrigin-RevId: 461724457

Diff: diff --git a/tensorflow/core/kernels/avgpooling_op.cc b/tensorflow/core/kernels/avgpooling_op.cc
index 0429d50cdec23..7e302a7e52e65 100644
--- a/tensorflow/core/kernels/avgpooling_op.cc
+++ b/tensorflow/core/kernels/avgpooling_op.cc
@@ -296,11 +296,11 @@ class AvgPoolingGradOp : public OpKernel {
     const int64_t out_backprop_depth = out_backprop.dim_size(3);
 
     TensorShape output_shape;
     auto shape_vec = tensor_in_shape.vec<int32>();
     for (int64_t i = 0; i < tensor_in_shape.NumElements(); ++i) {
-      output_shape.AddDim(shape_vec(i));
+      OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
     }
     const int64_t in_rows = output_shape.dim_size(1);
     const int64_t in_cols = output_shape.dim_size(2);
 
     Tensor* output = nullptr;
@@ -455,11 +455,11 @@ class AvgPoolingGradOp<GPUDevice, T> : public OpKernel {
                 errors::InvalidArgument("out_backprop must be 4-dimensional"));
 
     TensorShape output_shape;
     auto shape_vec = tensor_in_shape.vec<int32>();
     for (int64_t i = 0; i < tensor_in_shape.NumElements(); ++i) {
-      output_shape.AddDim(shape_vec(i));
+      OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
     }
 
     if (output_shape.num_elements() == 0) {
       Tensor* output = nullptr;
       OP_REQUIRES_OK(context,
@@ -541,11 +541,11 @@ class AvgPoolingGradOpCustomGPUKernel : public OpKernel {
     OP_REQUIRES(context, out_backprop.dims() == 4,
                 errors::InvalidArgument("out_backprop must be 4-dimensional"));
     TensorShape output_shape;
     auto shape_vec = tensor_in_shape.vec<int32>();
     for (int64_t i = 0; i < tensor_in_shape.NumElements(); ++i) {
-      output_shape.AddDim(shape_vec(i));
+      OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
     }
     if (output_shape.num_elements() == 0) {
       Tensor* output = nullptr;
       OP_REQUIRES_OK(context,
                      context->allocate_output(0, output_shape, &output));
diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
index aca29d05a3686..391f8ec025f3c 100644
--- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
@@ -2468,10 +2468,26 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
             r"Expected argmax shape to be \[1,1,1,1\], but got \[1,1,1,2\]"):
           gen_nn_ops.max_pool_grad_grad_with_argmax(
               inp, grad, argmax, ksize=[1, 1, 1, 1], strides=[1, 1, 1, 1],
               padding="VALID")
 
+  def testAvgPoolGradInvalidInputShapeRaiseError(self):
+    with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
+      with self.cached_session():
+        orig_input_shape = constant_op.constant(
+            -536870912, shape=[4], dtype=dtypes.int32)
+        grad = constant_op.constant(
+            .0890338004362538, shape=[1, 5, 7, 1], dtype=dtypes.float64)
+        t = gen_nn_ops.AvgPoolGrad(
+            orig_input_shape=orig_input_shape,
+            grad=grad,
+            ksize=[1, 2, 2, 1],
+            strides=[1, 2, 2, 1],
+            padding="VALID",
+            data_format="NHWC")
+        self.evaluate(t)
+
 
 def GetMaxPoolFwdTest(input_size, filter_size, strides, padding):
 
   def Test(self):
     # MaxPoolWithArgMax is implemented only on CUDA.

Vulnerable code:       output_shape.AddDim(shape_vec(i));
      output_shape.AddDim(shape_vec(i));
      output_shape.AddDim(shape_vec(i));
Fixed code:       OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
      OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
      OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
  def testAvgPoolGradInvalidInputShapeRaiseError(self):
    with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
      with self.cached_session():
        orig_input_shape = constant_op.constant(
            -536870912, shape=[4], dtype=dtypes.int32)
        grad = constant_op.constant(
            .0890338004362538, shape=[1, 5, 7, 1], dtype=dtypes.float64)
        t = gen_nn_ops.AvgPoolGrad(
            orig_input_shape=orig_input_shape,
            grad=grad,
            ksize=[1, 2, 2, 1],
            strides=[1, 2, 2, 1],
            padding="VALID",
            data_format="NHWC")
        self.evaluate(t)