Skip to content
Permalink
Browse files Browse the repository at this point in the history
Add security vulnerability test for raw_ops.Conv2DBackpropInput
PiperOrigin-RevId: 463395218
  • Loading branch information
tensorflower-gardener committed Jul 26, 2022
1 parent 8606329 commit 50156d5
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 4 deletions.
9 changes: 9 additions & 0 deletions tensorflow/core/kernels/conv_grad_input_ops.h
Expand Up @@ -422,6 +422,11 @@ class Conv2DBackpropInputOp : public OpKernel {
const Tensor& filter = context->input(1);
const Tensor& out_backprop = context->input(2);

OP_REQUIRES(
context, out_backprop.dims() == 4,
errors::InvalidArgument("input_sizes must be 4-dimensional, got: ",
out_backprop.dims()));

TensorShape input_shape;
OP_REQUIRES_OK(context,
Conv2DBackpropComputeInputShape(input_sizes, filter.shape(),
Expand Down Expand Up @@ -527,6 +532,10 @@ class Conv2DCustomBackpropInputOp : public OpKernel {
const Tensor& input_sizes = context->input(0);
const Tensor& filter = context->input(1);
const Tensor& out_backprop = context->input(2);
OP_REQUIRES(
context, out_backprop.dims() == 4,
errors::InvalidArgument("input_sizes must be 4-dimensional, got: ",
out_backprop.dims()));

TensorShape input_shape;
OP_REQUIRES_OK(context,
Expand Down
30 changes: 26 additions & 4 deletions tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
Expand Up @@ -32,6 +32,7 @@
from tensorflow.python.layers import convolutional
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_nn_ops
from tensorflow.python.ops import gradient_checker
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import math_ops
Expand Down Expand Up @@ -1319,7 +1320,7 @@ def _RunAndVerifyBackpropInputDilation(self, input_sizes, filter_sizes,
x2 = self._CreateNumpyTensor(filter_sizes)
default_dilations = (dilations[0] == 1 and dilations[1] == 1)
if default_dilations or use_gpu:
with self.cached_session(use_gpu=use_gpu) as sess:
with self.cached_session(use_gpu=use_gpu):
if data_format == "NCHW":
input_sizes = test_util.NHWCToNCHW(input_sizes)
t1 = constant_op.constant(x1, shape=input_sizes)
Expand Down Expand Up @@ -1365,7 +1366,7 @@ def _RunAndVerifyBackpropFilterDilation(self, input_sizes, filter_sizes,
x2 = self._CreateNumpyTensor(filter_sizes)
default_dilations = (dilations[0] == 1 and dilations[1] == 1)
if default_dilations or use_gpu:
with self.cached_session(use_gpu=use_gpu) as sess:
with self.cached_session(use_gpu=use_gpu):
if data_format == "NCHW":
input_sizes = test_util.NHWCToNCHW(input_sizes)
t1 = constant_op.constant(x1, shape=input_sizes)
Expand Down Expand Up @@ -2628,6 +2629,27 @@ def testOpEdgeCases(self):
strides=[1, 1, 1, 1],
padding=[[0, 0], [-1, 0], [0, 0], [0, 0]]))

def testConv2DBackpropInputInvalidOutBackpropRaiseError(self):
with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
with self.cached_session():
input_sizes = constant_op.constant([65534, 65534],
shape=[2],
dtype=dtypes.int32)
filters = constant_op.constant(
0.159749106, shape=[3, 3, 2, 2], dtype=dtypes.float32)
out_backprop = constant_op.constant(0, shape=[], dtype=dtypes.float32)
t = gen_nn_ops.conv2d_backprop_input(
input_sizes=input_sizes,
filter=filters,
out_backprop=out_backprop,
strides=[1, 1, 1, 1],
padding="SAME",
use_cudnn_on_gpu=True,
explicit_paddings=[],
data_format="NHWC",
dilations=[1, 1, 1, 1])
self.evaluate(t)


@test_util.run_all_without_tensor_float_32("Avoid TF32 conv on GPU")
class DepthwiseConv2DTest(test.TestCase):
Expand Down Expand Up @@ -2655,7 +2677,7 @@ def _VerifyValues(self, tensor_in_sizes, filter_in_sizes, stride, padding,
# numbers from 1.
x1 = [f * 1.0 for f in range(1, total_size_1 + 1)]
x2 = [f * 1.0 for f in range(1, total_size_2 + 1)]
with self.cached_session() as sess:
with self.cached_session():
t1 = constant_op.constant(x1, shape=tensor_in_sizes)
t1.set_shape(tensor_in_sizes)
t2 = constant_op.constant(x2, shape=filter_in_sizes)
Expand Down Expand Up @@ -2926,7 +2948,7 @@ def _CompareFwdConv2D(self, tensor_in_sizes, filter_in_sizes, conv_strides,
x1 = np.random.rand(*tensor_in_sizes).astype(np.float32)
x2 = np.random.rand(*filter_in_sizes).astype(np.float32)

with self.cached_session(use_gpu=False) as sess:
with self.cached_session(use_gpu=False):
t1 = constant_op.constant(x1, shape=tensor_in_sizes)
t2 = constant_op.constant(x2, shape=filter_in_sizes)
strides = [1] + conv_strides + [1]
Expand Down

0 comments on commit 50156d5

Please sign in to comment.