Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix heap buffer overflow in tf.raw_ops.SparseFillEmptyRowsGrad.
Also add tests as they were lacking

PiperOrigin-RevId: 332566071
Change-Id: I44277578e26ff5fb3fdb0dcbba6e91b2ec3e7859
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Sep 19, 2020
1 parent a633411 commit 390611e
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 2 deletions.
12 changes: 10 additions & 2 deletions tensorflow/core/kernels/sparse_fill_empty_rows_op.cc
Expand Up @@ -236,6 +236,9 @@ class SparseFillEmptyRowsGradOp : public OpKernel {
context, TensorShapeUtils::IsVector(reverse_index_map_t->shape()),
errors::InvalidArgument("reverse_index_map must be a vector, saw: ",
reverse_index_map_t->shape().DebugString()));
OP_REQUIRES(context, TensorShapeUtils::IsVector(grad_values_t->shape()),
errors::InvalidArgument("grad_values must be a vector, saw: ",
grad_values_t->shape().DebugString()));

const auto reverse_index_map = reverse_index_map_t->vec<int64>();
const auto grad_values = grad_values_t->vec<T>();
Expand Down Expand Up @@ -264,8 +267,13 @@ class SparseFillEmptyRowsGradOp : public OpKernel {
// Locate the index of the output of the forward prop associated
// with this location in the input of the forward prop. Copy
// the gradient into it. Mark it as visited.
d_values(i) = grad_values(reverse_index_map(i));
visited(reverse_index_map(i)) = true;
int64 reverse_index = reverse_index_map(i);
OP_REQUIRES(
context, 0 <= reverse_index && reverse_index < N_full,
errors::InvalidArgument("Elements in reverse index must be in [0, ",
N_full, ") but got ", reverse_index));
d_values(i) = grad_values(reverse_index);
visited(reverse_index) = true;
}
for (int j = 0; j < N_full; ++j) {
// The default value gradient gets the accumulated remainder of
Expand Down
54 changes: 54 additions & 0 deletions tensorflow/python/ops/sparse_ops_test.py
Expand Up @@ -21,6 +21,7 @@
from absl.testing import parameterized
import numpy as np

from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
Expand All @@ -30,6 +31,7 @@
# Need array_grad to register gradient for Identity.
from tensorflow.python.ops import array_grad # pylint: disable=unused-import
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_sparse_ops
from tensorflow.python.ops import gradient_checker_v2 as gradient_checker
from tensorflow.python.ops import math_ops
# Need sparse_grad to register gradient for SparseToDense.
Expand Down Expand Up @@ -234,5 +236,57 @@ def testConstantStringToSparse(self):
self.assertAllEqual([5], result.dense_shape)


@test_util.run_all_in_graph_and_eager_modes
class RawOpsTest(test_util.TensorFlowTestCase, parameterized.TestCase):

def testSparseFillEmptyRowsGrad(self):
reverse_index_map = [2, 1]
grad_values = [0, 1, 2, 3]
d_values, d_default_value = self.evaluate(
gen_sparse_ops.SparseFillEmptyRowsGrad(
reverse_index_map=reverse_index_map, grad_values=grad_values))
self.assertAllEqual([2, 1], d_values)
self.assertEqual(3, d_default_value)

def testSparseFillEmptyRowsGradNegativeIndexMapValue(self):
reverse_index_map = [2, -1]
grad_values = [0, 1, 2, 3]
with self.assertRaisesRegex(
errors.InvalidArgumentError,
r'Elements in reverse index must be in \[0, 4\)'):
self.evaluate(
gen_sparse_ops.SparseFillEmptyRowsGrad(
reverse_index_map=reverse_index_map, grad_values=grad_values))

def testSparseFillEmptyRowsGradLargeIndexMapValue(self):
reverse_index_map = [2, 10]
grad_values = [0, 1, 2, 3]
with self.assertRaisesRegex(
errors.InvalidArgumentError,
r'Elements in reverse index must be in \[0, 4\)'):
self.evaluate(
gen_sparse_ops.SparseFillEmptyRowsGrad(
reverse_index_map=reverse_index_map, grad_values=grad_values))

def testSparseFillEmptyRowsGradMatrix(self):
reverse_index_map = [0, 1]
grad_values = [[0, 1], [2, 3]]
# Note: Eager mode and graph mode throw different errors here. Graph mode
# will fail with a ValueError from the shape checking logic, while Eager
# will fail with an InvalidArgumentError from the kernel itself.
if context.executing_eagerly():
with self.assertRaisesRegex(errors.InvalidArgumentError,
r'grad_values must be a vector'):
self.evaluate(
gen_sparse_ops.SparseFillEmptyRowsGrad(
reverse_index_map=reverse_index_map, grad_values=grad_values))
else:
with self.assertRaisesRegex(ValueError,
r'Shape must be rank 1 but is rank 2'):
self.evaluate(
gen_sparse_ops.SparseFillEmptyRowsGrad(
reverse_index_map=reverse_index_map, grad_values=grad_values))


if __name__ == '__main__':
googletest.main()

0 comments on commit 390611e

Please sign in to comment.