42 lines
2.1 KiB
Diff
42 lines
2.1 KiB
Diff
|
From 63c6a29d0f2d692b247f7bf81f8732d6442fad09 Mon Sep 17 00:00:00 2001
|
||
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
||
|
|
Date: Wed, 5 May 2021 18:07:02 -0700
|
||
|
|
Subject: [PATCH] Add missing validation, prevent heap OOB
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 372246723
|
||
|
|
Change-Id: I1a454a643810e77d7d14821b342098c56a09fbbf
|
||
|
|
---
|
||
|
|
tensorflow/core/kernels/pooling_ops_3d.cc | 12 ++++++++++++
|
||
|
|
1 file changed, 12 insertions(+)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/pooling_ops_3d.cc b/tensorflow/core/kernels/pooling_ops_3d.cc
|
||
|
|
index 7d133b66a1ebd..9da2d62b0a21d 100644
|
||
|
|
--- a/tensorflow/core/kernels/pooling_ops_3d.cc
|
||
|
|
+++ b/tensorflow/core/kernels/pooling_ops_3d.cc
|
||
|
|
@@ -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(
|
||
|
|
@@ -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);
|