CVE-2022-29203 (Other, published 2022-05-20T23:15Z)
Description: value: TensorFlow is an open source platform for machine learning. Prior to versions 2.9.0, 2.8.1, 2.7.2, and 2.6.4, the implementation of `tf.raw_ops.SpaceToBatchND` (in all backends such as XLA and handwritten kernels) is vulnerable to an integer overflow: The result of this integer overflow is used to allocate the output tensor, hence we get a denial of service via a `CHECK`-failure (assertion failure), as in TFSA-2021-198. Versions 2.9.0, 2.8.1, 2.7.2, and 2.6.4 contain a patch for this issue.
Severity: LOW (CVSS v3 base score: 5.5)
Weakness: CWE-190 \u2014 Integer Overflow or Wraparound
The product performs a calculation that can
         produce an integer overflow or wraparound when the logic
         assumes that the resulting value will always be larger than
         the original value. This occurs when an integer value is
         incremented to a value that is too large to store in the
         associated representation. When this occurs, the value may
         become a very small or negative number.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/compiler/tests/spacetobatch_op_test.py
tensorflow/compiler/tf2xla/kernels/BUILD
tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc
tensorflow/core/framework/BUILD
tensorflow/core/framework/shape_inference.cc
tensorflow/core/kernels/BUILD
tensorflow/core/kernels/spacetobatch_op.cc
tensorflow/core/util/BUILD
tensorflow/python/kernel_tests/array_ops/spacetobatch_op_test.py
Commit message: Fix security vulnerability with SpaceToBatchNDOp.

PiperOrigin-RevId: 445527615


Vulnerable code:
    block_num_elems *= block_shape[i];
  const int64_t batch_size = input_shape[0];
  output_shape[0] = batch_size * block_num_elems;
    const int64_t product = first_value * second_value;
load(
    "//third_party/mkl:build_defs.bzl",
    "mkl_deps",
)
#include "tensorflow/core/kernels/spacetobatch_functor.h"

    block_shape_product *= block_shape[block_dim];
  external_output_shape.AddDim(orig_input_tensor.dim_size(0) *
                               block_shape_product);

Fixed code:
from tensorflow.python.framework import constant_op
class SpaceToBatchNDErrorHandlingTest(xla_test.XLATestCase):

  def testInvalidBlockShape(self):
    with self.assertRaisesRegex(ValueError, "block_shape must be positive"):
      with self.session() as sess, self.test_scope():
        tf_in = constant_op.constant(
            -3.5e+35, shape=[10, 20, 20], dtype=dtypes.float32)
        block_shape = constant_op.constant(-10, shape=[2], dtype=dtypes.int64)
        paddings = constant_op.constant(0, shape=[2, 2], dtype=dtypes.int32)
        sess.run(array_ops.space_to_batch_nd(tf_in, block_shape, paddings))

  def testOutputSizeOutOfBounds(self):
    with self.assertRaisesRegex(ValueError,
                                "Negative.* dimension size caused by overflow"):
      with self.session() as sess, self.test_scope():
        tf_in = constant_op.constant(
            -3.5e+35, shape=[10, 19, 22], dtype=dtypes.float32)
        block_shape = constant_op.constant(
            1879048192, shape=[2], dtype=dtypes.int64)
        paddings = constant_op.constant(0, shape=[2, 2], dtype=dtypes.int32)
        sess.run(array_ops.space_to_batch_nd(tf_in, block_shape, paddings))


        "//tensorflow/core/util:overflow",
#include "tensorflow/core/util/overflow.h"
    OP_REQUIRES(ctx, block_shape[i] >= 1,
                errors::InvalidArgument(
                    "All values in block_shape must be positive, got value, ",
                    block_shape[i], " at index ", i, "."));
    block_num_elems = MultiplyWithoutOverflow(block_num_elems, block_shape[i]);
  const int64_t batch_size = input_shape[0];
  const int64_t output_dim =
      MultiplyWithoutOverflow(batch_size, block_num_elems);
  if (output_dim < 0) {
    OP_REQUIRES(
        ctx, output_dim >= 0,
        errors::InvalidArgument("Negative output dimension size caused by "
                                "overflow when multiplying ",
                                batch_size, " and ", block_num_elems));
  }
  output_shape[0] = output_dim;
        "//tensorflow/core/util:overflow",
#include "tensorflow/core/util/overflow.h"
    const int64_t product = MultiplyWithoutOverflow(first_value, second_value);
    "mkl_deps",
        "//tensorflow/core/util:overflow",
#include "tensorflow/core/kernels/spacetobatch_functor.h"
#include "tensorflow/core/util/overflow.h"
    if (block_shape[block_dim] < 1) {
      return errors::InvalidArgument(
          "All values in block_shape must be positive, got value, ",
          block_shape[block_dim], " at index ", block_dim, ".");
    }
    block_shape_product =
        MultiplyWithoutOverflow(block_shape_product, block_shape[block_dim]);
  const int64_t output_shape = MultiplyWithoutOverflow(
      orig_input_tensor.dim_size(0), block_shape_product);
  if (output_shape < 0) {
    return errors::InvalidArgument(
        "Negative output dimension size caused by overflow when multiplying ",
        orig_input_tensor.dim_size(0), " and ", block_shape_product);
  }
  external_output_shape.AddDim(output_shape);
    visibility = [
        "//tensorflow:internal",
    ],
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import errors
  @test_util.run_in_graph_and_eager_modes
  def testInvalidBlockShape(self):
    tf_in = constant_op.constant(
        -3.5e+35, shape=[10, 20, 20], dtype=dtypes.float32)
    block_shape = constant_op.constant(-10, shape=[2], dtype=dtypes.int64)
    paddings = constant_op.constant(0, shape=[2, 2], dtype=dtypes.int32)
    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "block_shape must be positive"):
      array_ops.space_to_batch_nd(tf_in, block_shape, paddings)

  @test_util.run_in_graph_and_eager_modes
  def testOutputSizeOutOfBounds(self):
    tf_in = constant_op.constant(
        -3.5e+35, shape=[10, 19, 22], dtype=dtypes.float32)
    block_shape = constant_op.constant(
        1879048192, shape=[2], dtype=dtypes.int64)
    paddings = constant_op.constant(0, shape=[2, 2], dtype=dtypes.int32)
    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "Negative.* dimension size caused by overflow"):
      array_ops.space_to_batch_nd(tf_in, block_shape, paddings)