CVE-2022-21729: Integer Overflow or Wraparound (CWE-190)
value: Tensorflow is an Open Source Machine Learning Framework. The implementation of `UnravelIndex` is vulnerable to a division by zero caused by an integer overflow bug. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.
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: 6.5
Published: 2022-02-03T13:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/unravel_index_op.cc, tensorflow/python/kernel_tests/array_ops/array_ops_test.py
Commit date: 2021-11-30 14:55:43 -0800
Commit message: Fix integer overflow leading to divide by zero error in Unravel index kernel when dimensions product exceeds max int value.

PiperOrigin-RevId: 413250052
Change-Id: I9450b6e8acecd2e881a64b882e2b7c70e8e9289a

Diff: diff --git a/tensorflow/core/kernels/unravel_index_op.cc b/tensorflow/core/kernels/unravel_index_op.cc
index cff04387d6024..b8524e87598ce 100644
--- a/tensorflow/core/kernels/unravel_index_op.cc
+++ b/tensorflow/core/kernels/unravel_index_op.cc
@@ -11,10 +11,14 @@ distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <cstdint>
+
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/types.h"
 #define EIGEN_USE_THREADS
 
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
@@ -33,11 +37,12 @@ struct mod_op {
 typedef Eigen::ThreadPoolDevice CPUDevice;
 
 template <typename Tidx>
 class UnravelIndexOp : public OpKernel {
  public:
-  explicit UnravelIndexOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  explicit UnravelIndexOp(OpKernelConstruction* ctx)
+      : OpKernel(ctx), dtidx_(DataTypeToEnum<Tidx>::v()) {}
 
   void Compute(OpKernelContext* ctx) override {
     const Tensor& indices_tensor = ctx->input(0);
     OP_REQUIRES(ctx,
                 TensorShapeUtils::IsVector(indices_tensor.shape()) ||
@@ -52,16 +57,35 @@ class UnravelIndexOp : public OpKernel {
         errors::InvalidArgument("The indices can only be 1-D, got \"",
                                 dims_tensor.shape().DebugString(), "\""));
 
     auto dims = dims_tensor.vec<Tidx>();
     // Make sure dims does not contain a zero
+    double prod = 1;
+    uint64_t limit;
+    if (dtidx_ == DataType::DT_INT64) {
+      limit = kint64max;
+    } else {
+      limit = kint32max;
+    }
+
     for (int i = 0; i < dims.size(); i++) {
       OP_REQUIRES(
           ctx, dims(i) != 0,
           errors::InvalidArgument("Input dims cannot contain a dim of zero, "
                                   "but dims contains zero at index ",
                                   i));
+      OP_REQUIRES(ctx, dims(i) > 0,
+                  errors::InvalidArgument(
+                      "Input dims cannot be negative. Got dim = ", dims(i),
+                      " at index ", i));
+      // Check interger overflow
+      OP_REQUIRES(
+          ctx, prod <= limit / dims(i),
+          errors::InvalidArgument("Input dims product is causing integer "
+                                  "overflow: (",
+                                  dims, ")"));
+      prod = (prod * dims(i));
     }
 
     // Check to make sure indices is not out of boundary
     Eigen::Tensor<Tidx, 0, Eigen::RowMajor> dims_prod_eigen = dims.prod();
     Tidx dims_prod = dims_prod_eigen();
@@ -130,10 +154,11 @@ class UnravelIndexOp : public OpKernel {
       output = output.binaryExpr(strides.reshape(reshape).broadcast(bcast),
                                  mod_op<Tidx>()) /
                strides_shifted.reshape(reshape).broadcast(bcast);
     }
   }
+  const DataType dtidx_;
 };
 
 #define REGISTER_KERNEL(type)                                               \
   REGISTER_KERNEL_BUILDER(                                                  \
       Name("UnravelIndex").Device(DEVICE_CPU).TypeConstraint<type>("Tidx"), \
diff --git a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
index 9634184039e1d..e50dabc8120ab 100644
--- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
@@ -1578,10 +1578,24 @@ class UnravelIndexTest(test_util.TensorFlowTestCase):
                                     "dims cannot contain a dim of zero"):
           indices = constant_op.constant([2, 5, 7], dtype=dtype)
           dims = constant_op.constant([3, 0], dtype=dtype)
           self.evaluate(array_ops.unravel_index(indices=indices, dims=dims))
 
+  def testUnravelIndexIntegerOverflow(self):
+    with self.cached_session():
+      for dtype in [dtypes.int32, dtypes.int64]:
+        with self.assertRaisesRegex(
+            errors.InvalidArgumentError,
+            r"Input dims product is causing integer overflow"):
+          indices = constant_op.constant(-0x100000, dtype=dtype)
+          if dtype == dtypes.int32:
+            value = 0x10000000
+          else:
+            value = 0x7FFFFFFFFFFFFFFF
+          dims = constant_op.constant([value, value], dtype=dtype)
+          self.evaluate(array_ops.unravel_index(indices=indices, dims=dims))
+
 
 class GuaranteeConstOpTest(test_util.TensorFlowTestCase):
 
   def testSimple(self):
     a = array_ops.constant(10)

Vulnerable code:   explicit UnravelIndexOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
Fixed code: #include <cstdint>

#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/types.h"
  explicit UnravelIndexOp(OpKernelConstruction* ctx)
      : OpKernel(ctx), dtidx_(DataTypeToEnum<Tidx>::v()) {}
    double prod = 1;
    uint64_t limit;
    if (dtidx_ == DataType::DT_INT64) {
      limit = kint64max;
    } else {
      limit = kint32max;
    }

      OP_REQUIRES(ctx, dims(i) > 0,
                  errors::InvalidArgument(
                      "Input dims cannot be negative. Got dim = ", dims(i),
                      " at index ", i));
      // Check interger overflow
      OP_REQUIRES(
          ctx, prod <= limit / dims(i),
          errors::InvalidArgument("Input dims product is causing integer "
                                  "overflow: (",
                                  dims, ")"));
      prod = (prod * dims(i));
  const DataType dtidx_;
  def testUnravelIndexIntegerOverflow(self):
    with self.cached_session():
      for dtype in [dtypes.int32, dtypes.int64]:
        with self.assertRaisesRegex(
            errors.InvalidArgumentError,
            r"Input dims product is causing integer overflow"):
          indices = constant_op.constant(-0x100000, dtype=dtype)
          if dtype == dtypes.int32:
            value = 0x10000000
          else:
            value = 0x7FFFFFFFFFFFFFFF
          dims = constant_op.constant([value, value], dtype=dtype)
          self.evaluate(array_ops.unravel_index(indices=indices, dims=dims))