CVE-2021-37642: Divide By Zero (CWE-369)
value: TensorFlow is an end-to-end open source platform for machine learning. In affected versions the implementation of `tf.raw_ops.ResourceScatterDiv` is vulnerable to a division by 0 error. The [implementation](https://github.com/tensorflow/tensorflow/blob/8d72537c6abf5a44103b57b9c2e22c14f5f49698/tensorflow/core/kernels/resource_variable_ops.cc#L865) uses a common class for all binary operations but fails to treat the division by 0 case separately. We have patched the issue in GitHub commit 4aacb30888638da75023e6601149415b39763d76. The fix will be included in TensorFlow 2.6.0. We will also cherrypick this commit on TensorFlow 2.5.1, TensorFlow 2.4.3, and TensorFlow 2.3.4, as these are also affected and still in supported range.
CWE description: The product divides a value by zero.
CVSS v3 base score: 5.5
Published: 2021-08-12T18:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/resource_variable_ops.cc, tensorflow/python/distribute/sharded_variable_test.py
Commit date: 2021-08-03 12:33:45 -0700
Commit message: Disallow division by zero FPE in `tf.raw_ops.ResourceScatterDiv`

Had to update a test that was broken.

PiperOrigin-RevId: 388516976
Change-Id: Ic358e6bf0559e011539974d453fc7aa18b427e9c

Diff: diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index b81a7a517ea6d..1a74774785f94 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -871,10 +871,39 @@ TF_CALL_GPU_NUMBER_TYPES(REGISTER_GATHER_ND_GPU);
 #undef REGISTER_GATHER_ND_CPU
 #undef REGISTER_GATHER_ND_GPU
 #undef REGISTER_GATHER_ND_ALL_INDICES
 #undef REGISTER_GATHER_ND_FULL
 
+namespace {
+
+template <typename Device>
+bool isCPUDevice() {
+  return false;
+}
+
+template <>
+bool isCPUDevice<CPUDevice>() {
+  return true;
+}
+
+template <typename T>
+bool ValidateInput(const Tensor& updates) {
+  const auto updates_flat = updates.flat<T>();
+  const T zero(0);
+  for (int i = 0; i < updates.NumElements(); i++) {
+    if (updates_flat(i) == zero) return false;
+  }
+  return true;
+}
+
+template <>
+bool ValidateInput<Variant>(const Tensor& updates) {
+  return true;
+}
+
+}  // namespace
+
 template <typename Device, typename T, typename Index, scatter_op::UpdateOp op>
 class ResourceScatterUpdateOp : public OpKernel {
  public:
   explicit ResourceScatterUpdateOp(OpKernelConstruction* c) : OpKernel(c) {
     // We use the same kernel for many operations.
@@ -937,10 +966,16 @@ class ResourceScatterUpdateOp : public OpKernel {
         errors::InvalidArgument("params.shape[0] too large for ",
                                 DataTypeString(DataTypeToEnum<Index>::v()),
                                 " indexing: ", params->dim_size(0), " > ",
                                 std::numeric_limits<Index>::max()));
 
+    // Prevent division by 0
+    if (isCPUDevice<Device>() && op == tensorflow::scatter_op::UpdateOp::DIV) {
+      OP_REQUIRES(c, ValidateInput<T>(updates),
+                  errors::InvalidArgument("updates must not contain 0"));
+    }
+
     if (N > 0) {
       auto indices_flat = indices.flat<Index>();
       auto params_flat = params->flat_outer_dims<T>();
       if (TensorShapeUtils::IsScalar(updates.shape())) {
         const auto update = updates.scalar<T>();
diff --git a/tensorflow/python/distribute/sharded_variable_test.py b/tensorflow/python/distribute/sharded_variable_test.py
index 9e1a32dc70286..cb61cd5715aca 100644
--- a/tensorflow/python/distribute/sharded_variable_test.py
+++ b/tensorflow/python/distribute/sharded_variable_test.py
@@ -173,12 +173,13 @@ class ShardedVariableTest(test.TestCase, parameterized.TestCase):
   @parameterized.parameters('scatter_add', 'scatter_div', 'scatter_max',
                             'scatter_min', 'scatter_mul', 'scatter_sub',
                             'scatter_update')
   def test_scatter_ops_even_partition(self, op):
     v = variables_lib.Variable(array_ops.zeros((30, 1)))
+    # Make sure values does not contain 0 due to testing `scatter_div`!
     sparse_delta = ops.IndexedSlices(
-        values=constant_op.constant([[0.], [1.], [2.], [3.], [4.]]),
+        values=constant_op.constant([[1.], [2.], [3.], [4.], [5.]]),
         indices=constant_op.constant([0, 10, 12, 21, 22]))
 
     v0 = variables_lib.Variable(array_ops.zeros((10, 1)))
     v1 = variables_lib.Variable(array_ops.zeros((10, 1)))
     v2 = variables_lib.Variable(array_ops.zeros((10, 1)))

Vulnerable code:         values=constant_op.constant([[0.], [1.], [2.], [3.], [4.]]),
Fixed code: namespace {

template <typename Device>
bool isCPUDevice() {
  return false;
}

template <>
bool isCPUDevice<CPUDevice>() {
  return true;
}

template <typename T>
bool ValidateInput(const Tensor& updates) {
  const auto updates_flat = updates.flat<T>();
  const T zero(0);
  for (int i = 0; i < updates.NumElements(); i++) {
    if (updates_flat(i) == zero) return false;
  }
  return true;
}

template <>
bool ValidateInput<Variant>(const Tensor& updates) {
  return true;
}

}  // namespace

    // Prevent division by 0
    if (isCPUDevice<Device>() && op == tensorflow::scatter_op::UpdateOp::DIV) {
      OP_REQUIRES(c, ValidateInput<T>(updates),
                  errors::InvalidArgument("updates must not contain 0"));
    }

    # Make sure values does not contain 0 due to testing `scatter_div`!
        values=constant_op.constant([[1.], [2.], [3.], [4.], [5.]]),