tensorflow/CVE-2021-29564.patch

From f4c364a5d6880557f6f5b6eb5cee2c407f0186b3 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Tue, 4 May 2021 18:06:03 -0700
Subject: [PATCH] Fix multiple issues in EditDistance

PiperOrigin-RevId: 372033948
Change-Id: Ieb957c29894af05bdfeb1a0402fced808dfcfd7b
---
 tensorflow/core/kernels/edit_distance_op.cc | 47 +++++++++++++++++++++
 1 file changed, 47 insertions(+)

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
@@ -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;
     }