Fix segfault/heap buffer overflow in {Experimental,}DatasetToTFRecord where dataset is numeric.

mihaimaruseac · tensorflower-gardener · commit e0b6e58c3280 · 2021-07-29T15:08:55.000-07:00
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 --git a/tensorflow/core/kernels/data/experimental/to_tf_record_op.cc b/tensorflow/core/kernels/data/experimental/to_tf_record_op.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #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"
@@ -91,8 +92,20 @@ class ToTFRecordOp : public AsyncOpKernel {
     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(
