46 lines
2.1 KiB
Diff
46 lines
2.1 KiB
Diff
|
From 3a7362750d5c372420aa8f0caf7bf5b5c3d0f52d Mon Sep 17 00:00:00 2001
|
||
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
||
|
|
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<int>()();
|
||
|
|
+ OP_REQUIRES(context, output_size >= 0,
|
||
|
|
+ errors::InvalidArgument("output size must be non-negative"));
|
||
|
|
|
||
|
|
std::vector<T> scores_data(num_boxes);
|
||
|
|
std::copy_n(scores.flat<T>().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;
|