From 87158f43f05f2720a374f3e6d22a7aaa3a33f750 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Fri, 30 Jul 2021 21:11:18 -0700
Subject: [PATCH] Prevent heap OOB in sparse reduction ops.

PiperOrigin-RevId: 387934524
Change-Id: I894aa30f1e454f09b471d565b4a325da49322c1a
---
 tensorflow/core/kernels/sparse_reduce_op.cc | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/tensorflow/core/kernels/sparse_reduce_op.cc b/tensorflow/core/kernels/sparse_reduce_op.cc
index b65f31e5..2bfa3299 100644
--- a/tensorflow/core/kernels/sparse_reduce_op.cc
+++ b/tensorflow/core/kernels/sparse_reduce_op.cc
@@ -219,7 +219,20 @@ class SparseReduceOp : public OpKernel {
     sp.Reorder<T>(reduction.reorder_dims);
     for (const auto &g : sp.group(reduction.group_by_dims)) {
       Op::template Run<T>(ctx, reduced_val, g.template values<T>());
+      OP_REQUIRES(ctx,
+		  output_strides.empty() ||
+		  (g.group().size() == output_strides.size()),
+		  errors::Internal(
+		      "Expected group size and output_strides size to match",
+		      ", but got ", g.group().size(), " and ",
+		      output_strides.size()));
       const int64 idx = CoordinatesToFlatIndex(g.group(), output_strides);
+      OP_REQUIRES(ctx,
+		  idx >= 0 && idx < out_flat.size(),
+		  errors::Internal(
+		      "Obtained a write index of ", idx,
+		      " which is outside of bounds of [0, ",
+		      out_flat.size(), ")"));
       out_flat(idx) = reduced_val();
       VLOG(2) << "coords: " << absl::StrJoin(g.group(), ",")
               << "; idx: " << idx << "; group " << Op::Name() << ": "
-- 
2.27.0