CVE-2020-15200 (Python, published 2020-09-25T19:15Z)
Description: value: In Tensorflow before version 2.3.1, the `RaggedCountSparseOutput` implementation does not validate that the input arguments form a valid ragged tensor. In particular, there is no validation that the values in the `splits` tensor generate a valid partitioning of the `values` tensor. Thus, the code sets up conditions to cause a heap buffer overflow. A `BatchedMap` is equivalent to a vector where each element is a hashmap. However, if the first element of `splits_values` is not 0, `batch_idx` will never be 1, hence there will be no hashmap at index 0 in `per_batch_counts`. Trying to access that in the user code results in a segmentation fault. The issue is patched in commit 3cbb917b4714766030b28eba9fb41bb97ce9ee02 and is released in TensorFlow version 2.3.1.
Severity: MEDIUM (CVSS v3 base score: 5.9)
Weakness: CWE-787 \u2014 Out-of-bounds Write
The product writes data past the end, or before the beginning, of the intended buffer.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/kernels/count_ops.cc
tensorflow/python/ops/bincount_ops_test.py
Commit message: Fix multiple vulnerabilities in `tf.raw_ops.*CountSparseOutput`.

Also add tests for these API points, both for the happy paths and for the vulnerable ones.

PiperOrigin-RevId: 332563222
Change-Id: Ib3b52116a83a134c2e742a7c66e5e956db8fba05


Vulnerable code:

Fixed code:
    OP_REQUIRES(context, TensorShapeUtils::IsMatrix(indices.shape()),
                errors::InvalidArgument(
                    "Input indices must be a 2-dimensional tensor. Got: ",
                    indices.shape().DebugString()));

    if (use_weights) {
      OP_REQUIRES(
          context, weights.shape() == values.shape(),
          errors::InvalidArgument(
              "Weights and values must have the same shape. Weight shape: ",
              weights.shape().DebugString(),
              "; values shape: ", values.shape().DebugString()));
    }

    OP_REQUIRES(context, num_values == indices.shape().dim_size(0),
                errors::InvalidArgument(
                    "Number of values must match first dimension of indices.",
                    "Got ", num_values,
                    " values, indices shape: ", indices.shape().DebugString()));

    if (use_weights) {
      OP_REQUIRES(
          context, weights.shape() == values.shape(),
          errors::InvalidArgument(
              "Weights and values must have the same shape. Weight shape: ",
              weights.shape().DebugString(),
              "; values shape: ", values.shape().DebugString()));
    }

    OP_REQUIRES(
        context, num_batches > 0,
        errors::InvalidArgument(
            "Must provide at least 2 elements for the splits argument"));
    OP_REQUIRES(context, splits_values(0) == 0,
                errors::InvalidArgument("Splits must start with 0, not with ",
                                        splits_values(0)));
    OP_REQUIRES(context, splits_values(num_batches) == num_values,
                errors::InvalidArgument(
                    "Splits must end with the number of values, got ",
                    splits_values(num_batches), " instead of ", num_values));

from tensorflow.python.framework import test_util
from tensorflow.python.ops import gen_count_ops
@test_util.run_all_in_graph_and_eager_modes
@test_util.disable_tfrt
class RawOpsTest(test.TestCase, parameterized.TestCase):

  def testSparseCountSparseOutputBadIndicesShape(self):
    indices = [[[0], [0]], [[0], [1]], [[1], [0]], [[1], [2]]]
    values = [1, 1, 1, 10]
    weights = [1, 2, 4, 6]
    dense_shape = [2, 3]
    with self.assertRaisesRegex(errors.InvalidArgumentError,
                                "Input indices must be a 2-dimensional tensor"):
      self.evaluate(
          gen_count_ops.SparseCountSparseOutput(
              indices=indices,
              values=values,
              dense_shape=dense_shape,
              weights=weights,
              binary_output=False))

  def testSparseCountSparseOutputBadWeightsShape(self):
    indices = [[0, 0], [0, 1], [1, 0], [1, 2]]
    values = [1, 1, 1, 10]
    weights = [1, 2, 4]
    dense_shape = [2, 3]
    with self.assertRaisesRegex(errors.InvalidArgumentError,
                                "Weights and values must have the same shape"):
      self.evaluate(
          gen_count_ops.SparseCountSparseOutput(
              indices=indices,
              values=values,
              dense_shape=dense_shape,
              weights=weights,
              binary_output=False))

  def testSparseCountSparseOutputBadNumberOfValues(self):
    indices = [[0, 0], [0, 1], [1, 0]]
    values = [1, 1, 1, 10]
    weights = [1, 2, 4, 6]
    dense_shape = [2, 3]
    with self.assertRaisesRegex(
        errors.InvalidArgumentError,
        "Number of values must match first dimension of indices"):
      self.evaluate(
          gen_count_ops.SparseCountSparseOutput(
              indices=indices,
              values=values,
              dense_shape=dense_shape,
              weights=weights,
              binary_output=False))

  def testRaggedCountSparseOutput(self):
    splits = [0, 4, 7]
    values = [1, 1, 2, 1, 2, 10, 5]
    weights = [1, 2, 3, 4, 5, 6, 7]
    output_indices, output_values, output_shape = self.evaluate(
        gen_count_ops.RaggedCountSparseOutput(
            splits=splits, values=values, weights=weights, binary_output=False))
    self.assertAllEqual([[0, 1], [0, 2], [1, 2], [1, 5], [1, 10]],
                        output_indices)
    self.assertAllEqual([7, 3, 5, 7, 6], output_values)
    self.assertAllEqual([2, 11], output_shape)

  def testRaggedCountSparseOutputBadWeightsShape(self):
    splits = [0, 4, 7]
    values = [1, 1, 2, 1, 2, 10, 5]
    weights = [1, 2, 3, 4, 5, 6]
    with self.assertRaisesRegex(errors.InvalidArgumentError,
                                "Weights and values must have the same shape"):
      self.evaluate(
          gen_count_ops.RaggedCountSparseOutput(
              splits=splits,
              values=values,
              weights=weights,
              binary_output=False))

  def testRaggedCountSparseOutputEmptySplits(self):
    splits = []
    values = [1, 1, 2, 1, 2, 10, 5]
    weights = [1, 2, 3, 4, 5, 6, 7]
    with self.assertRaisesRegex(
        errors.InvalidArgumentError,
        "Must provide at least 2 elements for the splits argument"):
      self.evaluate(
          gen_count_ops.RaggedCountSparseOutput(
              splits=splits,
              values=values,
              weights=weights,
              binary_output=False))

  def testRaggedCountSparseOutputBadSplitsStart(self):
    splits = [1, 7]
    values = [1, 1, 2, 1, 2, 10, 5]
    weights = [1, 2, 3, 4, 5, 6, 7]
    with self.assertRaisesRegex(errors.InvalidArgumentError,
                                "Splits must start with 0"):
      self.evaluate(
          gen_count_ops.RaggedCountSparseOutput(
              splits=splits,
              values=values,
              weights=weights,
              binary_output=False))

  def testRaggedCountSparseOutputBadSplitsEnd(self):
    splits = [0, 5]
    values = [1, 1, 2, 1, 2, 10, 5]
    weights = [1, 2, 3, 4, 5, 6, 7]
    with self.assertRaisesRegex(errors.InvalidArgumentError,
                                "Splits must end with the number of values"):
      self.evaluate(
          gen_count_ops.RaggedCountSparseOutput(
              splits=splits,
              values=values,
              weights=weights,
              binary_output=False))