CVE-2022-41900: Out-of-bounds Write (CWE-787)
value: TensorFlow is an open source platform for machine learning. The security vulnerability results in FractionalMax(AVG)Pool with illegal pooling_ratio. Attackers using Tensorflow can exploit the vulnerability. They can access heap memory which is not in the control of user, leading to a crash or remote code execution. We have patched the issue in GitHub commit 216525144ee7c910296f5b05d214ca1327c9ce48. The fix will be included in TensorFlow 2.11.0. We will also cherry pick this commit on TensorFlow 2.10.1.
CWE description: The product writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 9.8
Published: 2022-11-18T22:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/fractional_avg_pool_op.cc, tensorflow/core/kernels/fractional_max_pool_op.cc, tensorflow/core/ops/nn_ops.cc, tensorflow/core/ops/nn_ops_test.cc, tensorflow/python/kernel_tests/nn_ops/fractional_avg_pool_op_test.py, tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py
Commit date: 2022-10-24 14:34:32 -0700
Commit message: Fix security vulnerability with FractionalMax(AVG)Pool with illegal pooling_ratio

PiperOrigin-RevId: 483486453

Diff: diff --git a/tensorflow/core/kernels/fractional_avg_pool_op.cc b/tensorflow/core/kernels/fractional_avg_pool_op.cc
index e1c536af7065d..3bb206866089a 100644
--- a/tensorflow/core/kernels/fractional_avg_pool_op.cc
+++ b/tensorflow/core/kernels/fractional_avg_pool_op.cc
@@ -42,10 +42,16 @@ class FractionalAvgPoolOp : public OpKernel {
     OP_REQUIRES_OK(context, context->GetAttr("pseudo_random", &pseudo_random_));
     OP_REQUIRES_OK(context, context->GetAttr("overlapping", &overlapping_));
     OP_REQUIRES(context, pooling_ratio_.size() == 4,
                 errors::InvalidArgument(
                     "pooling_ratio field must specify 4 dimensions"));
+    for (std::size_t i = 0; i < pooling_ratio_.size(); ++i) {
+      OP_REQUIRES(context, pooling_ratio_[i] >= 1,
+                  errors::InvalidArgument(
+                      "pooling_ratio cannot be smaller than 1, got: ",
+                      pooling_ratio_[i]));
+    }
     OP_REQUIRES(
         context, pooling_ratio_[0] == 1 || pooling_ratio_[3] == 1,
         errors::Unimplemented("Fractional average pooling is not yet "
                               "supported on the batch nor channel dimension."));
     OP_REQUIRES_OK(context, context->GetAttr("deterministic", &deterministic_));
@@ -80,13 +86,15 @@ class FractionalAvgPoolOp : public OpKernel {
     std::vector<int> input_size(tensor_in_and_out_dims);
     std::vector<int> output_size(tensor_in_and_out_dims);
     for (int i = 0; i < tensor_in_and_out_dims; ++i) {
       input_size[i] = tensor_in.dim_size(i);
       OP_REQUIRES(
-          context, pooling_ratio_[i] <= input_size[i],
-          errors::InvalidArgument(
-              "Pooling ratio cannot be bigger than input tensor dim size."));
+          context, input_size[i] >= pooling_ratio_[i],
+          errors::InvalidArgument("Pooling ratio is higher than input "
+                                  "dimension size for dimension ",
+                                  i, ". Input dim size: ", input_size[i],
+                                  " pooling ratio: ", pooling_ratio_[i]));
     }
     // Output size.
     for (int i = 0; i < tensor_in_and_out_dims; ++i) {
       output_size[i] =
           static_cast<int>(std::floor(input_size[i] / pooling_ratio_[i]));
diff --git a/tensorflow/core/kernels/fractional_max_pool_op.cc b/tensorflow/core/kernels/fractional_max_pool_op.cc
index 3ea45277bb5a4..ec08b5c502872 100644
--- a/tensorflow/core/kernels/fractional_max_pool_op.cc
+++ b/tensorflow/core/kernels/fractional_max_pool_op.cc
@@ -43,10 +43,16 @@ class FractionalMaxPoolOp : public OpKernel {
     OP_REQUIRES_OK(context, context->GetAttr("overlapping", &overlapping_));
 
     OP_REQUIRES(context, pooling_ratio_.size() == 4,
                 errors::InvalidArgument("pooling_ratio field must "
                                         "specify 4 dimensions"));
+    for (std::size_t i = 0; i < pooling_ratio_.size(); ++i) {
+      OP_REQUIRES(context, pooling_ratio_[i] >= 1,
+                  errors::InvalidArgument(
+                      "pooling_ratio cannot be smaller than 1, got: ",
+                      pooling_ratio_[i]));
+    }
 
     OP_REQUIRES(
         context, pooling_ratio_[0] == 1 || pooling_ratio_[3] == 1,
         errors::Unimplemented("Fractional max pooling is not yet "
                               "supported on the batch nor channel dimension."));
diff --git a/tensorflow/core/ops/nn_ops.cc b/tensorflow/core/ops/nn_ops.cc
index 417034477d1fc..c16563001e48c 100644
--- a/tensorflow/core/ops/nn_ops.cc
+++ b/tensorflow/core/ops/nn_ops.cc
@@ -61,10 +61,17 @@ Status FractionalPoolShapeFn(InferenceContext* c) {
     } else {
       output_dims.push_back(c->UnknownDim());
     }
   }
 
+  for (std::size_t i = 0; i < pooling_ratio.size(); ++i) {
+    if (pooling_ratio[i] < 1) {
+      return errors::InvalidArgument(
+          "pooling_ratio cannot be smaller than 1, got: ", pooling_ratio[i]);
+    }
+  }
+
   c->set_output(0, c->MakeShape(output_dims));
   c->set_output(1, c->Vector(output_dims[1]));
   c->set_output(2, c->Vector(output_dims[2]));
   return OkStatus();
 }
diff --git a/tensorflow/core/ops/nn_ops_test.cc b/tensorflow/core/ops/nn_ops_test.cc
index 469a9015a17c9..41940da69ef4c 100644
--- a/tensorflow/core/ops/nn_ops_test.cc
+++ b/tensorflow/core/ops/nn_ops_test.cc
@@ -521,24 +521,25 @@ TEST(NNOpsTest, FractionalPool_ShapeFn) {
                        .Input("input", 0, DT_FLOAT)
                        .Attr("pooling_ratio", pooling_ratio)
                        .Finalize(&op.node_def));
     };
 
-    set_op(std::vector<float>{2.0f, 1, 1 / 1.5f, 1 / 2.0f});
+    // pooling_ratio must >= 1.0
+    set_op(std::vector<float>{2.0f, 1, 1.5f, 4.0f});
 
     // Rank check.
     INFER_ERROR("must be rank 4", op, "[?,?,?]");
 
     // Unknown inputs.
     INFER_OK(op, "?", "[?,?,?,?];[?];[?]");
     INFER_OK(op, "[?,?,?,?]", "[?,?,?,?];[?];[?]");
 
-    INFER_OK(op, "[10,20,30,40]", "[5,20,45,80];[20];[45]");
-    INFER_OK(op, "[?,20,30,40]", "[?,20,45,80];[20];[45]");
-    INFER_OK(op, "[10,?,30,40]", "[5,?,45,80];[?];[45]");
-    INFER_OK(op, "[10,20,?,40]", "[5,20,?,80];[20];[?]");
-    INFER_OK(op, "[10,20,30,?]", "[5,20,45,?];[20];[45]");
+    INFER_OK(op, "[10,20,30,40]", "[5,20,20,10];[20];[20]");
+    INFER_OK(op, "[?,20,30,40]", "[?,20,20,10];[20];[20]");
+    INFER_OK(op, "[10,?,30,40]", "[5,?,20,10];[?];[20]");
+    INFER_OK(op, "[10,20,?,40]", "[5,20,?,10];[20];[?]");
+    INFER_OK(op, "[10,20,30,?]", "[5,20,20,?];[20];[20]");
 
     // Wrong number of values for pooling_ratio.
     set_op(std::vector<float>{.5, 1.0, 1.5});
     INFER_ERROR("pooling_ratio field", op, "?");
     set_op(std::vector<float>{1, 2, 3, 4, 5});
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 642a9d4a88002..59b20de84b4c4 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
@@ -331,10 +331,45 @@ class FractionalAvgTest(test.TestCase):
             col_pooling_sequence=[1, 2, -3, 4],
             overlapping=True)
 
         self.evaluate(z)
 
+  def testPoolingRatioHasMoreDimThanInput(self):
+    with self.cached_session() as _:
+      with self.assertRaisesRegex(
+          errors.InvalidArgumentError,
+          r"Pooling ratio is higher than input dimension size for dimension 1.*"
+      ):
+        result = nn_ops.gen_nn_ops.fractional_avg_pool(
+            value=constant_op.constant(
+                value=[[[[1, 4, 2, 3]]]], dtype=dtypes.int64),
+            pooling_ratio=[1.0, 1.44, 1.73, 1.0],
+            pseudo_random=False,
+            overlapping=False,
+            deterministic=False,
+            seed=0,
+            seed2=0,
+            name=None)
+        self.evaluate(result)
+
+  def testPoolingRatioValueOutOfRange(self):
+    with self.cached_session() as _:
+      # Whether turn on `TF2_BEHAVIOR` generates different error messages
+      with self.assertRaisesRegex(
+          (errors.InvalidArgumentError, ValueError),
+          r"(pooling_ratio cannot be smaller than 1, got: .*)|(is negative)"):
+        result = nn_ops.gen_nn_ops.fractional_avg_pool(
+            value=np.zeros([3, 30, 30, 3]),
+            pooling_ratio=[1, -1, 3, 1],
+            pseudo_random=False,
+            overlapping=False,
+            deterministic=False,
+            seed=0,
+            seed2=0,
+        )
+        self.evaluate(result)
+
 
 class FractionalAvgPoolGradTest(test.TestCase):
   """Tests for FractionalAvgPoolGrad.
 
   Two types of tests for FractionalAvgPoolGrad.
diff --git a/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py b/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py
index 79b2c799d58e3..9102973fa13b5 100644
--- a/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py
@@ -318,11 +318,11 @@ class FractionalMaxPoolTest(test.TestCase):
                                   'requires "seed" and "seed2" to be non-zero'):
         nn_ops.fractional_max_pool(rand_mat, [1, 1.5, 1.5, 1])
       nn_ops.fractional_max_pool(
           rand_mat, [1, 1.5, 1.5, 1], seed=1, seed2=1, deterministic=True)
 
-  def testPoolingRatio(self):
+  def testPoolingRatioHasMoreDimThanInput(self):
     with self.cached_session() as _:
       with self.assertRaisesRegex(
           errors.InvalidArgumentError,
           r"Pooling ratio is higher than input dimension size for dimension 1.*"
       ):
@@ -336,10 +336,27 @@ class FractionalMaxPoolTest(test.TestCase):
             seed=0,
             seed2=0,
             name=None)
         self.evaluate(result)
 
+  def testPoolingRatioValueOutOfRange(self):
+    with self.cached_session() as _:
+      # Whether turn on `TF2_BEHAVIOR` generates different error messages
+      with self.assertRaisesRegex(
+          (errors.InvalidArgumentError, ValueError),
+          r"(pooling_ratio cannot be smaller than 1, got: .*)|(is negative)"):
+        result = nn_ops.gen_nn_ops.fractional_max_pool(
+            value=np.zeros([3, 30, 30, 3]),
+            pooling_ratio=[1, -1, 3, 1],
+            pseudo_random=False,
+            overlapping=False,
+            deterministic=False,
+            seed=0,
+            seed2=0,
+        )
+        self.evaluate(result)
+
 
 class FractionalMaxPoolGradTest(test.TestCase):
   """Tests for FractionalMaxPoolGrad.
 
   Two types of tests for FractionalMaxPoolGrad.

Vulnerable code:           context, pooling_ratio_[i] <= input_size[i],
          errors::InvalidArgument(
              "Pooling ratio cannot be bigger than input tensor dim size."));
    set_op(std::vector<float>{2.0f, 1, 1 / 1.5f, 1 / 2.0f});
    INFER_OK(op, "[10,20,30,40]", "[5,20,45,80];[20];[45]");
    INFER_OK(op, "[?,20,30,40]", "[?,20,45,80];[20];[45]");
    INFER_OK(op, "[10,?,30,40]", "[5,?,45,80];[?];[45]");
    INFER_OK(op, "[10,20,?,40]", "[5,20,?,80];[20];[?]");
    INFER_OK(op, "[10,20,30,?]", "[5,20,45,?];[20];[45]");
  def testPoolingRatio(self):
Fixed code:     for (std::size_t i = 0; i < pooling_ratio_.size(); ++i) {
      OP_REQUIRES(context, pooling_ratio_[i] >= 1,
                  errors::InvalidArgument(
                      "pooling_ratio cannot be smaller than 1, got: ",
                      pooling_ratio_[i]));
    }
          context, input_size[i] >= pooling_ratio_[i],
          errors::InvalidArgument("Pooling ratio is higher than input "
                                  "dimension size for dimension ",
                                  i, ". Input dim size: ", input_size[i],
                                  " pooling ratio: ", pooling_ratio_[i]));
    for (std::size_t i = 0; i < pooling_ratio_.size(); ++i) {
      OP_REQUIRES(context, pooling_ratio_[i] >= 1,
                  errors::InvalidArgument(
                      "pooling_ratio cannot be smaller than 1, got: ",
                      pooling_ratio_[i]));
    }
  for (std::size_t i = 0; i < pooling_ratio.size(); ++i) {
    if (pooling_ratio[i] < 1) {
      return errors::InvalidArgument(
          "pooling_ratio cannot be smaller than 1, got: ", pooling_ratio[i]);
    }
  }

    // pooling_ratio must >= 1.0
    set_op(std::vector<float>{2.0f, 1, 1.5f, 4.0f});
    INFER_OK(op, "[10,20,30,40]", "[5,20,20,10];[20];[20]");
    INFER_OK(op, "[?,20,30,40]", "[?,20,20,10];[20];[20]");
    INFER_OK(op, "[10,?,30,40]", "[5,?,20,10];[?];[20]");
    INFER_OK(op, "[10,20,?,40]", "[5,20,?,10];[20];[?]");
    INFER_OK(op, "[10,20,30,?]", "[5,20,20,?];[20];[20]");
  def testPoolingRatioHasMoreDimThanInput(self):
    with self.cached_session() as _:
      with self.assertRaisesRegex(
          errors.InvalidArgumentError,
          r"Pooling ratio is higher than input dimension size for dimension 1.*"
      ):
        result = nn_ops.gen_nn_ops.fractional_avg_pool(
            value=constant_op.constant(
                value=[[[[1, 4, 2, 3]]]], dtype=dtypes.int64),
            pooling_ratio=[1.0, 1.44, 1.73, 1.0],
            pseudo_random=False,
            overlapping=False,
            deterministic=False,
            seed=0,
            seed2=0,
            name=None)
        self.evaluate(result)

  def testPoolingRatioValueOutOfRange(self):
    with self.cached_session() as _:
      # Whether turn on `TF2_BEHAVIOR` generates different error messages
      with self.assertRaisesRegex(
          (errors.InvalidArgumentError, ValueError),
          r"(pooling_ratio cannot be smaller than 1, got: .*)|(is negative)"):
        result = nn_ops.gen_nn_ops.fractional_avg_pool(
            value=np.zeros([3, 30, 30, 3]),
            pooling_ratio=[1, -1, 3, 1],
            pseudo_random=False,
            overlapping=False,
            deterministic=False,
            seed=0,
            seed2=0,
        )
        self.evaluate(result)

  def testPoolingRatioHasMoreDimThanInput(self):
  def testPoolingRatioValueOutOfRange(self):
    with self.cached_session() as _:
      # Whether turn on `TF2_BEHAVIOR` generates different error messages
      with self.assertRaisesRegex(
          (errors.InvalidArgumentError, ValueError),
          r"(pooling_ratio cannot be smaller than 1, got: .*)|(is negative)"):
        result = nn_ops.gen_nn_ops.fractional_max_pool(
            value=np.zeros([3, 30, 30, 3]),
            pooling_ratio=[1, -1, 3, 1],
            pseudo_random=False,
            overlapping=False,
            deterministic=False,
            seed=0,
            seed2=0,
        )
        self.evaluate(result)