Vulnerability: Heap buffer overflow in `QuantizedReshape` ### Impact
An attacker can cause a heap buffer overflow in `QuantizedReshape` by passing in invalid thresholds for the quantization:

```python
import tensorflow as tf

tensor = tf.constant([], dtype=tf.qint32)
shape = tf.constant([], dtype=tf.int32)
input_min = tf.constant([], dtype=tf.float32)
input_max = tf.constant([], dtype=tf.float32)

tf.raw_ops.QuantizedReshape(tensor=tensor, shape=shape, input_min=input_min, input_max=input_max)
```

This is because the [implementation](https://github.com/tensorflow/tensorflow/blob/a324ac84e573fba362a5e53d4e74d5de6729933e/tensorflow/core/kernels/quantized_reshape_op.cc#L38-L55) assumes that the 2 arguments are always valid scalars and tries to access the numeric value directly:

```cc
const auto& input_min_float_tensor = ctx->input(2);
...
const float input_min_float = input_min_float_tensor.flat<float>()(0);
const auto& input_max_float_tensor = ctx->input(3);
...
const float input_max_float = input_max_float_tensor.flat<float>()(0);
```

However, if any of these tensors is empty, then `.flat<T>()` is an empty buffer and accessing the element at position 0 results in overflow.

### Patches
We have patched the issue in GitHub commit [a324ac84e573fba362a5e53d4e74d5de6729933e](https://github.com/tensorflow/tensorflow/commit/a324ac84e573fba362a5e53d4e74d5de6729933e).

The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range.

### For more information
Please consult [our security guide](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for more information regarding the security model and how to contact us with issues and questions.

### Attribution
This vulnerability has been reported by Ying Wang and Yakun Zhang of Baidu X-Team.
Vulnerability ID: GHSA-2gfx-95x2-5v3x
CWE: CWE-787
Severity score: 2.5
Project: https://github.com/tensorflow/tensorflow
Published: 2021-05-21
File: tensorflow/core/kernels/quantized_reshape_op.cc
Commit message: Validate arguments to `QuantizedReshape`.

Ensure that validations from `Reshape` also terminate `QuantizedReshape` on failure.

PiperOrigin-RevId: 369775421
Change-Id: If8c5342267aceea65b7cb83a4b183304886f1ce8
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 limitations under the License.  #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/kernels/reshape_op.h" class QuantizedReshapeOp : public ReshapeOp {   void Compute(OpKernelContext* ctx) override {     // This call processes inputs 1 and 2 to write output 0.     ReshapeOp::Compute(ctx);      const float input_min_float = ctx->input(2).flat<float>()(0);     const float input_max_float = ctx->input(3).flat<float>()(0);     Tensor* output_min = nullptr;     OP_REQUIRES_OK(ctx, ctx->allocate_output(1, TensorShape({}), &output_min));     output_min->flat<float>()(0) = input_min_float;

Fixed version (after fix):
 limitations under the License.  #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/kernels/reshape_op.h" class QuantizedReshapeOp : public ReshapeOp {   void Compute(OpKernelContext* ctx) override {     // This call processes inputs 1 and 2 to write output 0.     ReshapeOp::Compute(ctx);     if (!ctx->status().ok()) {       return;     }      const auto& input_min_float_tensor = ctx->input(2);     const auto& input_min_float_shape = input_min_float_tensor.shape();     OP_REQUIRES(ctx,                 TensorShapeUtils::IsScalar(input_min_float_shape) ||                     (TensorShapeUtils::IsVector(input_min_float_shape) &&                      (input_min_float_shape.dim_size(0) == 1)),                 errors::InvalidArgument(                     "input_min must be a scalar or a vector of 1 element"));     const float input_min_float = input_min_float_tensor.flat<float>()(0);     const auto& input_max_float_tensor = ctx->input(3);     const auto& input_max_float_shape = input_max_float_tensor.shape();     OP_REQUIRES(ctx,                 TensorShapeUtils::IsScalar(input_max_float_shape) ||                     (TensorShapeUtils::IsVector(input_max_float_shape) &&                      (input_max_float_shape.dim_size(0) == 1)),                 errors::InvalidArgument(                     "input_max must be a scalar or a vector of 1 element"));     const float input_max_float = input_max_float_tensor.flat<float>()(0);      Tensor* output_min = nullptr;     OP_REQUIRES_OK(ctx, ctx->allocate_output(1, TensorShape({}), &output_min));     output_min->flat<float>()(0) = input_min_float;

Patch:
patch: @@ -17,6 +17,7 @@ limitations under the License.
 
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/reshape_op.h"
@@ -30,9 +31,29 @@ class QuantizedReshapeOp : public ReshapeOp {
   void Compute(OpKernelContext* ctx) override {
     // This call processes inputs 1 and 2 to write output 0.
     ReshapeOp::Compute(ctx);
+    if (!ctx->status().ok()) {
+      return;
+    }
+
+    const auto& input_min_float_tensor = ctx->input(2);
+    const auto& input_min_float_shape = input_min_float_tensor.shape();
+    OP_REQUIRES(ctx,
+                TensorShapeUtils::IsScalar(input_min_float_shape) ||
+                    (TensorShapeUtils::IsVector(input_min_float_shape) &&
+                     (input_min_float_shape.dim_size(0) == 1)),
+                errors::InvalidArgument(
+                    "input_min must be a scalar or a vector of 1 element"));
+    const float input_min_float = input_min_float_tensor.flat<float>()(0);
+    const auto& input_max_float_tensor = ctx->input(3);
+    const auto& input_max_float_shape = input_max_float_tensor.shape();
+    OP_REQUIRES(ctx,
+                TensorShapeUtils::IsScalar(input_max_float_shape) ||
+                    (TensorShapeUtils::IsVector(input_max_float_shape) &&
+                     (input_max_float_shape.dim_size(0) == 1)),
+                errors::InvalidArgument(
+                    "input_max must be a scalar or a vector of 1 element"));
+    const float input_max_float = input_max_float_tensor.flat<float>()(0);
 
-    const float input_min_float = ctx->input(2).flat<float>()(0);
-    const float input_max_float = ctx->input(3).flat<float>()(0);
     Tensor* output_min = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(1, TensorShape({}), &output_min));
     output_min->flat<float>()(0) = input_min_float;
path: tensorflow/core/kernels/quantized_reshape_op.cc