CVE-2021-29608 (C++, published 2021-05-14T20:15Z)
Description: 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.
Severity: MEDIUM (CVSS v3 base score: 7.8)
Weakness: CWE-131 \u2014 Incorrect Calculation of Buffer Size
The product does not correctly calculate the size to be used when allocating a buffer, which could lead to a buffer overflow.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
Commit message: Fix heap OOB / undefined behavior in `RaggedTensorToTensor`

PiperOrigin-RevId: 373244623
Change-Id: I2d6cbbc8c67b238a8815bf58097f7586d87c54f2


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);