From c93b92a43f9b4bc66a5773e86e47f09bc1764807 Mon Sep 17 00:00:00 2001 From: Xiang Zhang Date: Wed, 7 Aug 2024 16:27:18 -0700 Subject: [PATCH] fix wrong check for tree ensemble regressor (#21595) Fix missed ORT_ENFORCE check which caused heap buffer overflow because of out of bound access. --- onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h index 8f847fe66a..df27f888bb 100644 --- a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h +++ b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h @@ -172,7 +172,7 @@ Status TreeEnsembleCommon::Init( nodes_falsenodeids.size() == nodes_values_as_tensor.size()); ORT_ENFORCE(target_class_ids.size() == target_class_nodeids.size()); ORT_ENFORCE(target_class_ids.size() == target_class_treeids.size()); - ORT_ENFORCE(target_class_ids.size() == target_class_treeids.size()); + ORT_ENFORCE(target_class_weights.empty() || target_class_ids.size() == target_class_weights.size()); ORT_ENFORCE(base_values.empty() || base_values_as_tensor.empty()); ORT_ENFORCE(nodes_hitrates.empty() || nodes_hitrates_as_tensor.empty()); ORT_ENFORCE(nodes_values.empty() || nodes_values_as_tensor.empty());