## TFSA-2020-002: Out of bounds write in TFLite implementation of segment sum

### CVE Number
CVE-2020-15214

### Impact
In TensorFlow Lite models using segment sum can trigger a write out bounds /
segmentation fault if the segment ids are not sorted. Code assumes that the
segment ids are in increasing order, [using the last element of the tensor
holding them to determine the dimensionality of output
tensor](https://github.com/tensorflow/tensorflow/blob/0e68f4d3295eb0281a517c3662f6698992b7b2cf/tensorflow/lite/kernels/segment_sum.cc#L39-L44):
```cc
  if (segment_id_size > 0) {
    max_index = segment_ids->data.i32[segment_id_size - 1];
  }
  TfLiteIntArray* output_shape = TfLiteIntArrayCreate(NumDimensions(data));
  output_shape->data[0] = max_index + 1;
```

This results in allocating insufficient memory for the output tensor and in a
[write outside the bounds of the output
array](https://github.com/tensorflow/tensorflow/blob/0e68f4d3295eb0281a517c3662f6698992b7b2cf/tensorflow/lite/kernels/internal/reference/reference_ops.h#L2625-L2631):
```cc
  memset(output_data, 0, sizeof(T) * output_shape.FlatSize());
  for (int i = 0; i < input_shape.Dims(0); i++) {
    int output_index = segment_ids_data[i];
    for (int j = 0; j < segment_flat_size; ++j) {
      output_data[output_index * segment_flat_size + j] +=
          input_data[i * segment_flat_size + j];
    }
  }
```

This usually results in a segmentation fault, but depending on runtime
conditions it can provide for a write gadget to be used in future memory
corruption-based exploits.

### Vulnerable Versions
TensorFlow 2.2.0, 2.3.0.

### Patches
We have patched the issue in
[204945b](https://github.com/tensorflow/tensorflow/commit/204945b) and will
release patch releases for all affected versions.

We recommend users to upgrade to TensorFlow 2.2.1, or 2.3.1.

### Workarounds
A potential workaround would be to add a custom `Verifier` to the model loading
code to ensure that the segment ids are sorted, although this only handles the
case when the segment ids are stored statically in the model.

A similar validation could be done if the segment ids are generated at runtime
between inference steps.

If the segment ids are generated as outputs of a tensor during inference steps,
then there are no possible workaround and users are advised to upgrade to
patched code.

### For more information
Please consult [our security
guide](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for
more information regarding the security model and how to contact us with issues
and questions.

### Attribution
This vulnerability has been reported by members of the Aivul Team from Qihoo
360.