CVE-2021-37665: 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 MKL implementation of requantization, 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/460e000de3a83278fb00b61a16d161b1964f15f4/tensorflow/core/kernels/mkl/mkl_requantization_range_per_channel_op.cc) does not validate the dimensions of the `input` tensor. A similar issue occurs in `MklRequantizePerChannelOp`. The [implementation](https://github.com/tensorflow/tensorflow/blob/460e000de3a83278fb00b61a16d161b1964f15f4/tensorflow/core/kernels/mkl/mkl_requantize_per_channel_op.cc) does not perform full validation for all the input arguments. We have patched the issue in GitHub commit 9e62869465573cb2d9b5053f1fa02a81fce21d69 and in the Github commit 203214568f5bc237603dbab6e1fd389f1572f5c9. 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/mkl/mkl_requantize_per_channel_op.cc
Commit date: 2021-07-30 16:11:14 -0700
Commit message: Reorganize and add more validation to MKL requantization

PiperOrigin-RevId: 387901341
Change-Id: I2515b9034c64e113db0bcec8337d30643ab0a0f1

Diff: diff --git a/tensorflow/core/kernels/mkl/mkl_requantize_per_channel_op.cc b/tensorflow/core/kernels/mkl/mkl_requantize_per_channel_op.cc
index c0f9845cd4b08..6ffbd09b44f54 100644
--- a/tensorflow/core/kernels/mkl/mkl_requantize_per_channel_op.cc
+++ b/tensorflow/core/kernels/mkl/mkl_requantize_per_channel_op.cc
@@ -47,39 +47,49 @@ class MklRequantizePerChannelOp : public OpKernel {
   }
   virtual ~MklRequantizePerChannelOp() {}
   void Compute(OpKernelContext* ctx) override {
     try {
       const Tensor& input = ctx->input(kInputTensorIndex);
+      OP_REQUIRES(
+          ctx, input.dims() == 4,
+          errors::InvalidArgument("Current RequantizePerChannel operator"
+                                  "supports 4D tensors only."));
+
       const Tensor& input_min_vec = ctx->input(kInputMinVecIndex);
+      size_t depth = input_min_vec.NumElements();
       float* input_min_vec_data = (float*)const_cast<void*>(
           static_cast<const void*>(input_min_vec.flat<float>().data()));
+
       const Tensor& input_max_vec = ctx->input(kInputMaxVecIndex);
+      OP_REQUIRES(
+          ctx, input_max_vec.NumElements() == depth,
+          errors::InvalidArgument("input_max has incorrect size, expected ",
+                                  depth, " was ", input_max_vec.NumElements()));
       float* input_max_vec_data = (float*)const_cast<void*>(
           static_cast<const void*>(input_max_vec.flat<float>().data()));
 
       const Tensor& input_requested_min = ctx->input(this->kRequestMinIndex);
+      OP_REQUIRES(
+          ctx, input_requested_min.NumElements() == 1,
+          errors::InvalidArgument("requested_output_min must be a scalar"));
       const float input_requested_min_float =
           input_requested_min.flat<float>()(0);
+
       const Tensor& input_requested_max = ctx->input(this->kRequestMaxIndex);
+      OP_REQUIRES(
+          ctx, input_requested_min.NumElements() == 1,
+          errors::InvalidArgument("requested_output_max must be a scalar"));
       const float input_requested_max_float =
           input_requested_max.flat<float>()(0);
 
-      size_t depth = input_min_vec.NumElements();
-      OP_REQUIRES(
-          ctx, input.dims() == 4,
-          errors::InvalidArgument("Current RequantizePerChannel operator"
-                                  "supports 4D tensors only."));
-      OP_REQUIRES(
-          ctx, input_min_vec.dim_size(0) == depth,
-          errors::InvalidArgument("input_min has incorrect size, expected ",
-                                  depth, " was ", input_min_vec.dim_size(0)));
-      OP_REQUIRES(
-          ctx, input_max_vec.dim_size(0) == depth,
-          errors::InvalidArgument("input_max has incorrect size, expected ",
-                                  depth, " was ", input_max_vec.dim_size(0)));
-
-      if (out_type_ == DT_QINT8) DCHECK(input_requested_min_float < 0.0f);
+      if (out_type_ == DT_QINT8) {
+        OP_REQUIRES(ctx, input_requested_min_float < 0.0f,
+                    errors::InvalidArgument(
+                        "If out_type is QINT8, requested_output_max must be "
+                        "non negative, got ",
+                        input_requested_min_float));
+      }
 
       const float factor = (out_type_ == DT_QINT8) ? 127.0f : 255.0f;
       const float requested_min_max =
           std::max(std::abs(input_requested_min_float),
                    std::abs(input_requested_max_float));

Vulnerable code:       size_t depth = input_min_vec.NumElements();
      OP_REQUIRES(
          ctx, input.dims() == 4,
          errors::InvalidArgument("Current RequantizePerChannel operator"
                                  "supports 4D tensors only."));
      OP_REQUIRES(
          ctx, input_min_vec.dim_size(0) == depth,
          errors::InvalidArgument("input_min has incorrect size, expected ",
                                  depth, " was ", input_min_vec.dim_size(0)));
      OP_REQUIRES(
          ctx, input_max_vec.dim_size(0) == depth,
          errors::InvalidArgument("input_max has incorrect size, expected ",
                                  depth, " was ", input_max_vec.dim_size(0)));

      if (out_type_ == DT_QINT8) DCHECK(input_requested_min_float < 0.0f);
Fixed code:       OP_REQUIRES(
          ctx, input.dims() == 4,
          errors::InvalidArgument("Current RequantizePerChannel operator"
                                  "supports 4D tensors only."));

      size_t depth = input_min_vec.NumElements();

      OP_REQUIRES(
          ctx, input_max_vec.NumElements() == depth,
          errors::InvalidArgument("input_max has incorrect size, expected ",
                                  depth, " was ", input_max_vec.NumElements()));
      OP_REQUIRES(
          ctx, input_requested_min.NumElements() == 1,
          errors::InvalidArgument("requested_output_min must be a scalar"));

      OP_REQUIRES(
          ctx, input_requested_min.NumElements() == 1,
          errors::InvalidArgument("requested_output_max must be a scalar"));
      if (out_type_ == DT_QINT8) {
        OP_REQUIRES(ctx, input_requested_min_float < 0.0f,
                    errors::InvalidArgument(
                        "If out_type is QINT8, requested_output_max must be "
                        "non negative, got ",
                        input_requested_min_float));
      }