Fix crash, add assert and test

diff --git a/modules/imgproc/src/connectedcomponents.cpp b/modules/imgproc/src/connectedcomponents.cpp
index 1009584..9241c6c 100644 (file)
--- a/modules/imgproc/src/connectedcomponents.cpp
+++ b/modules/imgproc/src/connectedcomponents.cpp
@@ -2542,7 +2542,8 @@ namespace cv{
 
             //Array used to store info and labeled pixel by each thread.
             //Different threads affect different memory location of chunksSizeAndLabels
-            int *chunksSizeAndLabels = (int *)cv::fastMalloc(h * sizeof(int));
+            const int chunksSizeAndLabelsSize = h + 1;
+            int *chunksSizeAndLabels = (int *)cv::fastMalloc(chunksSizeAndLabelsSize * sizeof(int));
 
             //Tree of labels
             LabelT *P = (LabelT *)cv::fastMalloc(Plength * sizeof(LabelT));
@@ -2561,6 +2562,7 @@ namespace cv{
 
             LabelT nLabels = 1;
             for (int i = 0; i < h; i = chunksSizeAndLabels[i]){
+                CV_Assert(i + 1 < chunksSizeAndLabelsSize);
                 flattenL(P, LabelT((i + 1) / 2) * LabelT((w + 1) / 2) + 1, chunksSizeAndLabels[i + 1], nLabels);
             }
 

diff --git a/modules/imgproc/test/test_connectedcomponents.cpp b/modules/imgproc/test/test_connectedcomponents.cpp
index abd6fd4..3817f6d 100644 (file)
--- a/modules/imgproc/test/test_connectedcomponents.cpp
+++ b/modules/imgproc/test/test_connectedcomponents.cpp
@@ -136,4 +136,18 @@ void CV_ConnectedComponentsTest::run( int /* start_from */)
 
 TEST(Imgproc_ConnectedComponents, regression) { CV_ConnectedComponentsTest test; test.safe_run(); }
 
+TEST(Imgproc_ConnectedComponents, grana_buffer_overflow)
+{
+    cv::Mat darkMask;
+    darkMask.create(31, 87, CV_8U);
+    darkMask = 0;
+
+    cv::Mat labels;
+    cv::Mat stats;
+    cv::Mat centroids;
+
+    int nbComponents = cv::connectedComponentsWithStats(darkMask, labels, stats, centroids, 8, CV_32S, cv::CCL_GRANA);
+    EXPECT_EQ(1, nbComponents);
+}
+
 }} // namespace