From 44b7f486c0143f68b56c34e2d01e146ee445134a Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Wed, 21 Apr 2021 16:19:54 -0700
Subject: [PATCH] Fix out of bounds read in `ragged_cross_op.cc`.

PiperOrigin-RevId: 369757702
Change-Id: Ie6e5d2c21513a8d56bf41fcf35960caf76e890f9
---
 tensorflow/core/kernels/ragged_cross_op.cc | 30 ++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/tensorflow/core/kernels/ragged_cross_op.cc b/tensorflow/core/kernels/ragged_cross_op.cc
index ea65c0ee2b5b2..5dfe93f416659 100644
--- a/tensorflow/core/kernels/ragged_cross_op.cc
+++ b/tensorflow/core/kernels/ragged_cross_op.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/fingerprint.h"
 #include "tensorflow/core/util/util.h"
 #include "tensorflow/core/util/work_sharder.h"
@@ -466,16 +467,45 @@ class RaggedCrossOp : public OpKernel {
     int next_dense = 0;
     for (char c : input_order_) {
       if (c == 'R') {
+        if (next_ragged >= ragged_values_list.size())
+          return errors::InvalidArgument(
+              "input_order \"", input_order_,
+              "\" specifies reading a ragged tensor value at index ",
+              next_ragged, " from a list of ", ragged_values_list.size(),
+              " values.");
+        if (next_ragged >= ragged_splits_list.size())
+          return errors::InvalidArgument(
+              "input_order \"", input_order_,
+              "\" specifies reading a ragged tensor split at index ",
+              next_ragged, " from a list of ", ragged_splits_list.size(),
+              " splits.");
         TF_RETURN_IF_ERROR(BuildRaggedFeatureReader(
             ragged_values_list[next_ragged], ragged_splits_list[next_ragged],
             features));
         next_ragged++;
       } else if (c == 'S') {
+        if (next_sparse >= sparse_values_list.size())
+          return errors::InvalidArgument(
+              "input_order \"", input_order_,
+              "\" specifies reading a sparse tensor value at index ",
+              next_sparse, " from a list of ", sparse_values_list.size(),
+              " values.");
+        if (next_sparse >= sparse_indices_list.size())
+          return errors::InvalidArgument(
+              "input_order \"", input_order_,
+              "\" specifies reading a sparse tensor index at index ",
+              next_sparse, " from a list of ", sparse_indices_list.size(),
+              " indices.");
         TF_RETURN_IF_ERROR(BuildSparseFeatureReader(
             sparse_indices_list[next_sparse], sparse_values_list[next_sparse],
             batch_size, features));
         next_sparse++;
       } else if (c == 'D') {
+        if (next_dense >= dense_list.size())
+          return errors::InvalidArgument(
+              "input_order \"", input_order_,
+              "\" specifies reading a dense tensor at index ", next_dense,
+              " from a list of ", dense_list.size(), " tensors.");
         TF_RETURN_IF_ERROR(
             BuildDenseFeatureReader(dense_list[next_dense++], features));
       } else {