diff options
author | Mihai Maruseac <mihaimaruseac@google.com> | 2022-08-19 13:16:20 -0700 |
---|---|---|
committer | GitHub <noreply@github.com> | 2022-08-19 13:16:20 -0700 |
commit | 50f90f4ce5c2fb06c8c478e2692a8ccc2a9b718f (patch) | |
tree | c86cc018903467eb3049b366543951fe5351a5ff | |
parent | be8b86aff1e2233af7323ae04dba521930dbd41c (diff) | |
parent | da45c48c412fb9585ee0d99cd13b28acb1a4f20a (diff) | |
download | tensorflow-50f90f4ce5c2fb06c8c478e2692a8ccc2a9b718f.tar.gz |
Merge pull request #57273 from tensorflow/r2.7-8741e57d163
r2.7 cherry-pick: 8741e57d163 "Fix security vulnerability with FractionalMaxPoolGrad"
-rw-r--r-- | tensorflow/core/kernels/fractional_max_pool_op.cc | 20 | ||||
-rw-r--r-- | tensorflow/python/kernel_tests/fractional_max_pool_op_test.py | 27 |
2 files changed, 37 insertions, 10 deletions
diff --git a/tensorflow/core/kernels/fractional_max_pool_op.cc b/tensorflow/core/kernels/fractional_max_pool_op.cc index 0722c408fba..375786615eb 100644 --- a/tensorflow/core/kernels/fractional_max_pool_op.cc +++ b/tensorflow/core/kernels/fractional_max_pool_op.cc @@ -19,12 +19,13 @@ limitations under the License. #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/framework/op_requires.h" +#include "tensorflow/core/kernels/fractional_pool_common.h" #include "tensorflow/core/lib/random/random.h" +#include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/util/guarded_philox_random.h" @@ -352,7 +353,9 @@ class FractionalMaxPoolGradOp : public OpKernel { output_size[2] * output_size[1] * output_size[0]; for (int64_t i = 0; i < num_reshaped_cols; ++i) { for (int64_t j = 0; j < output_size[3]; ++j) { - DCHECK_EQ(tensor_out_dup_mat(j, i), tensor_out_mat(j, i)); + OP_REQUIRES(context, tensor_out_dup_mat(j, i) == tensor_out_mat(j, i), + errors::InvalidArgument( + "tensor_out_dup is not the same as tensor_out")); } } @@ -369,11 +372,12 @@ class FractionalMaxPoolGradOp : public OpKernel { for (int index = 0; index < num_total_outputs; ++index) { int input_backprop_index = out_arg_max_flat(index); - // According to maxpooling_op.cc, the performance impact below is small. - CHECK(input_backprop_index >= 0 && - input_backprop_index < num_total_inputs) - << "Invalid input backprop index: " << input_backprop_index << ", " - << num_total_inputs; + OP_REQUIRES( + context, + input_backprop_index >= 0 && input_backprop_index < num_total_inputs, + errors::InvalidArgument( + "Invalid input backprop index: ", input_backprop_index, ", ", + num_total_inputs)); input_backprop_flat(input_backprop_index) += out_backprop_flat(index); } } diff --git a/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py b/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py index f395fefeb22..b8daa8bb919 100644 --- a/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py +++ b/tensorflow/python/kernel_tests/fractional_max_pool_op_test.py @@ -128,7 +128,7 @@ class FractionalMaxPoolTest(test.TestCase): Returns: None """ - with self.cached_session() as sess: + with self.cached_session(): p, r, c = nn_ops.fractional_max_pool_v2( input_tensor, pooling_ratio, @@ -159,7 +159,7 @@ class FractionalMaxPoolTest(test.TestCase): overlapping)) rand_mat = self._PRNG.randint(10, size=tensor_shape) pooling_ratio = [1, math.sqrt(2), math.sqrt(2), 1] - with self.cached_session() as sess: + with self.cached_session(): p, r, c = nn_ops.fractional_max_pool_v2( rand_mat, pooling_ratio, @@ -616,6 +616,29 @@ class FractionalMaxPoolGradTest(test.TestCase): self.assertAllClose(expected_input_backprop_overlapping, input_backprop_overlapping) + def testInvalidSeqRaiseErrorForFractionalMaxPoolGrad(self): + with self.assertRaises(errors.InvalidArgumentError): + with self.cached_session() as _: + overlapping = True + orig_input = constant_op.constant( + .453409232, shape=[1, 7, 13, 1], dtype=dtypes.float32) + orig_output = constant_op.constant( + .453409232, shape=[1, 7, 13, 1], dtype=dtypes.float32) + out_backprop = constant_op.constant( + .453409232, shape=[1, 7, 13, 1], dtype=dtypes.float32) + row_pooling_sequence = constant_op.constant( + 0, shape=[5], dtype=dtypes.int64) + col_pooling_sequence = constant_op.constant( + 0, shape=[5], dtype=dtypes.int64) + 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() |