1## TFSA-2021-009: Segfault in `SparseCountSparseOutput` 2 3### CVE Number 4CVE-2021-29521 5 6### Impact 7Specifying a negative dense shape in `tf.raw_ops.SparseCountSparseOutput` 8results in a segmentation fault being thrown out from the standard library as 9`std::vector` invariants are broken. 10 11```python 12import tensorflow as tf 13 14indices = tf.constant([], shape=[0, 0], dtype=tf.int64) 15values = tf.constant([], shape=[0, 0], dtype=tf.int64) 16dense_shape = tf.constant([-100, -100, -100], shape=[3], dtype=tf.int64) 17weights = tf.constant([], shape=[0, 0], dtype=tf.int64) 18 19tf.raw_ops.SparseCountSparseOutput(indices=indices, values=values, dense_shape=dense_shape, weights=weights, minlength=79, maxlength=96, binary_output=False) 20``` 21 22This is because the 23[implementation](https://github.com/tensorflow/tensorflow/blob/8f7b60ee8c0206a2c99802e3a4d1bb55d2bc0624/tensorflow/core/kernels/count_ops.cc#L199-L213) 24assumes the first element of the dense shape is always positive and uses it to 25initialize a `BatchedMap<T>` (i.e., 26[`std::vector<absl::flat_hash_map<int64,T>>`](https://github.com/tensorflow/tensorflow/blob/8f7b60ee8c0206a2c99802e3a4d1bb55d2bc0624/tensorflow/core/kernels/count_ops.cc#L27)) 27data structure. 28 29```cc 30 bool is_1d = shape.NumElements() == 1; 31 int num_batches = is_1d ? 1 : shape.flat<int64>()(0); 32 ... 33 auto per_batch_counts = BatchedMap<W>(num_batches); 34``` 35 36If the `shape` tensor has more than one element, `num_batches` is the first 37value in `shape`. 38 39Ensuring that the `dense_shape` argument is a valid tensor shape (that is, all 40elements are non-negative) solves this issue. 41 42### Patches 43We have patched the issue in GitHub commit 44[c57c0b9f3a4f8684f3489dd9a9ec627ad8b599f5](https://github.com/tensorflow/tensorflow/commit/c57c0b9f3a4f8684f3489dd9a9ec627ad8b599f5). 45 46The fix will be included in TensorFlow 2.5.0. We will also cherrypick this 47commit on TensorFlow 2.4.2 and TensorFlow 2.3.3. 48 49### For more information 50Please consult [our security 51guide](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for 52more information regarding the security model and how to contact us with issues 53and questions. 54 55### Attribution 56This vulnerability has been reported by Yakun Zhang and Ying Wang of Baidu 57X-Team. 58