Vulnerability: Heap out of bounds in `QuantizedBatchNormWithGlobalNormalization` ### Impact
An attacker can cause a segfault and denial of service via accessing data outside of bounds in `tf.raw_ops.QuantizedBatchNormWithGlobalNormalization`:

```python
import tensorflow as tf

t = tf.constant([1], shape=[1, 1, 1, 1], dtype=tf.quint8)
t_min = tf.constant([], shape=[0], dtype=tf.float32)
t_max = tf.constant([], shape=[0], dtype=tf.float32)
m = tf.constant([1], shape=[1], dtype=tf.quint8)
m_min = tf.constant([], shape=[0], dtype=tf.float32)
m_max = tf.constant([], shape=[0], dtype=tf.float32)
v = tf.constant([1], shape=[1], dtype=tf.quint8)
v_min = tf.constant([], shape=[0], dtype=tf.float32)
v_max = tf.constant([], shape=[0], dtype=tf.float32)
beta = tf.constant([1], shape=[1], dtype=tf.quint8)
beta_min = tf.constant([], shape=[0], dtype=tf.float32)
beta_max = tf.constant([], shape=[0], dtype=tf.float32)
gamma = tf.constant([1], shape=[1], dtype=tf.quint8)
gamma_min = tf.constant([], shape=[0], dtype=tf.float32)
gamma_max = tf.constant([], shape=[0], dtype=tf.float32) 

tf.raw_ops.QuantizedBatchNormWithGlobalNormalization(
  t=t, t_min=t_min, t_max=t_max, m=m, m_min=m_min, m_max=m_max,
  v=v, v_min=v_min, v_max=v_max, beta=beta, beta_min=beta_min,
  beta_max=beta_max, gamma=gamma, gamma_min=gamma_min,
  gamma_max=gamma_max, out_type=tf.qint32,
  variance_epsilon=0.1, scale_after_normalization=True)
```                         
                            
This is because the [implementation](https://github.com/tensorflow/tensorflow/blob/55a97caa9e99c7f37a0bbbeb414dc55553d3ae7f/tensorflow/core/kernels/quantized_batch_norm_op.cc#L176-L189) assumes the inputs are not empty: 
  
```cc
const float input_min = context->input(1).flat<float>()(0);
const float input_max = context->input(2).flat<float>()(0);
...
const float mean_min = context->input(4).flat<float>()(0);
const float mean_max = context->input(5).flat<float>()(0);
...
const float var_min = context->input(7).flat<float>()(0);
const float var_max = context->input(8).flat<float>()(0);
...
const float beta_min = context->input(10).flat<float>()(0);
const float beta_max = context->input(11).flat<float>()(0);
...
const float gamma_min = context->input(13).flat<float>()(0);
const float gamma_max = context->input(14).flat<float>()(0);
```

If any of these inputs is empty, `.flat<T>()` is an empty buffer, so accessing the element at index 0 is accessing data outside of bounds.

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

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 Yakun Zhang and Ying Wang of Baidu X-Team.
Vulnerability ID: GHSA-4fg4-p75j-w5xj
CWE: CWE-125
Severity score: 2.5
Project: https://github.com/tensorflow/tensorflow
Published: 2021-05-21
File: tensorflow/core/kernels/quantized_batch_norm_op.cc
Commit message: Add missing validation in `QuantizedBatchNormWithGlobalNormalization`

PiperOrigin-RevId: 370123451
Change-Id: Id234d6dab1ec21230bb8e503dba30f899af87f33
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 class QuantizedBatchNormOp : public OpKernel {    void Compute(OpKernelContext* context) override {     const Tensor& input = context->input(0);     const float input_min = context->input(1).flat<float>()(0);     const float input_max = context->input(2).flat<float>()(0);     const Tensor& mean = context->input(3);     const float mean_min = context->input(4).flat<float>()(0);     const float mean_max = context->input(5).flat<float>()(0);     const Tensor& var = context->input(6);     const float var_min = context->input(7).flat<float>()(0);     const float var_max = context->input(8).flat<float>()(0);     const Tensor& beta = context->input(9);     const float beta_min = context->input(10).flat<float>()(0);     const float beta_max = context->input(11).flat<float>()(0);     const Tensor& gamma = context->input(12);     const float gamma_min = context->input(13).flat<float>()(0);     const float gamma_max = context->input(14).flat<float>()(0);      OP_REQUIRES(context, input.dims() == 4,                 errors::InvalidArgument("input must be 4-dimensional", class QuantizedBatchNormOp : public OpKernel {     OP_REQUIRES(context, gamma.dims() == 1,                 errors::InvalidArgument("gamma must be 1-dimensional",                                         gamma.shape().DebugString()));      Tensor* output = nullptr;     OP_REQUIRES_OK(context,

Fixed version (after fix):
 class QuantizedBatchNormOp : public OpKernel {    void Compute(OpKernelContext* context) override {     const Tensor& input = context->input(0);     const auto& input_min_tensor = context->input(1);     OP_REQUIRES(context, input_min_tensor.NumElements() == 1,                 errors::InvalidArgument("input_min must have 1 element"));     const float input_min = input_min_tensor.flat<float>()(0);     const auto& input_max_tensor = context->input(2);     OP_REQUIRES(context, input_max_tensor.NumElements() == 1,                 errors::InvalidArgument("input_max must have 1 element"));     const float input_max = input_max_tensor.flat<float>()(0);     const Tensor& mean = context->input(3);     const auto& mean_min_tensor = context->input(4);     OP_REQUIRES(context, mean_min_tensor.NumElements() == 1,                 errors::InvalidArgument("mean_min must have 1 element"));     const float mean_min = mean_min_tensor.flat<float>()(0);     const auto& mean_max_tensor = context->input(5);     OP_REQUIRES(context, mean_max_tensor.NumElements() == 1,                 errors::InvalidArgument("mean_max must have 1 element"));     const float mean_max = mean_max_tensor.flat<float>()(0);     const Tensor& var = context->input(6);     const auto& var_min_tensor = context->input(7);     OP_REQUIRES(context, var_min_tensor.NumElements() == 1,                 errors::InvalidArgument("var_min must have 1 element"));     const float var_min = var_min_tensor.flat<float>()(0);     const auto& var_max_tensor = context->input(8);     OP_REQUIRES(context, var_max_tensor.NumElements() == 1,                 errors::InvalidArgument("var_max must have 1 element"));     const float var_max = var_max_tensor.flat<float>()(0);     const Tensor& beta = context->input(9);     const auto& beta_min_tensor = context->input(10);     OP_REQUIRES(context, beta_min_tensor.NumElements() == 1,                 errors::InvalidArgument("beta_min must have 1 element"));     const float beta_min = beta_min_tensor.flat<float>()(0);     const auto& beta_max_tensor = context->input(11);     OP_REQUIRES(context, beta_max_tensor.NumElements() == 1,                 errors::InvalidArgument("beta_max must have 1 element"));     const float beta_max = beta_max_tensor.flat<float>()(0);     const Tensor& gamma = context->input(12);     const auto& gamma_min_tensor = context->input(13);     OP_REQUIRES(context, gamma_min_tensor.NumElements() == 1,                 errors::InvalidArgument("gamma_min must have 1 element"));     const float gamma_min = gamma_min_tensor.flat<float>()(0);     const auto& gamma_max_tensor = context->input(14);     OP_REQUIRES(context, gamma_max_tensor.NumElements() == 1,                 errors::InvalidArgument("gamma_max must have 1 element"));     const float gamma_max = gamma_max_tensor.flat<float>()(0);      OP_REQUIRES(context, input.dims() == 4,                 errors::InvalidArgument("input must be 4-dimensional", class QuantizedBatchNormOp : public OpKernel {     OP_REQUIRES(context, gamma.dims() == 1,                 errors::InvalidArgument("gamma must be 1-dimensional",                                         gamma.shape().DebugString()));     OP_REQUIRES(context, mean.NumElements() > 1,                 errors::InvalidArgument("Must have at least a mean value",                                         gamma.shape().DebugString()));     OP_REQUIRES(context, mean.NumElements() > 1,                 errors::InvalidArgument("Must have at least a mean value"));     const auto last_dim = input.shape().dims() - 1;     OP_REQUIRES(context,                 mean.shape().dim_size(0) == input.shape().dim_size(last_dim),                 errors::InvalidArgument("Must provide as many means as the "                                         "last dimension of the input tensor: ",                                         mean.shape().DebugString(), " vs. ",                                         input.shape().DebugString()));     OP_REQUIRES(         context, mean.shape().dim_size(0) == var.shape().dim_size(0),         errors::InvalidArgument(             "Mean and variance tensors must have the same shape: ",             mean.shape().DebugString(), " vs. ", var.shape().DebugString()));     OP_REQUIRES(         context, mean.shape().dim_size(0) == beta.shape().dim_size(0),         errors::InvalidArgument(             "Mean and beta tensors must have the same shape: ",             mean.shape().DebugString(), " vs. ", beta.shape().DebugString()));     OP_REQUIRES(         context, mean.shape().dim_size(0) == gamma.shape().dim_size(0),         errors::InvalidArgument(             "Mean and gamma tensors must have the same shape: ",             mean.shape().DebugString(), " vs. ", gamma.shape().DebugString()));      Tensor* output = nullptr;     OP_REQUIRES_OK(context,

Patch:
patch: @@ -173,20 +173,50 @@ class QuantizedBatchNormOp : public OpKernel {
 
   void Compute(OpKernelContext* context) override {
     const Tensor& input = context->input(0);
-    const float input_min = context->input(1).flat<float>()(0);
-    const float input_max = context->input(2).flat<float>()(0);
+    const auto& input_min_tensor = context->input(1);
+    OP_REQUIRES(context, input_min_tensor.NumElements() == 1,
+                errors::InvalidArgument("input_min must have 1 element"));
+    const float input_min = input_min_tensor.flat<float>()(0);
+    const auto& input_max_tensor = context->input(2);
+    OP_REQUIRES(context, input_max_tensor.NumElements() == 1,
+                errors::InvalidArgument("input_max must have 1 element"));
+    const float input_max = input_max_tensor.flat<float>()(0);
     const Tensor& mean = context->input(3);
-    const float mean_min = context->input(4).flat<float>()(0);
-    const float mean_max = context->input(5).flat<float>()(0);
+    const auto& mean_min_tensor = context->input(4);
+    OP_REQUIRES(context, mean_min_tensor.NumElements() == 1,
+                errors::InvalidArgument("mean_min must have 1 element"));
+    const float mean_min = mean_min_tensor.flat<float>()(0);
+    const auto& mean_max_tensor = context->input(5);
+    OP_REQUIRES(context, mean_max_tensor.NumElements() == 1,
+                errors::InvalidArgument("mean_max must have 1 element"));
+    const float mean_max = mean_max_tensor.flat<float>()(0);
     const Tensor& var = context->input(6);
-    const float var_min = context->input(7).flat<float>()(0);
-    const float var_max = context->input(8).flat<float>()(0);
+    const auto& var_min_tensor = context->input(7);
+    OP_REQUIRES(context, var_min_tensor.NumElements() == 1,
+                errors::InvalidArgument("var_min must have 1 element"));
+    const float var_min = var_min_tensor.flat<float>()(0);
+    const auto& var_max_tensor = context->input(8);
+    OP_REQUIRES(context, var_max_tensor.NumElements() == 1,
+                errors::InvalidArgument("var_max must have 1 element"));
+    const float var_max = var_max_tensor.flat<float>()(0);
     const Tensor& beta = context->input(9);
-    const float beta_min = context->input(10).flat<float>()(0);
-    const float beta_max = context->input(11).flat<float>()(0);
+    const auto& beta_min_tensor = context->input(10);
+    OP_REQUIRES(context, beta_min_tensor.NumElements() == 1,
+                errors::InvalidArgument("beta_min must have 1 element"));
+    const float beta_min = beta_min_tensor.flat<float>()(0);
+    const auto& beta_max_tensor = context->input(11);
+    OP_REQUIRES(context, beta_max_tensor.NumElements() == 1,
+                errors::InvalidArgument("beta_max must have 1 element"));
+    const float beta_max = beta_max_tensor.flat<float>()(0);
     const Tensor& gamma = context->input(12);
-    const float gamma_min = context->input(13).flat<float>()(0);
-    const float gamma_max = context->input(14).flat<float>()(0);
+    const auto& gamma_min_tensor = context->input(13);
+    OP_REQUIRES(context, gamma_min_tensor.NumElements() == 1,
+                errors::InvalidArgument("gamma_min must have 1 element"));
+    const float gamma_min = gamma_min_tensor.flat<float>()(0);
+    const auto& gamma_max_tensor = context->input(14);
+    OP_REQUIRES(context, gamma_max_tensor.NumElements() == 1,
+                errors::InvalidArgument("gamma_max must have 1 element"));
+    const float gamma_max = gamma_max_tensor.flat<float>()(0);
 
     OP_REQUIRES(context, input.dims() == 4,
                 errors::InvalidArgument("input must be 4-dimensional",
@@ -203,6 +233,33 @@ class QuantizedBatchNormOp : public OpKernel {
     OP_REQUIRES(context, gamma.dims() == 1,
                 errors::InvalidArgument("gamma must be 1-dimensional",
                                         gamma.shape().DebugString()));
+    OP_REQUIRES(context, mean.NumElements() > 1,
+                errors::InvalidArgument("Must have at least a mean value",
+                                        gamma.shape().DebugString()));
+    OP_REQUIRES(context, mean.NumElements() > 1,
+                errors::InvalidArgument("Must have at least a mean value"));
+    const auto last_dim = input.shape().dims() - 1;
+    OP_REQUIRES(context,
+                mean.shape().dim_size(0) == input.shape().dim_size(last_dim),
+                errors::InvalidArgument("Must provide as many means as the "
+                                        "last dimension of the input tensor: ",
+                                        mean.shape().DebugString(), " vs. ",
+                                        input.shape().DebugString()));
+    OP_REQUIRES(
+        context, mean.shape().dim_size(0) == var.shape().dim_size(0),
+        errors::InvalidArgument(
+            "Mean and variance tensors must have the same shape: ",
+            mean.shape().DebugString(), " vs. ", var.shape().DebugString()));
+    OP_REQUIRES(
+        context, mean.shape().dim_size(0) == beta.shape().dim_size(0),
+        errors::InvalidArgument(
+            "Mean and beta tensors must have the same shape: ",
+            mean.shape().DebugString(), " vs. ", beta.shape().DebugString()));
+    OP_REQUIRES(
+        context, mean.shape().dim_size(0) == gamma.shape().dim_size(0),
+        errors::InvalidArgument(
+            "Mean and gamma tensors must have the same shape: ",
+            mean.shape().DebugString(), " vs. ", gamma.shape().DebugString()));
 
     Tensor* output = nullptr;
     OP_REQUIRES_OK(context,
path: tensorflow/core/kernels/quantized_batch_norm_op.cc