Add input shapes validation to prevent heap overflow in Concat layer

Author: Deepdive543443

src/layer/arm/concat_arm.cpp

@@ -197,10 +197,13 @@ int Concat_arm::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>&
         // total channels
         size_t elemsize = bottom_blobs[0].elemsize;
         int elempack = bottom_blobs[0].elempack;
-        int top_channels = 0;
-        for (size_t b = 0; b < bottom_blobs.size(); b++)
+        int top_channels = bottom_blobs[0].c * bottom_blobs[0].elempack;
+        for (size_t b = 1; b < bottom_blobs.size(); b++)
         {
             const Mat& bottom_blob = bottom_blobs[b];
+            if (bottom_blob.w != w || bottom_blob.h != h || bottom_blob.d != d)
+                return -100;
+
             elemsize = std::min(elemsize, bottom_blob.elemsize);
             elempack = std::min(elempack, bottom_blob.elempack);
             top_channels += bottom_blob.c * bottom_blob.elempack;

src/layer/concat.cpp

@@ -128,10 +128,13 @@ int Concat::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_
         int d = bottom_blobs[0].d;
 
         // total channels
-        int top_channels = 0;
-        for (size_t b = 0; b < bottom_blobs.size(); b++)
+        int top_channels = bottom_blobs[0].c;
+        for (size_t b = 1; b < bottom_blobs.size(); b++)
         {
             const Mat& bottom_blob = bottom_blobs[b];
+            if (bottom_blob.w != w || bottom_blob.h != h || bottom_blob.d != d)
+                return -100;
+
             top_channels += bottom_blob.c;
         }
 

src/layer/loongarch/concat_loongarch.cpp

@@ -176,10 +176,13 @@ int Concat_loongarch::forward(const std::vector<Mat>& bottom_blobs, std::vector<
         // total channels
         size_t elemsize = bottom_blobs[0].elemsize;
         int elempack = bottom_blobs[0].elempack;
-        int top_channels = 0;
-        for (size_t b = 0; b < bottom_blobs.size(); b++)
+        int top_channels = bottom_blobs[0].c * bottom_blobs[0].elempack;
+        for (size_t b = 1; b < bottom_blobs.size(); b++)
         {
             const Mat& bottom_blob = bottom_blobs[b];
+            if (bottom_blob.w != w || bottom_blob.h != h || bottom_blob.d != d)
+                return -100;
+
             elemsize = std::min(elemsize, bottom_blob.elemsize);
             elempack = std::min(elempack, bottom_blob.elempack);
             top_channels += bottom_blob.c * bottom_blob.elempack;

src/layer/mips/concat_mips.cpp

@@ -176,10 +176,13 @@ int Concat_mips::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>&
         // total channels
         size_t elemsize = bottom_blobs[0].elemsize;
         int elempack = bottom_blobs[0].elempack;
-        int top_channels = 0;
-        for (size_t b = 0; b < bottom_blobs.size(); b++)
+        int top_channels = bottom_blobs[0].c * bottom_blobs[0].elempack;
+        for (size_t b = 1; b < bottom_blobs.size(); b++)
         {
             const Mat& bottom_blob = bottom_blobs[b];
+            if (bottom_blob.w != w || bottom_blob.h != h || bottom_blob.d != d)
+                return -100;
+
             elemsize = std::min(elemsize, bottom_blob.elemsize);
             elempack = std::min(elempack, bottom_blob.elempack);
             top_channels += bottom_blob.c * bottom_blob.elempack;

src/layer/riscv/concat_riscv.cpp

@@ -218,10 +218,13 @@ int Concat_riscv::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>
         // total channels
         size_t elemsize = bottom_blobs[0].elemsize;
         int elempack = bottom_blobs[0].elempack;
-        int top_channels = 0;
-        for (size_t b = 0; b < bottom_blobs.size(); b++)
+        int top_channels = bottom_blobs[0].c * bottom_blobs[0].elempack;
+        for (size_t b = 1; b < bottom_blobs.size(); b++)
         {
             const Mat& bottom_blob = bottom_blobs[b];
+            if (bottom_blob.w != w || bottom_blob.h != h || bottom_blob.d != d)
+                return -100;
+
             elemsize = std::min(elemsize, bottom_blob.elemsize);
             elempack = std::min(elempack, bottom_blob.elempack);
             top_channels += bottom_blob.c * bottom_blob.elempack;

src/layer/x86/concat_x86.cpp

@@ -388,10 +388,13 @@ int Concat_x86::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>&
         // total channels
         size_t elemsize = bottom_blobs[0].elemsize;
         int elempack = bottom_blobs[0].elempack;
-        int top_channels = 0;
-        for (size_t b = 0; b < bottom_blobs.size(); b++)
+        int top_channels = bottom_blobs[0].c * bottom_blobs[0].elempack;
+        for (size_t b = 1; b < bottom_blobs.size(); b++)
         {
             const Mat& bottom_blob = bottom_blobs[b];
+            if (bottom_blob.w != w || bottom_blob.h != h || bottom_blob.d != d)
+                return -100;
+
             elemsize = std::min(elemsize, bottom_blob.elemsize);
             elempack = std::min(elempack, bottom_blob.elempack);
             top_channels += bottom_blob.c * bottom_blob.elempack;