CVE-2022-35974: Insufficient Information (NVD-CWE-noinfo)
value: TensorFlow is an open source platform for machine learning. If `QuantizeDownAndShrinkRange` is given nonscalar inputs for `input_min` or `input_max`, it results in a segfault that can be used to trigger a denial of service attack. We have patched the issue in GitHub commit 73ad1815ebcfeb7c051f9c2f7ab5024380ca8613. 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/quantize_down_and_shrink_range.cc, tensorflow/core/kernels/quantize_down_and_shrink_range_op_test.cc, tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
Commit date: 2022-07-21 09:26:51 -0700
Commit message: Add IsScalar (rank == 0) check to input_min/max tensors for QuantizeDownAndShrinkRangeOp.

PiperOrigin-RevId: 462401306

Diff: diff --git a/tensorflow/core/kernels/quantize_down_and_shrink_range.cc b/tensorflow/core/kernels/quantize_down_and_shrink_range.cc
index 1b948c8108de8..83f8996b4cc74 100644
--- a/tensorflow/core/kernels/quantize_down_and_shrink_range.cc
+++ b/tensorflow/core/kernels/quantize_down_and_shrink_range.cc
@@ -38,12 +38,24 @@ class QuantizeDownAndShrinkRangeOp : public OpKernel {
   explicit QuantizeDownAndShrinkRangeOp(OpKernelConstruction* ctx)
       : OpKernel(ctx) {}
 
   void Compute(OpKernelContext* ctx) override {
     const Tensor& input = ctx->input(0);
-    const float input_min_float = ctx->input(1).flat<float>()(0);
-    const float input_max_float = ctx->input(2).flat<float>()(0);
+    const Tensor& input_min = ctx->input(1);
+    const Tensor& input_max = ctx->input(2);
+
+    OP_REQUIRES(
+        ctx, TensorShapeUtils::IsScalar(input_min.shape()),
+        errors::InvalidArgument("`input_min` must be rank 0 but is rank ",
+                                input_min.dims()));
+    OP_REQUIRES(
+        ctx, TensorShapeUtils::IsScalar(input_max.shape()),
+        errors::InvalidArgument("`input_max` must be rank 0 but is rank ",
+                                input_max.dims()));
+
+    const float input_min_float = input_min.scalar<float>()();
+    const float input_max_float = input_max.scalar<float>()();
     Tensor* output = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input.shape(), &output));
     Tensor* output_min = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(1, TensorShape({}), &output_min));
     Tensor* output_max = nullptr;
diff --git a/tensorflow/core/kernels/quantize_down_and_shrink_range_op_test.cc b/tensorflow/core/kernels/quantize_down_and_shrink_range_op_test.cc
index 583ba1fded79e..f3d777e8e0e3c 100644
--- a/tensorflow/core/kernels/quantize_down_and_shrink_range_op_test.cc
+++ b/tensorflow/core/kernels/quantize_down_and_shrink_range_op_test.cc
@@ -51,12 +51,12 @@ TEST_F(QuantizeDownAndShrinkRangeTest, HandCrafted) {
   // the available range in the lower bit depth, and update to the true min and
   // max ranges.
   const int value_count = 3;
   AddInputFromArray<qint32>(TensorShape({value_count}),
                             {-(1 << 23), 0, (1 << 23)});
-  AddInputFromArray<float>(TensorShape({1}), {-256.0f});
-  AddInputFromArray<float>(TensorShape({1}), {256.0f});
+  AddInputFromArray<float>(TensorShape({}), {-256.0f});
+  AddInputFromArray<float>(TensorShape({}), {256.0f});
   TF_ASSERT_OK(RunOpKernel());
   Tensor expected(allocator(), DT_QUINT8, TensorShape({value_count}));
   test::FillValues<quint8>(&expected, {0, 128, 255});
   test::ExpectTensorEqual<quint8>(expected, *GetOutput(0));
   Tensor expected_min(allocator(), DT_FLOAT, TensorShape({}));
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 8be78a0e0b21f..708edf71a6bce 100644
--- a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
+++ b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
@@ -259,7 +259,23 @@ class QuantizedRelu6OpTest(test_util.TensorFlowTestCase):
               min_features=[],
               max_features=127.0,
               out_type=dtypes.quint8))
 
 
+class QuantizeDownAndShrinkRangeOpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    inputs = constant_op.constant(
+        np.int32(0), shape=[3, 3, 3, 3], dtype=dtypes.qint32)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          math_ops.quantize_down_and_shrink_range(input=inputs,
+                                                  input_min=[],
+                                                  input_max=4.0,
+                                                  out_type=dtypes.quint8))
+
+
 if __name__ == "__main__":
   googletest.main()

Vulnerable code:     const float input_min_float = ctx->input(1).flat<float>()(0);
    const float input_max_float = ctx->input(2).flat<float>()(0);
  AddInputFromArray<float>(TensorShape({1}), {-256.0f});
  AddInputFromArray<float>(TensorShape({1}), {256.0f});
Fixed code:     const Tensor& input_min = ctx->input(1);
    const Tensor& input_max = ctx->input(2);

    OP_REQUIRES(
        ctx, TensorShapeUtils::IsScalar(input_min.shape()),
        errors::InvalidArgument("`input_min` must be rank 0 but is rank ",
                                input_min.dims()));
    OP_REQUIRES(
        ctx, TensorShapeUtils::IsScalar(input_max.shape()),
        errors::InvalidArgument("`input_max` must be rank 0 but is rank ",
                                input_max.dims()));

    const float input_min_float = input_min.scalar<float>()();
    const float input_max_float = input_max.scalar<float>()();
  AddInputFromArray<float>(TensorShape({}), {-256.0f});
  AddInputFromArray<float>(TensorShape({}), {256.0f});
class QuantizeDownAndShrinkRangeOpTest(test_util.TensorFlowTestCase):

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

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          math_ops.quantize_down_and_shrink_range(input=inputs,
                                                  input_min=[],
                                                  input_max=4.0,
                                                  out_type=dtypes.quint8))