Verify w and h input multiplication overflow to rleEncode() (#1544)

openvinotoolkit · Jun 27, 2024 · d00d0cf · d00d0cf
1 parent c296000
commit d00d0cf
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/datumaro/capi/pybind.cpp b/src/datumaro/capi/pybind.cpp
@@ -50,7 +50,12 @@ struct RLE
 
 RLE rleEncode(const byte *M, siz h, siz w)
 {
+
     siz j, k, a = w * h;
+
+    if (a / w != h)
+        throw std::runtime_error("result overflow. input size is too big.");
+
     uint c;
     byte p;
     auto cnts = std::make_unique<std::vector<uint>>(a + 1);
@@ -82,6 +87,7 @@ py::dict pyRleEncode(py::array_t<std::uint8_t, py::array::f_style | py::array::f
 
     const siz h = buf.shape[0];
     const siz w = buf.shape[1];
+
     const auto rle = rleEncode(static_cast<const byte *>(buf.ptr), h, w);
 
     py::array_t<uint> cnts(rle.m);