CVE-2022-35966 (published 2022-09-16T21:15Z)

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.

CVSS v3 base score: 7.5, HIGH (CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H)

Weakness information:
cwe_id: NVD-CWE-noinfo
cwe_name: Insufficient Information
description: There is insufficient information about the issue to classify it; details are unkown or unspecified.
extended_description: Insufficient Information

Fix information:
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Antonio Sanchez
author_date: 2022-07-21 10:32:19-07:00
msg: 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
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: quantized_pooling_ops.cc
change_type: ModificationType.MODIFY
diff: @@ -15,18 +15,18 @@ 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"
@@ -67,8 +67,20 @@ class QuantizedAvgPoolingOp : public OpKernel {
       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 "
@@ -119,20 +131,20 @@ class QuantizedMaxPoolingOp : public MaxPoolingOp<Device, T> {
       : 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));

programming_language: C++
filename: quantized_pooling_ops_test.cc
change_type: ModificationType.MODIFY
diff: @@ -69,8 +69,8 @@ TEST_F(QuantizedPoolingTest, SmallAveragePooling) {
 
   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);
@@ -114,8 +114,8 @@ TEST_F(QuantizedPoolingTest, SmallMaxPooling) {
 
   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);

programming_language: C++
filename: quantization_ops_test.py
change_type: ModificationType.MODIFY
diff: @@ -154,6 +154,72 @@ class QuantizedInstanceNormOpTest(test_util.TensorFlowTestCase):
               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

programming_language: Python