Add input shapes validation to prevent heap overflow in Concat layer

Author: Deepdive543443

src/layer/arm/concat_arm.cpp

@@ -197,10 +197,20 @@ 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)
+            {
+                NCNN_LOGE("Sizes of bottom blobs must match except in axis 0. "
+                          "Expected size (w:%d h: %d d: %d) "
+                          "but got size (w:%d h: %d d: %d) "
+                          "for blob number %d in the list.",
+                          w, h, d, bottom_blob.w, bottom_blob.h, bottom_blob.d, (int)b);
+                return -1;
+            }
+
             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,20 @@ 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)
+            {
+                NCNN_LOGE("Sizes of bottom blobs must match except in axis 0. "
+                          "Expected size (w:%d h: %d d: %d) "
+                          "but got size (w:%d h: %d d: %d) "
+                          "for blob number %d in the list.",
+                          w, h, d, bottom_blob.w, bottom_blob.h, bottom_blob.d, (int)b);
+                return -1;
+            }
+
             top_channels += bottom_blob.c;
         }
 

src/layer/loongarch/concat_loongarch.cpp

@@ -176,10 +176,20 @@ 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)
+            {
+                NCNN_LOGE("Sizes of bottom blobs must match except in axis 0. "
+                          "Expected size (w:%d h: %d d: %d) "
+                          "but got size (w:%d h: %d d: %d) "
+                          "for blob number %d in the list.",
+                          w, h, d, bottom_blob.w, bottom_blob.h, bottom_blob.d, (int)b);
+                return -1;
+            }
+
             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,20 @@ 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)
+            {
+                NCNN_LOGE("Sizes of bottom blobs must match except in axis 0. "
+                          "Expected size (w:%d h: %d d: %d) "
+                          "but got size (w:%d h: %d d: %d) "
+                          "for blob number %d in the list.",
+                          w, h, d, bottom_blob.w, bottom_blob.h, bottom_blob.d, (int)b);
+                return -1;
+            }
+
             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,20 @@ 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)
+            {
+                NCNN_LOGE("Sizes of bottom blobs must match except in axis 0. "
+                          "Expected size (w:%d h: %d d: %d) "
+                          "but got size (w:%d h: %d d: %d) "
+                          "for blob number %d in the list.",
+                          w, h, d, bottom_blob.w, bottom_blob.h, bottom_blob.d, (int)b);
+                return -1;
+            }
+
             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,20 @@ 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)
+            {
+                NCNN_LOGE("Sizes of bottom blobs must match except in axis 0. "
+                          "Expected size (w:%d h: %d d: %d) "
+                          "but got size (w:%d h: %d d: %d) "
+                          "for blob number %d in the list.",
+                          w, h, d, bottom_blob.w, bottom_blob.h, bottom_blob.d, (int)b);
+                return -1;
+            }
+
             elemsize = std::min(elemsize, bottom_blob.elemsize);
             elempack = std::min(elempack, bottom_blob.elempack);
             top_channels += bottom_blob.c * bottom_blob.elempack;