tensorflow/CVE-2021-37651.patch
2021-08-31 15:42:43 +08:00

39 lines
1.5 KiB
Diff

From 0f931751fb20f565c4e94aa6df58d54a003cdb30 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Mon, 2 Aug 2021 13:03:44 -0700
Subject: [PATCH] Validate dimensions of input tensor in
`FractionalAvgPoolGrad`
PiperOrigin-RevId: 388286227
Change-Id: Ieb7566155e92acc8993a2212c76deacadc0edc8a
---
tensorflow/core/kernels/fractional_avg_pool_op.cc | 12 ++++++++++++
1 file changed, 12 insertions(+)
diff --git a/tensorflow/core/kernels/fractional_avg_pool_op.cc b/tensorflow/core/kernels/fractional_avg_pool_op.cc
index 0452638a..a338dd95 100644
--- a/tensorflow/core/kernels/fractional_avg_pool_op.cc
+++ b/tensorflow/core/kernels/fractional_avg_pool_op.cc
@@ -271,6 +271,18 @@ class FractionalAvgPoolGradOp : public OpKernel {
const int64 in_rows = orig_input_tensor_shape_flat(1);
const int64 in_cols = orig_input_tensor_shape_flat(2);
const int64 in_depth = orig_input_tensor_shape_flat(3);
+ OP_REQUIRES(
+ context, in_batch != 0,
+ errors::InvalidArgument("Batch dimension of input must not be 0"));
+ OP_REQUIRES(
+ context, in_rows != 0,
+ errors::InvalidArgument("Rows dimension of input must not be 0"));
+ OP_REQUIRES(
+ context, in_cols != 0,
+ errors::InvalidArgument("Columns dimension of input must not be 0"));
+ OP_REQUIRES(
+ context, in_depth != 0,
+ errors::InvalidArgument("Depth dimension of input must not be 0"));
constexpr int tensor_in_and_out_dims = 4;
// Transform orig_input_tensor_shape into TensorShape
--
2.27.0