CVE-2021-29576: Out-of-bounds Write (CWE-787)
value: TensorFlow is an end-to-end open source platform for machine learning. The implementation of `tf.raw_ops.MaxPool3DGradGrad` is vulnerable to a heap buffer overflow. 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. 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. 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.
CWE description: The product writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 7.8
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/pooling_ops_3d.cc
Commit date: 2021-05-05 18:18:17 -0700
Commit message: Add missing validation, prevent heap OOB

PiperOrigin-RevId: 372246723
Change-Id: I1a454a643810e77d7d14821b342098c56a09fbbf

Diff: 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
@@ -691,10 +691,11 @@ class MaxPooling3dGradGradOp : public OpKernel {
         context, out_grad_backprop.dims() == 5,
         errors::InvalidArgument("out_grad_backprop must be 5-dimensional"));
 
     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(
                                 {2}, 0, tensor_out.shape(), &output));
 
@@ -708,10 +709,21 @@ class MaxPooling3dGradGradOp : public OpKernel {
                                         tensor_out.DebugString()));
     OP_REQUIRES(
         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);
   }
 

Fixed code:     if (!context->status().ok()) return;  // params is invalid
    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(), ">"));