CVE-2022-36026 (Python, published 2022-09-16T22:15Z)
Description: value: TensorFlow is an open source platform for machine learning. If `QuantizeAndDequantizeV3` is given a nonscalar `num_bits` input tensor, it results in a `CHECK` fail that can be used to trigger a denial of service attack. We have patched the issue in GitHub commit f3f9cb38ecfe5a8a703f2c4a8fead434ef291713. 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: CWE-617 \u2014 Reachable Assertion
The product contains an assert() or similar statement that can be triggered by an attacker, which leads to an application exit or other behavior that is more severe than necessary.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/kernels/quantize_and_dequantize_op.cc
tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
Commit message: Validate the rank and number of elements of the `num_bits` tensor for QuantizeAndDequantizeV3.

QuantizeAndDequantizeV3Op, which accepts `num_bits` as a tensor, has a precondition that it should be rank <= 1 and the number of elements should be 1.
This change adds a validation for the Compute() method for this condition.

PiperOrigin-RevId: 463755293


Vulnerable code:
#include "tensorflow/core/kernels/quantize_and_dequantize_op.h"

typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
                errors::InvalidArgument("num_bits is out of range: ", num_bits_,
                                        " with signed_input_ ", signed_input_));
        errors::InvalidArgument("Round mode string must be "
                                "'HALF_UP' or "
                                "'HALF_TO_EVEN', is '" +
                                round_mode_string + "'"));
    OP_REQUIRES(
        ctx, axis_ >= -1,
        errors::InvalidArgument("Axis must be at least -1. Found ", axis_));
    OP_REQUIRES(
        ctx, (axis_ == -1 || axis_ < input.shape().dims()),
        errors::InvalidArgument("Shape must be at least rank ", axis_ + 1,
                    errors::InvalidArgument("Invalid range: input_min ",
                                            min_val, " > input_max ", max_val));
        OP_REQUIRES(ctx, input_min_tensor.dim_size(0) == depth,
                    errors::InvalidArgument(
                        "input_min_tensor has incorrect size, was ",
                        input_min_tensor.dim_size(0), " expected ", depth,
                        " to match dim ", axis_, " of the input ",
                        input_min_tensor.shape()));
        OP_REQUIRES(ctx, input_max_tensor.dim_size(0) == depth,
                    errors::InvalidArgument(
                        "input_max_tensor has incorrect size, was ",
                        input_max_tensor.dim_size(0), " expected ", depth,
                        " to match dim ", axis_, " of the input ",
                        input_max_tensor.shape()));
    OP_REQUIRES(
        ctx, axis_ >= -1,
        errors::InvalidArgument("Axis must be at least -1. Found ", axis_));
                errors::InvalidArgument(
    OP_REQUIRES(
        ctx, input.IsSameSize(gradient),
        errors::InvalidArgument("gradient and input must be the same size"));
                errors::InvalidArgument(
                errors::InvalidArgument(
      OP_REQUIRES(
          ctx, input_min_tensor.dim_size(0) == depth,
          errors::InvalidArgument("min has incorrect size, expected ", depth,
      OP_REQUIRES(
          ctx, input_max_tensor.dim_size(0) == depth,
          errors::InvalidArgument("max has incorrect size, expected ", depth,
      OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(input_min_tensor.shape()),
                  errors::InvalidArgument(
                      "input_min must be a scalar if axis is unspecified"));
      OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(input_max_tensor.shape()),
                  errors::InvalidArgument(
                      "input_max must be a scalar if axis is unspecified"));
                errors::InvalidArgument(
    Tensor num_bits_tensor;
    num_bits_tensor = ctx->input(3);
    int num_bits_val = num_bits_tensor.scalar<int32>()();
    OP_REQUIRES(
        ctx, num_bits_val > 0 && num_bits_val < (signed_input_ ? 62 : 63),
        errors::InvalidArgument("num_bits is out of range: ", num_bits_val,
                                " with signed_input_ ", signed_input_));
        auto min_val = input_min_tensor.scalar<T>()();
        auto max_val = input_max_tensor.scalar<T>()();
                    errors::InvalidArgument("Invalid range: input_min ",
                                            min_val, " > input_max ", max_val));
        OP_REQUIRES(ctx, input_min_tensor.dim_size(0) == depth,
                    errors::InvalidArgument(
                        "input_min_tensor has incorrect size, was ",
                        input_min_tensor.dim_size(0), " expected ", depth,
                        " to match dim ", axis_, " of the input ",
                        input_min_tensor.shape()));
        OP_REQUIRES(ctx, input_max_tensor.dim_size(0) == depth,
                    errors::InvalidArgument(
                        "input_max_tensor has incorrect size, was ",
                        input_max_tensor.dim_size(0), " expected ", depth,
                        " to match dim ", axis_, " of the input ",
                        input_max_tensor.shape()));
                errors::InvalidArgument("num_bits is out of range: ", num_bits_,
                                        " with signed_input_ ", signed_input_));
      OP_REQUIRES(
          ctx, input_min_ <= input_max_,
          errors::InvalidArgument("Invalid range: input_min ", input_min_,
// Specializations for CPUDevice.
struct QuantizeAndDequantizeOneScaleFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d, typename TTypes<T>::ConstVec input,
    QuantizeAndDequantizeOneScaleImpl<CPUDevice, T>::Compute(
struct QuantizeAndDequantizePerChannelFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d, typename TTypes<T, 3>::ConstTensor input,
    QuantizeAndDequantizePerChannelImpl<CPUDevice, T>::Compute(
struct QuantizeAndDequantizeOneScaleGradientFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d, typename TTypes<T>::ConstFlat gradient,
    QuantizeAndDequantizeOneScaleGradientImpl<CPUDevice, T>::Compute(
struct QuantizeAndDequantizePerChannelGradientFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d,
    QuantizeAndDequantizePerChannelGradientImpl<CPUDevice, T>::Compute(
template struct functor::QuantizeAndDequantizeOneScaleGradientFunctor<CPUDevice,
    CPUDevice, double>;
                          QuantizeAndDequantizeV2Op<CPUDevice, T>);            \
                          QuantizeAndDequantizeV3Op<CPUDevice, T>);            \
                          QuantizeAndDequantizeV2Op<CPUDevice, T>);            \
                          QuantizeAndDequantizeV4GradientOp<CPUDevice, T>);    \
      QuantizeAndDequantizeOp<CPUDevice, T>);
                          QuantizeAndDequantizeV2Op<GPUDevice, T>);            \
                          QuantizeAndDequantizeV3Op<GPUDevice, T>);            \
                          QuantizeAndDequantizeV2Op<GPUDevice, T>);            \
                          QuantizeAndDequantizeV4GradientOp<GPUDevice, T>);    \
      QuantizeAndDequantizeOp<GPUDevice, T>);

Fixed code:
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/kernels/quantize_and_dequantize_op.h"
namespace {
using CpuDevice = ::Eigen::ThreadPoolDevice;
using GpuDevice = ::Eigen::GpuDevice;
using ::tensorflow::errors::InvalidArgument;

}  // namespace
                InvalidArgument("num_bits is out of range: ", num_bits_,
                                " with signed_input_ ", signed_input_));
        InvalidArgument("Round mode string must be "
                        "'HALF_UP' or "
                        "'HALF_TO_EVEN', is '" +
                        round_mode_string + "'"));
    OP_REQUIRES(ctx, axis_ >= -1,
                InvalidArgument("Axis must be at least -1. Found ", axis_));
    OP_REQUIRES(ctx, (axis_ == -1 || axis_ < input.shape().dims()),
                InvalidArgument("Shape must be at least rank ", axis_ + 1,
                    InvalidArgument("Invalid range: input_min ", min_val,
                                    " > input_max ", max_val));
        OP_REQUIRES(
            ctx, input_min_tensor.dim_size(0) == depth,
            InvalidArgument("input_min_tensor has incorrect size, was ",
                            input_min_tensor.dim_size(0), " expected ", depth,
                            " to match dim ", axis_, " of the input ",
                            input_min_tensor.shape()));
        OP_REQUIRES(
            ctx, input_max_tensor.dim_size(0) == depth,
            InvalidArgument("input_max_tensor has incorrect size, was ",
                            input_max_tensor.dim_size(0), " expected ", depth,
                            " to match dim ", axis_, " of the input ",
                            input_max_tensor.shape()));
    OP_REQUIRES(ctx, axis_ >= -1,
                InvalidArgument("Axis must be at least -1. Found ", axis_));
                InvalidArgument(
    OP_REQUIRES(ctx, input.IsSameSize(gradient),
                InvalidArgument("gradient and input must be the same size"));
                InvalidArgument(
                InvalidArgument(
      OP_REQUIRES(ctx, input_min_tensor.dim_size(0) == depth,
                  InvalidArgument("min has incorrect size, expected ", depth,
      OP_REQUIRES(ctx, input_max_tensor.dim_size(0) == depth,
                  InvalidArgument("max has incorrect size, expected ", depth,
      OP_REQUIRES(
          ctx, TensorShapeUtils::IsScalar(input_min_tensor.shape()),
          InvalidArgument("input_min must be a scalar if axis is unspecified"));
      OP_REQUIRES(
          ctx, TensorShapeUtils::IsScalar(input_max_tensor.shape()),
          InvalidArgument("input_max must be a scalar if axis is unspecified"));
                InvalidArgument(
    // Get num_bits and validate.
    const Tensor num_bits_tensor = ctx->input(3);
    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(num_bits_tensor.shape()),
                InvalidArgument("Invalid shape. The `num_bits` tensor should "
                                "be a scalar. Got dimensions: ",
                                num_bits_tensor.dims()));
    const int num_bits_val = num_bits_tensor.scalar<int32>()();
    OP_REQUIRES(ctx,
                num_bits_val > 0 && num_bits_val < (signed_input_ ? 62 : 63),
                InvalidArgument("num_bits is out of range: ", num_bits_val,
                                " with `signed_input_` ", signed_input_));
        const auto min_val = input_min_tensor.scalar<T>()();
        const auto max_val = input_max_tensor.scalar<T>()();
                    InvalidArgument("Invalid range: input_min ", min_val,
                                    " > input_max ", max_val));
        OP_REQUIRES(
            ctx, input_min_tensor.dim_size(0) == depth,
            InvalidArgument("input_min_tensor has incorrect size, was ",
                            input_min_tensor.dim_size(0), " expected ", depth,
                            " to match dim ", axis_, " of the input ",
                            input_min_tensor.shape()));
        OP_REQUIRES(
            ctx, input_max_tensor.dim_size(0) == depth,
            InvalidArgument("input_max_tensor has incorrect size, was ",
                            input_max_tensor.dim_size(0), " expected ", depth,
                            " to match dim ", axis_, " of the input ",
                            input_max_tensor.shape()));
                InvalidArgument("num_bits is out of range: ", num_bits_,
                                " with signed_input_ ", signed_input_));
      OP_REQUIRES(ctx, input_min_ <= input_max_,
                  InvalidArgument("Invalid range: input_min ", input_min_,
// Specializations for CpuDevice.
struct QuantizeAndDequantizeOneScaleFunctor<CpuDevice, T> {
  void operator()(const CpuDevice& d, typename TTypes<T>::ConstVec input,
    QuantizeAndDequantizeOneScaleImpl<CpuDevice, T>::Compute(
struct QuantizeAndDequantizePerChannelFunctor<CpuDevice, T> {
  void operator()(const CpuDevice& d, typename TTypes<T, 3>::ConstTensor input,
    QuantizeAndDequantizePerChannelImpl<CpuDevice, T>::Compute(
struct QuantizeAndDequantizeOneScaleGradientFunctor<CpuDevice, T> {
  void operator()(const CpuDevice& d, typename TTypes<T>::ConstFlat gradient,
    QuantizeAndDequantizeOneScaleGradientImpl<CpuDevice, T>::Compute(
struct QuantizeAndDequantizePerChannelGradientFunctor<CpuDevice, T> {
  void operator()(const CpuDevice& d,
    QuantizeAndDequantizePerChannelGradientImpl<CpuDevice, T>::Compute(
template struct functor::QuantizeAndDequantizeOneScaleGradientFunctor<CpuDevice,
    CpuDevice, double>;
                          QuantizeAndDequantizeV2Op<CpuDevice, T>);            \
                          QuantizeAndDequantizeV3Op<CpuDevice, T>);            \
                          QuantizeAndDequantizeV2Op<CpuDevice, T>);            \
                          QuantizeAndDequantizeV4GradientOp<CpuDevice, T>);    \
      QuantizeAndDequantizeOp<CpuDevice, T>);
                          QuantizeAndDequantizeV2Op<GpuDevice, T>);            \
                          QuantizeAndDequantizeV3Op<GpuDevice, T>);            \
                          QuantizeAndDequantizeV2Op<GpuDevice, T>);            \
                          QuantizeAndDequantizeV4GradientOp<GpuDevice, T>);    \
      QuantizeAndDequantizeOp<GpuDevice, T>);
from tensorflow.python.eager import context
from tensorflow.python.framework import ops
class QuantizeAndDequantizeV3OpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_valid(self):
    with ops.Graph().as_default(), context.eager_mode():
      input_value = constant_op.constant([-0.8, -0.5, 0, 0.3, 0.8, -2.0],
                                         shape=(6,),
                                         dtype=dtypes.float32),
      input_min = constant_op.constant(-127, shape=(), dtype=dtypes.float32)
      input_max = constant_op.constant(127, shape=(), dtype=dtypes.float32)
      num_bits = constant_op.constant(8, shape=(), dtype=dtypes.int32)

      quantized = array_ops.quantize_and_dequantize_v3(
          input_value,
          input_min,
          input_max,
          num_bits,
          signed_input=True,
          range_given=False)
      self.assertSequenceAlmostEqual(
          input_value[0].numpy(), quantized.numpy()[0], delta=0.05)

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    input_value = constant_op.constant([-0.8, -0.5, 0, 0.3, 0.8, -2.0],
                                       shape=(6,),
                                       dtype=dtypes.float32),
    input_min = constant_op.constant(-127, shape=(), dtype=dtypes.float32)
    input_max = constant_op.constant(127, shape=(), dtype=dtypes.float32)
    # Tensor with invalid shape and invalid number of elements.
    num_bits = constant_op.constant([], shape=(0,), dtype=dtypes.int32)

    # Test that running the op raises error. It raises different errors
    # depending on whether the shape inference is run first or the op's
    # Compute() is run first.
    try:
      array_ops.quantize_and_dequantize_v3(
          input_value, input_min, input_max, num_bits, signed_input=True)
    except Exception as ex:  # pylint: disable=broad-except
      if isinstance(ex, errors.InvalidArgumentError):
        self.assertRegex(str(ex), "The `num_bits` tensor should be a scalar.")
      elif isinstance(ex, ValueError):
        self.assertRegex(str(ex), "Shape must be rank 0")
      else:
        self.fail(
            "Raised exception other than expected: %s. "
            "Expected exceptions are errors.InvalidArgumentError or ValueError",
            ex.__name__)
    else:
      self.fail(
          "Did not raise an exception where it is expected to raise either "
          "a ValueError or errors.InvalidArgumentError.")