Vulnerability: Out of bounds write in TFLite ### Impact 
An attacker can craft a TFLite model that would cause a write outside of bounds of an array in TFLite. In fact, the attacker can override the linked list used by the memory allocator. This can be leveraged for an arbitrary write primitive under certain conditions.

### Patches
We have patched the issue in GitHub commit [6c0b2b70eeee588591680f5b7d5d38175fd7cdf6](https://github.com/tensorflow/tensorflow/commit/6c0b2b70eeee588591680f5b7d5d38175fd7cdf6).
  
The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, 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 Wang Xuan of Qihoo 360 AIVul Team.
Vulnerability ID: GHSA-9c78-vcq7-7vxq
CWE: CWE-787
Severity score: 8.8
Project: https://github.com/tensorflow/tensorflow
Published: 2022-02-09
File: tensorflow/lite/kernels/fully_connected.cc
Commit message: [lite] add validation check for sparse fully connected

PiperOrigin-RevId: 417629354
Change-Id: If96171c4bd4f5fdb01d6368d6deab19d1c9beca7
Commit author: Karim Nosir

Vulnerable version (before fix):
 TfLiteStatus EvalShuffledQuantized(TfLiteContext* context, TfLiteNode* node,   return kTfLiteOk; }  template <KernelType kernel_type> TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,                        TfLiteFullyConnectedParams* params, OpData* data, TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,                            "Unsupported sparse fully-connected weight format.");         return kTfLiteError;       }        if (sparsity.dim_metadata_size == kDimMetadataSizeRandomSparse) {         // Random sparse.         optimized_ops::FullyConnectedSparseWeight(             sparsity, op_params, GetTensorShape(input),             GetTensorData<float>(input), GetTensorShape(filter),             GetTensorData<float>(filter), GetTensorShape(bias),             GetTensorData<float>(bias), GetTensorShape(output),             GetTensorData<float>(output));       } else if (sparsity.dim_metadata_size == kDimMetadataSizeBlockSparse &&                  sparsity.dim_metadata[2].dense_size == 4) {         // Block sparse with block size of 1x4.         optimized_ops::FullyConnectedSparseWeight1x4(             sparsity, op_params, GetTensorShape(input),             GetTensorData<float>(input), GetTensorShape(filter),             GetTensorData<float>(filter), GetTensorShape(bias),             GetTensorData<float>(bias), GetTensorShape(output),             GetTensorData<float>(output),             CpuBackendContext::GetFromContext(context));       } else {         TF_LITE_KERNEL_LOG(context,

Fixed version (after fix):
 TfLiteStatus EvalShuffledQuantized(TfLiteContext* context, TfLiteNode* node,   return kTfLiteOk; }  // Verifies that sparsity values are valid given input/weight/output. bool VerifySparsity(const RuntimeShape& weights_shape,                     const RuntimeShape& input_shape,                     const RuntimeShape& output_shape,                     const TfLiteSparsity* sparsity) {   const int weights_dims_count = weights_shape.DimensionsCount();   const int output_dims_count = output_shape.DimensionsCount();   const int w0_size = sparsity->dim_metadata[0].dense_size;   const int accum_depth = weights_shape.Dims(weights_dims_count - 1);   const int output_elements = output_shape.FlatSize();   const int input_elements = input_shape.FlatSize();   const int batches = FlatSizeSkipDim(output_shape, output_dims_count - 1);   const int output_depth = MatchingDim(weights_shape, weights_dims_count - 2,                                        output_shape, output_dims_count - 1);   const int max_batch_index = batches - 1;   const int max_output = max_batch_index * output_depth + w0_size;   const int max_batch_depth = accum_depth * max_batch_index;    // Verify output size is enough.   if (output_elements < max_output) return false;    // Verify index from sparse in input is valid.   for (int i = 0; i < sparsity->dim_metadata[1].array_indices->size; ++i) {     if (input_elements <=         max_batch_depth + sparsity->dim_metadata[1].array_indices->data[i])       return false;   }   return true; }  template <KernelType kernel_type> TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,                        TfLiteFullyConnectedParams* params, OpData* data, TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,                            "Unsupported sparse fully-connected weight format.");         return kTfLiteError;       }       const auto& input_shape = GetTensorShape(input);       const auto& filter_shape = GetTensorShape(filter);       const auto& output_shape = GetTensorShape(output);       const auto& bias_shape = GetTensorShape(bias);       if (!VerifySparsity(filter_shape, input_shape, output_shape, &sparsity)) {         TF_LITE_KERNEL_LOG(context, "Invalid sparse fully-connected format.");         return kTfLiteError;       }        if (sparsity.dim_metadata_size == kDimMetadataSizeRandomSparse) {         // Random sparse.         optimized_ops::FullyConnectedSparseWeight(             sparsity, op_params,                         // Disable formatting             input_shape, GetTensorData<float>(input),    // Disable formatting             filter_shape, GetTensorData<float>(filter),  // Disable formatting             bias_shape, GetTensorData<float>(bias),      // Disable formatting             output_shape, GetTensorData<float>(output));       } else if (sparsity.dim_metadata_size == kDimMetadataSizeBlockSparse &&                  sparsity.dim_metadata[2].dense_size == 4) {         // Block sparse with block size of 1x4.         optimized_ops::FullyConnectedSparseWeight1x4(             sparsity, op_params,                         // Disable formatting             input_shape, GetTensorData<float>(input),    // Disable formatting             filter_shape, GetTensorData<float>(filter),  // Disable formatting             bias_shape, GetTensorData<float>(bias),      // Disable formatting             output_shape, GetTensorData<float>(output),             CpuBackendContext::GetFromContext(context));       } else {         TF_LITE_KERNEL_LOG(context,

Patch:
patch: @@ -928,6 +928,36 @@ TfLiteStatus EvalShuffledQuantized(TfLiteContext* context, TfLiteNode* node,
   return kTfLiteOk;
 }
 
+// Verifies that sparsity values are valid given input/weight/output.
+bool VerifySparsity(const RuntimeShape& weights_shape,
+                    const RuntimeShape& input_shape,
+                    const RuntimeShape& output_shape,
+                    const TfLiteSparsity* sparsity) {
+  const int weights_dims_count = weights_shape.DimensionsCount();
+  const int output_dims_count = output_shape.DimensionsCount();
+  const int w0_size = sparsity->dim_metadata[0].dense_size;
+  const int accum_depth = weights_shape.Dims(weights_dims_count - 1);
+  const int output_elements = output_shape.FlatSize();
+  const int input_elements = input_shape.FlatSize();
+  const int batches = FlatSizeSkipDim(output_shape, output_dims_count - 1);
+  const int output_depth = MatchingDim(weights_shape, weights_dims_count - 2,
+                                       output_shape, output_dims_count - 1);
+  const int max_batch_index = batches - 1;
+  const int max_output = max_batch_index * output_depth + w0_size;
+  const int max_batch_depth = accum_depth * max_batch_index;
+
+  // Verify output size is enough.
+  if (output_elements < max_output) return false;
+
+  // Verify index from sparse in input is valid.
+  for (int i = 0; i < sparsity->dim_metadata[1].array_indices->size; ++i) {
+    if (input_elements <=
+        max_batch_depth + sparsity->dim_metadata[1].array_indices->data[i])
+      return false;
+  }
+  return true;
+}
+
 template <KernelType kernel_type>
 TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,
                        TfLiteFullyConnectedParams* params, OpData* data,
@@ -968,24 +998,32 @@ TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,
                            "Unsupported sparse fully-connected weight format.");
         return kTfLiteError;
       }
+      const auto& input_shape = GetTensorShape(input);
+      const auto& filter_shape = GetTensorShape(filter);
+      const auto& output_shape = GetTensorShape(output);
+      const auto& bias_shape = GetTensorShape(bias);
+      if (!VerifySparsity(filter_shape, input_shape, output_shape, &sparsity)) {
+        TF_LITE_KERNEL_LOG(context, "Invalid sparse fully-connected format.");
+        return kTfLiteError;
+      }
 
       if (sparsity.dim_metadata_size == kDimMetadataSizeRandomSparse) {
         // Random sparse.
         optimized_ops::FullyConnectedSparseWeight(
-            sparsity, op_params, GetTensorShape(input),
-            GetTensorData<float>(input), GetTensorShape(filter),
-            GetTensorData<float>(filter), GetTensorShape(bias),
-            GetTensorData<float>(bias), GetTensorShape(output),
-            GetTensorData<float>(output));
+            sparsity, op_params,                         // Disable formatting
+            input_shape, GetTensorData<float>(input),    // Disable formatting
+            filter_shape, GetTensorData<float>(filter),  // Disable formatting
+            bias_shape, GetTensorData<float>(bias),      // Disable formatting
+            output_shape, GetTensorData<float>(output));
       } else if (sparsity.dim_metadata_size == kDimMetadataSizeBlockSparse &&
                  sparsity.dim_metadata[2].dense_size == 4) {
         // Block sparse with block size of 1x4.
         optimized_ops::FullyConnectedSparseWeight1x4(
-            sparsity, op_params, GetTensorShape(input),
-            GetTensorData<float>(input), GetTensorShape(filter),
-            GetTensorData<float>(filter), GetTensorShape(bias),
-            GetTensorData<float>(bias), GetTensorShape(output),
-            GetTensorData<float>(output),
+            sparsity, op_params,                         // Disable formatting
+            input_shape, GetTensorData<float>(input),    // Disable formatting
+            filter_shape, GetTensorData<float>(filter),  // Disable formatting
+            bias_shape, GetTensorData<float>(bias),      // Disable formatting
+            output_shape, GetTensorData<float>(output),
             CpuBackendContext::GetFromContext(context));
       } else {
         TF_LITE_KERNEL_LOG(context,
path: tensorflow/lite/kernels/fully_connected.cc