CVE-2021-29583: NULL Pointer Dereference (CWE-476)
value: TensorFlow is an end-to-end open source platform for machine learning. The implementation of `tf.raw_ops.FusedBatchNorm` is vulnerable to a heap buffer overflow. If the tensors are empty, the same implementation can trigger undefined behavior by dereferencing null pointers. The implementation(https://github.com/tensorflow/tensorflow/blob/57d86e0db5d1365f19adcce848dfc1bf89fdd4c7/tensorflow/core/kernels/fused_batch_norm_op.cc) fails to validate that `scale`, `offset`, `mean` and `variance` (the last two only when required) all have the same number of elements as the number of channels of `x`. This results in heap out of bounds reads when the buffers backing these tensors are indexed past their boundary. If the tensors are empty, the validation mentioned in the above paragraph would also trigger and prevent the undefined behavior. 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 dereferences a pointer that it expects to be valid but is NULL.
CVSS v3 base score: 7.8
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/fused_batch_norm_op.cc
Commit date: 2021-05-06 17:50:14 -0700
Commit message: Add missing valuidation to FusedBatchNorm.

PiperOrigin-RevId: 372460336
Change-Id: Ic8c4e4de67c58a741bd87f2e182bed07247d1126

Diff: diff --git a/tensorflow/core/kernels/fused_batch_norm_op.cc b/tensorflow/core/kernels/fused_batch_norm_op.cc
index e564b19857c38..7b0932d953261 100644
--- a/tensorflow/core/kernels/fused_batch_norm_op.cc
+++ b/tensorflow/core/kernels/fused_batch_norm_op.cc
@@ -1280,10 +1280,36 @@ class FusedBatchNormOpBase : public OpKernel {
                                    {{in_planes, in_rows * in_cols}}, in_depth);
       OP_REQUIRES(context, x.CopyFrom(x, dest_shape),
                   errors::InvalidArgument("Error during tensor copy."));
     }
 
+    const auto num_channels = GetTensorDim(x, tensor_format_, 'C');
+    OP_REQUIRES(
+        context, scale.NumElements() == num_channels,
+        errors::InvalidArgument("scale must have the same number of elements "
+                                "as the channels of x, got ",
+                                scale.NumElements(), " and ", num_channels));
+    OP_REQUIRES(
+        context, offset.NumElements() == num_channels,
+        errors::InvalidArgument("offset must have the same number of elements "
+                                "as the channels of x, got ",
+                                offset.NumElements(), " and ", num_channels));
+    if (estimated_mean.NumElements() != 0) {
+      OP_REQUIRES(context, estimated_mean.NumElements() == num_channels,
+                  errors::InvalidArgument(
+                      "mean must be empty or have the same number of "
+                      "elements as the channels of x, got ",
+                      estimated_mean.NumElements(), " and ", num_channels));
+    }
+    if (estimated_variance.NumElements() != 0) {
+      OP_REQUIRES(context, estimated_variance.NumElements() == num_channels,
+                  errors::InvalidArgument(
+                      "variance must be empty or have the same number of "
+                      "elements as the channels of x, got ",
+                      estimated_variance.NumElements(), " and ", num_channels));
+    }
+
     if (has_side_input_) {
       OP_REQUIRES(context, side_input->shape() == x.shape(),
                   errors::InvalidArgument(
                       "side_input shape must be equal to input shape: ",
                       side_input->shape().DebugString(),
@@ -1292,11 +1318,11 @@ class FusedBatchNormOpBase : public OpKernel {
 
     if (activation_mode_ != FbnActivationMode::kIdentity) {
       // NOTE(ezhulenev): This requirement is coming from implementation
       // details of cudnnBatchNormalizationForwardTrainingEx.
       OP_REQUIRES(
-          context, !is_training_ || x.dim_size(3) % 4 == 0,
+          context, !is_training_ || num_channels % 4 == 0,
           errors::InvalidArgument("FusedBatchNorm with activation requires "
                                   "channel dimension to be a multiple of 4."));
     }
 
     Tensor* y = nullptr;

Vulnerable code:           context, !is_training_ || x.dim_size(3) % 4 == 0,
Fixed code:     const auto num_channels = GetTensorDim(x, tensor_format_, 'C');
    OP_REQUIRES(
        context, scale.NumElements() == num_channels,
        errors::InvalidArgument("scale must have the same number of elements "
                                "as the channels of x, got ",
                                scale.NumElements(), " and ", num_channels));
    OP_REQUIRES(
        context, offset.NumElements() == num_channels,
        errors::InvalidArgument("offset must have the same number of elements "
                                "as the channels of x, got ",
                                offset.NumElements(), " and ", num_channels));
    if (estimated_mean.NumElements() != 0) {
      OP_REQUIRES(context, estimated_mean.NumElements() == num_channels,
                  errors::InvalidArgument(
                      "mean must be empty or have the same number of "
                      "elements as the channels of x, got ",
                      estimated_mean.NumElements(), " and ", num_channels));
    }
    if (estimated_variance.NumElements() != 0) {
      OP_REQUIRES(context, estimated_variance.NumElements() == num_channels,
                  errors::InvalidArgument(
                      "variance must be empty or have the same number of "
                      "elements as the channels of x, got ",
                      estimated_variance.NumElements(), " and ", num_channels));
    }

          context, !is_training_ || num_channels % 4 == 0,