CVE-2021-29608 (published 2021-05-14T20:15Z)

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.

Severity: MEDIUM
CVSS v2 base score: 4.6 (AV:L/AC:L/Au:N/C:P/I:P/A:P)
CVSS v3 base score: 7.8, HIGH (CVSS:3.1/AV:L/AC:L/PR:L/UI:N/S:U/C:H/I:H/A:H)

Weakness information:
cwe_id: CWE-131
cwe_name: Incorrect Calculation of Buffer Size
description: The product does not correctly calculate the size to be used when allocating a buffer, which could lead to a buffer overflow.
extended_description: 
            
            
             
			 
            
         
         

Fix information:
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Mihai Maruseac
author_date: 2021-05-11 15:22:49-07:00
msg: Fix heap OOB / undefined behavior in `RaggedTensorToTensor`

PiperOrigin-RevId: 373244623
Change-Id: I2d6cbbc8c67b238a8815bf58097f7586d87c54f2
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: ragged_tensor_to_tensor_op.cc
change_type: ModificationType.MODIFY
diff: @@ -207,8 +207,8 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
     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) {
@@ -232,10 +232,11 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
         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.
@@ -259,20 +260,26 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
   // 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) {
@@ -289,13 +296,23 @@ 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,
@@ -308,10 +325,9 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
     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(
@@ -319,10 +335,9 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
               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:",

programming_language: C++
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Amit Patankar
author_date: 2021-04-13 14:54:18-07:00
msg: Fix `tf.raw_ops.RaggedTensorToTensor` failing CHECK in `tensor.cc`.

PiperOrigin-RevId: 368300502
Change-Id: I91255d23c4bfd3aa3c029aac773937c09daf3c64
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: ragged_tensor_to_tensor_op.cc
change_type: ModificationType.MODIFY
diff: @@ -345,6 +345,11 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
 
   void Compute(OpKernelContext* context) override {
     INDEX_TYPE first_dimension;
+    const Tensor first_partition_tensor =
+        context->input(kFirstPartitionInputIndex);
+    OP_REQUIRES(context, first_partition_tensor.NumElements() > 0,
+                errors::InvalidArgument("Invalid first partition input. Tensor "
+                                        "requires at least one element."));
     OP_REQUIRES_OK(context, GetFirstDimensionSize(context, &first_dimension));
     vector<INDEX_TYPE> output_size;
     OP_REQUIRES_OK(context,

programming_language: C++
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Amit Patankar
author_date: 2021-04-15 13:28:49-07:00
msg: Fix `tf.raw_ops.RaggedTensorToTensor` failing CHECK.

PiperOrigin-RevId: 368706628
Change-Id: I5c9ea4833f38835ee183ca50d63251dc89c9f3bc
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: ragged_tensor_to_tensor_op.cc
change_type: ModificationType.MODIFY
diff: @@ -208,7 +208,7 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
   }
 
   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) {
@@ -233,7 +233,8 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
       }
     }
     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."));
     }
   }
 
@@ -259,7 +260,7 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
   // 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) {
@@ -293,7 +294,8 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
       }
       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,
@@ -307,13 +309,13 @@ class RaggedTensorToTensorBaseOp : public OpKernel {
     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(

programming_language: C++