40 lines
1.6 KiB
Diff
40 lines
1.6 KiB
Diff
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
|
|
|