CVE-2022-35979: Insufficient Information (NVD-CWE-noinfo)
value: TensorFlow is an open source platform for machine learning. If `QuantizedRelu` or `QuantizedRelu6` are given nonscalar inputs for `min_features` or `max_features`, it results in a segfault that can be used to trigger a denial of service attack. We have patched the issue in GitHub commit 49b3824d83af706df0ad07e4e677d88659756d89. 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-16T22:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/quantized_activation_ops.cc, tensorflow/core/kernels/quantized_activation_ops_test.cc, tensorflow/core/kernels/quantized_add_op.cc, tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
Commit date: 2022-07-19 09:36:21 -0700
Commit message: Add IsScalar (rank == 0) check to min/max input tensors for QuantizedAdd/Relu/Relu6 op.

PiperOrigin-RevId: 461902847

Diff: diff --git a/tensorflow/core/kernels/quantized_activation_ops.cc b/tensorflow/core/kernels/quantized_activation_ops.cc
index 2896c3d45a702..36d321a8e1713 100644
--- a/tensorflow/core/kernels/quantized_activation_ops.cc
+++ b/tensorflow/core/kernels/quantized_activation_ops.cc
@@ -30,12 +30,25 @@ class QuantizedReluOp : public OpKernel {
  public:
   explicit QuantizedReluOp(OpKernelConstruction* context) : OpKernel(context) {}
 
   void Compute(OpKernelContext* context) override {
     const Tensor& input = context->input(0);
-    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` must be rank 0 but is rank ",
+                                min_input_tensor.dims()));
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsScalar(max_input_tensor.shape()),
+        errors::InvalidArgument("`max_input` must be rank 0 but is rank ",
+                                max_input_tensor.dims()));
+
+    const float min_input = min_input_tensor.scalar<float>()();
+    const float max_input = max_input_tensor.scalar<float>()();
+
     Tensor* output = nullptr;
     OP_REQUIRES_OK(context,
                    context->allocate_output(0, input.shape(), &output));
     const T min_as_quantized = FloatToQuantized<T>(0.0f, min_input, max_input);
 
@@ -63,12 +76,25 @@ class QuantizedRelu6Op : public OpKernel {
   explicit QuantizedRelu6Op(OpKernelConstruction* context)
       : OpKernel(context) {}
 
   void Compute(OpKernelContext* context) override {
     const Tensor& input = context->input(0);
-    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` must be rank 0 but is rank ",
+                                min_input_tensor.dims()));
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsScalar(max_input_tensor.shape()),
+        errors::InvalidArgument("`max_input` must be rank 0 but is rank ",
+                                max_input_tensor.dims()));
+
+    const float min_input = min_input_tensor.scalar<float>()();
+    const float max_input = max_input_tensor.scalar<float>()();
+
     Tensor* output = nullptr;
     OP_REQUIRES_OK(context,
                    context->allocate_output(0, input.shape(), &output));
     const T min_as_quantized = FloatToQuantized<T>(0.0f, min_input, max_input);
     const T max_as_quantized = FloatToQuantized<T>(6.0f, min_input, max_input);
diff --git a/tensorflow/core/kernels/quantized_activation_ops_test.cc b/tensorflow/core/kernels/quantized_activation_ops_test.cc
index b3b7cb58b9a45..34c5130f4759b 100644
--- a/tensorflow/core/kernels/quantized_activation_ops_test.cc
+++ b/tensorflow/core/kernels/quantized_activation_ops_test.cc
@@ -53,12 +53,12 @@ TEST_F(QuantizedActivationsTest, TestRelu) {
   Tensor expected_float(DT_FLOAT, {input_height, input_width});
   test::FillValues<float>(&expected_float, {0, 0, 0, 1, 3, 6, 7, 100});
 
   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 =
@@ -84,12 +84,12 @@ TEST_F(QuantizedActivationsTest, TestRelu6) {
   Tensor expected_float(DT_FLOAT, {input_height, input_width});
   test::FillValues<float>(&expected_float, {0, 0, 0, 1, 3, 6, 6, 6});
 
   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/core/kernels/quantized_add_op.cc b/tensorflow/core/kernels/quantized_add_op.cc
index 1f9897b9b61cd..5cf7ed1456034 100644
--- a/tensorflow/core/kernels/quantized_add_op.cc
+++ b/tensorflow/core/kernels/quantized_add_op.cc
@@ -23,10 +23,11 @@ limitations under the License.
 #include <arm_neon.h>
 #endif
 
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/kernels/meta_support.h"
 #include "tensorflow/core/kernels/quantization_utils.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/util/bcast.h"
 
@@ -455,14 +456,32 @@ class QuantizedAddOp : public OpKernel {
   explicit QuantizedAddOp(OpKernelConstruction* context) : OpKernel(context) {}
 
   void Compute(OpKernelContext* context) override {
     const Tensor& x = context->input(0);
     const Tensor& y = context->input(1);
-    const float min_x = context->input(2).flat<float>()(0);
-    const float max_x = context->input(3).flat<float>()(0);
-    const float min_y = context->input(4).flat<float>()(0);
-    const float max_y = context->input(5).flat<float>()(0);
+    const Tensor& min_x_tensor = context->input(2);
+    const Tensor& max_x_tensor = context->input(3);
+    const Tensor& min_y_tensor = context->input(4);
+    const Tensor& max_y_tensor = context->input(5);
+
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_x_tensor.shape()),
+                errors::InvalidArgument("`min_x` must be rank 0 but is rank ",
+                                        min_x_tensor.dims()));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_x_tensor.shape()),
+                errors::InvalidArgument("`max_x` must be rank 0 but is rank ",
+                                        max_x_tensor.dims()));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_y_tensor.shape()),
+                errors::InvalidArgument("`min_y` must be rank 0 but is rank ",
+                                        min_y_tensor.dims()));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_y_tensor.shape()),
+                errors::InvalidArgument("`max_y` must be rank 0 but is rank ",
+                                        max_y_tensor.dims()));
+
+    const float min_x = min_x_tensor.scalar<float>()();
+    const float max_x = max_x_tensor.scalar<float>()();
+    const float min_y = min_y_tensor.scalar<float>()();
+    const float max_y = max_y_tensor.scalar<float>()();
 
     BCast bcast(BCast::FromShape(x.shape()), BCast::FromShape(y.shape()));
     if (!bcast.IsValid()) {
       context->SetStatus(errors::InvalidArgument(
           "Incompatible shapes: ", x.shape().DebugString(), " vs. ",
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 af7cecf2ffb02..8be78a0e0b21f 100644
--- a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
+++ b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
@@ -204,7 +204,62 @@ class RequantizeOpTest(test_util.TensorFlowTestCase):
               requested_output_min=0.0,
               requested_output_max=[],
               out_type=dtypes.qint8))
 
 
+class QuantizedAddOpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    x = constant_op.constant(
+        np.int8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)
+    y = constant_op.constant(np.int8(0), shape=[3], dtype=dtypes.quint8)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          math_ops.quantized_add(
+              x=x,
+              y=y,
+              min_x=[],
+              max_x=1.0,
+              min_y=0.0,
+              max_y=1.0,
+              Toutput=dtypes.qint32))
+
+
+class QuantizedReluOpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    inputs = constant_op.constant(
+        np.int8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          nn_ops.quantized_relu(
+              features=inputs,
+              min_features=[],
+              max_features=127.0,
+              out_type=dtypes.quint8))
+
+
+class QuantizedRelu6OpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    inputs = constant_op.constant(
+        np.int8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          nn_ops.quantized_relu6(
+              features=inputs,
+              min_features=[],
+              max_features=127.0,
+              out_type=dtypes.quint8))
+
+
 if __name__ == "__main__":
   googletest.main()

Vulnerable code:     const float min_input = context->input(1).flat<float>()(0);
    const float max_input = context->input(2).flat<float>()(0);
    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});
    const float min_x = context->input(2).flat<float>()(0);
    const float max_x = context->input(3).flat<float>()(0);
    const float min_y = context->input(4).flat<float>()(0);
    const float max_y = context->input(5).flat<float>()(0);
Fixed code:     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` must be rank 0 but is rank ",
                                min_input_tensor.dims()));
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(max_input_tensor.shape()),
        errors::InvalidArgument("`max_input` must be rank 0 but is rank ",
                                max_input_tensor.dims()));

    const float min_input = min_input_tensor.scalar<float>()();
    const float max_input = max_input_tensor.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` must be rank 0 but is rank ",
                                min_input_tensor.dims()));
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(max_input_tensor.shape()),
        errors::InvalidArgument("`max_input` must be rank 0 but is rank ",
                                max_input_tensor.dims()));

    const float min_input = min_input_tensor.scalar<float>()();
    const float max_input = max_input_tensor.scalar<float>()();

  AddInputFromArray<float>(TensorShape({}), {input_min});
  AddInputFromArray<float>(TensorShape({}), {input_max});
  AddInputFromArray<float>(TensorShape({}), {input_min});
  AddInputFromArray<float>(TensorShape({}), {input_max});
#include "tensorflow/core/framework/tensor_shape.h"
    const Tensor& min_x_tensor = context->input(2);
    const Tensor& max_x_tensor = context->input(3);
    const Tensor& min_y_tensor = context->input(4);
    const Tensor& max_y_tensor = context->input(5);

    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_x_tensor.shape()),
                errors::InvalidArgument("`min_x` must be rank 0 but is rank ",
                                        min_x_tensor.dims()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_x_tensor.shape()),
                errors::InvalidArgument("`max_x` must be rank 0 but is rank ",
                                        max_x_tensor.dims()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_y_tensor.shape()),
                errors::InvalidArgument("`min_y` must be rank 0 but is rank ",
                                        min_y_tensor.dims()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_y_tensor.shape()),
                errors::InvalidArgument("`max_y` must be rank 0 but is rank ",
                                        max_y_tensor.dims()));

    const float min_x = min_x_tensor.scalar<float>()();
    const float max_x = max_x_tensor.scalar<float>()();
    const float min_y = min_y_tensor.scalar<float>()();
    const float max_y = max_y_tensor.scalar<float>()();
class QuantizedAddOpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    x = constant_op.constant(
        np.int8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)
    y = constant_op.constant(np.int8(0), shape=[3], dtype=dtypes.quint8)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          math_ops.quantized_add(
              x=x,
              y=y,
              min_x=[],
              max_x=1.0,
              min_y=0.0,
              max_y=1.0,
              Toutput=dtypes.qint32))


class QuantizedReluOpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    inputs = constant_op.constant(
        np.int8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          nn_ops.quantized_relu(
              features=inputs,
              min_features=[],
              max_features=127.0,
              out_type=dtypes.quint8))


class QuantizedRelu6OpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    inputs = constant_op.constant(
        np.int8(0), shape=[3, 3, 3, 3], dtype=dtypes.quint8)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          nn_ops.quantized_relu6(
              features=inputs,
              min_features=[],
              max_features=127.0,
              out_type=dtypes.quint8))