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
CVE: CVE-2021-37664
CWE: CWE-369
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. In affected versions an attacker can read from outside of bounds of heap allocated data by sending specially crafted illegal arguments to `BoostedTreesSparseCalculateBestFeatureSplit`. 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. We have patched the issue in GitHub 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.

Project: tensorflow
File: stats_ops.cc

  void CalculateBestEqualitySplit(
      TTypes<float, 4>::ConstTensor stats_summary,
      const Eigen::VectorXf& total_grad, const Eigen::VectorXf& total_hess,
      const int32_t node_id, const int32_t feature_dims,
      const int32_t logits_dim, const int32_t hessian_dim,
      const int32_t num_buckets, const float l1, const float l2,
      float* best_gain, int32* best_bucket, int32* best_f_dim,
      string* best_split_type, Eigen::VectorXf* best_contrib_for_left,
      Eigen::VectorXf* best_contrib_for_right) {
    const string kEqualityDefaultRight =
        boosted_trees::SplitTypeWithDefault_Name(
            boosted_trees::EQUALITY_DEFAULT_RIGHT);
    for (int f_dim = 0; f_dim < feature_dims; ++f_dim) {
      for (int bucket = 0; bucket < num_buckets; ++bucket) {
        ConstVectorMap stats_vec(&stats_summary(node_id, f_dim, bucket, 0),
                                 logits_dim + hessian_dim);
        Eigen::VectorXf curr_grad = stats_vec.head(logits_dim);
        Eigen::VectorXf curr_hess = stats_vec.tail(hessian_dim);
        MaybeUpdateBestSplit(curr_grad, total_grad - curr_grad, curr_hess,
                             total_hess - curr_hess, logits_dim, bucket, f_dim,
                             l1, l2, kEqualityDefaultRight, best_gain,
                             best_bucket, best_f_dim, best_split_type,
                             best_contrib_for_left, best_contrib_for_right);
      }
    }
  }

Project: tensorflow
File: stats_ops.cc

static void AddInstanceStatsToMap(
    const int32_t instance, const int32_t feature_dim, const int32_t bucket_id,
    const int32_t logits_dims, const int32_t stats_dims,
    StatsPartitionMap* stats_map, const TTypes<float>::ConstMatrix& gradients,
    const TTypes<float>::ConstMatrix& hessians,
    const TTypes<int32>::ConstVec& node_ids) {
  const int32_t node_id = node_ids(instance);
  const auto key = StatsPartitionKey(node_id, feature_dim, bucket_id);
  std::pair<StatsPartitionIterator, bool> const& insert_result =
      stats_map->insert(StatsPartitionIterator::value_type(
          key, std::vector<float>(stats_dims, 0.0f)));
  auto& stats = insert_result.first->second;
  for (int stat_dim = 0; stat_dim < logits_dims; ++stat_dim) {
    stats[stat_dim] += gradients(instance, stat_dim);
  }
  for (int stat_dim = logits_dims; stat_dim < stats_dims; ++stat_dim) {
    stats[stat_dim] += hessians(instance, stat_dim - logits_dims);
  }
}

Project: tensorflow
File: stats_ops.cc

  bool operator==(const StatsPartitionKey& other) const {
    return (node_id == other.node_id) && (feature_dim == other.feature_dim) &&
           (bucket_id == other.bucket_id);
  }

Project: tensorflow
File: stats_ops.cc

  explicit BoostedTreesAggregateStatsOp(OpKernelConstruction* const context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("max_splits", &max_splits_));
    OP_REQUIRES_OK(context, context->GetAttr("num_buckets", &num_buckets_));
  }

Project: tensorflow
File: stats_ops.cc

  explicit BoostedTreesCalculateBestGainsPerFeatureOp(
      OpKernelConstruction* const context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("max_splits", &max_splits_));
    OP_REQUIRES_OK(context, context->GetAttr("num_features", &num_features_));
  }

Project: tensorflow
File: stats_ops.cc

  InstanceFeatureDimKey(const int32_t instance, const int32_t feature_dim)
      : instance(instance), feature_dim(feature_dim) {}

Project: tensorflow
File: stats_ops.cc

  InstanceFeatureDimKey() : instance(-1), feature_dim(-1) {}

Project: tensorflow
File: stats_ops.cc

  explicit BoostedTreesMakeStatsSummaryOp(OpKernelConstruction* const context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("max_splits", &max_splits_));
    OP_REQUIRES_OK(context, context->GetAttr("num_buckets", &num_buckets_));
    OP_REQUIRES_OK(context, context->GetAttr("num_features", &num_features_));
  }

Project: tensorflow
File: stats_ops.cc

  explicit BoostedTreesSparseAggregateStatsOp(
      OpKernelConstruction* const context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("max_splits", &max_splits_));
    OP_REQUIRES_OK(context, context->GetAttr("num_buckets", &num_buckets_));
  }

Project: tensorflow
File: stats_ops.cc

    bool operator()(const StatsPartitionKey& a,
                    const StatsPartitionKey& b) const {
      if (a.node_id < b.node_id) {
        return true;
      }
      if ((a.node_id == b.node_id) && (a.feature_dim < b.feature_dim)) {
        return true;
      }
      if ((a.node_id == b.node_id) && (a.feature_dim == b.feature_dim) &&
          (a.bucket_id < b.bucket_id)) {
        return true;
      }
      return false;
    }

Project: tensorflow
File: stats_ops.cc

  void process_node(const FeatureMap& f_map,
                    std::vector<int32>* output_node_ids,
                    std::vector<float>* output_gains,
                    std::vector<int32>* output_feature_dimensions,
                    std::vector<int32>* output_thresholds,
                    std::vector<float>* output_left_node_contribs,
                    std::vector<float>* output_right_node_contribs,
                    std::vector<string>* output_split_types,
                    const int32_t node_id, const float min_node_weight,
                    const float l1, const float l2, const int32_t num_buckets) {
    float parent_gain;
    Eigen::VectorXf unused(logits_dim_);
    Eigen::MatrixXf identity;
    identity.setIdentity(1, 1);

    // start processing for previous node id.
    float best_gain = std::numeric_limits<float>::lowest();
    float best_bucket = 0;
    float best_f_dim = 0;
    string best_split_type = boosted_trees::SplitTypeWithDefault_Name(
        boosted_trees::INEQUALITY_DEFAULT_LEFT);
    float best_contrib_for_left = 0.0;
    float best_contrib_for_right = 0.0;
    // the sum of gradients including default bucket.
    float total_grad = 0;
    // the sum of hessians including default bucket.
    float total_hess = 0;

    for (auto f_iter = f_map.begin(); f_iter != f_map.end(); ++f_iter) {
      const int32_t feature_dim = f_iter->first;
      const auto buckets_to_stats_map = f_iter->second;

      // The very last bucket contains stats for missing values.
      // TODO(crawles): use vector for multi-class.
      const float default_grad =
          (buckets_to_stats_map.find(num_buckets) == buckets_to_stats_map.end()
               ? 0
               : buckets_to_stats_map.at(num_buckets)[0]);
      const float default_hess =
          (buckets_to_stats_map.find(num_buckets) == buckets_to_stats_map.end()
               ? 0
               : buckets_to_stats_map.at(num_buckets)[1]);

      if (f_iter == f_map.begin()) {
        // first get the sum of grads, including default bucket.
        for (auto b_iter = buckets_to_stats_map.begin();
             b_iter != buckets_to_stats_map.end(); ++b_iter) {
          total_grad += b_iter->second[0];
          total_hess += b_iter->second[1];
        }
        if (total_hess < min_node_weight) {
          // Do not split the node because not enough avg hessian.
          break;
        }
        CalculateWeightsAndGains(total_grad * identity, total_hess * identity,
                                 l1, l2, &unused, &parent_gain);
      }

      float total_left_grad = 0;
      float total_left_hess = 0;
      for (auto b_iter = buckets_to_stats_map.begin();
           b_iter != buckets_to_stats_map.end(); ++b_iter) {
        const int32_t bucket_id = b_iter->first;
        // total_left_stats should exclude stats from default bucket.
        if (bucket_id == num_buckets) {
          break;
        }
        // TODO(crawles): vector for multi-class.
        total_left_grad += b_iter->second[0];
        total_left_hess += b_iter->second[1];
        // From left to right, default right.
        // Left child.
        Eigen::VectorXf contrib_for_left(1);
        float gain_for_left;
        CalculateWeightsAndGains(total_left_grad * identity,
                                 total_left_hess * identity, l1, l2,
                                 &contrib_for_left, &gain_for_left);
        // Right child.
        Eigen::VectorXf contrib_for_right(1);
        float gain_for_right;
        CalculateWeightsAndGains((total_grad - total_left_grad) * identity,
                                 (total_hess - total_left_hess) * identity, l1,
                                 l2, &contrib_for_right, &gain_for_right);
        if (GainIsLarger(gain_for_left + gain_for_right, best_gain)) {
          best_gain = gain_for_left + gain_for_right;
          best_bucket = bucket_id;
          best_f_dim = feature_dim;
          best_split_type = boosted_trees::SplitTypeWithDefault_Name(
              boosted_trees::INEQUALITY_DEFAULT_RIGHT);
          best_contrib_for_left = contrib_for_left[0];
          best_contrib_for_right = contrib_for_right[0];
        }

        // From right to left, default left.
        CalculateWeightsAndGains((total_left_grad + default_grad) * identity,
                                 (total_left_hess + default_hess) * identity,
                                 l1, l2, &contrib_for_left, &gain_for_left);
        CalculateWeightsAndGains(
            (total_grad - default_grad - total_left_grad) * identity,
            (total_hess - default_hess - total_left_hess) * identity, l1, l2,
            &contrib_for_right, &gain_for_right);
        if (GainIsLarger(gain_for_left + gain_for_right, best_gain)) {
          best_gain = gain_for_left + gain_for_right;
          best_bucket = bucket_id;
          best_f_dim = feature_dim;
          best_split_type = boosted_trees::SplitTypeWithDefault_Name(
              boosted_trees::INEQUALITY_DEFAULT_LEFT);
          best_contrib_for_left = contrib_for_left[0];
          best_contrib_for_right = contrib_for_right[0];
        }
      }  // for bucket_id
    }    // for feature_dim
    if (best_gain != std::numeric_limits<float>::lowest()) {
      output_node_ids->push_back(node_id);
      // Remove the parent gain.
      output_gains->push_back(best_gain - parent_gain);
      output_feature_dimensions->push_back(best_f_dim);
      output_split_types->push_back(best_split_type);
      output_thresholds->push_back(best_bucket);
      output_left_node_contribs->push_back(best_contrib_for_left);
      output_right_node_contribs->push_back(best_contrib_for_right);
    }
  }

Project: tensorflow
File: stats_ops.cc

static void AddRangeStats(const int start_instance, const int end_instance,
                          const int start_feature_dim,
                          const int end_feature_dim,
                          StatsPartitionMap* stats_map,
                          const TTypes<float>::ConstMatrix& gradients,
                          const TTypes<float>::ConstMatrix& hessians,
                          const TTypes<int32>::ConstVec& node_ids,
                          const int32_t feature_dims, const int32_t bucket_id,
                          const int32_t logits_dims, const int32_t stats_dims) {
  DCHECK_LE(start_instance, end_instance);
  if (start_instance == end_instance) {
    DCHECK_LT(start_feature_dim, end_feature_dim);
  }
  for (int32_t instance = start_instance; instance <= end_instance;
       ++instance) {
    const int32_t start_f_dim =
        (instance == start_instance) ? start_feature_dim + 1 : 0;
    const int32_t end_f_dim =
        (instance == end_instance) ? end_feature_dim : feature_dims;
    for (int32_t f_dim = start_f_dim; f_dim < end_f_dim; ++f_dim) {
      AddInstanceStatsToMap(instance, f_dim, bucket_id, logits_dims, stats_dims,
                            stats_map, gradients, hessians, node_ids);
    }
  }
}

Project: tensorflow
File: stats_ops.cc

  void MaybeUpdateBestSplit(const Eigen::VectorXf& grad_for_left,
                            const Eigen::VectorXf& grad_for_right,
                            const Eigen::VectorXf& hess_for_left,
                            const Eigen::VectorXf& hess_for_right,
                            const int32_t logits_dim, const int32_t bucket,
                            const int32_t f_dim, const float l1, const float l2,
                            const string split_type, float* best_gain,
                            int32* best_bucket, int32* best_f_dim,
                            string* best_split_type,
                            Eigen::VectorXf* best_contrib_for_left,
                            Eigen::VectorXf* best_contrib_for_right) {
    // Left child.
    Eigen::VectorXf contrib_for_left(logits_dim);
    float gain_for_left;
    CalculateWeightsAndGains(grad_for_left, hess_for_left, l1, l2,
                             &contrib_for_left, &gain_for_left);
    Eigen::VectorXf contrib_for_right(logits_dim);
    float gain_for_right;
    CalculateWeightsAndGains(grad_for_right, hess_for_right, l1, l2,
                             &contrib_for_right, &gain_for_right);
    if (GainIsLarger(gain_for_left + gain_for_right, *best_gain)) {
      *best_gain = gain_for_left + gain_for_right;
      *best_bucket = bucket;
      *best_f_dim = f_dim;
      *best_contrib_for_left = contrib_for_left;
      *best_contrib_for_right = contrib_for_right;
      *best_split_type = split_type;
    }
  }

Project: tensorflow
File: stats_ops.cc

  explicit BoostedTreesCalculateBestFeatureSplitV2(
      OpKernelConstruction* const context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("logits_dimension", &logits_dim_));
    OP_REQUIRES_OK(context, context->GetAttr("num_features", &num_features_));
  }

Project: tensorflow
File: stats_ops.cc

  bool operator==(const InstanceFeatureDimKey& other) const {
    return (instance == other.instance) && (feature_dim == other.feature_dim);
  }

Project: tensorflow
File: stats_ops.cc

    bool operator()(const InstanceFeatureDimKey& a,
                    const InstanceFeatureDimKey& b) const {
      if (a.instance < b.instance) {
        return true;
      }
      if ((a.instance == b.instance) && (a.feature_dim < b.feature_dim)) {
        return true;
      }
      return false;
    }

Project: tensorflow
File: stats_ops.cc

  StatsPartitionKey(const int32_t node_id, const int32_t feature_dim,
                    const int32_t bucket_id)
      : node_id(node_id), feature_dim(feature_dim), bucket_id(bucket_id) {}

Project: tensorflow
File: stats_ops.cc

  void CalculateBestInequalitySplit(
      TTypes<float, 4>::ConstTensor stats_summary, const int32_t node_id,
      const int32_t feature_dims, const int32_t logits_dim,
      const int32_t hessian_dim, const int32_t num_buckets,
      const float min_node_weight, const float l1, const float l2,
      float* best_gain, int32* best_bucket, int32* best_f_dim,
      string* best_split_type, Eigen::VectorXf* best_contrib_for_left,
      Eigen::VectorXf* best_contrib_for_right) {
    std::vector<Eigen::VectorXf> cum_grad;
    std::vector<Eigen::VectorXf> cum_hess;
    // get all cumulative gradients including default bucket.
    cum_grad.reserve(num_buckets);
    cum_hess.reserve(num_buckets);

    for (int f_dim = 0; f_dim < feature_dims; ++f_dim) {
      ConstVectorMap default_stats_vec(
          &stats_summary(node_id, f_dim, num_buckets, 0),
          logits_dim + hessian_dim);
      Eigen::VectorXf missing_bucket_grad = default_stats_vec.head(logits_dim);
      Eigen::VectorXf missing_bucket_hess = default_stats_vec.tail(hessian_dim);
      cum_grad.clear();
      cum_hess.clear();
      Eigen::VectorXf total_grad = Eigen::VectorXf::Zero(logits_dim);
      Eigen::VectorXf total_hess = Eigen::VectorXf::Zero(hessian_dim);
      // sum all the gradients including default bucket.
      for (int bucket = 0; bucket <= num_buckets; ++bucket) {
        for (int i = 0; i < logits_dim; ++i) {
          total_grad[i] += stats_summary(node_id, f_dim, bucket, i);
        }
        for (int i = 0; i < hessian_dim; ++i) {
          // Full hessian.
          total_hess[i] +=
              stats_summary(node_id, f_dim, bucket, logits_dim + i);
        }
        if (bucket < num_buckets) {
          cum_grad.push_back(total_grad);
          cum_hess.push_back(total_hess);
        }
      }
      const string kInequalityDefaultLeft =
          boosted_trees::SplitTypeWithDefault_Name(
              boosted_trees::INEQUALITY_DEFAULT_LEFT);
      const string kInequalityDefaultRight =
          boosted_trees::SplitTypeWithDefault_Name(
              boosted_trees::INEQUALITY_DEFAULT_RIGHT);

      // Iterate from left to right, excluding default bucket.
      for (int bucket = 0; bucket < num_buckets; ++bucket) {
        // default value goes to left node.
        const Eigen::VectorXf total_left_grad =
            cum_grad[bucket] + missing_bucket_grad;
        const Eigen::VectorXf total_left_hess =
            cum_hess[bucket] + missing_bucket_hess;
        MaybeUpdateBestSplit(
            total_left_grad, total_grad - total_left_grad, total_left_hess,
            total_hess - total_left_hess, logits_dim, bucket, f_dim, l1, l2,
            kInequalityDefaultLeft, best_gain, best_bucket, best_f_dim,
            best_split_type, best_contrib_for_left, best_contrib_for_right);
        // default value goes to right node.
        MaybeUpdateBestSplit(
            cum_grad[bucket], total_grad - cum_grad[bucket], cum_hess[bucket],
            total_hess - cum_hess[bucket], logits_dim, bucket, f_dim, l1, l2,
            kInequalityDefaultRight, best_gain, best_bucket, best_f_dim,
            best_split_type, best_contrib_for_left, best_contrib_for_right);
      }  // for bucket
    }
  }

Project: tensorflow
File: stats_ops.cc

  explicit BoostedTreesCalculateBestFeatureSplitOp(
      OpKernelConstruction* const context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("logits_dimension", &logits_dim_));
    OP_REQUIRES_OK(context, context->GetAttr("split_type", &split_type_));
  }

Project: tensorflow
File: stats_ops.cc

  explicit BoostedTreesSparseCalculateBestFeatureSplitOp(
      OpKernelConstruction* const context)
      : OpKernel(context) {
    // TODO(crawles): Using logits_dim_ for multi-class split.
    OP_REQUIRES_OK(context, context->GetAttr("logits_dimension", &logits_dim_));
    // TODO(tanzheny): Using this for equality split.
    OP_REQUIRES_OK(context, context->GetAttr("split_type", &split_type_));
  }