From 3a7362750d5c372420aa8f0caf7bf5b5c3d0f52d Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Fri, 30 Jul 2021 22:02:22 -0700 Subject: [PATCH] Prevent crash/heap OOB due to integer conversion to unsigned in NMS kernels PiperOrigin-RevId: 387938262 Change-Id: Id361a715307e7179977cf5c64391c199a966f2ad --- tensorflow/core/kernels/non_max_suppression_op.cc | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tensorflow/core/kernels/non_max_suppression_op.cc b/tensorflow/core/kernels/non_max_suppression_op.cc index 69b05cc9d84f8..1ec4c853f5f5b 100644 --- a/tensorflow/core/kernels/non_max_suppression_op.cc +++ b/tensorflow/core/kernels/non_max_suppression_op.cc @@ -161,6 +161,8 @@ void DoNonMaxSuppressionOp(OpKernelContext* context, const Tensor& scores, bool pad_to_max_output_size = false, int* ptr_num_valid_outputs = nullptr) { const int output_size = max_output_size.scalar()(); + OP_REQUIRES(context, output_size >= 0, + errors::InvalidArgument("output size must be non-negative")); std::vector scores_data(num_boxes); std::copy_n(scores.flat().data(), num_boxes, scores_data.begin()); @@ -759,6 +761,9 @@ class NonMaxSuppressionV4Op : public OpKernel { context, scores, num_boxes, max_output_size, iou_threshold_val, score_threshold_val, dummy_soft_nms_sigma, similarity_fn, return_scores_tensor_, pad_to_max_output_size_, &num_valid_outputs); + if (!context->status().ok()) { + return; + } // Allocate scalar output tensor for number of indices computed. Tensor* num_outputs_t = nullptr; @@ -836,6 +841,9 @@ class NonMaxSuppressionV5Op : public OpKernel { context, scores, num_boxes, max_output_size, iou_threshold_val, score_threshold_val, soft_nms_sigma_val, similarity_fn, return_scores_tensor_, pad_to_max_output_size_, &num_valid_outputs); + if (!context->status().ok()) { + return; + } // Allocate scalar output tensor for number of indices computed. Tensor* num_outputs_t = nullptr;