Vulnerability: Null pointer dereference in `EditDistance` ### Impact
An attacker can trigger a null pointer dereference in the implementation of `tf.raw_ops.EditDistance`: 
    
```python
import tensorflow as tf

hypothesis_indices = tf.constant([247, 247, 247], shape=[1, 3], dtype=tf.int64)
hypothesis_values = tf.constant([-9.9999], shape=[1], dtype=tf.float32)
hypothesis_shape = tf.constant([0, 0, 0], shape=[3], dtype=tf.int64)
truth_indices = tf.constant([], shape=[0, 3], dtype=tf.int64)
truth_values = tf.constant([], shape=[0], dtype=tf.float32)
truth_shape = tf.constant([0, 0, 0], shape=[3], dtype=tf.int64)

tf.raw_ops.EditDistance(
    hypothesis_indices=hypothesis_indices, hypothesis_values=hypothesis_values,
    hypothesis_shape=hypothesis_shape, truth_indices=truth_indices,
    truth_values=truth_values, truth_shape=truth_shape, normalize=True)
```

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.

In the above scenario, an attacker causes an allocation of an empty tensor for the output:

```cc
OP_REQUIRES_OK(ctx, ctx->allocate_output("output", output_shape, &output));
auto output_t = output->flat<float>();
output_t.setZero();
```

Because `output_shape` has 0 elements, the result of `output->flat<T>()` has an empty buffer, so calling `setZero` would result in a null dereference.

### Patches
We have patched the issue in GitHub commit [f4c364a5d6880557f6f5b6eb5cee2c407f0186b3](https://github.com/tensorflow/tensorflow/commit/f4c364a5d6880557f6f5b6eb5cee2c407f0186b3).

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.

### For more information
Please consult [our security guide](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for more information regarding the security model and how to contact us with issues and questions.

### Attribution
This vulnerability has been reported by Yakun Zhang and Ying Wang of Baidu X-Team.
Vulnerability ID: GHSA-75f6-78jr-4656
CWE: CWE-476
Severity score: 2.5
Project: https://github.com/tensorflow/tensorflow
Published: 2021-05-21
File: tensorflow/core/kernels/edit_distance_op.cc
Commit message: Fix multiple issues in EditDistance

PiperOrigin-RevId: 372033948
Change-Id: Ieb957c29894af05bdfeb1a0402fced808dfcfd7b
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 Status ValidateShapes(OpKernelContext* ctx, const Tensor& hypothesis_indices,     return errors::InvalidArgument(         "truth_shape should be a vector, but got shape: ",         truth_shape.shape().DebugString());   if (hypothesis_shape.NumElements() != hypothesis_indices.dim_size(1))     return errors::InvalidArgument(         "Expected hypothesis_shape.NumElements == " Status ValidateShapes(OpKernelContext* ctx, const Tensor& hypothesis_indices,         "Input SparseTensors must have rank at least 2, but truth_shape "         "rank is: ",         truth_shape.NumElements());   if (truth_shape.NumElements() != truth_indices.dim_size(1))     return errors::InvalidArgument(         "Expected truth_shape.NumElements == " class EditDistanceOp : public OpKernel {       output_shape.AddDim(std::max(hypothesis_st_shape.dim_size(d),                                    truth_st_shape.dim_size(d)));     }      Tensor* output = nullptr;     OP_REQUIRES_OK(ctx, ctx->allocate_output("output", output_shape, &output)); class EditDistanceOp : public OpKernel {       if (g_truth == g_hypothesis) {         auto loc = std::inner_product(g_truth.begin(), g_truth.end(),                                       output_strides.begin(), int64{0});         output_t(loc) =             gtl::LevenshteinDistance<T>(truth_seq, hypothesis_seq, cmp);         if (normalize_) output_t(loc) /= truth_seq.size(); class EditDistanceOp : public OpKernel {       } 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});         output_t(loc) = hypothesis_seq.size();         if (normalize_ && output_t(loc) != 0.0f) {           output_t(loc) = std::numeric_limits<float>::infinity(); class EditDistanceOp : public OpKernel {       } else {  // zero-length hypothesis         auto loc = std::inner_product(g_truth.begin(), g_truth.end(),                                       output_strides.begin(), int64{0});         output_t(loc) = (normalize_) ? 1.0 : truth_seq.size();         ++truth_iter;       } class EditDistanceOp : public OpKernel {       auto hypothesis_seq = hypothesis_j.values<T>();       auto loc = std::inner_product(g_hypothesis.begin(), g_hypothesis.end(),                                     output_strides.begin(), int64{0});       output_t(loc) = hypothesis_seq.size();       if (normalize_ && output_t(loc) != 0.0f) {         output_t(loc) = std::numeric_limits<float>::infinity(); class EditDistanceOp : public OpKernel {       auto truth_seq = truth_i.values<T>();       auto loc = std::inner_product(g_truth.begin(), g_truth.end(),                                     output_strides.begin(), int64{0});       output_t(loc) = (normalize_) ? 1.0 : truth_seq.size();       ++truth_iter;     }

Fixed version (after fix):
 Status ValidateShapes(OpKernelContext* ctx, const Tensor& hypothesis_indices,     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 == " Status ValidateShapes(OpKernelContext* ctx, const Tensor& hypothesis_indices,         "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 == " class EditDistanceOp : public OpKernel {       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)); class EditDistanceOp : public OpKernel {       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(); class EditDistanceOp : public OpKernel {       } 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(); class EditDistanceOp : public OpKernel {       } 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;       } class EditDistanceOp : public OpKernel {       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(); class EditDistanceOp : public OpKernel {       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;     }

Patch:
patch: @@ -64,6 +64,12 @@ Status ValidateShapes(OpKernelContext* ctx, const Tensor& hypothesis_indices,
     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 == "
@@ -75,6 +81,12 @@ Status ValidateShapes(OpKernelContext* ctx, const Tensor& hypothesis_indices,
         "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 == "
@@ -153,6 +165,11 @@ class EditDistanceOp : public OpKernel {
       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));
@@ -185,6 +202,12 @@ class EditDistanceOp : public OpKernel {
       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();
@@ -194,6 +217,12 @@ class EditDistanceOp : public OpKernel {
       } 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();
@@ -202,6 +231,12 @@ class EditDistanceOp : public OpKernel {
       } 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;
       }
@@ -212,6 +247,12 @@ class EditDistanceOp : public OpKernel {
       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();
@@ -224,6 +265,12 @@ class EditDistanceOp : public OpKernel {
       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;
     }
path: tensorflow/core/kernels/edit_distance_op.cc