From 87158f43f05f2720a374f3e6d22a7aaa3a33f750 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac 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(reduction.reorder_dims); for (const auto &g : sp.group(reduction.group_by_dims)) { Op::template Run(ctx, reduced_val, g.template values()); + 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