CVE-2021-29564: NULL Pointer Dereference (CWE-476)
value: TensorFlow is an end-to-end open source platform for machine learning. An attacker can trigger a null pointer dereference in the implementation of `tf.raw_ops.EditDistance`. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/79865b542f9ffdc9caeb255631f7c56f1d4b6517/tensorflow/core/kernels/edit_distance_op.cc#L103-L159) has incomplete validation of the input parameters. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit 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 dereferences a pointer that it expects to be valid but is NULL.
CVSS v3 base score: 5.5
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/edit_distance_op.cc
Commit date: 2021-05-04 18:18:01 -0700
Commit message: Fix multiple issues in EditDistance

PiperOrigin-RevId: 372033948
Change-Id: Ieb957c29894af05bdfeb1a0402fced808dfcfd7b

Diff: diff --git a/tensorflow/core/kernels/edit_distance_op.cc b/tensorflow/core/kernels/edit_distance_op.cc
index 4aecdc9e414d3..386a1af08409f 100644
--- a/tensorflow/core/kernels/edit_distance_op.cc
+++ b/tensorflow/core/kernels/edit_distance_op.cc
@@ -62,10 +62,16 @@ Status ValidateShapes(OpKernelContext* ctx, const Tensor& hypothesis_indices,
         hypothesis_shape.shape().DebugString());
   if (!TensorShapeUtils::IsVector(truth_shape.shape()))
     return errors::InvalidArgument(
         "truth_shape should be a vector, but got shape: ",
         truth_shape.shape().DebugString());
+  if (hypothesis_values.NumElements() != hypothesis_indices.dim_size(0))
+    return errors::InvalidArgument(
+        "Expected hypothesis_values.NumElements == "
+        "#rows(hypothesis_indices), their shapes are: ",
+        hypothesis_values.shape().DebugString(), " and ",
+        hypothesis_indices.shape().DebugString());
   if (hypothesis_shape.NumElements() != hypothesis_indices.dim_size(1))
     return errors::InvalidArgument(
         "Expected hypothesis_shape.NumElements == "
         "#cols(hypothesis_indices), their shapes are: ",
         hypothesis_shape.shape().DebugString(), " and ",
@@ -73,10 +79,16 @@ Status ValidateShapes(OpKernelContext* ctx, const Tensor& hypothesis_indices,
   if (truth_shape.NumElements() < 2)
     return errors::InvalidArgument(
         "Input SparseTensors must have rank at least 2, but truth_shape "
         "rank is: ",
         truth_shape.NumElements());
+  if (truth_values.NumElements() != truth_indices.dim_size(0))
+    return errors::InvalidArgument(
+        "Expected truth_values.NumElements == "
+        "#rows(truth_indices), their shapes are: ",
+        truth_values.shape().DebugString(), " and ",
+        truth_indices.shape().DebugString());
   if (truth_shape.NumElements() != truth_indices.dim_size(1))
     return errors::InvalidArgument(
         "Expected truth_shape.NumElements == "
         "#cols(truth_indices), their shapes are: ",
         truth_shape.shape().DebugString(), " and ",
@@ -151,10 +163,15 @@ class EditDistanceOp : public OpKernel {
     TensorShape output_shape;
     for (int d = 0; d < static_cast<int>(group_dims.size()); ++d) {
       output_shape.AddDim(std::max(hypothesis_st_shape.dim_size(d),
                                    truth_st_shape.dim_size(d)));
     }
+    const auto output_elements = output_shape.num_elements();
+    OP_REQUIRES(
+        ctx, output_elements > 0,
+        errors::InvalidArgument("Got output shape ", output_shape.DebugString(),
+                                " which has 0 elements"));
 
     Tensor* output = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output("output", output_shape, &output));
     auto output_t = output->flat<float>();
     output_t.setZero();
@@ -183,37 +200,61 @@ class EditDistanceOp : public OpKernel {
       auto hypothesis_seq = hypothesis_j.values<T>();
 
       if (g_truth == g_hypothesis) {
         auto loc = std::inner_product(g_truth.begin(), g_truth.end(),
                                       output_strides.begin(), int64{0});
+        OP_REQUIRES(
+            ctx, loc < output_elements,
+            errors::Internal("Got an inner product ", loc,
+                             " which would require in writing to outside of "
+                             "the buffer for the output tensor (max elements ",
+                             output_elements, ")"));
         output_t(loc) =
             gtl::LevenshteinDistance<T>(truth_seq, hypothesis_seq, cmp);
         if (normalize_) output_t(loc) /= truth_seq.size();
 
         ++hypothesis_iter;
         ++truth_iter;
       } else if (g_truth > g_hypothesis) {  // zero-length truth
         auto loc = std::inner_product(g_hypothesis.begin(), g_hypothesis.end(),
                                       output_strides.begin(), int64{0});
+        OP_REQUIRES(
+            ctx, loc < output_elements,
+            errors::Internal("Got an inner product ", loc,
+                             " which would require in writing to outside of "
+                             "the buffer for the output tensor (max elements ",
+                             output_elements, ")"));
         output_t(loc) = hypothesis_seq.size();
         if (normalize_ && output_t(loc) != 0.0f) {
           output_t(loc) = std::numeric_limits<float>::infinity();
         }
         ++hypothesis_iter;
       } else {  // zero-length hypothesis
         auto loc = std::inner_product(g_truth.begin(), g_truth.end(),
                                       output_strides.begin(), int64{0});
+        OP_REQUIRES(
+            ctx, loc < output_elements,
+            errors::Internal("Got an inner product ", loc,
+                             " which would require in writing to outside of "
+                             "the buffer for the output tensor (max elements ",
+                             output_elements, ")"));
         output_t(loc) = (normalize_) ? 1.0 : truth_seq.size();
         ++truth_iter;
       }
     }
     while (hypothesis_iter != hypothesis_grouper.end()) {  // zero-length truths
       sparse::Group hypothesis_j = *hypothesis_iter;
       std::vector<int64> g_hypothesis = hypothesis_j.group();
       auto hypothesis_seq = hypothesis_j.values<T>();
       auto loc = std::inner_product(g_hypothesis.begin(), g_hypothesis.end(),
                                     output_strides.begin(), int64{0});
+      OP_REQUIRES(
+          ctx, loc < output_elements,
+          errors::Internal("Got an inner product ", loc,
+                           " which would require in writing to outside of the "
+                           "buffer for the output tensor (max elements ",
+                           output_elements, ")"));
       output_t(loc) = hypothesis_seq.size();
       if (normalize_ && output_t(loc) != 0.0f) {
         output_t(loc) = std::numeric_limits<float>::infinity();
       }
       ++hypothesis_iter;
@@ -222,10 +263,16 @@ class EditDistanceOp : public OpKernel {
       sparse::Group truth_i = *truth_iter;
       std::vector<int64> g_truth = truth_i.group();
       auto truth_seq = truth_i.values<T>();
       auto loc = std::inner_product(g_truth.begin(), g_truth.end(),
                                     output_strides.begin(), int64{0});
+      OP_REQUIRES(
+          ctx, loc < output_elements,
+          errors::Internal("Got an inner product ", loc,
+                           " which would require in writing to outside of the "
+                           "buffer for the output tensor (max elements ",
+                           output_elements, ")"));
       output_t(loc) = (normalize_) ? 1.0 : truth_seq.size();
       ++truth_iter;
     }
   }
 

Fixed code:   if (hypothesis_values.NumElements() != hypothesis_indices.dim_size(0))
    return errors::InvalidArgument(
        "Expected hypothesis_values.NumElements == "
        "#rows(hypothesis_indices), their shapes are: ",
        hypothesis_values.shape().DebugString(), " and ",
        hypothesis_indices.shape().DebugString());
  if (truth_values.NumElements() != truth_indices.dim_size(0))
    return errors::InvalidArgument(
        "Expected truth_values.NumElements == "
        "#rows(truth_indices), their shapes are: ",
        truth_values.shape().DebugString(), " and ",
        truth_indices.shape().DebugString());
    const auto output_elements = output_shape.num_elements();
    OP_REQUIRES(
        ctx, output_elements > 0,
        errors::InvalidArgument("Got output shape ", output_shape.DebugString(),
                                " which has 0 elements"));
        OP_REQUIRES(
            ctx, loc < output_elements,
            errors::Internal("Got an inner product ", loc,
                             " which would require in writing to outside of "
                             "the buffer for the output tensor (max elements ",
                             output_elements, ")"));
        OP_REQUIRES(
            ctx, loc < output_elements,
            errors::Internal("Got an inner product ", loc,
                             " which would require in writing to outside of "
                             "the buffer for the output tensor (max elements ",
                             output_elements, ")"));
        OP_REQUIRES(
            ctx, loc < output_elements,
            errors::Internal("Got an inner product ", loc,
                             " which would require in writing to outside of "
                             "the buffer for the output tensor (max elements ",
                             output_elements, ")"));
      OP_REQUIRES(
          ctx, loc < output_elements,
          errors::Internal("Got an inner product ", loc,
                           " which would require in writing to outside of the "
                           "buffer for the output tensor (max elements ",
                           output_elements, ")"));
      OP_REQUIRES(
          ctx, loc < output_elements,
          errors::Internal("Got an inner product ", loc,
                           " which would require in writing to outside of the "
                           "buffer for the output tensor (max elements ",
                           output_elements, ")"));