[FIL] Validate Treelite model input to prevent integer overflow and OOB memory access

cpp/include/cuml/fil/detail/bitset.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -22,6 +22,9 @@ struct bitset {
   using storage_type = storage_t;
   using index_type   = index_t;
 
+  // Ensrue that index_t is unsigned. Bound checks below rely on index_t being unsigned
+  static_assert(std::is_unsigned_v<index_t>, "index_t must be unsigned");
+
   auto constexpr static const bin_width = index_type(sizeof(storage_type) * 8);
 
   HOST DEVICE bitset() : data_{nullptr}, num_bits_{0} {}
@@ -39,12 +42,13 @@ struct bitset {
   // Standard bit-wise mutators and accessor
   HOST DEVICE auto& set(index_type index)
   {
-    data_[bin_from_index(index)] |= mask_in_bin(index);
+    // Guard against OOB writes; silently ignored to preserve memory safety
+    if (index < num_bits_) { data_[bin_from_index(index)] |= mask_in_bin(index); }
     return *this;
   }
   HOST DEVICE auto& clear(index_type index)
   {
-    data_[bin_from_index(index)] &= ~mask_in_bin(index);
+    if (index < num_bits_) { data_[bin_from_index(index)] &= ~mask_in_bin(index); }
     return *this;
   }
   HOST DEVICE auto test(index_type index) const

cpp/include/cuml/fil/detail/decision_forest_builder.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -18,26 +18,58 @@
 #include <algorithm>
 #include <cmath>
 #include <cstddef>
+#include <iterator>
 #include <numeric>
 #include <optional>
+#include <sstream>
+#include <string>
+#include <type_traits>
+#include <utility>
 #include <vector>
 
 namespace ML {
 namespace fil {
 namespace detail {
 
-/*
- * Exception indicating that FIL model could not be built from given input
- */
-struct model_builder_error : std::exception {
-  model_builder_error() : model_builder_error("Error while building model") {}
-  model_builder_error(char const* msg) : msg_{msg} {}
-  virtual char const* what() const noexcept { return msg_; }
+struct floating_point_truncation_error : std::exception {
+  floating_point_truncation_error() {}
+  floating_point_truncation_error(std::string msg) : msg_{msg} {}
+  floating_point_truncation_error(char const* msg) : msg_{msg} {}
+  virtual char const* what() const noexcept { return msg_.c_str(); }
 
  private:
-  char const* msg_;
+  std::string msg_;
 };
 
+template <typename To, typename From>
+To safe_cast_floating_point(From x)
+{
+  static_assert(std::is_floating_point_v<From> && std::is_floating_point_v<To>,
+                "Source and destination types must be both floating-point types.");
+  if constexpr (sizeof(To) >= sizeof(From)) {
+    // Widening cast
+    return static_cast<To>(x);
+  } else {
+    // Narrowing cast: should be checked
+    if (!std::isfinite(x)) {
+      throw floating_point_truncation_error{"Cannot cast an INF or NaN value"};
+    }
+    auto constexpr lower_limit = static_cast<From>(std::numeric_limits<To>::lowest());
+    auto constexpr upper_limit = static_cast<From>(std::numeric_limits<To>::max());
+    if (x < lower_limit) {
+      std::ostringstream ss;
+      ss << "Input must be at least " << lower_limit << ".";
+      throw floating_point_truncation_error{ss.str()};
+    }
+    if (x > upper_limit) {
+      std::ostringstream ss;
+      ss << "Input must be at most " << upper_limit << ".";
+      throw floating_point_truncation_error{ss.str()};
+    }
+    return static_cast<To>(x);
+  }
+}
+
 /*
  * Struct used to build FIL forests
  */
@@ -57,20 +89,36 @@ struct decision_forest_builder {
     typename node_type::metadata_storage_type feature = typename node_type::metadata_storage_type{},
     typename node_type::offset_type offset            = typename node_type::offset_type{})
   {
-    auto constexpr const bin_width = index_type(sizeof(typename node_type::index_type) * 8);
-    auto node_value                = typename node_type::index_type{};
-    auto set_storage               = &node_value;
-    auto max_node_categories =
-      (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) + 1 : 1;
+    auto constexpr const bin_width =
+      typename node_type::index_type{sizeof(typename node_type::index_type) * 8};
+    auto node_value  = typename node_type::index_type{};
+    auto set_storage = &node_value;
+
+    // Check invariants for data types
+    using cat_t   = typename std::iterator_traits<iter_t>::value_type;
+    using index_t = typename node_type::index_type;
+    static_assert(std::is_same_v<cat_t, std::uint32_t>, "Category value must be uint32_t");
+    static_assert(std::is_same_v<index_t, std::uint32_t> || std::is_same_v<index_t, std::uint64_t>,
+                  "Index type in tree node must be either uint32_t or uint64_t");
+
+    // Ensure that (max_cat + 1) can be represented as index_t to prevent integer overflow.
+    auto max_cat = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) : cat_t{0};
+    if constexpr (std::is_same_v<cat_t, index_t>) {
+      if (max_cat == std::numeric_limits<index_t>::max()) {
+        throw model_import_error{std::string{"Category index must be at most "} +
+                                 std::to_string(std::numeric_limits<index_t>::max() - 1)};
+      }
+    }
+    auto max_cat_plus_one = static_cast<index_t>(max_cat) + index_t{1};
+
     if (max_num_categories_ > bin_width) {
-      // TODO(wphicks): Check for overflow here
       node_value         = categorical_storage_.size();
-      auto bins_required = raft_proto::ceildiv(max_node_categories, bin_width);
-      categorical_storage_.push_back(max_node_categories);
+      auto bins_required = raft_proto::ceildiv(max_cat_plus_one, bin_width);
+      categorical_storage_.push_back(max_cat_plus_one);
       categorical_storage_.resize(categorical_storage_.size() + bins_required);
       set_storage = &(categorical_storage_[node_value + 1]);
     }
-    auto set = bitset{set_storage, max_node_categories};
+    auto set = bitset{set_storage, max_cat_plus_one};
     std::for_each(vec_begin, vec_end, [&set](auto&& cat_index) { set.set(cat_index); });
 
     add_node(
@@ -153,7 +201,7 @@ struct decision_forest_builder {
   void set_output_size(index_type val)
   {
     if (output_size_ != index_type{1} && output_size_ != val) {
-      throw model_import_error("Inconsistent leaf vector size");
+      throw unusable_model_exception("Inconsistent leaf vector size");
     }
     output_size_ = val;
   }
@@ -182,11 +230,78 @@ struct decision_forest_builder {
                            int device                       = 0,
                            raft_proto::cuda_stream stream   = raft_proto::cuda_stream{})
   {
-    // Allow narrowing for preprocessing constants. They are stored as doubles
-    // for consistency in the builder but must be converted to the proper types
-    // for the concrete forest model.
-#pragma GCC diagnostic push
-#pragma GCC diagnostic ignored "-Wnarrowing"
+    // Validate forest invariants the inference kernel relies on. After this
+    // function returns, the forest is treated as trusted by the kernel.
+
+    // tree_index arithmetic in the kernel uses index_type, so the tree count
+    // must fit without narrowing.
+    if (root_node_indexes_.size() > std::numeric_limits<index_type>::max()) {
+      throw model_import_error{std::string{"Forest has "} +
+                               std::to_string(root_node_indexes_.size()) +
+                               " trees, which exceeds the maximum representable in index_type (" +
+                               std::to_string(std::numeric_limits<index_type>::max()) + ")"};
+    }
+
+    // forest::get_tree_root(tree_index) dereferences nodes_ + root_index.
+    // Ensure each root index points into the nodes buffer.
+    for (auto i = std::size_t{0}; i < root_node_indexes_.size(); ++i) {
+      if (root_node_indexes_[i] >= nodes_.size()) {
+        throw model_import_error{
+          std::string{"Tree "} + std::to_string(i) + ": root node index out of bounds (" +
+          std::to_string(root_node_indexes_[i]) + " >= " + std::to_string(nodes_.size()) + ")"};
+      }
+    }
+
+    auto constexpr const cat_bin_width =
+      typename node_type::index_type{sizeof(typename node_type::index_type) * 8};
+    if (max_num_categories_ > cat_bin_width) {
+      auto const storage_size = categorical_storage_.size();
+      for (auto i = std::size_t{0}; i < nodes_.size(); ++i) {
+        auto const& n = nodes_[i];
+        if (n.is_leaf() || !n.is_categorical()) { continue; }
+        auto const offset = n.index();
+
+        // evaluate_tree_impl() reads categorical_storage[offset] as the number
+        // of categories for this node; offset must be in-range.
+        if (offset >= storage_size) {
+          throw model_import_error{std::string{"Categorical node "} + std::to_string(i) +
+                                   ": storage offset out of bounds (" + std::to_string(offset) +
+                                   " >= " + std::to_string(storage_size) + ")"};
+        }
+        auto const stored_num_cats = categorical_storage_[offset];
+        auto const bins_required   = raft_proto::ceildiv(stored_num_cats, cat_bin_width);
+
+        // evaluate_tree_impl() reconstructs a bitset from
+        // [offset + 1, offset + 1 + bins_required). Compute this range using
+        // size_t to keep the arithmetic explicit and overflow-safe.
+        auto const bits_begin = static_cast<std::size_t>(offset) + std::size_t{1};
+        auto const bits_end   = bits_begin + static_cast<std::size_t>(bins_required);
+        if (bits_end > storage_size) {
+          throw model_import_error{std::string{"Categorical node "} + std::to_string(i) +
+                                   ": bitset extends past categorical_storage end"};
+        }
+      }
+    }
+
+    // Safely cast average_factor_ and postproc_constant_ to node_type::threshold_type
+    auto average_factor_casted    = typename node_type::threshold_type{};
+    auto postproc_constant_casted = typename node_type::threshold_type{};
+    try {
+      average_factor_casted =
+        safe_cast_floating_point<typename node_type::threshold_type>(average_factor_);
+      // We can't use cuda::narrow here, because it throws for imprecise conversion, i.e. casting
+      // double{3.1} to float.
+    } catch (const floating_point_truncation_error& e) {
+      throw unusable_model_exception{std::string{"Found an invalid value for averaging factor: "} +
+                                     e.what()};
+    }
+    try {
+      postproc_constant_casted =
+        safe_cast_floating_point<typename node_type::threshold_type>(postproc_constant_);
+    } catch (const floating_point_truncation_error& e) {
+      throw unusable_model_exception{
+        std::string{"Found an invalid value for postprocessing constant: "} + e.what()};
+    }
     return decision_forest_t{
       raft_proto::buffer{
         raft_proto::buffer{nodes_.data(), nodes_.size()}, mem_type, device, stream},
@@ -219,9 +334,8 @@ struct decision_forest_builder {
       output_size_,
       row_postproc_,
       element_postproc_,
-      static_cast<typename node_type::threshold_type>(average_factor_),
-      static_cast<typename node_type::threshold_type>(postproc_constant_)};
-#pragma GCC diagnostic pop
+      average_factor_casted,
+      postproc_constant_casted};
   }
 
  private:

cpp/include/cuml/fil/detail/degenerate_trees.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -17,7 +17,7 @@ namespace ML::fil::detail {
 // This function returns a modified copy of a given Treelite model if it contains
 // at least one degenerate tree (a single root node with no child).
 // If the model contains no degenerate tree, then the function returns nullptr.
-std::unique_ptr<treelite::Model> convert_degenerate_trees(treelite::Model const& tl_model)
+inline std::unique_ptr<treelite::Model> convert_degenerate_trees(treelite::Model const& tl_model)
 {
   bool contains_degenerate =
     ML::forest::tree_accumulate(tl_model, false, [](auto&& contains, auto&& tree) {

cpp/include/cuml/fil/detail/node.hpp

@@ -1,11 +1,12 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
 
 #include <cuml/fil/detail/index_type.hpp>
 #include <cuml/fil/detail/raft_proto/gpu_support.hpp>
+#include <cuml/fil/exceptions.hpp>
 #include <cuml/fil/tree_layout.hpp>
 
 #include <iostream>
@@ -104,14 +105,15 @@ struct alignas(detail::get_node_alignment<threshold_t, index_t, metadata_storage
 
   // TODO(wphicks): Add custom type to ensure given child offset is at least
   // one
-#pragma GCC diagnostic push
-#pragma GCC diagnostic ignored "-Wnarrowing"
-  HOST DEVICE constexpr node(threshold_type value             = threshold_type{},
-                             bool is_leaf_node                = true,
-                             bool default_to_distant_child    = false,
-                             bool is_categorical_node         = false,
-                             metadata_storage_type feature    = metadata_storage_type{},
-                             offset_type distant_child_offset = offset_type{})
+
+  // Assumption: Node construction occurs on the host. This allows us to perform
+  // bound check on the 'feature' parameter.
+  constexpr node(threshold_type value             = threshold_type{},
+                 bool is_leaf_node                = true,
+                 bool default_to_distant_child    = false,
+                 bool is_categorical_node         = false,
+                 metadata_storage_type feature    = metadata_storage_type{},
+                 offset_type distant_child_offset = offset_type{})
     : aligned_data{
         .inner_data = {
           {.value = value},
@@ -120,20 +122,19 @@ struct alignas(detail::get_node_alignment<threshold_t, index_t, metadata_storage
   {
   }
 
-  HOST DEVICE constexpr node(index_type index,
-                             bool is_leaf_node                = true,
-                             bool default_to_distant_child    = false,
-                             bool is_categorical_node         = false,
-                             metadata_storage_type feature    = metadata_storage_type{},
-                             offset_type distant_child_offset = offset_type{})
+  constexpr node(index_type index,
+                 bool is_leaf_node                = true,
+                 bool default_to_distant_child    = false,
+                 bool is_categorical_node         = false,
+                 metadata_storage_type feature    = metadata_storage_type{},
+                 offset_type distant_child_offset = offset_type{})
     : aligned_data{
         .inner_data = {
           {.index = index},
           distant_child_offset,
           construct_metadata(is_leaf_node, default_to_distant_child, is_categorical_node, feature)}}
   {
   }
-#pragma GCC diagnostic pop
 
   /* The index of the feature for this node */
   HOST DEVICE auto constexpr feature_index() const
@@ -213,6 +214,13 @@ struct alignas(detail::get_node_alignment<threshold_t, index_t, metadata_storage
                                            bool is_categorical_node      = false,
                                            metadata_storage_type feature = metadata_storage_type{})
   {
+    // Ensure that 'feature' is not truncated.
+    static_assert(std::is_unsigned_v<metadata_storage_type>, "metadata storage must be unsigned");
+    if (feature > FEATURE_MASK) {
+      throw model_import_error{std::string{"The 'feature' value in the node must be at most "} +
+                               std::to_string(FEATURE_MASK) + "."};
+    }
+
     return metadata_storage_type(
       (is_leaf_node << LEAF_BIT) + (default_to_distant_child << DEFAULT_DISTANT_BIT) +
       (is_categorical_node << CATEGORICAL_BIT) + (feature & FEATURE_MASK));

cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp

@@ -1,14 +1,17 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
 #include <cuml/fil/detail/raft_proto/gpu_support.hpp>
 
+#include <type_traits>
+
 namespace raft_proto {
 template <typename T, typename U>
 HOST DEVICE auto constexpr ceildiv(T dividend, U divisor)
 {
-  return (dividend + divisor - T{1}) / divisor;
+  static_assert(std::is_integral_v<T> && std::is_integral_v<U>, "Arguments must be integers");
+  return dividend / divisor + (dividend % divisor != 0);
 }
 }  // namespace raft_proto

cpp/include/cuml/fil/exceptions.hpp

@@ -23,11 +23,12 @@ struct unusable_model_exception : std::exception {
 /** Exception indicating model import failed */
 struct model_import_error : std::exception {
   model_import_error() : model_import_error("Error while importing model") {}
+  model_import_error(std::string msg) : msg_{msg} {}
   model_import_error(char const* msg) : msg_{msg} {}
-  virtual char const* what() const noexcept { return msg_; }
+  virtual char const* what() const noexcept { return msg_.c_str(); }
 
  private:
-  char const* msg_;
+  std::string msg_;
 };
 
 /**