Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix heap out of bounds error in tf.raw_ops.SparseCountSparseOutput sh…
…ape inference when it is called with invalid inputs, and add a test for it.

PiperOrigin-RevId: 405766415
Change-Id: I77d244ef35f351ef7b6f821efd959cac2c66db24
  • Loading branch information
broken authored and tensorflower-gardener committed Oct 26, 2021
1 parent 3e77733 commit 701cfac
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 0 deletions.
2 changes: 2 additions & 0 deletions tensorflow/core/ops/count_ops.cc
Expand Up @@ -41,6 +41,8 @@ Status DenseCountSparseOutputShapeFn(InferenceContext *c) {
}

Status SparseCountSparseOutputShapeFn(InferenceContext *c) {
ShapeHandle unused;
TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &unused));
auto rank = c->Dim(c->input(0), 1);
auto nvals = c->UnknownDim();
c->set_output(0, c->Matrix(nvals, rank)); // out.indices
Expand Down
19 changes: 19 additions & 0 deletions tensorflow/python/ops/bincount_ops_test.py
Expand Up @@ -831,6 +831,25 @@ def test_ragged_input_different_shape_fails(self):
self.evaluate(bincount_ops.sparse_bincount(x, weights=weights, axis=-1))


class RawOpsHeapOobTest(test.TestCase, parameterized.TestCase):

@test_util.run_v1_only("Test security error")
def testSparseCountSparseOutputBadIndicesShapeTooSmall(self):
indices = [1]
values = [[1]]
weights = []
dense_shape = [10]
with self.assertRaisesRegex(ValueError,
"Shape must be rank 2 but is rank 1 for"):
self.evaluate(
gen_count_ops.SparseCountSparseOutput(
indices=indices,
values=values,
dense_shape=dense_shape,
weights=weights,
binary_output=True))


@test_util.run_all_in_graph_and_eager_modes
@test_util.disable_tfrt
class RawOpsTest(test.TestCase, parameterized.TestCase):
Expand Down

0 comments on commit 701cfac

Please sign in to comment.