CVE-2021-37663: Improper Input Validation (CWE-20)
value: TensorFlow is an end-to-end open source platform for machine learning. In affected versions due to incomplete validation in `tf.raw_ops.QuantizeV2`, an attacker can trigger undefined behavior via binding a reference to a null pointer or can access data outside the bounds of heap allocated arrays. The [implementation](https://github.com/tensorflow/tensorflow/blob/84d053187cb80d975ef2b9684d4b61981bca0c41/tensorflow/core/kernels/quantize_op.cc#L59) has some validation but does not check that `min_range` and `max_range` both have the same non-zero number of elements. If `axis` is provided (i.e., not `-1`), then validation should check that it is a value in range for the rank of `input` tensor and then the lengths of `min_range` and `max_range` inputs match the `axis` dimension of the `input` tensor. We have patched the issue in GitHub commit 6da6620efad397c85493b8f8667b821403516708. The fix will be included in TensorFlow 2.6.0. We will also cherrypick this commit on TensorFlow 2.5.1, TensorFlow 2.4.3, and TensorFlow 2.3.4, as these are also affected and still in supported range.
CWE description: The product receives input or data, but it does
        not validate or incorrectly validates that the input has the
        properties that are required to process the data safely and
        correctly.
CVSS v3 base score: 7.8
Published: 2021-08-12T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/quantize_op.cc
Commit date: 2021-07-27 17:27:17 -0700
Commit message: Secure tf.raw_ops.QuantizeV2

Validate size and shape of min_range and max_range
Ensure axis is within input dims limits

PiperOrigin-RevId: 387232799
Change-Id: I36975281f7b5758e9e31a8dcc73fe610ef456318

Diff: diff --git a/tensorflow/core/kernels/quantize_op.cc b/tensorflow/core/kernels/quantize_op.cc
index f64a2188fa954..be73d4f8291f7 100644
--- a/tensorflow/core/kernels/quantize_op.cc
+++ b/tensorflow/core/kernels/quantize_op.cc
@@ -111,11 +111,54 @@ class QuantizeV2Op : public OpKernel {
     const Tensor& input_min_range = ctx->input(1);
     const Tensor& input_max_range = ctx->input(2);
 
     int num_slices = 1;
     if (axis_ > -1) {
+      OP_REQUIRES(
+          ctx, input.dims() > axis_,
+          errors::InvalidArgument(
+              "Axis is on a zero-based index, so its value must always be less "
+              "than number of input's dims, but given axis value was ",
+              axis_, " and input's dims was ", input.dims()));
       num_slices = input.dim_size(axis_);
+      OP_REQUIRES(ctx, input_min_range.dims() == 1,
+                  errors::InvalidArgument(
+                      "If axis is specified, min_range must be a 1-D tensor "
+                      "whose size matches the axis dimension of the input and "
+                      "output tensors, but min_range dims are ",
+                      input_min_range.dims()));
+      OP_REQUIRES(ctx, input_min_range.dim_size(0) == num_slices,
+                  errors::InvalidArgument(
+                      "If axis is specified, min_range must be a 1-D tensor "
+                      "whose size matches the axis dimension of the input and "
+                      "output tensors, but min_range is a 1-D tensor of size ",
+                      input_min_range.dim_size(0),
+                      " and input's axis dimension is of size ", num_slices));
+      OP_REQUIRES(ctx, input_max_range.dims() == 1,
+                  errors::InvalidArgument(
+                      "If axis is specified, max_range must be a 1-D tensor "
+                      "whose size matches the axis dimension of the input and "
+                      "output tensors, but max_range dims are ",
+                      input_max_range.dims()));
+      OP_REQUIRES(ctx, input_max_range.dim_size(0) == num_slices,
+                  errors::InvalidArgument(
+                      "If axis is specified, max_range must be a 1-D tensor "
+                      "whose size matches the axis dimension of the input and "
+                      "output tensors, but max_range is a 1-D tensor of size ",
+                      input_max_range.dim_size(0),
+                      " and input's axis dimension is of size ", num_slices));
+    } else {
+      OP_REQUIRES(ctx, input_min_range.NumElements() == 1,
+                  errors::InvalidArgument(
+                      "If axis is not specified, min_range must contain a "
+                      "single float element, but it contains ",
+                      input_min_range.NumElements(), " elements"));
+      OP_REQUIRES(ctx, input_max_range.NumElements() == 1,
+                  errors::InvalidArgument(
+                      "If axis is not specified, max_range must contain a "
+                      "single float element, but it contains ",
+                      input_max_range.NumElements(), " elements"));
     }
 
     const TensorShape& minmax_shape = ctx->input(1).shape();
     Tensor* output = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input.shape(), &output));

Fixed code:       OP_REQUIRES(
          ctx, input.dims() > axis_,
          errors::InvalidArgument(
              "Axis is on a zero-based index, so its value must always be less "
              "than number of input's dims, but given axis value was ",
              axis_, " and input's dims was ", input.dims()));
      OP_REQUIRES(ctx, input_min_range.dims() == 1,
                  errors::InvalidArgument(
                      "If axis is specified, min_range must be a 1-D tensor "
                      "whose size matches the axis dimension of the input and "
                      "output tensors, but min_range dims are ",
                      input_min_range.dims()));
      OP_REQUIRES(ctx, input_min_range.dim_size(0) == num_slices,
                  errors::InvalidArgument(
                      "If axis is specified, min_range must be a 1-D tensor "
                      "whose size matches the axis dimension of the input and "
                      "output tensors, but min_range is a 1-D tensor of size ",
                      input_min_range.dim_size(0),
                      " and input's axis dimension is of size ", num_slices));
      OP_REQUIRES(ctx, input_max_range.dims() == 1,
                  errors::InvalidArgument(
                      "If axis is specified, max_range must be a 1-D tensor "
                      "whose size matches the axis dimension of the input and "
                      "output tensors, but max_range dims are ",
                      input_max_range.dims()));
      OP_REQUIRES(ctx, input_max_range.dim_size(0) == num_slices,
                  errors::InvalidArgument(
                      "If axis is specified, max_range must be a 1-D tensor "
                      "whose size matches the axis dimension of the input and "
                      "output tensors, but max_range is a 1-D tensor of size ",
                      input_max_range.dim_size(0),
                      " and input's axis dimension is of size ", num_slices));
    } else {
      OP_REQUIRES(ctx, input_min_range.NumElements() == 1,
                  errors::InvalidArgument(
                      "If axis is not specified, min_range must contain a "
                      "single float element, but it contains ",
                      input_min_range.NumElements(), " elements"));
      OP_REQUIRES(ctx, input_max_range.NumElements() == 1,
                  errors::InvalidArgument(
                      "If axis is not specified, max_range must contain a "
                      "single float element, but it contains ",
                      input_max_range.NumElements(), " elements"));