Skip to content
Permalink
Browse files Browse the repository at this point in the history
Add negative bound check for row and column pooling_sequence in Fract…
…ionalAvgPoolGrad op to avoid out of bound heap access

PiperOrigin-RevId: 413837346
Change-Id: I2b86034101df31bee161abcb781755e236c7bccd
  • Loading branch information
ishark authored and tensorflower-gardener committed Dec 3, 2021
1 parent e29199a commit 002408c
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 0 deletions.
11 changes: 11 additions & 0 deletions tensorflow/core/kernels/fractional_avg_pool_op.cc
Expand Up @@ -311,15 +311,26 @@ class FractionalAvgPoolGradOp : public OpKernel {
for (int64_t b = 0; b < out_batch; ++b) {
for (int64_t r = 0; r < out_rows; ++r) {
const int64_t in_row_start = row_seq_tensor_flat(r);

int64_t in_row_end = overlapping_ ? row_seq_tensor_flat(r + 1)
: row_seq_tensor_flat(r + 1) - 1;
in_row_end = std::min(in_row_end, in_max_row_index);
OP_REQUIRES(context, in_row_start >= 0 && in_row_end >= 0,
errors::InvalidArgument(
"Row sequence tensor values must not be negative, got ",
row_seq_tensor_flat));

for (int64_t c = 0; c < out_cols; ++c) {
const int64_t in_col_start = col_seq_tensor_flat(c);
int64_t in_col_end = overlapping_ ? col_seq_tensor_flat(c + 1)
: col_seq_tensor_flat(c + 1) - 1;
in_col_end = std::min(in_col_end, in_max_col_index);

OP_REQUIRES(
context, in_col_start >= 0 && in_col_end >= 0,
errors::InvalidArgument(
"Column sequence tensor values must not be negative, got ",
col_seq_tensor_flat));
const int64_t num_elements_in_pooling_cell =
(in_row_end - in_row_start + 1) * (in_col_end - in_col_start + 1);
const int64_t out_index = (b * out_rows + r) * out_cols + c;
Expand Down
Expand Up @@ -20,6 +20,7 @@

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_nn_ops
Expand Down Expand Up @@ -306,6 +307,32 @@ def testDifferentInputTensorShape(self):
input_b, row_seq, col_seq, overlapping)
self.assertSequenceEqual(expected.shape, actual.shape)

def testNegativeSeqValuesForGradOp(self):
with self.assertRaisesRegex(
errors.InvalidArgumentError,
r"Row sequence tensor values must not be negative.*"):
y = nn_ops.gen_nn_ops.fractional_avg_pool_grad(
orig_input_tensor_shape=[2, 2, 2, 2],
out_backprop=[[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11,
12]]]],
row_pooling_sequence=[-10, 1, 2, 3],
col_pooling_sequence=[1, 2, 3, 4],
overlapping=True)

self.evaluate(y)
with self.assertRaisesRegex(
errors.InvalidArgumentError,
r"Column sequence tensor values must not be negative.*"):
z = nn_ops.gen_nn_ops.fractional_avg_pool_grad(
orig_input_tensor_shape=[2, 2, 2, 2],
out_backprop=[[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11,
12]]]],
row_pooling_sequence=[10, 1, 2, 3],
col_pooling_sequence=[1, 2, -3, 4],
overlapping=True)

self.evaluate(z)


class FractionalAvgPoolGradTest(test.TestCase):
"""Tests for FractionalAvgPoolGrad.
Expand Down

0 comments on commit 002408c

Please sign in to comment.