CVE-2021-29514

TensorFlow RaggedBincount heap buffer overflow via splits argument

• CVSS 7.8
• CWE-787 Out-of-bounds Write
• Memory Corruption
• Local

TensorFlow is an end-to-end open source platform for machine learning. If the `splits` argument of `RaggedBincount` does not specify a valid `SparseTensor`(https://www.tensorflow.org/api_docs/python/tf/sparse/SparseTensor), then an attacker can trigger a heap buffer overflow. This will cause a read from outside the bounds of the `splits` tensor buffer in the implementation of the `RaggedBincount` op(https://github.com/tensorflow/tensorflow/blob/8b677d79167799f71c42fd3fa074476e0295413a/tensorflow/core/kernels/bincount_op.cc#L430-L446). Before the `for` loop, `batch_idx` is set to 0. The attacker sets `splits(0)` to be 7, hence the `while` loop does not execute and `batch_idx` remains 0. This then results in writing to `out(-1, bin)`, which is before the heap allocated buffer for the output tensor. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2 and TensorFlow 2.3.3, as these are also affected.

CVE

CVE-2021-29514

CWE

CWE-787 Out-of-bounds Write

CVSS base score

7.8

Published

2021-05-14

OWASP

A08 Software and Data Integrity Failures

Orthogonal defect classification

Checking

Code defect classification

Incorrect Check

Category

Memory Corruption

Subcategory

Buffer Overflows

Accessibility scope

Local

Impact

Denial of Service (DoS)

Affected component

TensorFlow

Fixed by upgrading

Yes

Solution

Update TensorFlow to version 2.5.0 or higher.

import tensorflow as tf

def ragged_bincount(splits, values, num_bins):
    out = tf.zeros([num_bins], dtype=tf.int32)

    for i in range(tf.shape(splits)[0]):
        if splits[i] >= num_bins:
            print(f"Access with index {splits[i]} at position {i}")
        else:
            out[splits[i]] += 1

    return out

splits = tf.constant([7])
values = tf.constant([0, 1, 2])
num_bins = 3

output = ragged_bincount(splits, values, num_bins)
print(output)

import tensorflow as tf

def ragged_bincount(splits, values, num_bins):
    if tf.reduce_any(splits >= num_bins) or tf.reduce_any(splits < 0):
        raise ValueError("Splits values are out of bounds")

    out = tf.zeros([num_bins], dtype=tf.int32)

    for i in range(tf.shape(splits)[0]):
        out = tf.tensor_scatter_nd_add(out, [[splits[i]]], [1])

    return out

splits = tf.constant([7])
values = tf.constant([0, 1, 2])
num_bins = 3

try:
    output = ragged_bincount(splits, values, num_bins)
    print(output)
except ValueError as e:
    print(e)

@misc{vaitp:cve202129514,
  title        = {{TensorFlow RaggedBincount heap buffer overflow via splits argument}},
  author       = {Bogaerts, Fr\'ed\'eric and Ivaki, Naghmeh and Fonseca, Jos\'e},
  year         = {2021},
  note         = {VAITP Python Vulnerability Dataset, entry CVE-2021-29514},
  howpublished = {\url{https://netpack.pt/vaitp/vulnerability/CVE-2021-29514/}}
}