CVE-2021-41223: Out-of-bounds Read (CWE-125)
value: TensorFlow is an open source platform for machine learning. In affected versions the implementation of `FusedBatchNorm` kernels is vulnerable to a heap OOB access. The fix will be included in TensorFlow 2.7.0. We will also cherrypick this commit on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, as these are also affected and still in supported range.
CWE description: The product reads data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 7.1
Published: 2021-11-05T21:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/fused_batch_norm_op.cc, tensorflow/python/ops/nn_fused_batchnorm_test.py
Commit date: 2021-09-29 13:17:15 -0700
Commit message: Add shape checks to FusedBatchNorm kernels.

PiperOrigin-RevId: 399755576
Change-Id: If8049fde109cc33badb5509d174b9b95aee1ea5e

Diff: diff --git a/tensorflow/core/kernels/fused_batch_norm_op.cc b/tensorflow/core/kernels/fused_batch_norm_op.cc
index f1488e7407725..1278365c67f40 100644
--- a/tensorflow/core/kernels/fused_batch_norm_op.cc
+++ b/tensorflow/core/kernels/fused_batch_norm_op.cc
@@ -1338,22 +1338,24 @@ class FusedBatchNormOpBase : public OpKernel {
     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) {
+    if (!is_training_ || exponential_avg_factor_ != 1.) {
+      std::string prefix_msg = is_training_ ? "When exponential_avg_factor != 1"
+                                            : "When is_training=false";
       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 ",
+                      prefix_msg,
+                      ", mean must 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 ",
+                      prefix_msg,
+                      ", variance must 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(),
@@ -1541,10 +1543,15 @@ class FusedBatchNormGradOpBase : public OpKernel {
                                 saved_mean_or_pop_mean.shape().DebugString()));
     OP_REQUIRES(context, saved_maybe_inv_var_or_pop_var.dims() == 1,
                 errors::InvalidArgument(
                     "saved variance must be 1-dimensional",
                     saved_maybe_inv_var_or_pop_var.shape().DebugString()));
+    OP_REQUIRES(
+        context, x.shape() == y_backprop.shape(),
+        errors::InvalidArgument(
+            "x and y_backprop must have same shape, but x has shape ",
+            x.shape(), " and y_backprop has shape ", y_backprop.shape()));
     if (use_activation) {
       OP_REQUIRES(
           context, x.dim_size(3) % 4 == 0,
           errors::InvalidArgument("FusedBatchNormGrad with activation requires "
                                   "channel dimension to be a multiple of 4."));
@@ -1567,10 +1574,27 @@ class FusedBatchNormGradOpBase : public OpKernel {
                   errors::InvalidArgument("Error during tensor copy."));
       OP_REQUIRES(context, y_backprop.CopyFrom(y_backprop, 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, saved_mean_or_pop_mean.NumElements() == num_channels,
+        errors::InvalidArgument("reserve_space_1 must have the same number of "
+                                "elements as the channels of x, got ",
+                                scale.NumElements(), " and ", num_channels));
+    OP_REQUIRES(
+        context, saved_maybe_inv_var_or_pop_var.NumElements() == num_channels,
+        errors::InvalidArgument("reserve_space_2 must have the same number of "
+                                "elements as the channels of x, got ",
+                                scale.NumElements(), " and ", num_channels));
+
     Tensor* x_backprop = nullptr;
     auto alloc_shape = use_reshape ? dest_shape : x_shape;
     OP_REQUIRES_OK(context,
                    context->allocate_output(0, alloc_shape, &x_backprop));
 
diff --git a/tensorflow/python/ops/nn_fused_batchnorm_test.py b/tensorflow/python/ops/nn_fused_batchnorm_test.py
index fb3400fd49b99..5004e8625be5f 100644
--- a/tensorflow/python/ops/nn_fused_batchnorm_test.py
+++ b/tensorflow/python/ops/nn_fused_batchnorm_test.py
@@ -14,14 +14,17 @@
 # ==============================================================================
 """Tests for fused_batch_norm related functionality in tensorflow.ops.nn."""
 
 import numpy as np
 
+from tensorflow.python.eager import context
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_nn_ops
 from tensorflow.python.ops import gradient_checker
 from tensorflow.python.ops import gradients_impl
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import nn_grad
 from tensorflow.python.ops import nn_impl
@@ -692,8 +695,128 @@ class BatchNormalizationTest(test.TestCase):
     y_ref = self._inference_ref(x, scale, offset, mean, var, epsilon,
                                 'NCDHW')
     y_ref = np.maximum(y_ref, 0.)
     self.assertAllClose(y_ref, y_val, atol=1e-3)
 
+  def testEagerShapeErrors(self):
+    with context.eager_mode():
+      x = array_ops.ones((2, 2, 2, 2))
+      scale = array_ops.ones((3,))
+      offset = array_ops.ones((2,))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          'scale must have the same number of elements'):
+        nn_impl.fused_batch_norm(x, scale, offset)
+
+      x = array_ops.ones((2, 2, 2, 2))
+      scale = array_ops.ones((2,))
+      offset = array_ops.ones((3,))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          'offset must have the same number of elements'):
+        nn_impl.fused_batch_norm(x, scale, offset)
+
+      x = array_ops.ones((2, 2, 2, 2))
+      scale = array_ops.ones((2,))
+      offset = array_ops.ones((2,))
+      mean = array_ops.ones((0,))
+      variance = array_ops.ones((2,))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          'When is_training=false, mean must have the same number of elements'):
+        nn_impl.fused_batch_norm(
+            x, scale, offset, mean=mean, variance=variance, is_training=False)
+
+      x = array_ops.ones((2, 2, 2, 2))
+      scale = array_ops.ones((2,))
+      offset = array_ops.ones((2,))
+      mean = array_ops.ones((2,))
+      variance = array_ops.ones((0,))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          'When is_training=false, variance must have the same number of '
+          'elements'):
+        nn_impl.fused_batch_norm(
+            x, scale, offset, mean=mean, variance=variance, is_training=False)
+
+      x = array_ops.ones((2, 2, 2, 2))
+      scale = array_ops.ones((2,))
+      offset = array_ops.ones((2,))
+      mean = array_ops.ones((0,))
+      variance = array_ops.ones((2,))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          'When exponential_avg_factor != 1, mean must have the same number of '
+          'elements'):
+        nn_impl.fused_batch_norm(
+            x,
+            scale,
+            offset,
+            mean=mean,
+            variance=variance,
+            exponential_avg_factor=0.5)
+
+      x = array_ops.ones((2, 2, 2, 2))
+      scale = array_ops.ones((2,))
+      offset = array_ops.ones((2,))
+      mean = array_ops.ones((2,))
+      variance = array_ops.ones((0,))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          'When exponential_avg_factor != 1, variance must have the same '
+          'number of elements'):
+        nn_impl.fused_batch_norm(
+            x,
+            scale,
+            offset,
+            mean=mean,
+            variance=variance,
+            exponential_avg_factor=0.5)
+
+  def testEagerShapeGradErrors(self):
+    with context.eager_mode():
+      y_backprop = array_ops.ones((2, 2, 2, 3))
+      x = array_ops.ones((2, 2, 2, 2))
+      scale = array_ops.ones((2,))
+      reserve_space_1 = array_ops.ones((2,))
+      reserve_space_2 = array_ops.ones((2,))
+      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                  'x and y_backprop must have same shape,'):
+        gen_nn_ops.fused_batch_norm_grad_v2(y_backprop, x, scale,
+                                            reserve_space_1, reserve_space_2)
+
+      y_backprop = array_ops.ones((2, 2, 2, 2))
+      x = array_ops.ones((2, 2, 2, 2))
+      scale = array_ops.ones((3,))
+      reserve_space_1 = array_ops.ones((2,))
+      reserve_space_2 = array_ops.ones((2,))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          'scale must have the same number of elements'):
+        gen_nn_ops.fused_batch_norm_grad_v2(y_backprop, x, scale,
+                                            reserve_space_1, reserve_space_2)
+
+      y_backprop = array_ops.ones((2, 2, 2, 2))
+      x = array_ops.ones((2, 2, 2, 2))
+      scale = array_ops.ones((2,))
+      reserve_space_1 = array_ops.ones((3,))
+      reserve_space_2 = array_ops.ones((2,))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          'reserve_space_1 must have the same number of elements'):
+        gen_nn_ops.fused_batch_norm_grad_v2(y_backprop, x, scale,
+                                            reserve_space_1, reserve_space_2)
+
+      y_backprop = array_ops.ones((2, 2, 2, 2))
+      x = array_ops.ones((2, 2, 2, 2))
+      scale = array_ops.ones((2,))
+      reserve_space_1 = array_ops.ones((2,))
+      reserve_space_2 = array_ops.ones((3,))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          'reserve_space_2 must have the same number of elements'):
+        gen_nn_ops.fused_batch_norm_grad_v2(y_backprop, x, scale,
+                                            reserve_space_1, reserve_space_2)
+
 
 if __name__ == '__main__':
   test.main()

Vulnerable code:     if (estimated_mean.NumElements() != 0) {
                      "mean must be empty or have the same number of "
                      "elements as the channels of x, got ",
    }
    if (estimated_variance.NumElements() != 0) {
                      "variance must be empty or have the same number of "
                      "elements as the channels of x, got ",
Fixed code:     if (!is_training_ || exponential_avg_factor_ != 1.) {
      std::string prefix_msg = is_training_ ? "When exponential_avg_factor != 1"
                                            : "When is_training=false";
                      prefix_msg,
                      ", mean must have the same number "
                      "of elements as the channels of x, got ",
                      prefix_msg,
                      ", variance must have the same "
                      "number of elements as the channels of x, got ",
    OP_REQUIRES(
        context, x.shape() == y_backprop.shape(),
        errors::InvalidArgument(
            "x and y_backprop must have same shape, but x has shape ",
            x.shape(), " and y_backprop has shape ", y_backprop.shape()));
    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, saved_mean_or_pop_mean.NumElements() == num_channels,
        errors::InvalidArgument("reserve_space_1 must have the same number of "
                                "elements as the channels of x, got ",
                                scale.NumElements(), " and ", num_channels));
    OP_REQUIRES(
        context, saved_maybe_inv_var_or_pop_var.NumElements() == num_channels,
        errors::InvalidArgument("reserve_space_2 must have the same number of "
                                "elements as the channels of x, got ",
                                scale.NumElements(), " and ", num_channels));

from tensorflow.python.eager import context
from tensorflow.python.framework import errors_impl
from tensorflow.python.ops import gen_nn_ops
  def testEagerShapeErrors(self):
    with context.eager_mode():
      x = array_ops.ones((2, 2, 2, 2))
      scale = array_ops.ones((3,))
      offset = array_ops.ones((2,))
      with self.assertRaisesRegex(
          errors_impl.InvalidArgumentError,
          'scale must have the same number of elements'):
        nn_impl.fused_batch_norm(x, scale, offset)

      x = array_ops.ones((2, 2, 2, 2))
      scale = array_ops.ones((2,))
      offset = array_ops.ones((3,))
      with self.assertRaisesRegex(
          errors_impl.InvalidArgumentError,
          'offset must have the same number of elements'):
        nn_impl.fused_batch_norm(x, scale, offset)

      x = array_ops.ones((2, 2, 2, 2))
      scale = array_ops.ones((2,))
      offset = array_ops.ones((2,))
      mean = array_ops.ones((0,))
      variance = array_ops.ones((2,))
      with self.assertRaisesRegex(
          errors_impl.InvalidArgumentError,
          'When is_training=false, mean must have the same number of elements'):
        nn_impl.fused_batch_norm(
            x, scale, offset, mean=mean, variance=variance, is_training=False)

      x = array_ops.ones((2, 2, 2, 2))
      scale = array_ops.ones((2,))
      offset = array_ops.ones((2,))
      mean = array_ops.ones((2,))
      variance = array_ops.ones((0,))
      with self.assertRaisesRegex(
          errors_impl.InvalidArgumentError,
          'When is_training=false, variance must have the same number of '
          'elements'):
        nn_impl.fused_batch_norm(
            x, scale, offset, mean=mean, variance=variance, is_training=False)

      x = array_ops.ones((2, 2, 2, 2))
      scale = array_ops.ones((2,))
      offset = array_ops.ones((2,))
      mean = array_ops.ones((0,))
      variance = array_ops.ones((2,))
      with self.assertRaisesRegex(
          errors_impl.InvalidArgumentError,
          'When exponential_avg_factor != 1, mean must have the same number of '
          'elements'):
        nn_impl.fused_batch_norm(
            x,
            scale,
            offset,
            mean=mean,
            variance=variance,
            exponential_avg_factor=0.5)

      x = array_ops.ones((2, 2, 2, 2))
      scale = array_ops.ones((2,))
      offset = array_ops.ones((2,))
      mean = array_ops.ones((2,))
      variance = array_ops.ones((0,))
      with self.assertRaisesRegex(
          errors_impl.InvalidArgumentError,
          'When exponential_avg_factor != 1, variance must have the same '
          'number of elements'):
        nn_impl.fused_batch_norm(
            x,
            scale,
            offset,
            mean=mean,
            variance=variance,
            exponential_avg_factor=0.5)

  def testEagerShapeGradErrors(self):
    with context.eager_mode():
      y_backprop = array_ops.ones((2, 2, 2, 3))
      x = array_ops.ones((2, 2, 2, 2))
      scale = array_ops.ones((2,))
      reserve_space_1 = array_ops.ones((2,))
      reserve_space_2 = array_ops.ones((2,))
      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
                                  'x and y_backprop must have same shape,'):
        gen_nn_ops.fused_batch_norm_grad_v2(y_backprop, x, scale,
                                            reserve_space_1, reserve_space_2)

      y_backprop = array_ops.ones((2, 2, 2, 2))
      x = array_ops.ones((2, 2, 2, 2))
      scale = array_ops.ones((3,))
      reserve_space_1 = array_ops.ones((2,))
      reserve_space_2 = array_ops.ones((2,))
      with self.assertRaisesRegex(
          errors_impl.InvalidArgumentError,
          'scale must have the same number of elements'):
        gen_nn_ops.fused_batch_norm_grad_v2(y_backprop, x, scale,
                                            reserve_space_1, reserve_space_2)

      y_backprop = array_ops.ones((2, 2, 2, 2))
      x = array_ops.ones((2, 2, 2, 2))
      scale = array_ops.ones((2,))
      reserve_space_1 = array_ops.ones((3,))
      reserve_space_2 = array_ops.ones((2,))
      with self.assertRaisesRegex(
          errors_impl.InvalidArgumentError,
          'reserve_space_1 must have the same number of elements'):
        gen_nn_ops.fused_batch_norm_grad_v2(y_backprop, x, scale,
                                            reserve_space_1, reserve_space_2)

      y_backprop = array_ops.ones((2, 2, 2, 2))
      x = array_ops.ones((2, 2, 2, 2))
      scale = array_ops.ones((2,))
      reserve_space_1 = array_ops.ones((2,))
      reserve_space_2 = array_ops.ones((3,))
      with self.assertRaisesRegex(
          errors_impl.InvalidArgumentError,
          'reserve_space_2 must have the same number of elements'):
        gen_nn_ops.fused_batch_norm_grad_v2(y_backprop, x, scale,
                                            reserve_space_1, reserve_space_2)