Vulnerability: Heap buffer overflow and undefined behavior in `FusedBatchNorm` ### Impact
The implementation of `tf.raw_ops.FusedBatchNorm` is vulnerable to a heap buffer overflow:
      
```python
import tensorflow as tf

x = tf.zeros([10, 10, 10, 6], dtype=tf.float32)
scale = tf.constant([0.0], shape=[1], dtype=tf.float32)
offset = tf.constant([0.0], shape=[1], dtype=tf.float32)
mean = tf.constant([0.0], shape=[1], dtype=tf.float32)
variance = tf.constant([0.0], shape=[1], dtype=tf.float32)
epsilon = 0.0
exponential_avg_factor = 0.0
data_format = "NHWC"
is_training = False
    
tf.raw_ops.FusedBatchNorm(
  x=x, scale=scale, offset=offset, mean=mean, variance=variance,
  epsilon=epsilon, exponential_avg_factor=exponential_avg_factor,
  data_format=data_format, is_training=is_training)
```
  
If the tensors are empty, the same implementation can trigger undefined behavior by dereferencing null pointers:

```python 
import tensorflow as tf
import numpy as np

x = tf.zeros([10, 10, 10, 1], dtype=tf.float32)
scale = tf.constant([], shape=[0], dtype=tf.float32)
offset = tf.constant([], shape=[0], dtype=tf.float32)
mean = tf.constant([], shape=[0], dtype=tf.float32)
variance = tf.constant([], shape=[0], dtype=tf.float32)
epsilon = 0.0
exponential_avg_factor = 0.0
data_format = "NHWC"
is_training = False

tf.raw_ops.FusedBatchNorm(
  x=x, scale=scale, offset=offset, mean=mean, variance=variance, 
  epsilon=epsilon, exponential_avg_factor=exponential_avg_factor,
  data_format=data_format, is_training=is_training)
``` 

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.

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

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-9xh4-23q4-v6wr
CWE: CWE-476
Severity score: 2.5
Project: https://github.com/tensorflow/tensorflow
Published: 2021-05-21
File: tensorflow/core/kernels/fused_batch_norm_op.cc
Commit message: Add missing valuidation to FusedBatchNorm.

PiperOrigin-RevId: 372460336
Change-Id: Ic8c4e4de67c58a741bd87f2e182bed07247d1126
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 class FusedBatchNormOpBase : public OpKernel {                   errors::InvalidArgument("Error during tensor copy."));     }      if (has_side_input_) {       OP_REQUIRES(context, side_input->shape() == x.shape(),                   errors::InvalidArgument( class FusedBatchNormOpBase : public OpKernel {       // NOTE(ezhulenev): This requirement is coming from implementation       // details of cudnnBatchNormalizationForwardTrainingEx.       OP_REQUIRES(           context, !is_training_ || x.dim_size(3) % 4 == 0,           errors::InvalidArgument("FusedBatchNorm with activation requires "                                   "channel dimension to be a multiple of 4."));     }

Fixed version (after fix):
 class FusedBatchNormOpBase : public OpKernel {                   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( class FusedBatchNormOpBase : public OpKernel {       // NOTE(ezhulenev): This requirement is coming from implementation       // details of cudnnBatchNormalizationForwardTrainingEx.       OP_REQUIRES(           context, !is_training_ || num_channels % 4 == 0,           errors::InvalidArgument("FusedBatchNorm with activation requires "                                   "channel dimension to be a multiple of 4."));     }

Patch:
patch: @@ -1282,6 +1282,32 @@ class FusedBatchNormOpBase : public OpKernel {
                   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(
@@ -1294,7 +1320,7 @@ class FusedBatchNormOpBase : public OpKernel {
       // 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."));
     }
path: tensorflow/core/kernels/fused_batch_norm_op.cc