Vulnerability: Heap OOB in boosted trees ### Impact
An attacker can read from outside of bounds of heap allocated data by sending specially crafted illegal arguments to `BoostedTreesSparseCalculateBestFeatureSplit`:

```python
import tensorflow as tf

tf.raw_ops.BoostedTreesSparseCalculateBestFeatureSplit(
  node_id_range=[0,10],
  stats_summary_indices=[[1, 2, 3, 0x1000000]],
  stats_summary_values=[1.0],
  stats_summary_shape=[1,1,1,1],
  l1=l2=[1.0],
  tree_complexity=[0.5],
  min_node_weight=[1.0],
  logits_dimension=3,
  split_type='inequality')                                                                                                                                                                                                                                                                
```

The [implementation](https://github.com/tensorflow/tensorflow/blob/84d053187cb80d975ef2b9684d4b61981bca0c41/tensorflow/core/kernels/boosted_trees/stats_ops.cc) needs to validate that each value in `stats_summary_indices` is in range.
  
### Patches
We have patched the issue in GitHub commit [e84c975313e8e8e38bb2ea118196369c45c51378](https://github.com/tensorflow/tensorflow/commit/e84c975313e8e8e38bb2ea118196369c45c51378).
  
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.
  
### 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-r4c4-5fpq-56wg
CWE: CWE-125
Severity score: 7.3
Project: https://github.com/tensorflow/tensorflow
Published: 2021-08-25
File: tensorflow/core/kernels/boosted_trees/stats_ops.cc
Commit message: In tf.raw_ops.BoostedTreesSparseCalculateBestFeatureSplit, limit stat_dim in stats_summary_indices to under stats_dims in stats_summary_shape

PiperOrigin-RevId: 387171191
Change-Id: I83ca8a75b22aa78c037e8b98779da6cced16bfaa
Commit author: Laura Pak

Vulnerable version (before fix):
 class BoostedTreesSparseCalculateBestFeatureSplitOp : public OpKernel {       const int32_t feature_dim = stats_summary_indices(idx, 1);       const int32_t bucket_id = stats_summary_indices(idx, 2);       const int32_t stat_dim = stats_summary_indices(idx, 3);       std::pair<FeatureMapIterator, bool> const& f_insert_result = f_map.insert(           FeatureMapIterator::value_type(feature_dim, BucketMap()));       auto& b_map = f_insert_result.first->second;

Fixed version (after fix):
 class BoostedTreesSparseCalculateBestFeatureSplitOp : public OpKernel {       const int32_t feature_dim = stats_summary_indices(idx, 1);       const int32_t bucket_id = stats_summary_indices(idx, 2);       const int32_t stat_dim = stats_summary_indices(idx, 3);       OP_REQUIRES(context, stat_dim < stats_dims,                   errors::InvalidArgument(                       "Stat dim, the sum of logits dim and hessian dim in "                       "stats_summary_indices, cannot be greater than stats "                       "dims, the last value in stats_summary_shape, which was ",                       stats_dims, ". At index (", idx,                       ", 4), stats_summary_indices contains value ", stat_dim));       std::pair<FeatureMapIterator, bool> const& f_insert_result = f_map.insert(           FeatureMapIterator::value_type(feature_dim, BucketMap()));       auto& b_map = f_insert_result.first->second;

Patch:
patch: @@ -1050,6 +1050,13 @@ class BoostedTreesSparseCalculateBestFeatureSplitOp : public OpKernel {
       const int32_t feature_dim = stats_summary_indices(idx, 1);
       const int32_t bucket_id = stats_summary_indices(idx, 2);
       const int32_t stat_dim = stats_summary_indices(idx, 3);
+      OP_REQUIRES(context, stat_dim < stats_dims,
+                  errors::InvalidArgument(
+                      "Stat dim, the sum of logits dim and hessian dim in "
+                      "stats_summary_indices, cannot be greater than stats "
+                      "dims, the last value in stats_summary_shape, which was ",
+                      stats_dims, ". At index (", idx,
+                      ", 4), stats_summary_indices contains value ", stat_dim));
       std::pair<FeatureMapIterator, bool> const& f_insert_result = f_map.insert(
           FeatureMapIterator::value_type(feature_dim, BucketMap()));
       auto& b_map = f_insert_result.first->second;
path: tensorflow/core/kernels/boosted_trees/stats_ops.cc