CVE-2023-25669: Incorrect Comparison (CWE-697)
value: TensorFlow is an open source platform for machine learning. Prior to versions 2.12.0 and 2.11.1, if the stride and window size are not positive for `tf.raw_ops.AvgPoolGrad`, it can give a floating point exception. A fix is included in TensorFlow version 2.12.0 and version 2.11.1.

CWE description: The product compares two entities in a security-relevant context, but the comparison is incorrect, which may lead to resultant weaknesses.
CVSS v3 base score: 7.5
Published: 2023-03-25T00:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/compiler/tests/pooling_ops_test.py, tensorflow/compiler/tf2xla/kernels/pooling_ops.cc, tensorflow/compiler/xla/client/padding.cc
Commit date: 2023-01-26 09:49:53 -0800
Commit message: [tf2xla] Validate that stride and window size are positive

PiperOrigin-RevId: 504866231

Diff: diff --git a/tensorflow/compiler/tests/pooling_ops_test.py b/tensorflow/compiler/tests/pooling_ops_test.py
index 3d2695b15e9b0..3a7e22c02e54c 100644
--- a/tensorflow/compiler/tests/pooling_ops_test.py
+++ b/tensorflow/compiler/tests/pooling_ops_test.py
@@ -16,11 +16,13 @@
 
 import numpy as np
 
 from tensorflow.compiler.tests import xla_test
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gen_nn_ops
 from tensorflow.python.ops import nn_ops
 from tensorflow.python.platform import googletest
 
@@ -558,10 +560,38 @@ class PoolGradTest(xla_test.XLATestCase):
           padding=padding,
           data_format=data_format)
 
     self._TestPooling(nn_ops.avg_pool, AvgPoolGrad)
 
+  @test_util.disable_mlir_bridge(
+      "TODO(b/266613412): investigate FPE in AvgPoolGrad for TPU"
+  )
+  def testAvgPoolGradSamePaddingZeroStrideZeroSize(self):
+    output_gradient_vals = np.array([0.39117979], dtype=np.float32)
+    output_gradient_vals = output_gradient_vals.reshape([1, 1, 1, 1])
+    with self.session() as sess:
+      with self.test_scope():
+        output_gradients = array_ops.placeholder(
+            dtypes.float32, shape=output_gradient_vals.shape
+        )
+        t = gen_nn_ops.avg_pool_grad(
+            orig_input_shape=[1, 0, 0, 0],
+            grad=output_gradients,
+            ksize=[1, 0, 0, 0],
+            strides=[1, 0, 0, 0],
+            padding="SAME",
+            data_format="NCHW",
+        )
+      with self.assertRaisesRegex(
+          errors.InvalidArgumentError,
+          (
+              "Sliding window ksize field for dimension 1 must be positive but"
+              " is 0"
+          ),
+      ):
+        sess.run(t, {output_gradients: output_gradient_vals})
+
   # The CPU implementation of AvgPoolGrad doesn't accept kernels smaller than
   # the stride size, so we only run the following tests on MaxPoolGrad.
 
   def testMaxPoolKernelSmallerThanStrideValid(self):
     self._VerifyValues(
diff --git a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc
index 7f0d712f13df1..64351c6a7416b 100644
--- a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc
@@ -31,19 +31,45 @@ limitations under the License.
 #include "tensorflow/compiler/xla/client/xla_computation.h"
 #include "tensorflow/compiler/xla/literal.h"
 #include "tensorflow/compiler/xla/util.h"
 #include "tensorflow/core/framework/bounds_check.h"
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/util/determinism.h"
 #include "tensorflow/core/util/tensor_format.h"
+#include "tensorflow/tsl/platform/errors.h"
 
 namespace tensorflow {
 namespace {
 
+template <typename T>
+static Status ValidateKernelSizes(const T& ksizes) {
+  for (size_t i = 0; i < ksizes.size(); ++i) {
+    if (ksizes[i] <= 0) {
+      return errors::InvalidArgument(
+          "Sliding window ksize field for dimension ", i,
+          " must be positive but is ", ksizes[i]);
+    }
+  }
+  return OkStatus();
+}
+
+template <typename T>
+static Status ValidateStrides(const T& strides) {
+  for (size_t i = 0; i < strides.size(); ++i) {
+    if (strides[i] <= 0) {
+      return errors::InvalidArgument(
+          "Sliding window stride field for dimension ", i,
+          " must be positive but is ", strides[i]);
+    }
+  }
+  return OkStatus();
+}
+
 // Superclass of pooling ops.
 class PoolingOp : public XlaOpKernel {
  public:
   PoolingOp(OpKernelConstruction* ctx, int num_spatial_dims,
             const DataType reduction_type)
@@ -81,54 +107,58 @@ class PoolingOp : public XlaOpKernel {
 
   int num_dims() const { return num_spatial_dims_ + 2; }
 
  protected:
   StatusOr<std::vector<int64_t>> GetKernelSize(XlaOpKernelContext* ctx) {
-    if (ctx->num_inputs() == 1) {
-      return ksize_;
-    }
-    const TensorShape ksize_shape = ctx->InputShape(1);
-    // Validate input sizes.
-    if (!TensorShapeUtils::IsVector(ksize_shape)) {
-      return errors::InvalidArgument("ksize must be a vector, not shape ",
-                                     ksize_shape.DebugString());
-    }
-    if (ksize_shape.num_elements() != num_dims()) {
-      return errors::InvalidArgument(
-          "Sliding window ksize field must "
-          "specify ",
-          num_dims(), " dimensions");
-    }
     std::vector<int64_t> ksize;
-    auto status = ctx->ConstantInputAsIntVector(1, &ksize);
-    if (!status.ok()) {
-      return status;
+    if (ctx->num_inputs() == 1) {
+      ksize = ksize_;
+    } else {
+      const TensorShape ksize_shape = ctx->InputShape(1);
+      // Validate input sizes.
+      if (!TensorShapeUtils::IsVector(ksize_shape)) {
+        return errors::InvalidArgument("ksize must be a vector, not shape ",
+                                       ksize_shape.DebugString());
+      }
+      if (ksize_shape.num_elements() != num_dims()) {
+        return errors::InvalidArgument(
+            "Sliding window ksize field must "
+            "specify ",
+            num_dims(), " dimensions");
+      }
+      auto status = ctx->ConstantInputAsIntVector(1, &ksize);
+      if (!status.ok()) {
+        return status;
+      }
     }
+    TF_RETURN_IF_ERROR(ValidateKernelSizes(ksize));
     return ksize;
   }
 
   StatusOr<std::vector<int64_t>> GetStride(XlaOpKernelContext* ctx) {
-    if (ctx->num_inputs() == 1) {
-      return stride_;
-    }
-    const TensorShape stride_shape = ctx->InputShape(2);
-    // Validate input sizes.
-    if (!TensorShapeUtils::IsVector(stride_shape)) {
-      return errors::InvalidArgument("stride must be a vector, not shape ",
-                                     stride_shape.DebugString());
-    }
-    if (stride_shape.num_elements() != num_dims()) {
-      return errors::InvalidArgument(
-          "Sliding window stride field must "
-          "specify ",
-          num_dims(), " dimensions");
-    }
     std::vector<int64_t> stride;
-    auto status = ctx->ConstantInputAsIntVector(2, &stride);
-    if (!status.ok()) {
-      return status;
+    if (ctx->num_inputs() == 1) {
+      stride = stride_;
+    } else {
+      const TensorShape stride_shape = ctx->InputShape(2);
+      // Validate input sizes.
+      if (!TensorShapeUtils::IsVector(stride_shape)) {
+        return errors::InvalidArgument("stride must be a vector, not shape ",
+                                       stride_shape.DebugString());
+      }
+      if (stride_shape.num_elements() != num_dims()) {
+        return errors::InvalidArgument(
+            "Sliding window stride field must "
+            "specify ",
+            num_dims(), " dimensions");
+      }
+      auto status = ctx->ConstantInputAsIntVector(2, &stride);
+      if (!status.ok()) {
+        return status;
+      }
     }
+    TF_RETURN_IF_ERROR(ValidateStrides(stride));
     return stride;
   }
 
  protected:
   const int num_spatial_dims_;
@@ -353,14 +383,16 @@ class MaxPoolGradOp : public XlaOpKernel {
 
     OP_REQUIRES(ctx, ksize_.size() == num_dims(),
                 errors::InvalidArgument("Sliding window ksize field must "
                                         "specify ",
                                         num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateKernelSizes(ksize_));
     OP_REQUIRES(ctx, stride_.size() == num_dims(),
                 errors::InvalidArgument("Sliding window strides field must "
                                         "specify ",
                                         num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateStrides(stride_));
 
     const TensorShape tensor_in_shape = ctx->InputShape(0);
     const TensorShape tensor_out_shape = ctx->InputShape(1);
     const TensorShape out_backprop_shape = ctx->InputShape(2);
 
@@ -444,15 +476,17 @@ class AvgPoolGradOp : public XlaOpKernel {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("ksize", &ksize_));
     OP_REQUIRES(ctx, ksize_.size() == num_dims(),
                 errors::InvalidArgument("Sliding window ksize field must "
                                         "specify ",
                                         num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateKernelSizes(ksize_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &stride_));
     OP_REQUIRES(ctx, stride_.size() == num_dims(),
                 errors::InvalidArgument("Sliding window strides field must "
                                         "specify ",
                                         num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateStrides(stride_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_));
     OP_REQUIRES(ctx, padding_ != EXPLICIT,
                 errors::Unimplemented(
                     "XLA does not support avgpoolgrad with explicit padding."));
     OP_REQUIRES(ctx, ksize_[0] == 1 && stride_[0] == 1,
@@ -577,14 +611,16 @@ class MaxPoolGradGradOp : public XlaOpKernel {
 
     OP_REQUIRES(ctx, ksize_.size() == num_dims(),
                 errors::InvalidArgument("Sliding window ksize field must "
                                         "specify ",
                                         num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateKernelSizes(ksize_));
     OP_REQUIRES(ctx, stride_.size() == num_dims(),
                 errors::InvalidArgument("Sliding window strides field must "
                                         "specify ",
                                         num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateStrides(stride_));
 
     const TensorShape tensor_in_shape = ctx->InputShape(0);
     const TensorShape tensor_out_shape = ctx->InputShape(1);
     const TensorShape out_backprop_shape = ctx->InputShape(2);
 
diff --git a/tensorflow/compiler/xla/client/padding.cc b/tensorflow/compiler/xla/client/padding.cc
index 1fe8d556d1d19..7fec04e2ac540 100644
--- a/tensorflow/compiler/xla/client/padding.cc
+++ b/tensorflow/compiler/xla/client/padding.cc
@@ -33,10 +33,20 @@ Status ValidatePaddingValues(absl::Span<const int64_t> input_dimensions,
         "Want input dimensions size %u = window dimensions size %u = window "
         "strides size %u",
         input_dimensions.size(), window_dimensions.size(),
         window_strides.size());
   }
+  for (size_t i = 0; i < input_dimensions.size(); ++i) {
+    if (window_dimensions[i] <= 0) {
+      return InvalidArgument("Window dimension %u has non-positive size %d", i,
+                             window_dimensions[i]);
+    }
+    if (window_strides[i] <= 0) {
+      return InvalidArgument("Window dimension %u has non-positive stride %d",
+                             i, window_strides[i]);
+    }
+  }
   return OkStatus();
 }
 
 std::vector<std::pair<int64_t, int64_t>> MakePadding(
     absl::Span<const int64_t> input_dimensions,

Vulnerable code:     if (ctx->num_inputs() == 1) {
      return ksize_;
    }
    const TensorShape ksize_shape = ctx->InputShape(1);
    // Validate input sizes.
    if (!TensorShapeUtils::IsVector(ksize_shape)) {
      return errors::InvalidArgument("ksize must be a vector, not shape ",
                                     ksize_shape.DebugString());
    }
    if (ksize_shape.num_elements() != num_dims()) {
      return errors::InvalidArgument(
          "Sliding window ksize field must "
          "specify ",
          num_dims(), " dimensions");
    }
    auto status = ctx->ConstantInputAsIntVector(1, &ksize);
    if (!status.ok()) {
      return status;
    if (ctx->num_inputs() == 1) {
      return stride_;
    }
    const TensorShape stride_shape = ctx->InputShape(2);
    // Validate input sizes.
    if (!TensorShapeUtils::IsVector(stride_shape)) {
      return errors::InvalidArgument("stride must be a vector, not shape ",
                                     stride_shape.DebugString());
    }
    if (stride_shape.num_elements() != num_dims()) {
      return errors::InvalidArgument(
          "Sliding window stride field must "
          "specify ",
          num_dims(), " dimensions");
    }
    auto status = ctx->ConstantInputAsIntVector(2, &stride);
    if (!status.ok()) {
      return status;
Fixed code: from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
  @test_util.disable_mlir_bridge(
      "TODO(b/266613412): investigate FPE in AvgPoolGrad for TPU"
  )
  def testAvgPoolGradSamePaddingZeroStrideZeroSize(self):
    output_gradient_vals = np.array([0.39117979], dtype=np.float32)
    output_gradient_vals = output_gradient_vals.reshape([1, 1, 1, 1])
    with self.session() as sess:
      with self.test_scope():
        output_gradients = array_ops.placeholder(
            dtypes.float32, shape=output_gradient_vals.shape
        )
        t = gen_nn_ops.avg_pool_grad(
            orig_input_shape=[1, 0, 0, 0],
            grad=output_gradients,
            ksize=[1, 0, 0, 0],
            strides=[1, 0, 0, 0],
            padding="SAME",
            data_format="NCHW",
        )
      with self.assertRaisesRegex(
          errors.InvalidArgumentError,
          (
              "Sliding window ksize field for dimension 1 must be positive but"
              " is 0"
          ),
      ):
        sess.run(t, {output_gradients: output_gradient_vals})

#include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/tsl/platform/errors.h"
template <typename T>
static Status ValidateKernelSizes(const T& ksizes) {
  for (size_t i = 0; i < ksizes.size(); ++i) {
    if (ksizes[i] <= 0) {
      return errors::InvalidArgument(
          "Sliding window ksize field for dimension ", i,
          " must be positive but is ", ksizes[i]);
    }
  }
  return OkStatus();
}

template <typename T>
static Status ValidateStrides(const T& strides) {
  for (size_t i = 0; i < strides.size(); ++i) {
    if (strides[i] <= 0) {
      return errors::InvalidArgument(
          "Sliding window stride field for dimension ", i,
          " must be positive but is ", strides[i]);
    }
  }
  return OkStatus();
}

    if (ctx->num_inputs() == 1) {
      ksize = ksize_;
    } else {
      const TensorShape ksize_shape = ctx->InputShape(1);
      // Validate input sizes.
      if (!TensorShapeUtils::IsVector(ksize_shape)) {
        return errors::InvalidArgument("ksize must be a vector, not shape ",
                                       ksize_shape.DebugString());
      }
      if (ksize_shape.num_elements() != num_dims()) {
        return errors::InvalidArgument(
            "Sliding window ksize field must "
            "specify ",
            num_dims(), " dimensions");
      }
      auto status = ctx->ConstantInputAsIntVector(1, &ksize);
      if (!status.ok()) {
        return status;
      }
    TF_RETURN_IF_ERROR(ValidateKernelSizes(ksize));
    if (ctx->num_inputs() == 1) {
      stride = stride_;
    } else {
      const TensorShape stride_shape = ctx->InputShape(2);
      // Validate input sizes.
      if (!TensorShapeUtils::IsVector(stride_shape)) {
        return errors::InvalidArgument("stride must be a vector, not shape ",
                                       stride_shape.DebugString());
      }
      if (stride_shape.num_elements() != num_dims()) {
        return errors::InvalidArgument(
            "Sliding window stride field must "
            "specify ",
            num_dims(), " dimensions");
      }
      auto status = ctx->ConstantInputAsIntVector(2, &stride);
      if (!status.ok()) {
        return status;
      }
    TF_RETURN_IF_ERROR(ValidateStrides(stride));
    OP_REQUIRES_OK(ctx, ValidateKernelSizes(ksize_));
    OP_REQUIRES_OK(ctx, ValidateStrides(stride_));
    OP_REQUIRES_OK(ctx, ValidateKernelSizes(ksize_));
    OP_REQUIRES_OK(ctx, ValidateStrides(stride_));
    OP_REQUIRES_OK(ctx, ValidateKernelSizes(ksize_));
    OP_REQUIRES_OK(ctx, ValidateStrides(stride_));
  for (size_t i = 0; i < input_dimensions.size(); ++i) {
    if (window_dimensions[i] <= 0) {
      return InvalidArgument("Window dimension %u has non-positive size %d", i,
                             window_dimensions[i]);
    }
    if (window_strides[i] <= 0) {
      return InvalidArgument("Window dimension %u has non-positive stride %d",
                             i, window_strides[i]);
    }
  }