CVE-2023-25664: Heap-based Buffer Overflow (CWE-122)
value: TensorFlow is an open source platform for machine learning. Prior to versions 2.12.0 and 2.11.1, there is a heap buffer overflow in TAvgPoolGrad. A fix is included in TensorFlow 2.12.0 and 2.11.1.

CWE description: A heap overflow condition is a buffer overflow, where the buffer that can be overwritten is allocated in the heap portion of memory, generally meaning that the buffer was allocated using a routine such as malloc().
CVSS v3 base score: 9.8
Published: 2023-03-25T00:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/avgpooling_op.cc, tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
Commit date: 2022-11-16 10:27:28 -0800
Commit message: Add inputs check for AvgPoolGrad

PiperOrigin-RevId: 488975844

Diff: diff --git a/tensorflow/core/kernels/avgpooling_op.cc b/tensorflow/core/kernels/avgpooling_op.cc
index bc536d187512a..4ee4b8ff32474 100644
--- a/tensorflow/core/kernels/avgpooling_op.cc
+++ b/tensorflow/core/kernels/avgpooling_op.cc
@@ -340,10 +340,23 @@ class AvgPoolingGradOp : public OpKernel {
                                          padding_, &out_width, &pad_cols));
 
     const T* out_backprop_ptr = out_backprop.flat<T>().data();
     T* input_backprop_ptr = output->flat<T>().data();
 
+    for (int64_t r = 0; r < out_backprop_rows; ++r) {
+      int rindex, rsize;
+      OP_REQUIRES_OK(context,
+                     GetBroadcastSize(r, in_rows, window_rows, row_stride,
+                                      pad_rows, &rindex, &rsize));
+      for (int64_t c = 0; c < out_backprop_cols; ++c) {
+        int cindex, csize;
+        OP_REQUIRES_OK(context,
+                       GetBroadcastSize(c, in_cols, window_cols, col_stride,
+                                        pad_cols, &cindex, &csize));
+      }
+    }
+
     auto shard = [context, out_backprop_ptr, input_backprop_ptr,
                   out_backprop_rows, out_backprop_cols, out_backprop_depth,
                   in_rows, in_cols, window_rows, window_cols, row_stride,
                   col_stride, pad_rows,
                   pad_cols](int64_t start, int64_t limit) {
diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
index e279f6e0027da..3c6f974d78451 100644
--- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
@@ -2508,10 +2508,25 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
             strides=[1, 2, 2, 1],
             padding="VALID",
             data_format="NHWC")
         self.evaluate(t)
 
+  def testAvgPoolGradInvalidStrideRaiseErrorProperly(self):
+    with self.assertRaises(errors_impl.InvalidArgumentError):
+      with self.cached_session():
+        orig_input_shape = [11, 9, 78, 9]
+        grad = constant_op.constant(
+            0.1, shape=[16, 16, 16, 16], dtype=dtypes.float64)
+        t = gen_nn_ops.AvgPoolGrad(
+            orig_input_shape=orig_input_shape,
+            grad=grad,
+            ksize=[1, 40, 128, 1],
+            strides=[1, 128, 128, 30],
+            padding="SAME",
+            data_format="NHWC")
+        self.evaluate(t)
+
 
 def GetMaxPoolFwdTest(input_size, filter_size, strides, padding):
 
   def Test(self):
     # MaxPoolWithArgMax is implemented only on CUDA.

Fixed code:     for (int64_t r = 0; r < out_backprop_rows; ++r) {
      int rindex, rsize;
      OP_REQUIRES_OK(context,
                     GetBroadcastSize(r, in_rows, window_rows, row_stride,
                                      pad_rows, &rindex, &rsize));
      for (int64_t c = 0; c < out_backprop_cols; ++c) {
        int cindex, csize;
        OP_REQUIRES_OK(context,
                       GetBroadcastSize(c, in_cols, window_cols, col_stride,
                                        pad_cols, &cindex, &csize));
      }
    }

  def testAvgPoolGradInvalidStrideRaiseErrorProperly(self):
    with self.assertRaises(errors_impl.InvalidArgumentError):
      with self.cached_session():
        orig_input_shape = [11, 9, 78, 9]
        grad = constant_op.constant(
            0.1, shape=[16, 16, 16, 16], dtype=dtypes.float64)
        t = gen_nn_ops.AvgPoolGrad(
            orig_input_shape=orig_input_shape,
            grad=grad,
            ksize=[1, 40, 128, 1],
            strides=[1, 128, 128, 30],
            padding="SAME",
            data_format="NHWC")
        self.evaluate(t)