CVE-2022-35966: Insufficient Information (NVD-CWE-noinfo)
value: TensorFlow is an open source platform for machine learning. If `QuantizedAvgPool` is given `min_input` or `max_input` tensors of a nonzero rank, it results in a segfault that can be used to trigger a denial of service attack. We have patched the issue in GitHub commit 7cdf9d4d2083b739ec81cfdace546b0c99f50622. The fix will be included in TensorFlow 2.10.0. We will also cherrypick this commit on TensorFlow 2.9.1, TensorFlow 2.8.1, and TensorFlow 2.7.2, as these are also affected and still in supported range. There are no known workarounds for this issue.
CWE description: There is insufficient information about the issue to classify it; details are unkown or unspecified.
CVSS v3 base score: 7.5
Published: 2022-09-16T21:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/quantized_pooling_ops.cc, tensorflow/core/kernels/quantized_pooling_ops_test.cc, tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
Commit date: 2022-07-21 10:36:03 -0700
Commit message: Fix QuantizedAvgPool invalid rank issue.

Passing in an invalid rank in eager mode led to a runtime crash.
Updated the runtime error to match the one generated by the
shape function for better consistency.

PiperOrigin-RevId: 462419362

Diff: diff --git a/tensorflow/core/kernels/quantized_pooling_ops.cc b/tensorflow/core/kernels/quantized_pooling_ops.cc
index b512369b3c4dd..5673fb6ee00a5 100644
--- a/tensorflow/core/kernels/quantized_pooling_ops.cc
+++ b/tensorflow/core/kernels/quantized_pooling_ops.cc
@@ -13,22 +13,22 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
 // See docs in ../ops/nn_ops.cc.
 
-#include "tensorflow/core/framework/op_requires.h"
-#include "tensorflow/core/platform/errors.h"
 #define EIGEN_USE_THREADS
 
 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 #include "tensorflow/core/framework/numeric_op.h"
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/kernels/ops_util.h"
 #include "tensorflow/core/kernels/pooling_ops_common.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/util/padding.h"
 #include "tensorflow/core/util/tensor_format.h"
 
 namespace tensorflow {
@@ -65,12 +65,24 @@ class QuantizedAvgPoolingOp : public OpKernel {
                           tensor_in.shape()};
     if (!context->status().ok()) {
       return;
     }
 
-    const float min_input = context->input(1).flat<float>()(0);
-    const float max_input = context->input(2).flat<float>()(0);
+    const Tensor& min_input_tensor = context->input(1);
+    const Tensor& max_input_tensor = context->input(2);
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_input_tensor.shape()),
+                errors::InvalidArgument(
+                    "min_input shape must be rank 0 but is rank ",
+                    min_input_tensor.dims(),
+                    ", received shape: ", min_input_tensor.shape()));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_input_tensor.shape()),
+                errors::InvalidArgument(
+                    "max_input shape must be rank 0 but is rank ",
+                    max_input_tensor.dims(),
+                    ", received shape: ", max_input_tensor.shape()));
+    const float min_input = context->input(1).scalar<float>()();
+    const float max_input = context->input(2).scalar<float>()();
 
     OP_REQUIRES(context, params.depth_window == 1,
                 errors::Unimplemented("Non-spatial pooling is not "
                                       "yet supported. Volunteers? :)"));
 
@@ -117,24 +129,24 @@ class QuantizedMaxPoolingOp : public MaxPoolingOp<Device, T> {
  public:
   explicit QuantizedMaxPoolingOp(OpKernelConstruction* context)
       : MaxPoolingOp<Device, T>(context) {}
 
   void Compute(OpKernelContext* context) override {
-    auto min_input_tensor = context->input(1);
-    auto max_input_tensor = context->input(2);
-    OP_REQUIRES(
-        context, min_input_tensor.NumElements() == 1,
-        errors::InvalidArgument(
-            "min_input must be a scalar float value, got tensor with shape ",
-            min_input_tensor.shape()));
-    OP_REQUIRES(
-        context, max_input_tensor.NumElements() == 1,
-        errors::InvalidArgument(
-            "max_input must be a scalar float value, got tensor with shape ",
-            max_input_tensor.shape()));
-    const float min_input = context->input(1).flat<float>()(0);
-    const float max_input = context->input(2).flat<float>()(0);
+    const Tensor& min_input_tensor = context->input(1);
+    const Tensor& max_input_tensor = context->input(2);
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_input_tensor.shape()),
+                errors::InvalidArgument(
+                    "min_input shape must be rank 0 but is rank ",
+                    min_input_tensor.dims(),
+                    ", received shape: ", min_input_tensor.shape()));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_input_tensor.shape()),
+                errors::InvalidArgument(
+                    "max_input shape must be rank 0 but is rank ",
+                    max_input_tensor.dims(),
+                    ", received shape: ", max_input_tensor.shape()));
+    const float min_input = context->input(1).scalar<float>()();
+    const float max_input = context->input(2).scalar<float>()();
     MaxPoolingOp<Device, T>::Compute(context);
     Tensor* output_min = nullptr;
     OP_REQUIRES_OK(context, context->allocate_output(1, {}, &output_min));
     output_min->flat<float>()(0) = min_input;
     Tensor* output_max = nullptr;
diff --git a/tensorflow/core/kernels/quantized_pooling_ops_test.cc b/tensorflow/core/kernels/quantized_pooling_ops_test.cc
index fc0417e5431b2..9e56890478be2 100644
--- a/tensorflow/core/kernels/quantized_pooling_ops_test.cc
+++ b/tensorflow/core/kernels/quantized_pooling_ops_test.cc
@@ -67,12 +67,12 @@ TEST_F(QuantizedPoolingTest, SmallAveragePooling) {
                         {1, expected_height, expected_width, input_channels});
   test::FillValues<float>(&expected_float, {6, 7, 10, 11, 22, 23, 26, 27});
 
   AddInputFromArray<quint8>(input_quantized.shape(),
                             input_quantized.flat<quint8>());
-  AddInputFromArray<float>(TensorShape({1}), {input_min});
-  AddInputFromArray<float>(TensorShape({1}), {input_max});
+  AddInputFromArray<float>(TensorShape({}), {input_min});
+  AddInputFromArray<float>(TensorShape({}), {input_max});
   TF_ASSERT_OK(RunOpKernel());
   const Tensor& output_quantized = *GetOutput(0);
   const float output_min = GetOutput(1)->flat<float>()(0);
   const float output_max = GetOutput(2)->flat<float>()(0);
   Tensor output_float =
@@ -112,12 +112,12 @@ TEST_F(QuantizedPoolingTest, SmallMaxPooling) {
                         {1, expected_height, expected_width, input_channels});
   test::FillValues<float>(&expected_float, {11, 12, 15, 16, 27, 28, 31, 32});
 
   AddInputFromArray<quint8>(input_quantized.shape(),
                             input_quantized.flat<quint8>());
-  AddInputFromArray<float>(TensorShape({1}), {input_min});
-  AddInputFromArray<float>(TensorShape({1}), {input_max});
+  AddInputFromArray<float>(TensorShape({}), {input_min});
+  AddInputFromArray<float>(TensorShape({}), {input_max});
   TF_ASSERT_OK(RunOpKernel());
   const Tensor& output_quantized = *GetOutput(0);
   const float output_min = GetOutput(1)->flat<float>()(0);
   const float output_max = GetOutput(2)->flat<float>()(0);
   Tensor output_float =
diff --git a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
index 708edf71a6bce..4c6d9974c1fc9 100644
--- a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
+++ b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
@@ -152,10 +152,76 @@ class QuantizedInstanceNormOpTest(test_util.TensorFlowTestCase):
       self.evaluate(
           array_ops.quantized_instance_norm(
               x=inputs, x_min=[[1.0], [2.0], [4.0]], x_max=1.0))
 
 
+class QuantizedAvgPoolingOpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    inputs = constant_op.constant(
+        np.uint8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)
+    ksize = [1, 1, 1, 1]
+    strides = [1, 1, 1, 1]
+    padding = "SAME"
+
+    with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
+                                "must be.* rank 0"):
+      self.evaluate(
+          nn_ops.quantized_avg_pool(
+              input=inputs,
+              min_input=[],
+              max_input=1.0,
+              ksize=ksize,
+              strides=strides,
+              padding=padding))
+
+    with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
+                                "must be.* rank 0"):
+      self.evaluate(
+          nn_ops.quantized_avg_pool(
+              input=inputs,
+              min_input=0.0,
+              max_input=[],
+              ksize=ksize,
+              strides=strides,
+              padding=padding))
+
+
+class QuantizedMaxPoolingOpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    inputs = constant_op.constant(
+        np.uint8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)
+    ksize = [1, 1, 1, 1]
+    strides = [1, 1, 1, 1]
+    padding = "SAME"
+
+    with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
+                                "must be.* rank 0"):
+      self.evaluate(
+          nn_ops.quantized_max_pool(
+              input=inputs,
+              min_input=[],
+              max_input=1.0,
+              ksize=ksize,
+              strides=strides,
+              padding=padding))
+
+    with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
+                                "must be.* rank 0"):
+      self.evaluate(
+          nn_ops.quantized_max_pool(
+              input=inputs,
+              min_input=0.0,
+              max_input=[],
+              ksize=ksize,
+              strides=strides,
+              padding=padding))
+
+
 class RequantizeOpTest(test_util.TensorFlowTestCase):
 
   @test_util.run_in_graph_and_eager_modes
   def test_invalid_inputs(self):
     inputs = constant_op.constant(

Vulnerable code: #include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/platform/errors.h"
    const float min_input = context->input(1).flat<float>()(0);
    const float max_input = context->input(2).flat<float>()(0);
    auto min_input_tensor = context->input(1);
    auto max_input_tensor = context->input(2);
    OP_REQUIRES(
        context, min_input_tensor.NumElements() == 1,
        errors::InvalidArgument(
            "min_input must be a scalar float value, got tensor with shape ",
            min_input_tensor.shape()));
    OP_REQUIRES(
        context, max_input_tensor.NumElements() == 1,
        errors::InvalidArgument(
            "max_input must be a scalar float value, got tensor with shape ",
            max_input_tensor.shape()));
    const float min_input = context->input(1).flat<float>()(0);
    const float max_input = context->input(2).flat<float>()(0);
  AddInputFromArray<float>(TensorShape({1}), {input_min});
  AddInputFromArray<float>(TensorShape({1}), {input_max});
  AddInputFromArray<float>(TensorShape({1}), {input_min});
  AddInputFromArray<float>(TensorShape({1}), {input_max});
Fixed code: #include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/platform/errors.h"
    const Tensor& min_input_tensor = context->input(1);
    const Tensor& max_input_tensor = context->input(2);
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_input_tensor.shape()),
                errors::InvalidArgument(
                    "min_input shape must be rank 0 but is rank ",
                    min_input_tensor.dims(),
                    ", received shape: ", min_input_tensor.shape()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_input_tensor.shape()),
                errors::InvalidArgument(
                    "max_input shape must be rank 0 but is rank ",
                    max_input_tensor.dims(),
                    ", received shape: ", max_input_tensor.shape()));
    const float min_input = context->input(1).scalar<float>()();
    const float max_input = context->input(2).scalar<float>()();
    const Tensor& min_input_tensor = context->input(1);
    const Tensor& max_input_tensor = context->input(2);
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_input_tensor.shape()),
                errors::InvalidArgument(
                    "min_input shape must be rank 0 but is rank ",
                    min_input_tensor.dims(),
                    ", received shape: ", min_input_tensor.shape()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_input_tensor.shape()),
                errors::InvalidArgument(
                    "max_input shape must be rank 0 but is rank ",
                    max_input_tensor.dims(),
                    ", received shape: ", max_input_tensor.shape()));
    const float min_input = context->input(1).scalar<float>()();
    const float max_input = context->input(2).scalar<float>()();
  AddInputFromArray<float>(TensorShape({}), {input_min});
  AddInputFromArray<float>(TensorShape({}), {input_max});
  AddInputFromArray<float>(TensorShape({}), {input_min});
  AddInputFromArray<float>(TensorShape({}), {input_max});
class QuantizedAvgPoolingOpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    inputs = constant_op.constant(
        np.uint8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)
    ksize = [1, 1, 1, 1]
    strides = [1, 1, 1, 1]
    padding = "SAME"

    with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
                                "must be.* rank 0"):
      self.evaluate(
          nn_ops.quantized_avg_pool(
              input=inputs,
              min_input=[],
              max_input=1.0,
              ksize=ksize,
              strides=strides,
              padding=padding))

    with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
                                "must be.* rank 0"):
      self.evaluate(
          nn_ops.quantized_avg_pool(
              input=inputs,
              min_input=0.0,
              max_input=[],
              ksize=ksize,
              strides=strides,
              padding=padding))


class QuantizedMaxPoolingOpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    inputs = constant_op.constant(
        np.uint8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)
    ksize = [1, 1, 1, 1]
    strides = [1, 1, 1, 1]
    padding = "SAME"

    with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
                                "must be.* rank 0"):
      self.evaluate(
          nn_ops.quantized_max_pool(
              input=inputs,
              min_input=[],
              max_input=1.0,
              ksize=ksize,
              strides=strides,
              padding=padding))

    with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError),
                                "must be.* rank 0"):
      self.evaluate(
          nn_ops.quantized_max_pool(
              input=inputs,
              min_input=0.0,
              max_input=[],
              ksize=ksize,
              strides=strides,
              padding=padding))