Heap OOB read in `tf.raw_ops.SparseCountSparseOutput`

@@ -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

@@ -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):