Fix security vulnerability with FractionalMaxPoolGrad

PiperOrigin-RevId: 477500477
tensorflow · Sep 28, 2022 · d71090c · d71090c
1 parent 984128e
commit d71090c
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/tensorflow/core/kernels/fractional_max_pool_op.cc b/tensorflow/core/kernels/fractional_max_pool_op.cc
@@ -258,6 +258,18 @@ class FractionalMaxPoolGradOp : public OpKernel {
     OP_REQUIRES(context, tensor_out.NumElements() > 0,
                 errors::InvalidArgument("orig_output must not be empty, got ",
                                         tensor_out.DebugString()));
+    OP_REQUIRES(
+        context,
+        height_seq_tensor.NumElements() * width_seq_tensor.NumElements() <=
+            tensor_in.NumElements(),
+        errors::InvalidArgument(
+            "Pooling region has more elements than the input tensor. "
+            "row_pooling_sequence: ",
+            height_seq_tensor.DebugString(),
+            "col_pooling_sequence: ", width_seq_tensor.DebugString(),
+            "orig_input: ", tensor_in.DebugString()));
+
+    //
     std::vector<int64_t> input_size(tensor_in_and_out_dims);
     std::vector<int64_t> output_size(tensor_in_and_out_dims);
     for (int i = 0; i < tensor_in_and_out_dims; ++i) {

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
@@ -632,7 +632,7 @@ def testWhenRepeatedMaxValueInPoolingRegion(self):
 
   def testInvalidSeqRaiseErrorForFractionalMaxPoolGrad(self):
     with self.assertRaises(errors.InvalidArgumentError):
-      with self.cached_session() as _:
+      with self.cached_session():
         overlapping = True
         orig_input = constant_op.constant(
             .453409232, shape=[1, 7, 13, 1], dtype=dtypes.float32)
@@ -653,6 +653,24 @@ def testInvalidSeqRaiseErrorForFractionalMaxPoolGrad(self):
             overlapping=overlapping)
         self.evaluate(t)
 
+  def testOverLargeSeqRaiseErrorForFractionalMaxPoolGrad(self):
+    with self.assertRaises(errors.InvalidArgumentError):
+      with self.cached_session():
+        overlapping = False
+        orig_input = [[[[1, 1, 1, 1, 1]]]]
+        orig_output = [[[[1, 1, 1]]]]
+        out_backprop = [[[[3], [3], [6]]]]
+        row_pooling_sequence = [-0x4000000, 1, 1]
+        col_pooling_sequence = [-0x4000000, 1, 1]
+        t = gen_nn_ops.FractionalMaxPoolGrad(
+            orig_input=orig_input,
+            orig_output=orig_output,
+            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()