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-05-11 15:28:27 -0700
Commit message: Fix heap OOB / undefined behavior in `RaggedTensorToTensor`

PiperOrigin-RevId: 373244623
Change-Id: I2d6cbbc8c67b238a8815bf58097f7586d87c54f2

Diff: diff --git a/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc b/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
index 376d55945d2ce..b79a07e67ba91 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
@@ -205,12 +205,12 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
       result->push_back(-1);
     }
     DCHECK_EQ(result->size(), first_dimension);
   }
 
-  void CalculateOutputIndexRowSplit(
-      OpKernelContext* context, const RowPartitionTensor& row_split,
+  Status CalculateOutputIndexRowSplit(
+      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) {
@@ -230,14 +230,15 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
       }
       for (INDEX_TYPE j = 0; j < row_length - real_length; ++j) {
         result->push_back(-1);
       }
     }
-    if (row_split_size > 0) {
-      OP_REQUIRES(context, result->size() == row_split(row_split_size - 1),
-                  errors::InvalidArgument("Invalid row split size."));
+    if (row_split_size > 0 && result->size() != row_split(row_split_size - 1)) {
+      return errors::InvalidArgument("Invalid row split size.");
     }
+
+    return Status::OK();
   }
 
   // Calculate the output index of the first element of a list.
   // The parent_output_index is the same computation for the previous list.
   // -1 indicates an element or list that is out of range.
@@ -257,24 +258,30 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
   // result[4] = -1 because it is the third element the size is 2.
   // 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(
-      OpKernelContext* context, const RowPartitionTensor& value_rowids,
+  Status CalculateOutputIndexValueRowID(
+      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);
     if (index_size == 0) {
-      return;
+      return Status::OK();
     }
 
     INDEX_TYPE current_output_column = 0;
     INDEX_TYPE current_value_rowid = value_rowids(0);
-    DCHECK_LT(current_value_rowid, parent_output_index.size());
+
+    if (current_value_rowid >= parent_output_index.size()) {
+      return errors::InvalidArgument(
+          "Got current_value_rowid=", current_value_rowid,
+          " which is not less than ", parent_output_index.size());
+    }
+
     INDEX_TYPE current_output_index = parent_output_index[current_value_rowid];
     result->push_back(current_output_index);
     for (INDEX_TYPE i = 1; i < index_size; ++i) {
       INDEX_TYPE next_value_rowid = value_rowids(i);
       if (next_value_rowid == current_value_rowid) {
@@ -287,17 +294,27 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
           }
         }
       } else {
         current_output_column = 0;
         current_value_rowid = next_value_rowid;
-        DCHECK_LT(next_value_rowid, parent_output_index.size());
+
+        if (next_value_rowid >= parent_output_index.size()) {
+          return errors::InvalidArgument(
+              "Got next_value_rowid=", next_value_rowid,
+              " which is not less than ", parent_output_index.size());
+        }
+
         current_output_index = parent_output_index[next_value_rowid];
       }
       result->push_back(current_output_index);
     }
-    OP_REQUIRES(context, result->size() == value_rowids.size(),
-                errors::InvalidArgument("Invalid row ids."));
+
+    if (result->size() != value_rowids.size()) {
+      return errors::InvalidArgument("Invalid row ids.");
+    }
+
+    return Status::OK();
   }
 
   Status CalculateOutputIndex(OpKernelContext* context, int dimension,
                               const vector<INDEX_TYPE>& parent_output_index,
                               INDEX_TYPE output_index_multiplier,
@@ -306,25 +323,23 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
     const RowPartitionTensor row_partition_tensor =
         GetRowPartitionTensor(context, dimension);
     auto partition_type = GetRowPartitionTypeByDimension(dimension);
     switch (partition_type) {
       case RowPartitionType::VALUE_ROWIDS:
-        CalculateOutputIndexValueRowID(
-            context, row_partition_tensor, parent_output_index,
-            output_index_multiplier, output_size, result);
-        return tensorflow::Status::OK();
+        return CalculateOutputIndexValueRowID(
+            row_partition_tensor, parent_output_index, output_index_multiplier,
+            output_size, result);
       case RowPartitionType::ROW_SPLITS:
         if (row_partition_tensor.size() - 1 > parent_output_index.size()) {
           return errors::InvalidArgument(
               "Row partition size is greater than output size: ",
               row_partition_tensor.size() - 1, " > ",
               parent_output_index.size());
         }
-        CalculateOutputIndexRowSplit(
-            context, row_partition_tensor, parent_output_index,
-            output_index_multiplier, output_size, result);
-        return tensorflow::Status::OK();
+        return CalculateOutputIndexRowSplit(
+            row_partition_tensor, parent_output_index, output_index_multiplier,
+            output_size, result);
       default:
         return errors::InvalidArgument(
             "Unsupported partition type:",
             RowPartitionTypeToString(partition_type));
     }

Vulnerable code:   void CalculateOutputIndexRowSplit(
      OpKernelContext* context, const RowPartitionTensor& row_split,
    if (row_split_size > 0) {
      OP_REQUIRES(context, result->size() == row_split(row_split_size - 1),
                  errors::InvalidArgument("Invalid row split size."));
  void CalculateOutputIndexValueRowID(
      OpKernelContext* context, const RowPartitionTensor& value_rowids,
      return;
    DCHECK_LT(current_value_rowid, parent_output_index.size());
        DCHECK_LT(next_value_rowid, parent_output_index.size());
    OP_REQUIRES(context, result->size() == value_rowids.size(),
                errors::InvalidArgument("Invalid row ids."));
        CalculateOutputIndexValueRowID(
            context, row_partition_tensor, parent_output_index,
            output_index_multiplier, output_size, result);
        return tensorflow::Status::OK();
        CalculateOutputIndexRowSplit(
            context, row_partition_tensor, parent_output_index,
            output_index_multiplier, output_size, result);
        return tensorflow::Status::OK();
Fixed code:   Status CalculateOutputIndexRowSplit(
      const RowPartitionTensor& row_split,
    if (row_split_size > 0 && result->size() != row_split(row_split_size - 1)) {
      return errors::InvalidArgument("Invalid row split size.");

    return Status::OK();
  Status CalculateOutputIndexValueRowID(
      const RowPartitionTensor& value_rowids,
      return Status::OK();

    if (current_value_rowid >= parent_output_index.size()) {
      return errors::InvalidArgument(
          "Got current_value_rowid=", current_value_rowid,
          " which is not less than ", parent_output_index.size());
    }


        if (next_value_rowid >= parent_output_index.size()) {
          return errors::InvalidArgument(
              "Got next_value_rowid=", next_value_rowid,
              " which is not less than ", parent_output_index.size());
        }


    if (result->size() != value_rowids.size()) {
      return errors::InvalidArgument("Invalid row ids.");
    }

    return Status::OK();
        return CalculateOutputIndexValueRowID(
            row_partition_tensor, parent_output_index, output_index_multiplier,
            output_size, result);
        return CalculateOutputIndexRowSplit(
            row_partition_tensor, parent_output_index, output_index_multiplier,
            output_size, result);