CVE-2021-41223 (Python, published 2021-11-05T21:15Z)
Description: 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.
Severity: LOW (CVSS v3 base score: 7.1)
Weakness: CWE-125 \u2014 Out-of-bounds Read
The product reads data past the end, or before the beginning, of the intended buffer.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/kernels/fused_batch_norm_op.cc
tensorflow/python/ops/nn_fused_batchnorm_test.py
Commit message: Add shape checks to FusedBatchNorm kernels.

PiperOrigin-RevId: 399755576
Change-Id: If8049fde109cc33badb5509d174b9b95aee1ea5e


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)