Vulnerability: Heap buffer overflow in `MaxPool3DGradGrad` ### Impact
The implementation of `tf.raw_ops.MaxPool3DGradGrad` is vulnerable to a heap buffer overflow: 

```python
import tensorflow as tf

values = [0.01] * 11
orig_input = tf.constant(values, shape=[11, 1, 1, 1, 1], dtype=tf.float32)
orig_output = tf.constant([0.01], shape=[1, 1, 1, 1, 1], dtype=tf.float32)
grad = tf.constant([0.01], shape=[1, 1, 1, 1, 1], dtype=tf.float32)
ksize = [1, 1, 1, 1, 1]
strides = [1, 1, 1, 1, 1]
padding = "SAME"

tf.raw_ops.MaxPool3DGradGrad(
    orig_input=orig_input, orig_output=orig_output, grad=grad, ksize=ksize,
    strides=strides, padding=padding)
```

The [implementation](https://github.com/tensorflow/tensorflow/blob/596c05a159b6fbb9e39ca10b3f7753b7244fa1e9/tensorflow/core/kernels/pooling_ops_3d.cc#L694-L696) does not check that the initialization of `Pool3dParameters` completes successfully:

```cc
Pool3dParameters params{context,  ksize_,       stride_,
                        padding_, data_format_, tensor_in.shape()};
```

Since [the constructor](https://github.com/tensorflow/tensorflow/blob/596c05a159b6fbb9e39ca10b3f7753b7244fa1e9/tensorflow/core/kernels/pooling_ops_3d.cc#L48-L88) uses `OP_REQUIRES` to validate conditions, the first assertion that fails interrupts the initialization of `params`, making it contain invalid data. In turn, this might cause a heap buffer overflow, depending on default initialized values.

### Patches
We have patched the issue in GitHub commit [63c6a29d0f2d692b247f7bf81f8732d6442fad09](https://github.com/tensorflow/tensorflow/commit/63c6a29d0f2d692b247f7bf81f8732d6442fad09).

The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range.

### For more information
Please consult [our security guide](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for more information regarding the security model and how to contact us with issues and questions.

### Attribution
This vulnerability has been reported by Ying Wang and Yakun Zhang of Baidu X-Team.
Vulnerability ID: GHSA-7cqx-92hp-x6wh
CWE: CWE-787
Severity score: 2.5
Project: https://github.com/tensorflow/tensorflow
Published: 2021-05-21
File: tensorflow/core/kernels/pooling_ops_3d.cc
Commit message: Add missing validation, prevent heap OOB

PiperOrigin-RevId: 372246723
Change-Id: I1a454a643810e77d7d14821b342098c56a09fbbf
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 class MaxPooling3dGradGradOp : public OpKernel {      Pool3dParameters params{context,  ksize_,       stride_,                             padding_, data_format_, tensor_in.shape()};      Tensor* output = nullptr;     OP_REQUIRES_OK(context, context->forward_input_or_allocate_output( class MaxPooling3dGradGradOp : public OpKernel {         context, out_grad_backprop.NumElements() > 0,         errors::InvalidArgument("received empty tensor out_grad_backprop: ",                                 out_grad_backprop.DebugString()));      LaunchMaxPooling3dGradGradOp<Device, T>::launch(         context, params, tensor_in, tensor_out, out_grad_backprop, output);

Fixed version (after fix):
 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( 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);

Patch:
patch: @@ -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);
path: tensorflow/core/kernels/pooling_ops_3d.cc