Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix nn.conv3d_transpose security vulnerability with illegal input shape
PiperOrigin-RevId: 487000656
  • Loading branch information
vufg authored and tensorflower-gardener committed Nov 8, 2022
1 parent 0eb1ae1 commit 948fe63
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 1 deletion.
5 changes: 5 additions & 0 deletions tensorflow/core/kernels/conv_grad_ops_3d.cc
Expand Up @@ -1322,6 +1322,11 @@ class Conv3DBackpropInputOp<GPUDevice, T> : public OpKernel {
Tensor* in_backprop;
OP_REQUIRES_OK(context,
context->allocate_output(0, input_shape, &in_backprop));
for (std::size_t i = 0; i < input_shape.dims(); ++i) {
if (input_shape.dim_size(i) == 0) {
return;
}
}

auto* stream = context->op_device_context()->stream();
OP_REQUIRES(context, stream, errors::Internal("No GPU stream available."));
Expand Down
16 changes: 15 additions & 1 deletion tensorflow/python/kernel_tests/nn_ops/conv3d_transpose_test.py
Expand Up @@ -126,7 +126,7 @@ def testConv3DTransposeShapeMismatch(self):
x_value = np.random.random_sample(x_shape).astype(np.float64)
f_value = np.random.random_sample(f_shape).astype(np.float64)
nn_ops.conv3d_transpose(
x_value, f_value, y_shape, strides, data_format='NCDHW')
x_value, f_value, y_shape, strides, data_format="NCDHW")

def testConv3DTransposeOutputShapeType(self):
# Test case for GitHub issue 18887
Expand Down Expand Up @@ -218,6 +218,20 @@ def testGradient(self):
err_tolerance = 0.00055
self.assertLess(err, err_tolerance)

def testConv3DTransposeZeroShapeDoNotRaiseError(self):
with self.cached_session():
x_value = np.zeros([10, 0, 2, 3, 3])
f_value = np.ones((3, 3, 3, 3, 3))
y_shape = np.stack([10, 0, 2, 3, 3])
output = nn_ops.conv3d_transpose(
x_value,
f_value,
y_shape,
strides=(1, 1, 1),
data_format="NDHWC",
padding="SAME",
)
_ = self.evaluate(output)

if __name__ == "__main__":
test.main()

0 comments on commit 948fe63

Please sign in to comment.