diff --git a/paligemma/image.cc b/paligemma/image.cc index 821f9bf0..b2c54195 100644 --- a/paligemma/image.cc +++ b/paligemma/image.cc @@ -83,8 +83,11 @@ const char* ParseUnsigned(const char* pos, const char* end, size_t& num) { } num = 0; for (; pos < end && std::isdigit(*pos); ++pos) { - num *= 10; - num += *pos - '0'; + const size_t digit = *pos - '0'; + if (num > (SIZE_MAX - digit) / 10) { + return nullptr; // overflow + } + num = num * 10 + digit; } return pos; } @@ -136,6 +139,14 @@ bool Image::ReadPPM(const hwy::Span& buf) { return false; } ++pos; + if (width == 0 || height == 0) { + HWY_ABORT("Invalid zero dimension\n"); + return false; + } + if (width > SIZE_MAX / 3 || width * 3 > SIZE_MAX / height) { + HWY_ABORT("Image dimensions overflow\n"); + return false; + } const size_t data_size = width * height * 3; if (buf.cend() - pos < static_cast(data_size)) { std::cerr << "Insufficient data remaining\n"; diff --git a/util/basics.h b/util/basics.h index bfab3f89..28e8ee2b 100644 --- a/util/basics.h +++ b/util/basics.h @@ -97,7 +97,12 @@ struct Extents2D { constexpr Extents2D() : rows(0), cols(0) {} constexpr Extents2D(size_t rows, size_t cols) : rows(rows), cols(cols) {} - size_t Area() const { return rows * cols; } + size_t Area() const { + if (rows != 0 && cols > SIZE_MAX / rows) { + HWY_ABORT("Tensor dimension overflow: rows=%zu cols=%zu", rows, cols); + } + return rows * cols; + } size_t rows; size_t cols;