Merge pull request #29546 from uwezkhan/dtree-split-index-bound

bound categorical split indices in dtrees readSplit

modules/ml/src/tree.cpp

@@ -1839,12 +1839,14 @@ int DTreesImpl::readSplit( const FileNode& fn )
     Split split;
 
     int vi = (int)fn["var"];
-    CV_Assert( 0 <= vi && vi <= (int)varType.size() );
+    CV_Assert( 0 <= vi && vi < (int)varMapping.size() );
     vi = varMapping[vi]; // convert to varIdx if needed
+    CV_Assert( 0 <= vi && vi < (int)varType.size() );
     split.varIdx = vi;
 
     if( varType[vi] == VAR_CATEGORICAL ) // split on categorical var
     {
+        CV_Assert( vi < (int)catOfs.size() );
         int i, val, ssize = getSubsetSize(vi);
         split.subsetOfs = (int)subsets.size();
         for( i = 0; i < ssize; i++ )
@@ -1860,6 +1862,7 @@ int DTreesImpl::readSplit( const FileNode& fn )
         if( fns.isInt() )
         {
             val = (int)fns;
+            CV_Assert( 0 <= val && (val >> 5) < ssize );
             subset[val >> 5] |= 1 << (val & 31);
         }
         else
@@ -1869,6 +1872,7 @@ int DTreesImpl::readSplit( const FileNode& fn )
             for( i = 0; i < n; i++, ++it )
             {
                 val = (int)*it;
+                CV_Assert( 0 <= val && (val >> 5) < ssize );
                 subset[val >> 5] |= 1 << (val & 31);
             }
         }

modules/ml/test/test_save_load.cpp

@@ -96,6 +96,61 @@ ML_Legacy_Param param_list[] = {
 
 INSTANTIATE_TEST_CASE_P(/**/, ML_Legacy_Params, testing::ValuesIn(param_list));
 
+TEST(ML_DTrees, load_bad_categorical_split)
+{
+    // Train a tree with a categorical input so the model carries a
+    // categorical split serialized as an "in"/"not_in" value list.
+    const int n = 40;
+    Mat samples(n, 1, CV_32F);
+    Mat responses(n, 1, CV_32S);
+    for (int i = 0; i < n; i++)
+    {
+        int cat = i % 4;
+        samples.at<float>(i, 0) = (float)cat;
+        responses.at<int>(i, 0) = (cat == 1 || cat == 2) ? 1 : 0;
+    }
+    Mat varType(2, 1, CV_8U);
+    varType.at<uchar>(0) = ml::VAR_CATEGORICAL;
+    varType.at<uchar>(1) = ml::VAR_CATEGORICAL;
+
+    Ptr<ml::TrainData> td = ml::TrainData::create(samples, ml::ROW_SAMPLE, responses,
+                                                  noArray(), noArray(), noArray(), varType);
+    Ptr<ml::DTrees> dt = ml::DTrees::create();
+    dt->setMaxDepth(4);
+    dt->setCVFolds(0);
+    dt->setMaxCategories(4);
+    dt->setMinSampleCount(1);
+    dt->train(td);
+
+    const string filename = cv::tempfile(".yml");
+    dt->save(filename);
+
+    string model;
+    {
+        std::ifstream in(filename.c_str());
+        std::stringstream ss;
+        ss << in.rdbuf();
+        model = ss.str();
+    }
+
+    // Category values in the split list index a per-split subset bitmask.
+    // Injecting a value larger than the number of categories used to write
+    // past that bitmask; loading such a model must be rejected, not crash.
+    size_t pos = model.find("in:");
+    ASSERT_NE(pos, string::npos);
+    size_t br = model.find('[', pos);
+    ASSERT_NE(br, string::npos);
+    model.insert(br + 1, "1000000,");
+    {
+        std::ofstream out(filename.c_str());
+        out << model;
+    }
+
+    Ptr<ml::DTrees> bad;
+    EXPECT_THROW(bad = ml::DTrees::load(filename), Exception);
+    remove(filename.c_str());
+}
+
 /*TEST(ML_SVM, throw_exception_when_save_untrained_model)
 {
     Ptr<cv::ml::SVM> svm;