CVE-2022-35963: Reachable Assertion (CWE-617)
value: TensorFlow is an open source platform for machine learning. The implementation of `FractionalAvgPoolGrad` does not fully validate the input `orig_input_tensor_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 03a659d7be9a1154fdf5eeac221e5950fec07dad. 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: Python
Files changed: tensorflow/core/kernels/fractional_avg_pool_op.cc, tensorflow/python/kernel_tests/nn_ops/fractional_avg_pool_op_test.py
Commit date: 2022-07-20 20:59:08 -0700
Commit message: Fix security vulnerability with FractionalAvgPoolGrad

PiperOrigin-RevId: 462292194

Diff: diff --git a/tensorflow/core/kernels/fractional_avg_pool_op.cc b/tensorflow/core/kernels/fractional_avg_pool_op.cc
index b3e65aeaee22f..e1c536af7065d 100644
--- a/tensorflow/core/kernels/fractional_avg_pool_op.cc
+++ b/tensorflow/core/kernels/fractional_avg_pool_op.cc
@@ -10,26 +10,27 @@ Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
+
 #define EIGEN_USE_THREADS
 
 #include <algorithm>
 #include <cmath>
 #include <random>
 #include <vector>
 
-#include "tensorflow/core/kernels/fractional_pool_common.h"
-
 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 #include "tensorflow/core/framework/numeric_op.h"
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/kernels/fractional_pool_common.h"
 #include "tensorflow/core/lib/random/random.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/util/guarded_philox_random.h"
+#include "tensorflow/core/util/overflow.h"
 
 namespace tensorflow {
 typedef Eigen::ThreadPoolDevice CPUDevice;
 
 template <typename T>
@@ -239,11 +240,36 @@ class FractionalAvgPoolGradOp : public OpKernel {
     OP_REQUIRES(context,
                 orig_input_tensor_shape.dims() == 1 &&
                     orig_input_tensor_shape.NumElements() == 4,
                 errors::InvalidArgument("original input tensor shape must be"
                                         "1-dimensional and 4 elements"));
+    int64_t num_elements = 1;
+    for (int i = 0; i < orig_input_tensor_shape.dims(); i++) {
+      OP_REQUIRES(context, orig_input_tensor_shape.dim_size(i) > 0,
+                  errors::InvalidArgument(
+                      "orig_input_tensor_shape must be positive, got: ",
+                      orig_input_tensor_shape.dim_size(i)));
+      num_elements = MultiplyWithoutOverflow(
+          num_elements, orig_input_tensor_shape.dim_size(i));
+      OP_REQUIRES(
+          context, num_elements > 0,
+          errors::InvalidArgument(
+              "The total elements specified by orig_input_tensor_shape",
+              " is too large. Encountered overflow after multiplying ",
+              orig_input_tensor_shape.dim_size(i), ", result: ", num_elements));
+    }
+
     const Tensor& out_backprop = context->input(1);
+    OP_REQUIRES(context, out_backprop.dims() == 4,
+                errors::InvalidArgument("out_backprop must be 4-dimensional"));
+    for (int i = 0; i < out_backprop.dims(); i++) {
+      OP_REQUIRES(context, out_backprop.dim_size(i) > 0,
+                  errors::InvalidArgument(
+                      "out_backprop must be positive for all dimension, got:",
+                      out_backprop.dim_size(i)));
+    }
+
     const Tensor& row_seq_tensor = context->input(2);
     const Tensor& col_seq_tensor = context->input(3);
 
     const int64_t out_batch = out_backprop.dim_size(0);
     const int64_t out_rows = out_backprop.dim_size(1);
diff --git a/tensorflow/python/kernel_tests/nn_ops/fractional_avg_pool_op_test.py b/tensorflow/python/kernel_tests/nn_ops/fractional_avg_pool_op_test.py
index 7b153ae1ed708..642a9d4a88002 100644
--- a/tensorflow/python/kernel_tests/nn_ops/fractional_avg_pool_op_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/fractional_avg_pool_op_test.py
@@ -539,8 +539,29 @@ class FractionalAvgPoolGradTest(test.TestCase):
           output_shape,
           x_init_value=input_data.reshape(input_shape),
           delta=1e-2)
       self.assertLess(gradient_error, error_margin)
 
+  def testInvalidSeqRaiseErrorForFractionalAvgPoolGrad(self):
+    with self.assertRaises((errors.InvalidArgumentError, ValueError)):
+      with self.cached_session() as _:
+        overlapping = True
+        orig_input_tensor_shape = constant_op.constant(
+            -1879048192, shape=[4], dtype=dtypes.int64)
+        out_backprop = constant_op.constant([],
+                                            shape=[0, 0, 0, 0],
+                                            dtype=dtypes.float64)
+        row_pooling_sequence = constant_op.constant(
+            1, shape=[4], dtype=dtypes.int64)
+        col_pooling_sequence = constant_op.constant(
+            1, shape=[4], dtype=dtypes.int64)
+        t = gen_nn_ops.fractional_avg_pool_grad(
+            orig_input_tensor_shape=orig_input_tensor_shape,
+            out_backprop=out_backprop,
+            row_pooling_sequence=row_pooling_sequence,
+            col_pooling_sequence=col_pooling_sequence,
+            overlapping=overlapping)
+        self.evaluate(t)
+
 
 if __name__ == "__main__":
   test.main()

Vulnerable code: #include "tensorflow/core/kernels/fractional_pool_common.h"

Fixed code: 
#include "tensorflow/core/kernels/fractional_pool_common.h"
#include "tensorflow/core/util/overflow.h"
    int64_t num_elements = 1;
    for (int i = 0; i < orig_input_tensor_shape.dims(); i++) {
      OP_REQUIRES(context, orig_input_tensor_shape.dim_size(i) > 0,
                  errors::InvalidArgument(
                      "orig_input_tensor_shape must be positive, got: ",
                      orig_input_tensor_shape.dim_size(i)));
      num_elements = MultiplyWithoutOverflow(
          num_elements, orig_input_tensor_shape.dim_size(i));
      OP_REQUIRES(
          context, num_elements > 0,
          errors::InvalidArgument(
              "The total elements specified by orig_input_tensor_shape",
              " is too large. Encountered overflow after multiplying ",
              orig_input_tensor_shape.dim_size(i), ", result: ", num_elements));
    }

    OP_REQUIRES(context, out_backprop.dims() == 4,
                errors::InvalidArgument("out_backprop must be 4-dimensional"));
    for (int i = 0; i < out_backprop.dims(); i++) {
      OP_REQUIRES(context, out_backprop.dim_size(i) > 0,
                  errors::InvalidArgument(
                      "out_backprop must be positive for all dimension, got:",
                      out_backprop.dim_size(i)));
    }

  def testInvalidSeqRaiseErrorForFractionalAvgPoolGrad(self):
    with self.assertRaises((errors.InvalidArgumentError, ValueError)):
      with self.cached_session() as _:
        overlapping = True
        orig_input_tensor_shape = constant_op.constant(
            -1879048192, shape=[4], dtype=dtypes.int64)
        out_backprop = constant_op.constant([],
                                            shape=[0, 0, 0, 0],
                                            dtype=dtypes.float64)
        row_pooling_sequence = constant_op.constant(
            1, shape=[4], dtype=dtypes.int64)
        col_pooling_sequence = constant_op.constant(
            1, shape=[4], dtype=dtypes.int64)
        t = gen_nn_ops.fractional_avg_pool_grad(
            orig_input_tensor_shape=orig_input_tensor_shape,
            out_backprop=out_backprop,
            row_pooling_sequence=row_pooling_sequence,
            col_pooling_sequence=col_pooling_sequence,
            overlapping=overlapping)
        self.evaluate(t)