diff options
| author | Alan Liu <poulsbo@users.noreply.github.com> | 2022-10-06 10:42:17 -0700 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2022-10-06 10:42:17 -0700 |
| commit | 36df7362a299f10df62395915d1c1570783da54b (patch) | |
| tree | d2b0ea8d008473f39c8bf212dde5c7e4e333dd50 | |
| parent | e2332e082c588dc82fa1d497598607c6c23f3421 (diff) | |
| download | tensorflow-36df7362a299f10df62395915d1c1570783da54b.tar.gz | |
Fix tf.raw_ops.ResizeNearestNeighborGrad vulnerability with large dimensions. (#57997)
Note: This fix will have to be cherry picked in r2.10, r2.9, and r2.8.
PiperOrigin-RevId: 479141644
| -rw-r--r-- | tensorflow/core/kernels/image/resize_nearest_neighbor_op.cc | 10 | ||||
| -rw-r--r-- | tensorflow/python/ops/image_ops_test.py | 19 |
2 files changed, 24 insertions, 5 deletions
diff --git a/tensorflow/core/kernels/image/resize_nearest_neighbor_op.cc b/tensorflow/core/kernels/image/resize_nearest_neighbor_op.cc index a54b60f0099..0d0e0cbdbe6 100644 --- a/tensorflow/core/kernels/image/resize_nearest_neighbor_op.cc +++ b/tensorflow/core/kernels/image/resize_nearest_neighbor_op.cc @@ -257,11 +257,11 @@ class ResizeNearestNeighborOpGrad : public OpKernel { const int64_t out_width = sizes(1); Tensor* output = nullptr; - OP_REQUIRES_OK( - context, - context->allocate_output( - 0, TensorShape({batch_size, out_height, out_width, channels}), - &output)); + TensorShape shape; + OP_REQUIRES_OK(context, + TensorShape::BuildTensorShape( + {batch_size, out_height, out_width, channels}, &shape)); + OP_REQUIRES_OK(context, context->allocate_output(0, shape, &output)); // Return if the output is empty. if (output->NumElements() == 0) return; diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py index ec105dfaeb0..cf96ded1010 100644 --- a/tensorflow/python/ops/image_ops_test.py +++ b/tensorflow/python/ops/image_ops_test.py @@ -4152,6 +4152,25 @@ class ResizeImageWithPadV2Test(test_util.TensorFlowTestCase): self._assertReturns(x, x_shape, y, y_shape) +class ResizeNearestNeighborGrad(test_util.TensorFlowTestCase): + + def testSizeTooLarge(self): + align_corners = True + half_pixel_centers = False + grads = constant_op.constant(1, shape=[1, 8, 16, 3], dtype=dtypes.float16) + size = constant_op.constant([1879048192, 1879048192], + shape=[2], + dtype=dtypes.int32) + with self.assertRaisesRegex(errors.InvalidArgumentError, + r"Encountered overflow when multiplying"): + self.evaluate( + gen_image_ops.ResizeNearestNeighborGrad( + grads=grads, + size=size, + align_corners=align_corners, + half_pixel_centers=half_pixel_centers)) + + class ResizeImageWithCropOrPadTest(test_util.TensorFlowTestCase): def _ResizeImageWithCropOrPad(self, x, target_height, target_width, |