CVE-2022-36017 (C++, published 2022-09-16T23:15Z)
Description: value: TensorFlow is an open source platform for machine learning. If `Requantize` is given `input_min`, `input_max`, `requested_output_min`, `requested_output_max` 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 785d67a78a1d533759fcd2f5e8d6ef778de849e0. 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.
Severity: nan (CVSS v3 base score: 7.5)
Weakness: NVD-CWE-noinfo \u2014 Insufficient Information
There is insufficient information about the issue to classify it; details are unkown or unspecified.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/kernels/fake_quant_ops.cc
tensorflow/core/kernels/quantized_bias_add_op.cc
tensorflow/core/kernels/quantized_bias_add_op_test.cc
tensorflow/core/kernels/quantized_instance_norm.cc
tensorflow/core/kernels/requantize.cc
tensorflow/core/kernels/requantize_op_test.cc
tensorflow/python/kernel_tests/quantization_ops/BUILD
tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
Commit message: Fix quantize ops input validation issues.

The majority of these are just missing checks on min/max.

PiperOrigin-RevId: 461800665


Vulnerable code:
    const Tensor& max = context->input(2);
    const float input_min = context->input(2).flat<float>()(0);
    const float input_max = context->input(3).flat<float>()(0);
    const float bias_min = context->input(4).flat<float>()(0);
    const float bias_max = context->input(5).flat<float>()(0);
  AddInputFromArray<float>(TensorShape({1}), {input_min});
  AddInputFromArray<float>(TensorShape({1}), {input_max});
  AddInputFromArray<float>(TensorShape({1}), {bias_min});
  AddInputFromArray<float>(TensorShape({1}), {bias_max});
  AddInputFromArray<float>(TensorShape({1}), {input_min});
  AddInputFromArray<float>(TensorShape({1}), {input_max});
  AddInputFromArray<float>(TensorShape({1}), {bias_min});
  AddInputFromArray<float>(TensorShape({1}), {bias_max});

    float input_min = context->input(1).flat<float>()(0);
    float input_max = context->input(2).flat<float>()(0);
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
    const float input_min_float = ctx->input(1).flat<float>()(0);
    const float input_max_float = ctx->input(2).flat<float>()(0);
    const float requested_output_min_float = ctx->input(3).flat<float>()(0);
    const float requested_output_max_float = ctx->input(4).flat<float>()(0);
  AddInputFromArray<float>(TensorShape({1}), {-256.0f});
  AddInputFromArray<float>(TensorShape({1}), {256.0f});
  AddInputFromArray<float>(TensorShape({1}), {-1.0f});
  AddInputFromArray<float>(TensorShape({1}), {1.0f});
  AddInputFromArray<float>(TensorShape({1}), {-256.0f});
  AddInputFromArray<float>(TensorShape({1}), {256.0f});
  AddInputFromArray<float>(TensorShape({1}), {0.01f});
  AddInputFromArray<float>(TensorShape({1}), {1.0f});
  AddInputFromArray<float>(TensorShape({1}), {-256.0f});
  AddInputFromArray<float>(TensorShape({1}), {256.0f});
  AddInputFromArray<float>(TensorShape({1}), {-10.0f});
  AddInputFromArray<float>(TensorShape({1}), {-11.0f});

Fixed code:
#include "tensorflow/core/framework/tensor_shape.h"
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(min.shape()),
        InvalidArgument("`min` must be rank 0 but is rank ", min.dims()));
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(max.shape()),
        InvalidArgument("`max` must be rank 0 but is rank ", max.dims()));

    const Tensor& max = context->input(2);

    OP_REQUIRES(
        context, TensorShapeUtils::IsVector(min.shape()),
        InvalidArgument("`min` must be rank 1 but is rank ", min.dims()));
    OP_REQUIRES(
        context, TensorShapeUtils::IsVector(max.shape()),
        InvalidArgument("`max` must be rank 1 but is rank ", max.dims()));
#include "tensorflow/core/framework/tensor_shape.h"

    const Tensor& min_input = context->input(2);
    const Tensor& max_input = context->input(3);
    const Tensor& min_bias = context->input(4);
    const Tensor& max_bias = context->input(5);
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(min_input.shape()),
        errors::InvalidArgument("`min_input` must be rank 0 but is rank ",
                                min_input.dims()));
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(max_input.shape()),
        errors::InvalidArgument("`max_input` must be rank 0 but is rank ",
                                max_input.dims()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(min_bias.shape()),
                errors::InvalidArgument(
                    "`min_bias` must be rank 0 but is rank ", min_bias.dims()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(max_bias.shape()),
                errors::InvalidArgument(
                    "`max_bias` must be rank 0 but is rank ", max_bias.dims()));

    const float input_min = min_input.flat<float>()(0);
    const float input_max = max_input.flat<float>()(0);
    const float bias_min = min_bias.flat<float>()(0);
    const float bias_max = max_bias.flat<float>()(0);
  AddInputFromArray<float>(TensorShape({}), {input_min});
  AddInputFromArray<float>(TensorShape({}), {input_max});
  AddInputFromArray<float>(TensorShape({}), {bias_min});
  AddInputFromArray<float>(TensorShape({}), {bias_max});
  AddInputFromArray<float>(TensorShape({}), {input_min});
  AddInputFromArray<float>(TensorShape({}), {input_max});
  AddInputFromArray<float>(TensorShape({}), {bias_min});
  AddInputFromArray<float>(TensorShape({}), {bias_max});
#include "tensorflow/core/framework/tensor_shape.h"
    const Tensor& x_min = context->input(1);
    const Tensor& x_max = context->input(2);
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(x_min.shape()),
                errors::InvalidArgument("`x_min` must be rank 0 but is rank ",
                                        x_min.dims()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(x_max.shape()),
                errors::InvalidArgument("`x_max` must be rank 0 but is rank ",
                                        x_max.dims()));
    float input_min = x_min.scalar<float>()();
    float input_max = x_max.scalar<float>()();

#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"

    const Tensor& input_min = ctx->input(1);
    const Tensor& input_max = ctx->input(2);
    const Tensor& requested_output_min = ctx->input(3);
    const Tensor& requested_output_max = ctx->input(4);
    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()));
    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(requested_output_min.shape()),
                errors::InvalidArgument(
                    "`requested_output_min` must be rank 0 but is rank ",
                    requested_output_min.dims()));
    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(requested_output_max.shape()),
                errors::InvalidArgument(
                    "`requested_output_max` must be rank 0 but is rank ",
                    requested_output_max.dims()));

    const float input_min_float = input_min.flat<float>()(0);
    const float input_max_float = input_max.flat<float>()(0);
    const float requested_output_min_float =
        requested_output_min.flat<float>()(0);
    const float requested_output_max_float =
        requested_output_max.flat<float>()(0);
  AddInputFromArray<float>(TensorShape({}), {-256.0f});
  AddInputFromArray<float>(TensorShape({}), {256.0f});
  AddInputFromArray<float>(TensorShape({}), {-1.0f});
  AddInputFromArray<float>(TensorShape({}), {1.0f});
  AddInputFromArray<float>(TensorShape({}), {-256.0f});
  AddInputFromArray<float>(TensorShape({}), {256.0f});
  AddInputFromArray<float>(TensorShape({}), {0.01f});
  AddInputFromArray<float>(TensorShape({}), {1.0f});
  AddInputFromArray<float>(TensorShape({}), {-256.0f});
  AddInputFromArray<float>(TensorShape({}), {256.0f});
  AddInputFromArray<float>(TensorShape({}), {-10.0f});
  AddInputFromArray<float>(TensorShape({}), {-11.0f});
# Tests of TensorFlow quantization ops written using the Python API.

# buildifier: disable=same-origin-load
load("//tensorflow:tensorflow.bzl", "tf_py_test")

package(
    default_visibility = ["//tensorflow:internal"],
    licenses = ["notice"],
)

tf_py_test(
    name = "quantization_ops_test",
    size = "small",
    srcs = ["quantization_ops_test.py"],
    deps = [
        "//tensorflow/python:array_ops",
        "//tensorflow/python:client",
        "//tensorflow/python:client_testlib",
        "//tensorflow/python:framework",
        "//tensorflow/python:framework_for_generated_wrappers",
        "//tensorflow/python:math_ops",
        "//third_party/py/numpy",
    ],
)
# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for tf.quantize ops."""
import numpy as np

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.platform import googletest


class FakeQuantWithMinMaxVarsOpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    inputs = constant_op.constant(
        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars(
              inputs=inputs, min=0.0, max=[[1.0], [2.0], [4.0]]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars(
              inputs=inputs, min=[[1.0], [2.0], [4.0]], max=1.0))


class FakeQuantWithMinMaxVarsPerChannelOpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    inputs = constant_op.constant(
        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 1"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel(
              inputs=inputs, min=[[0.0]], max=[1.0]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "Dimensions must be equal|incorrect size"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel(
              inputs=inputs, min=[0.0, 0.1], max=[1.0]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 1"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel(
              inputs=inputs, min=[1.0], max=[[1.0]]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "Dimensions must be equal|incorrect size"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel(
              inputs=inputs, min=[0.0], max=[1.0, 1.1]))


class QuantizedBiasedAddTest(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.qint8)
    bias = constant_op.constant(np.int8(0), shape=[3], dtype=dtypes.qint8)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          nn_ops.quantized_bias_add(
              input=inputs,
              bias=bias,
              min_input=[],
              max_input=1.0,
              min_bias=0.0,
              max_bias=1.0,
              out_type=dtypes.qint32))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          nn_ops.quantized_bias_add(
              input=inputs,
              bias=bias,
              min_input=0.0,
              max_input=[],
              min_bias=0.0,
              max_bias=1.0,
              out_type=dtypes.qint32))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          nn_ops.quantized_bias_add(
              input=inputs,
              bias=bias,
              min_input=0.0,
              max_input=1.0,
              min_bias=[],
              max_bias=1.0,
              out_type=dtypes.qint32))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          nn_ops.quantized_bias_add(
              input=inputs,
              bias=bias,
              min_input=0.0,
              max_input=1.0,
              min_bias=0.0,
              max_bias=[],
              out_type=dtypes.qint32))


class QuantizedInstanceNormOpTest(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)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          array_ops.quantized_instance_norm(
              x=inputs, x_min=0.0, x_max=[[1.0], [2.0], [4.0]]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          array_ops.quantized_instance_norm(
              x=inputs, x_min=[[1.0], [2.0], [4.0]], x_max=1.0))


class RequantizeOpTest(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.requantize(
              input=inputs,
              input_min=[],
              input_max=1.0,
              requested_output_min=0.0,
              requested_output_max=1.0,
              out_type=dtypes.qint8))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          math_ops.requantize(
              input=inputs,
              input_min=0.0,
              input_max=[],
              requested_output_min=0.0,
              requested_output_max=1.0,
              out_type=dtypes.qint8))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          math_ops.requantize(
              input=inputs,
              input_min=0.0,
              input_max=1.0,
              requested_output_min=[],
              requested_output_max=1.0,
              out_type=dtypes.qint8))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          math_ops.requantize(
              input=inputs,
              input_min=0.0,
              input_max=1.0,
              requested_output_min=0.0,
              requested_output_max=[],
              out_type=dtypes.qint8))


if __name__ == "__main__":
  googletest.main()