Skip to content
Permalink
Browse files Browse the repository at this point in the history
Add missing validation, prevent heap OOB
PiperOrigin-RevId: 372246723
Change-Id: I1a454a643810e77d7d14821b342098c56a09fbbf
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed May 6, 2021
1 parent 596c05a commit 63c6a29
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions tensorflow/core/kernels/pooling_ops_3d.cc
Expand Up @@ -693,6 +693,7 @@ class MaxPooling3dGradGradOp : public OpKernel {

Pool3dParameters params{context, ksize_, stride_,
padding_, data_format_, tensor_in.shape()};
if (!context->status().ok()) return; // params is invalid

Tensor* output = nullptr;
OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
Expand All @@ -710,6 +711,17 @@ class MaxPooling3dGradGradOp : public OpKernel {
context, out_grad_backprop.NumElements() > 0,
errors::InvalidArgument("received empty tensor out_grad_backprop: ",
out_grad_backprop.DebugString()));
OP_REQUIRES(context,
tensor_in.NumElements() == out_grad_backprop.NumElements(),
errors::InvalidArgument("tensor_in and out_grad_backprop must "
"have same number of elements, got <",
tensor_in.DebugString(), "> and <",
out_grad_backprop.DebugString(), ">"));
OP_REQUIRES(
context, tensor_out.NumElements() == output->NumElements(),
errors::InvalidArgument(
"tensor_out and output must have same number of elements, got <",
tensor_out.DebugString(), "> and <", output->DebugString(), ">"));

LaunchMaxPooling3dGradGradOp<Device, T>::launch(
context, params, tensor_in, tensor_out, out_grad_backprop, output);
Expand Down

0 comments on commit 63c6a29

Please sign in to comment.