CVE-2021-29534: Improper Check for Unusual or Exceptional Conditions (CWE-754)
value: TensorFlow is an end-to-end open source platform for machine learning. An attacker can trigger a denial of service via a `CHECK`-fail in `tf.raw_ops.SparseConcat`. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/b432a38fe0e1b4b904a6c222cbce794c39703e87/tensorflow/core/kernels/sparse_concat_op.cc#L76) takes the values specified in `shapes[0]` as dimensions for the output shape. The `TensorShape` constructor(https://github.com/tensorflow/tensorflow/blob/6f9896890c4c703ae0a0845394086e2e1e523299/tensorflow/core/framework/tensor_shape.cc#L183-L188) uses a `CHECK` operation which triggers when `InitDims`(https://github.com/tensorflow/tensorflow/blob/6f9896890c4c703ae0a0845394086e2e1e523299/tensorflow/core/framework/tensor_shape.cc#L212-L296) returns a non-OK status. This is a legacy implementation of the constructor and operations should use `BuildTensorShapeBase` or `AddDimWithStatus` to prevent `CHECK`-failures in the presence of overflows. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range.
CWE description: The product does not check or incorrectly checks for unusual or exceptional conditions that are not expected to occur frequently during day to day operation of the product.
CVSS v3 base score: 5.5
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/sparse_tensors_map_ops.cc
Commit date: 2021-04-20 12:24:04 -0700
Commit message: Fix overflow CHECK issue with `tf.raw_ops.AddManySparseToTensorsMap`.

PiperOrigin-RevId: 369492969
Change-Id: I1d70d6c0c92e3d7a25bc3b3aa2a0c0ac9688bf81

Diff: diff --git a/tensorflow/core/kernels/sparse_tensors_map_ops.cc b/tensorflow/core/kernels/sparse_tensors_map_ops.cc
index c2c0e43ca2ba8..5ea5fca544d3e 100644
--- a/tensorflow/core/kernels/sparse_tensors_map_ops.cc
+++ b/tensorflow/core/kernels/sparse_tensors_map_ops.cc
@@ -19,20 +19,18 @@ limitations under the License.
 #include <numeric>
 #include <unordered_map>
 #include <utility>
 #include <vector>
 
-#include "tensorflow/core/framework/op_kernel.h"
-#include "tensorflow/core/framework/register_types.h"
-
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_util.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "tensorflow/core/util/overflow.h"
 #include "tensorflow/core/util/sparse/sparse_tensor.h"
 
 namespace tensorflow {
 
 typedef Eigen::ThreadPoolDevice CPUDevice;
@@ -252,31 +250,45 @@ class AddManySparseToTensorsMapOp : public SparseTensorAccessingOp {
     OP_REQUIRES(
         context, rank > 1,
         errors::InvalidArgument(
             "Rank of input SparseTensor should be > 1, but saw rank: ", rank));
 
-    TensorShape tensor_input_shape(input_shape->vec<int64>());
+    auto input_shape_vec = input_shape->vec<int64>();
+    int new_num_elements = 1;
+    bool overflow_ocurred = false;
+    for (int i = 0; i < input_shape_vec.size(); i++) {
+      new_num_elements =
+          MultiplyWithoutOverflow(new_num_elements, input_shape_vec(i));
+      if (new_num_elements < 0) {
+        overflow_ocurred = true;
+      }
+    }
+
+    OP_REQUIRES(
+        context, !overflow_ocurred,
+        errors::Internal("Encountered overflow from large input shape."));
+
+    TensorShape tensor_input_shape(input_shape_vec);
     gtl::InlinedVector<int64, 8> std_order(rank);
     std::iota(std_order.begin(), std_order.end(), 0);
     SparseTensor input_st;
     OP_REQUIRES_OK(context, SparseTensor::Create(*input_indices, *input_values,
                                                  tensor_input_shape, std_order,
                                                  &input_st));
 
-    auto input_shape_t = input_shape->vec<int64>();
-    const int64 N = input_shape_t(0);
+    const int64 N = input_shape_vec(0);
 
     Tensor sparse_handles(DT_INT64, TensorShape({N}));
     auto sparse_handles_t = sparse_handles.vec<int64>();
 
     OP_REQUIRES_OK(context, input_st.IndicesValid());
 
     // We can generate the output shape proto string now, for all
     // minibatch entries.
     TensorShape output_shape;
     OP_REQUIRES_OK(context, TensorShapeUtils::MakeShape(
-                                input_shape_t.data() + 1,
+                                input_shape_vec.data() + 1,
                                 input_shape->NumElements() - 1, &output_shape));
 
     // Get groups by minibatch dimension
     std::unordered_set<int64> visited;
     sparse::GroupIterable minibatch = input_st.group({0});

Vulnerable code: #include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"

    TensorShape tensor_input_shape(input_shape->vec<int64>());
    auto input_shape_t = input_shape->vec<int64>();
    const int64 N = input_shape_t(0);
                                input_shape_t.data() + 1,
Fixed code: #include "tensorflow/core/util/overflow.h"
    auto input_shape_vec = input_shape->vec<int64>();
    int new_num_elements = 1;
    bool overflow_ocurred = false;
    for (int i = 0; i < input_shape_vec.size(); i++) {
      new_num_elements =
          MultiplyWithoutOverflow(new_num_elements, input_shape_vec(i));
      if (new_num_elements < 0) {
        overflow_ocurred = true;
      }
    }

    OP_REQUIRES(
        context, !overflow_ocurred,
        errors::Internal("Encountered overflow from large input shape."));

    TensorShape tensor_input_shape(input_shape_vec);
    const int64 N = input_shape_vec(0);
                                input_shape_vec.data() + 1,