fix posible heap buffer overflow

dmlc · Oct 31, 2021 · fecf89d · fecf89d
1 parent 8ed8a91
commit fecf89d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/src/common/opt_partition_builder.h b/src/common/opt_partition_builder.h
@@ -204,6 +204,13 @@ class OptPartitionBuilder {
     if (!all_dense) {
       std::vector<size_t>& local_states = states[tid];
       std::vector<bool>& local_default_flags = default_flags[tid];
+      if (row_indices_begin >= row_indices_end) {
+        rows[0] = rows_count;
+        if (is_loss_guided) {
+          rows_left[0] = rows_left_count;
+        }
+        return;
+      }
       const uint32_t first_row_id = !is_loss_guided ? row_indices_begin :
                                                       row_indices_ptr[row_indices_begin];
       for (const auto& nid : split_nodes) {

diff --git a/tests/cpp/common/test_partition_builder.cc b/tests/cpp/common/test_partition_builder.cc
@@ -21,7 +21,6 @@ TEST(OptPartitionBuilder, BasicTest) {
     {0.27f, 0.29f}, {0.37f, 0.39f}, {-0.47f, 0.49f}, {0.57f, 0.59f} };
   ColumnMatrix column_matrix;
   column_matrix.Init(gmat, 0);
-  uint32_t total_bins = gmat.cut.Ptrs().back();
   RegTree tree;
   common::OptPartitionBuilder opt_partition_builder;