Merge pull request #30559 from uwezkhan:flann-kmeans-load-roots

flann: reallocate the kmeans tree roots in loadIndex 🤖🤖🤖 - #30559

`KMeansIndex::loadIndex` reads `trees_` from the saved file and then calls `load_tree(stream, root_[i], i)` for every tree, but it never reallocates `root_`. That array was sized by the constructor from the `trees` index parameter. `cv::flann::Index::load` builds the index with only `algorithm` set, so `root_` always has the default single slot. A saved KMeans index whose `trees_` field is 2 or more makes `load_tree` write node pointers past the end of `root_`, and `free_centers()` and the destructor later read the same slots. With ASan this is a heap-buffer-overflow WRITE of size 8 in `load_tree`, 0 bytes after the 8-byte region allocated in the `KMeansIndex` constructor. The same path is reached through the composite and autotuned loaders, which call `KMeansIndex::loadIndex`.

The fix frees `root_` and reallocates it with the loaded tree count, as `HierarchicalClusteringIndex::loadIndex` and `KDTreeIndex::loadIndex` already do. Both arrays are value-initialized and reset to NULL before the reads, so a truncated file that throws partway through only leaves NULL slots for the destructor. Files saved by OpenCV always store the `trees` value they were built with, so they load exactly as before. Added `Flann_KMeansIndex.load_more_trees_than_default` to `opencv_test_flann`, which stores the same tree twice under `trees_ = 2`. It fails under ASan without the fix. With the fix it passes, along with the rest of `opencv_test_flann`.

### Pull Request Readiness Checklist

See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request

- [x] I agree to contribute to the project under Apache 2 License.
- [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV
- [x] The PR is proposed to the proper branch
- [ ] There is a reference to the original bug report and related work
- [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable
      Patch to opencv_extra has the same branch name.
- [ ] The feature is well documented and sample code can be built with the project CMake

modules/flann/include/opencv2/flann/kmeans_index.h

@@ -485,7 +485,10 @@ public:
         }
         if (root_!=NULL) {
             free_centers();
+            delete[] root_;
         }
+        indices_ = NULL;
+        root_ = NULL;
 
         load_value(stream, branching_);
         load_value(stream, iterations_);
@@ -493,7 +496,15 @@ public:
         load_value(stream, cb_index_);
         load_value(stream, trees_);
 
-        indices_ = new int*[trees_];
+        if (branching_<2) {
+            FLANN_THROW(cv::Error::StsError, "Branching factor must be at least 2");
+        }
+        if (trees_<1) {
+            FLANN_THROW(cv::Error::StsError, "Number of trees must be at least 1");
+        }
+
+        indices_ = new int*[trees_]();
+        root_ = new KMeansNodePtr[trees_]();
         for (int i=0; i<trees_; ++i) {
             indices_[i] = new int[size_];
             load_value(stream, *indices_[i], size_);
@@ -655,17 +666,27 @@ private:
 
     void load_tree(FILE* stream, KMeansNodePtr& node, int num)
     {
-        node = pool_.allocate<KMeansNode>();
-        load_value(stream, *node);
+        KMeansNodePtr loaded = pool_.allocate<KMeansNode>();
+        load_value(stream, *loaded);
+        // pivot and childs hold pointers from the saving process, clear them before the node
+        // is attached so free_centers() only sees NULL or owned memory if a later read throws
+        bool is_leaf = loaded->childs==NULL;
+        loaded->pivot = NULL;
+        loaded->childs = NULL;
+        node = loaded;
         node->pivot = new CentersType[veclen_];
         load_value(stream, *(node->pivot), (int)veclen_);
-        if (node->childs==NULL) {
+        if (is_leaf) {
             int indices_offset;
             load_value(stream, indices_offset);
             node->indices = indices_[num] + indices_offset;
         }
         else {
-            node->childs = pool_.allocate<KMeansNodePtr>(branching_);
+            KMeansNodePtr* childs = pool_.allocate<KMeansNodePtr>(branching_);
+            for(int i=0; i<branching_; ++i) {
+                childs[i] = NULL;
+            }
+            node->childs = childs;
             for(int i=0; i<branching_; ++i) {
                 load_tree(stream, node->childs[i], num);
             }
@@ -681,7 +702,9 @@ private:
         delete[] node->pivot;
         if (node->childs!=NULL) {
             for (int k=0; k<branching_; ++k) {
-                free_centers(node->childs[k]);
+                if (node->childs[k]!=NULL) {
+                    free_centers(node->childs[k]);
+                }
             }
         }
     }

modules/flann/test/test_kmeans_index.cpp

@@ -0,0 +1,123 @@
+// This file is part of OpenCV project.
+// It is subject to the license terms in the LICENSE file found in the top-level directory
+// of this distribution and at http://opencv.org/license.html.
+
+#include "test_precomp.hpp"
+
+namespace opencv_test { namespace {
+
+// header, distance type, then branching_, iterations_, memoryCounter_, cb_index_, trees_
+static const size_t kmeans_branching_ofs = sizeof(cvflann::IndexHeader) + sizeof(int);
+static const size_t kmeans_trees_ofs = kmeans_branching_ofs + 4 * sizeof(int);
+
+static std::vector<char> saveKMeansIndex(const Mat& data, const std::string& filename)
+{
+    {
+        cv::flann::Index index(data, cv::flann::KMeansIndexParams(), cvflann::FLANN_DIST_L2);
+        index.save(filename);
+    }
+    std::ifstream f(filename.c_str(), std::ios::binary);
+    return std::vector<char>((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
+}
+
+static void writeFile(const std::string& filename, const std::vector<char>& buf)
+{
+    std::ofstream f(filename.c_str(), std::ios::binary);
+    f.write(&buf[0], buf.size());
+}
+
+static int getInt(const std::vector<char>& buf, size_t ofs)
+{
+    int value = 0;
+    memcpy(&value, &buf[ofs], sizeof(int));
+    return value;
+}
+
+static void setInt(std::vector<char>& buf, size_t ofs, int value)
+{
+    memcpy(&buf[ofs], &value, sizeof(int));
+}
+
+TEST(Flann_KMeansIndex, load_more_trees_than_default)
+{
+    Mat data(100, 8, CV_32FC1);
+    randu(data, 0.f, 1.f);
+    std::string filename = cv::tempfile(".fln");
+    std::vector<char> buf = saveKMeansIndex(data, filename);
+    ASSERT_GT(buf.size(), kmeans_trees_ofs + sizeof(int));
+    ASSERT_EQ(1, getInt(buf, kmeans_trees_ofs));
+
+    // store the same tree twice: valid content, but more trees than the default the index was created with
+    setInt(buf, kmeans_trees_ofs, 2);
+    std::vector<char> tree(buf.begin() + kmeans_trees_ofs + sizeof(int), buf.end());
+    buf.insert(buf.end(), tree.begin(), tree.end());
+    writeFile(filename, buf);
+
+    cv::flann::Index index;
+    ASSERT_TRUE(index.load(data, filename));
+    remove(filename.c_str());
+
+    Mat indices, dists;
+    index.knnSearch(data.row(7), indices, dists, 1, cv::flann::SearchParams(-1));
+    EXPECT_EQ(7, indices.at<int>(0));
+    // enough checks to go through every tree, not just root_[0]
+    index.knnSearch(data.row(7), indices, dists, 1, cv::flann::SearchParams(2 * data.rows));
+    EXPECT_EQ(7, indices.at<int>(0));
+}
+
+TEST(Flann_KMeansIndex, load_rejects_bad_branching_and_trees)
+{
+    Mat data(100, 8, CV_32FC1);
+    randu(data, 0.f, 1.f);
+    std::string filename = cv::tempfile(".fln");
+    const std::vector<char> saved = saveKMeansIndex(data, filename);
+    ASSERT_GT(saved.size(), kmeans_trees_ofs + sizeof(int));
+
+    std::vector<char> buf = saved;
+    setInt(buf, kmeans_trees_ofs, 0);
+    writeFile(filename, buf);
+    {
+        cv::flann::Index index;
+        EXPECT_THROW(index.load(data, filename), cv::Exception);
+    }
+
+    buf = saved;
+    setInt(buf, kmeans_branching_ofs, 1);
+    writeFile(filename, buf);
+    {
+        cv::flann::Index index;
+        EXPECT_THROW(index.load(data, filename), cv::Exception);
+    }
+    remove(filename.c_str());
+}
+
+TEST(Flann_KMeansIndex, destroy_after_truncated_load)
+{
+    Mat data(100, 8, CV_32FC1);
+    randu(data, 0.f, 1.f);
+    cvflann::Matrix<float> dataset((float*)data.data, data.rows, data.cols);
+    std::string filename = cv::tempfile(".fln");
+    const std::vector<char> saved = saveKMeansIndex(data, filename);
+    ASSERT_GT(saved.size(), kmeans_trees_ofs + sizeof(int));
+
+    // every cut after trees_ stops inside the indices, a node, a pivot or a leaf offset,
+    // and the destructor then frees whatever part of the tree was loaded
+    for (size_t len = kmeans_trees_ofs + sizeof(int); len < saved.size(); ++len)
+    {
+        SCOPED_TRACE(cv::format("File length: %d", (int)len));
+        writeFile(filename, std::vector<char>(saved.begin(), saved.begin() + len));
+        FILE* fin = fopen(filename.c_str(), "rb");
+        ASSERT_TRUE(fin != NULL);
+        cvflann::load_header(fin);
+        int distType = 0;
+        cvflann::load_value(fin, distType);
+        {
+            cvflann::KMeansIndex<cvflann::L2<float> > index(dataset, cvflann::KMeansIndexParams());
+            EXPECT_THROW(index.loadIndex(fin), cv::Exception);
+        }
+        fclose(fin);
+    }
+    remove(filename.c_str());
+}
+
+}} // namespace