CVE-2022-29203: Integer Overflow or Wraparound (CWE-190)
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.
CWE description: 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.
CVSS v3 base score: 5.5
Published: 2022-05-20T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Other
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 date: 2022-04-29 16:14:54 -0700
Commit message: Fix security vulnerability with SpaceToBatchNDOp.

PiperOrigin-RevId: 445527615

Diff: diff --git a/tensorflow/compiler/tests/spacetobatch_op_test.py b/tensorflow/compiler/tests/spacetobatch_op_test.py
index bb3d4b1812080..016c05f11e0c2 100644
--- a/tensorflow/compiler/tests/spacetobatch_op_test.py
+++ b/tensorflow/compiler/tests/spacetobatch_op_test.py
@@ -15,10 +15,11 @@
 """Functional tests for SpaceToBatch and BatchToSpace ops."""
 
 import numpy as np
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gen_array_ops
 from tensorflow.python.platform import test
 
@@ -143,10 +144,33 @@ class SpaceToBatchTest(xla_test.XLATestCase):
              [[[6], [8]], [[14], [16]]], [[[22], [24]], [[30], [32]]]]
     block_size = 2
     self._testOne(x_np, block_size, x_out)
 
 
+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))
+
+
 class SpaceToBatchNDTest(xla_test.XLATestCase):
   """Tests input-output pairs for the SpaceToBatchND and BatchToSpaceND ops."""
 
   def _testPad(self, inputs, block_shape, paddings, outputs):
     block_shape = np.array(block_shape)
diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD
index ca28459476e8d..873b892423f9f 100644
--- a/tensorflow/compiler/tf2xla/kernels/BUILD
+++ b/tensorflow/compiler/tf2xla/kernels/BUILD
@@ -209,10 +209,11 @@ tf_kernel_library(
         "//tensorflow/core/kernels:resource_variable_util",
         "//tensorflow/core/kernels:scatter_nd_util",
         "//tensorflow/core/kernels:stateful_random_ops_header",
         "//tensorflow/core/kernels:stateless_random_ops_v2_header",
         "//tensorflow/core/tpu:tpu_defs",
+        "//tensorflow/core/util:overflow",
         "//tensorflow/stream_executor/lib",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
diff --git a/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc
index a4e9aec1c9705..d6e38f1309f91 100644
--- a/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc
@@ -15,10 +15,11 @@ limitations under the License.
 
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "tensorflow/compiler/xla/client/xla_builder.h"
+#include "tensorflow/core/util/overflow.h"
 
 namespace tensorflow {
 namespace {
 
 void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp& input,
@@ -58,22 +59,36 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp& input,
     auto* dim = padding_config.add_dimensions();
     int64_t pad_start = paddings.Get<int64_t>({i, 0});
     int64_t pad_end = paddings.Get<int64_t>({i, 1});
     OP_REQUIRES(ctx, pad_start >= 0 && pad_end >= 0,
                 errors::InvalidArgument("Paddings must be non-negative"));
+    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, "."));
     dim->set_edge_padding_low(pad_start);
     dim->set_edge_padding_high(pad_end);
     padded_shape[1 + i] += pad_start + pad_end;
-    block_num_elems *= block_shape[i];
+    block_num_elems = MultiplyWithoutOverflow(block_num_elems, block_shape[i]);
   }
   // Don't pad the remainder dimensions.
   for (int i = 0; i < remainder_shape.size(); ++i) {
     padding_config.add_dimensions();
   }
   OP_REQUIRES(ctx, block_num_elems > 0,
               errors::InvalidArgument(
                   "The product of the block dimensions must be positive"));
+  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));
+  }
 
   xla::XlaOp padded =
       xla::Pad(input, XlaHelpers::Zero(b, input_dtype), padding_config);
 
   // 2. Reshape `padded` to `reshaped_padded` of shape:
@@ -83,11 +98,10 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp& input,
   //        block_shape[0],
   //       ...,
   //       padded_shape[M] / block_shape[M-1],
   //       block_shape[M-1]] +
   //      remaining_shape
-  const int64_t batch_size = input_shape[0];
   std::vector<int64_t> reshaped_padded_shape(input_rank + block_rank);
   reshaped_padded_shape[0] = batch_size;
   for (int i = 0; i < block_rank; ++i) {
     OP_REQUIRES(ctx, padded_shape[1 + i] % block_shape[i] == 0,
                 errors::InvalidArgument("padded_shape[", 1 + i,
@@ -132,11 +146,11 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp& input,
   //       padded_shape[M] / block_shape[M-1]] +
   //      remaining_shape
   // Determine the length of the prefix of block dims that can be combined
   // into the batch dimension due to having no padding and block_shape=1.
   std::vector<int64_t> output_shape(input_rank);
-  output_shape[0] = batch_size * block_num_elems;
+  output_shape[0] = output_dim;
   for (int i = 0; i < block_rank; ++i) {
     output_shape[1 + i] = padded_shape[1 + i] / block_shape[i];
   }
   std::copy(remainder_shape.begin(), remainder_shape.end(),
             output_shape.begin() + 1 + block_rank);
diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
index 4f927b10aa484..e3dc844c0343d 100644
--- a/tensorflow/core/framework/BUILD
+++ b/tensorflow/core/framework/BUILD
@@ -889,10 +889,11 @@ cc_library(
         "//tensorflow/core/lib/gtl:inlined_vector",
         "//tensorflow/core/lib/strings:numbers",
         "//tensorflow/core/lib/strings:scanner",
         "//tensorflow/core/lib/strings:str_util",
         "//tensorflow/core/platform:macros",
+        "//tensorflow/core/util:overflow",
         "@com_google_absl//absl/memory",
     ],
 )
 
 cc_library(
diff --git a/tensorflow/core/framework/shape_inference.cc b/tensorflow/core/framework/shape_inference.cc
index 9f854788b1aed..2dedd683d7c10 100644
--- a/tensorflow/core/framework/shape_inference.cc
+++ b/tensorflow/core/framework/shape_inference.cc
@@ -24,10 +24,11 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/numbers.h"
 #include "tensorflow/core/lib/strings/scanner.h"
 #include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/util/overflow.h"
 
 namespace tensorflow {
 namespace shape_inference {
 
 constexpr int32_t InferenceContext::kUnknownRank;
@@ -1109,11 +1110,11 @@ Status InferenceContext::Multiply(DimensionHandle first,
     *out = first;
   } else if (first_value == kUnknownDim || second_value == kUnknownDim) {
     *out = UnknownDim();
   } else {
     // Invariant: Both values are known and greater than 1.
-    const int64_t product = first_value * second_value;
+    const int64_t product = MultiplyWithoutOverflow(first_value, second_value);
     if (product < 0) {
       return errors::InvalidArgument(
           "Negative dimension size caused by overflow when multiplying ",
           first_value, " and ", second_value);
     }
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 94943cc04adfa..8c7ee12b4a8dd 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -27,10 +27,11 @@ load(
     "tf_kernel_tests_linkstatic",
 )
 load(
     "//third_party/mkl:build_defs.bzl",
     "if_mkl",
+    "mkl_deps",
 )
 
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "cc_header_only_library")
 
@@ -59,14 +60,10 @@ load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_tests")
 load("//tensorflow:tensorflow.bzl", "filegroup")
 load(
     "//tensorflow/core/platform:build_config_root.bzl",
     "tf_cuda_tests_tags",
 )
-load(
-    "//third_party/mkl:build_defs.bzl",
-    "mkl_deps",
-)
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
 load(
     "@local_config_rocm//rocm:build_defs.bzl",
     "if_rocm",
 )
@@ -4567,10 +4564,11 @@ tf_kernel_library(
     visibility = [":friends"],
     deps = [
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core/framework:bounds_check",
+        "//tensorflow/core/util:overflow",
         "//third_party/eigen3",
     ],
 )
 
 tf_cuda_cc_test(
diff --git a/tensorflow/core/kernels/spacetobatch_op.cc b/tensorflow/core/kernels/spacetobatch_op.cc
index 009500f79268f..e391529d852b0 100644
--- a/tensorflow/core/kernels/spacetobatch_op.cc
+++ b/tensorflow/core/kernels/spacetobatch_op.cc
@@ -19,22 +19,22 @@ limitations under the License.
 
 #include <memory>
 #include <string>
 #include <utility>
 
-#include "tensorflow/core/kernels/spacetobatch_functor.h"
-
 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/kernels/spacetobatch_functor.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/overflow.h"
 
 namespace tensorflow {
 
 typedef Eigen::ThreadPoolDevice CPUDevice;
 typedef Eigen::GpuDevice GPUDevice;
@@ -97,11 +97,17 @@ Status SpaceToBatchOpCompute(OpKernelContext* context,
   }
 
   // Compute the product of the block_shape values.
   int64_t block_shape_product = 1;
   for (int block_dim = 0; block_dim < block_dims; ++block_dim) {
-    block_shape_product *= block_shape[block_dim];
+    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]);
   }
   if (block_shape_product <= 0) {
     return errors::InvalidArgument(
         "Product of block sizes must be positive, got ", block_shape_product);
   }
@@ -129,12 +135,18 @@ Status SpaceToBatchOpCompute(OpKernelContext* context,
   TensorShape internal_output_shape;
 
   // The actual output shape exposed to callers.
   TensorShape external_output_shape;
 
-  external_output_shape.AddDim(orig_input_tensor.dim_size(0) *
-                               block_shape_product);
+  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);
 
   int64_t input_batch_size = orig_input_tensor.dim_size(0);
   for (int block_dim = 0; block_dim < removed_prefix_block_dims; ++block_dim) {
     const int64_t size = orig_input_tensor.dim_size(block_dim + 1);
     input_batch_size *= size;
diff --git a/tensorflow/core/util/BUILD b/tensorflow/core/util/BUILD
index 8c267148fb74f..8881f6fd5e914 100644
--- a/tensorflow/core/util/BUILD
+++ b/tensorflow/core/util/BUILD
@@ -531,10 +531,13 @@ tf_cuda_library(
 )
 
 cc_library(
     name = "overflow",
     hdrs = ["overflow.h"],
+    visibility = [
+        "//tensorflow:internal",
+    ],
     deps = [
         "//tensorflow/core/platform:logging",
         "//tensorflow/core/platform:macros",
         "//tensorflow/core/platform:types",
     ],
diff --git a/tensorflow/python/kernel_tests/array_ops/spacetobatch_op_test.py b/tensorflow/python/kernel_tests/array_ops/spacetobatch_op_test.py
index a095aced262ec..5e682364837e5 100644
--- a/tensorflow/python/kernel_tests/array_ops/spacetobatch_op_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/spacetobatch_op_test.py
@@ -14,11 +14,13 @@
 # ==============================================================================
 """Functional tests for SpaceToBatch and BatchToSpace ops."""
 
 import numpy as np
 
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_util
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gen_array_ops
@@ -514,10 +516,31 @@ class SpaceToBatchNDErrorHandlingTest(test.TestCase):
     t = array_ops.space_to_batch_nd(
         array_ops.placeholder(
             dtypes.float32, shape=(3, 2, 3, 2)), [2, 3], [[1, 1], [0, 0]])
     self.assertEqual([3 * 2 * 3, 2, 1, 2], t.get_shape().as_list())
 
+  @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)
+
 
 class SpaceToBatchGradientTest(test.TestCase, PythonOpImpl):
 
   # Check the gradients.
   def _checkGrad(self, x, paddings, block_size):

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)