CVE-2021-37650: Out-of-bounds Write (CWE-787)
value: TensorFlow is an end-to-end open source platform for machine learning. In affected versions the implementation for `tf.raw_ops.ExperimentalDatasetToTFRecord` and `tf.raw_ops.DatasetToTFRecord` can trigger heap buffer overflow and segmentation fault. The [implementation](https://github.com/tensorflow/tensorflow/blob/f24faa153ad31a4b51578f8181d3aaab77a1ddeb/tensorflow/core/kernels/data/experimental/to_tf_record_op.cc#L93-L102) assumes that all records in the dataset are of string type. However, there is no check for that, and the example given above uses numeric types. We have patched the issue in GitHub commit e0b6e58c328059829c3eb968136f17aa72b6c876. The fix will be included in TensorFlow 2.6.0. We will also cherrypick this commit on TensorFlow 2.5.1, TensorFlow 2.4.3, and TensorFlow 2.3.4, as these are also affected and still in supported range.
CWE description: The product writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 7.8
Published: 2021-08-12T21:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/data/experimental/to_tf_record_op.cc
Commit date: 2021-07-29 15:08:55 -0700
Commit message: Fix segfault/heap buffer overflow in `{Experimental,}DatasetToTFRecord` where dataset is numeric.

Code assumes only strings inputs and then interprets numbers as valid `tstring`s. Then, when trying to compute the CRC of the record this results in heap buffer overflow.

PiperOrigin-RevId: 387675909
Change-Id: I7396b9b8afc1ac744112af7c0b1cd7bb41e0f556

Diff: diff --git a/tensorflow/core/kernels/data/experimental/to_tf_record_op.cc b/tensorflow/core/kernels/data/experimental/to_tf_record_op.cc
index 0ba04d0bd94bd..4f759bede55b6 100644
--- a/tensorflow/core/kernels/data/experimental/to_tf_record_op.cc
+++ b/tensorflow/core/kernels/data/experimental/to_tf_record_op.cc
@@ -16,10 +16,11 @@ limitations under the License.
 #include "tensorflow/core/data/root_dataset.h"
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/function_handle_cache.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/ops_util.h"
 #include "tensorflow/core/lib/core/threadpool.h"
 #include "tensorflow/core/lib/io/record_writer.h"
 #include "tensorflow/core/platform/file_system.h"
 #include "tensorflow/core/platform/resource.h"
@@ -89,12 +90,24 @@ class ToTFRecordOp : public AsyncOpKernel {
 
     std::unique_ptr<IteratorBase> iterator;
     TF_RETURN_IF_ERROR(finalized_dataset->MakeIterator(
         &iter_ctx, /*parent=*/nullptr, "ToTFRecordOpIterator", &iterator));
 
+    const int num_output_dtypes = finalized_dataset->output_dtypes().size();
+    if (num_output_dtypes != 1) {
+      return errors::InvalidArgument(
+          "ToTFRecordOp currently only support datasets of 1 single column, ",
+          "but got ", num_output_dtypes);
+    }
+    const DataType dt = finalized_dataset->output_dtypes()[0];
+    if (dt != DT_STRING) {
+      return errors::InvalidArgument(
+          "ToTFRecordOp currently only supports DT_STRING dataypes, but got ",
+          DataTypeString(dt));
+    }
     std::vector<Tensor> components;
-    components.reserve(finalized_dataset->output_dtypes().size());
+    components.reserve(num_output_dtypes);
     bool end_of_sequence;
     do {
       TF_RETURN_IF_ERROR(
           iterator->GetNext(&iter_ctx, &components, &end_of_sequence));
 

Vulnerable code:     components.reserve(finalized_dataset->output_dtypes().size());
Fixed code: #include "tensorflow/core/framework/types.h"
    const int num_output_dtypes = finalized_dataset->output_dtypes().size();
    if (num_output_dtypes != 1) {
      return errors::InvalidArgument(
          "ToTFRecordOp currently only support datasets of 1 single column, ",
          "but got ", num_output_dtypes);
    }
    const DataType dt = finalized_dataset->output_dtypes()[0];
    if (dt != DT_STRING) {
      return errors::InvalidArgument(
          "ToTFRecordOp currently only supports DT_STRING dataypes, but got ",
          DataTypeString(dt));
    }
    components.reserve(num_output_dtypes);