CVE-2022-35959: Reachable Assertion (CWE-617)
value: TensorFlow is an open source platform for machine learning. The implementation of `AvgPool3DGradOp` does not fully validate the input `orig_input_shape`. This results in an overflow that results in a `CHECK` failure which can be used to trigger a denial of service attack. We have patched the issue in GitHub commit 9178ac9d6389bdc54638ab913ea0e419234d14eb. 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-16T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Other
Files changed: tensorflow/compiler/tf2xla/BUILD, tensorflow/compiler/tf2xla/xla_op_kernel.cc, tensorflow/core/kernels/pooling_ops_3d.cc, tensorflow/python/kernel_tests/nn_ops/BUILD, tensorflow/python/kernel_tests/nn_ops/pooling_ops_3d_test.py
Commit date: 2022-07-15 14:57:23 -0700
Commit message: Fix security vulnerability with AvgPool3DGrad.

PiperOrigin-RevId: 461244371

Diff: diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD
index bcc6f018afeb3..f5d2d1db2ab73 100644
--- a/tensorflow/compiler/tf2xla/BUILD
+++ b/tensorflow/compiler/tf2xla/BUILD
@@ -401,10 +401,11 @@ cc_library(
         "//tensorflow/compiler/xla/client:client_library",
         "//tensorflow/compiler/xla/client:local_client",
         "//tensorflow/compiler/xla/client:xla_builder",
         "//tensorflow/compiler/xla/client:xla_computation",
         "//tensorflow/compiler/xla/service:hlo",
+        "//tensorflow/core/util:overflow",
         "//tensorflow/core:core_cpu",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.cc b/tensorflow/compiler/tf2xla/xla_op_kernel.cc
index 3d70377d206b8..4d30d25aa07a1 100644
--- a/tensorflow/compiler/tf2xla/xla_op_kernel.cc
+++ b/tensorflow/compiler/tf2xla/xla_op_kernel.cc
@@ -28,10 +28,11 @@ limitations under the License.
 #include "tensorflow/compiler/xla/client/xla_builder.h"
 #include "tensorflow/compiler/xla/client/xla_computation.h"
 #include "tensorflow/compiler/xla/status_macros.h"
 #include "tensorflow/core/common_runtime/dma_helper.h"
 #include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/util/overflow.h"
 
 namespace tensorflow {
 
 XlaOpKernelContext::XlaOpKernelContext(OpKernelContext* context)
     : context_(context),
@@ -441,10 +442,20 @@ Status XlaOpKernelContext::ConstantInputAsShape(int index, TensorShape* shape,
                                                 xla::ValueInferenceMode mode) {
   xla::Literal literal;
   TF_RETURN_IF_ERROR(ConstantInput(index, &literal, mode));
   std::vector<int64_t> dims;
   TF_RETURN_IF_ERROR(LiteralToInt64Vector(literal, &dims));
+
+  int64_t num_elements = 1;
+  for (auto i = dims.begin(); i != dims.end(); ++i) {
+    num_elements = MultiplyWithoutOverflow(num_elements, *i);
+    if (num_elements < 0)
+      return errors::InvalidArgument(
+          "The total elements specified by orig_input_shape is too large.",
+          "Encountered overflow after multiplying", *i,
+          ", result: ", num_elements);
+  }
   *shape = TensorShape(dims);
   return OkStatus();
 }
 
 Status XlaOpKernelContext::ConstantInputAsPartialShape(
diff --git a/tensorflow/core/kernels/pooling_ops_3d.cc b/tensorflow/core/kernels/pooling_ops_3d.cc
index 57f4ac80ecab7..60f1718f240e2 100644
--- a/tensorflow/core/kernels/pooling_ops_3d.cc
+++ b/tensorflow/core/kernels/pooling_ops_3d.cc
@@ -521,11 +521,11 @@ class AvgPooling3dGradOp : public OpKernel {
                 errors::InvalidArgument("out_backprop must be 5-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)));
     }
 
     Tensor* output;
     OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
 
diff --git a/tensorflow/python/kernel_tests/nn_ops/BUILD b/tensorflow/python/kernel_tests/nn_ops/BUILD
index 4109a67d9ee1f..3c3f26456d891 100644
--- a/tensorflow/python/kernel_tests/nn_ops/BUILD
+++ b/tensorflow/python/kernel_tests/nn_ops/BUILD
@@ -498,10 +498,11 @@ cuda_py_test(
     name = "pooling_ops_3d_test",
     size = "medium",
     srcs = ["pooling_ops_3d_test.py"],
     deps = [
         "//tensorflow/python:client_testlib",
+        "//tensorflow/python:dtypes",
         "//tensorflow/python:framework_for_generated_wrappers",
         "//tensorflow/python:nn_grad",
         "//tensorflow/python:nn_ops",
         "//third_party/py/numpy",
     ],
diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_3d_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_3d_test.py
index 71bf5d02bcb56..9c8f11eb74309 100644
--- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_3d_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_3d_test.py
@@ -16,10 +16,11 @@
 
 import numpy as np
 
 from tensorflow.python.eager import context
 from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gen_nn_ops
@@ -65,11 +66,11 @@ class PoolingTest(test.TestCase):
     for s in input_sizes:
       total_size *= s
     # Initializes the input tensor with array containing incrementing
     # numbers from 1.
     x = [f * 1.0 for f in range(1, total_size + 1)]
-    with self.cached_session(use_gpu=use_gpu) as sess:
+    with self.cached_session(use_gpu=use_gpu):
       t = constant_op.constant(x, shape=input_sizes)
       window = [1] + list(window) + [1]
       strides = [1] + list(strides) + [1]
       if data_format == "NCDHW":
         t = test_util.NHWCToNCHW(t)
@@ -122,10 +123,27 @@ class PoolingTest(test.TestCase):
         window=(1, 2, 3),
         strides=(2, 3, 1),
         padding="SAME",
         expected=expected_output)
 
+  def testMaxPool3dGrad(self):
+    with self.assertRaises(
+        (errors.ResourceExhaustedError, errors.InvalidArgumentError)):
+      with self.cached_session():
+        orig_input_shape = constant_op.constant(
+            1879048192, shape=[5], dtype=dtypes.int32)
+        grad = constant_op.constant(
+            1, shape=[1, 3, 2, 4, 2], dtype=dtypes.float32)
+        t = gen_nn_ops.AvgPool3DGrad(
+            orig_input_shape=orig_input_shape,
+            grad=grad,
+            ksize=[1, 1, 1, 1, 1],
+            strides=[1, 1, 1, 1, 1],
+            padding="SAME",
+            data_format="NDHWC")
+        self.evaluate(t)
+
   def testMaxPool3dValidPadding(self):
     expected_output = [40.0, 41.0, 42.0]
     self._VerifyValues(
         nn_ops.max_pool3d,
         input_sizes=[1, 3, 3, 3, 3],

Vulnerable code:       output_shape.AddDim(shape_vec(i));
    with self.cached_session(use_gpu=use_gpu) as sess:
Fixed code:         "//tensorflow/core/util:overflow",
#include "tensorflow/core/util/overflow.h"

  int64_t num_elements = 1;
  for (auto i = dims.begin(); i != dims.end(); ++i) {
    num_elements = MultiplyWithoutOverflow(num_elements, *i);
    if (num_elements < 0)
      return errors::InvalidArgument(
          "The total elements specified by orig_input_shape is too large.",
          "Encountered overflow after multiplying", *i,
          ", result: ", num_elements);
  }
      OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
        "//tensorflow/python:dtypes",
from tensorflow.python.framework import dtypes
    with self.cached_session(use_gpu=use_gpu):
  def testMaxPool3dGrad(self):
    with self.assertRaises(
        (errors.ResourceExhaustedError, errors.InvalidArgumentError)):
      with self.cached_session():
        orig_input_shape = constant_op.constant(
            1879048192, shape=[5], dtype=dtypes.int32)
        grad = constant_op.constant(
            1, shape=[1, 3, 2, 4, 2], dtype=dtypes.float32)
        t = gen_nn_ops.AvgPool3DGrad(
            orig_input_shape=orig_input_shape,
            grad=grad,
            ksize=[1, 1, 1, 1, 1],
            strides=[1, 1, 1, 1, 1],
            padding="SAME",
            data_format="NDHWC")
        self.evaluate(t)