fix posible heap buffer overflow

dmlc · Oct 31, 2021 · a5ec1d6 · a5ec1d6
1 parent 8ed8a91
commit a5ec1d6
Showing 1 changed file with 7 additions and 0 deletions.
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) {