CVE-2021-29608: Incorrect Calculation of Buffer Size (CWE-131)
value: TensorFlow is an end-to-end open source platform for machine learning. Due to lack of validation in `tf.raw_ops.RaggedTensorToTensor`, an attacker can exploit an undefined behavior if input arguments are empty. The implementation(https://github.com/tensorflow/tensorflow/blob/656e7673b14acd7835dc778867f84916c6d1cac2/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc#L356-L360) only checks that one of the tensors is not empty, but does not check for the other ones. There are multiple `DCHECK` validations to prevent heap OOB, but these are no-op in release builds, hence they don't prevent anything. The fix will be included in TensorFlow 2.5.0. We will also cherrypick these commits 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 correctly calculate the size to be used when allocating a buffer, which could lead to a buffer overflow.
CVSS v3 base score: 7.8
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
Commit date: 2021-04-15 13:36:54 -0700
Commit message: Fix `tf.raw_ops.RaggedTensorToTensor` failing CHECK.

PiperOrigin-RevId: 368706628
Change-Id: I5c9ea4833f38835ee183ca50d63251dc89c9f3bc

Diff: diff --git a/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc b/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
index 433d910f6090c..434c853b63daa 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
@@ -206,11 +206,11 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
     }
     DCHECK_EQ(result->size(), first_dimension);
   }
 
   void CalculateOutputIndexRowSplit(
-      const RowPartitionTensor& row_split,
+      OpKernelContext* context, const RowPartitionTensor& row_split,
       const vector<INDEX_TYPE>& parent_output_index,
       INDEX_TYPE output_index_multiplier, INDEX_TYPE output_size,
       vector<INDEX_TYPE>* result) {
     INDEX_TYPE row_split_size = row_split.size();
     if (row_split_size > 0) {
@@ -231,11 +231,12 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
       for (INDEX_TYPE j = 0; j < row_length - real_length; ++j) {
         result->push_back(-1);
       }
     }
     if (row_split_size > 0) {
-      DCHECK_EQ(result->size(), row_split(row_split_size - 1));
+      OP_REQUIRES(context, result->size() == row_split(row_split_size - 1),
+                  errors::InvalidArgument("Invalid row split size."));
     }
   }
 
   // Calculate the output index of the first element of a list.
   // The parent_output_index is the same computation for the previous list.
@@ -257,11 +258,11 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
   // result[5] = parent_output_index[value_rowids[3]]
   // result[6] = -1 because parent_output_index[value_rowids[6]] == -1
   // result[7] = -1 because parent_output_index[value_rowids[6]] == -1
   // result[8] = parent_output_index[value_rowids[7]]
   void CalculateOutputIndexValueRowID(
-      const RowPartitionTensor& value_rowids,
+      OpKernelContext* context, const RowPartitionTensor& value_rowids,
       const vector<INDEX_TYPE>& parent_output_index,
       INDEX_TYPE output_index_multiplier, INDEX_TYPE output_size,
       vector<INDEX_TYPE>* result) {
     const INDEX_TYPE index_size = value_rowids.size();
     result->reserve(index_size);
@@ -291,11 +292,12 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
         DCHECK_LT(next_value_rowid, parent_output_index.size());
         current_output_index = parent_output_index[next_value_rowid];
       }
       result->push_back(current_output_index);
     }
-    DCHECK_EQ(result->size(), value_rowids.size());
+    OP_REQUIRES(context, result->size() == value_rowids.size(),
+                errors::InvalidArgument("Invalid row ids."));
   }
 
   Status CalculateOutputIndex(OpKernelContext* context, int dimension,
                               const vector<INDEX_TYPE>& parent_output_index,
                               INDEX_TYPE output_index_multiplier,
@@ -305,17 +307,17 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
         GetRowPartitionTensor(context, dimension);
     auto partition_type = GetRowPartitionTypeByDimension(dimension);
     switch (partition_type) {
       case RowPartitionType::VALUE_ROWIDS:
         CalculateOutputIndexValueRowID(
-            row_partition_tensor, parent_output_index, output_index_multiplier,
-            output_size, result);
+            context, row_partition_tensor, parent_output_index,
+            output_index_multiplier, output_size, result);
         return tensorflow::Status::OK();
       case RowPartitionType::ROW_SPLITS:
-        CalculateOutputIndexRowSplit(row_partition_tensor, parent_output_index,
-                                     output_index_multiplier, output_size,
-                                     result);
+        CalculateOutputIndexRowSplit(
+            context, row_partition_tensor, parent_output_index,
+            output_index_multiplier, output_size, result);
         return tensorflow::Status::OK();
       default:
         return errors::InvalidArgument(
             "Unsupported partition type:",
             RowPartitionTypeToString(partition_type));

Vulnerable code:       const RowPartitionTensor& row_split,
      DCHECK_EQ(result->size(), row_split(row_split_size - 1));
      const RowPartitionTensor& value_rowids,
    DCHECK_EQ(result->size(), value_rowids.size());
            row_partition_tensor, parent_output_index, output_index_multiplier,
            output_size, result);
        CalculateOutputIndexRowSplit(row_partition_tensor, parent_output_index,
                                     output_index_multiplier, output_size,
                                     result);
Fixed code:       OpKernelContext* context, const RowPartitionTensor& row_split,
      OP_REQUIRES(context, result->size() == row_split(row_split_size - 1),
                  errors::InvalidArgument("Invalid row split size."));
      OpKernelContext* context, const RowPartitionTensor& value_rowids,
    OP_REQUIRES(context, result->size() == value_rowids.size(),
                errors::InvalidArgument("Invalid row ids."));
            context, row_partition_tensor, parent_output_index,
            output_index_multiplier, output_size, result);
        CalculateOutputIndexRowSplit(
            context, row_partition_tensor, parent_output_index,
            output_index_multiplier, output_size, result);