Vulnerability: Division by zero in TFLite's implementation of `OneHot` ### Impact
The implementation of the `OneHot` TFLite operator is [vulnerable to a division by zero error](https://github.com/tensorflow/tensorflow/blob/f61c57bd425878be108ec787f4d96390579fb83e/tensorflow/lite/kernels/one_hot.cc#L68-L72):

```cc
int prefix_dim_size = 1;
for (int i = 0; i < op_context.axis; ++i) {
  prefix_dim_size *= op_context.indices->dims->data[i];
}
const int suffix_dim_size = NumElements(op_context.indices) / prefix_dim_size;
```

An attacker can craft a model such that at least one of the dimensions of `indices` would be 0. In turn, the `prefix_dim_size` value would become 0.

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

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 members of the Aivul Team from Qihoo 360.
Vulnerability ID: GHSA-j8qh-3xrq-c825
CWE: CWE-369
Severity score: 2.5
Project: https://github.com/tensorflow/tensorflow
Published: 2021-05-21
File: tensorflow/lite/kernels/one_hot.cc
Commit message: Prevent division by 0 in OneHot implementation

If input indices is degenerate, the implementation would do a divide by zero. See https://github.com/tensorflow/tensorflow/blob/745d57df6d5e9bc568666a2a48ed8dd629c27241/tensorflow/lite/kernels/one_hot.cc#L68-L72

PiperOrigin-RevId: 370966870
Change-Id: Ie018337811c8016b5a1d3a277d00d5f2e19a2058
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 void OneHotComputeImpl(const OneHotContext& op_context) {   for (int i = 0; i < op_context.axis; ++i) {     prefix_dim_size *= op_context.indices->dims->data[i];   }   const int suffix_dim_size = NumElements(op_context.indices) / prefix_dim_size;   const int depth = *op_context.depth->data.i32; 

Fixed version (after fix):
 void OneHotComputeImpl(const OneHotContext& op_context) {   for (int i = 0; i < op_context.axis; ++i) {     prefix_dim_size *= op_context.indices->dims->data[i];   }   if (prefix_dim_size == 0) {     // If indices tensor is degenerate, return a degenerate tensor, just like     // TensorFlow does.     return;   }   const int suffix_dim_size = NumElements(op_context.indices) / prefix_dim_size;   const int depth = *op_context.depth->data.i32; 

Patch:
patch: @@ -69,6 +69,11 @@ void OneHotComputeImpl(const OneHotContext& op_context) {
   for (int i = 0; i < op_context.axis; ++i) {
     prefix_dim_size *= op_context.indices->dims->data[i];
   }
+  if (prefix_dim_size == 0) {
+    // If indices tensor is degenerate, return a degenerate tensor, just like
+    // TensorFlow does.
+    return;
+  }
   const int suffix_dim_size = NumElements(op_context.indices) / prefix_dim_size;
   const int depth = *op_context.depth->data.i32;
path: tensorflow/lite/kernels/one_hot.cc