diff --git a/test/correctness/image_io.cpp b/test/correctness/image_io.cpp index 427c3739fe7c..9e8f9211aec1 100644 --- a/test/correctness/image_io.cpp +++ b/test/correctness/image_io.cpp @@ -294,6 +294,38 @@ void test_read_big_endian_row_channel_offset() { } } +// A .npy whose 'descr' is missing its closing quote used to walk the header +// parser's cursor off the end of the buffer (sscanf matched %c%c%d and returned +// 3, but the %n was never reached, leaving the byte count uninitialized). The +// loader should just reject the file. +void test_malformed_npy_header() { + std::ostringstream o; + o << Internal::get_test_tmp_dir() << "malformed_descr.npy"; + std::string filename = o.str(); + + // v1 header: magic, version 1.0, 2-byte header length, then the dict. + // (6 + 2 + 2 + header_len) must be a multiple of 64. + std::string dict = "{'descr': '> 8) & 0xff)}; + fs.write(len_le, 2); + fs.write(dict.data(), dict.size()); + const std::vector payload(64, 0); + fs.write(payload.data(), payload.size()); + fs.close(); + + Buffer<> im; + if (Tools::load>(filename, &im)) { + std::cout << "Malformed .npy header was accepted by the loader\n"; + exit(1); + } +} + #ifndef HALIDE_NO_PNG void test_png_unsupported_bit_depth() { // A 1-bit grayscale PNG is a valid file, but load_png only supports 8- and @@ -352,6 +384,7 @@ int main(int argc, char **argv) { do_test(); test_mat_header(); test_read_big_endian_row_channel_offset(); + test_malformed_npy_header(); #ifndef HALIDE_NO_PNG test_png_unsupported_bit_depth(); #endif diff --git a/tools/halide_image_io.h b/tools/halide_image_io.h index 32a923ef5c86..969d0cb00d48 100644 --- a/tools/halide_image_io.h +++ b/tools/halide_image_io.h @@ -1236,8 +1236,12 @@ struct NpyHeader { } while (true) { char endian; - int consumed; - if (std::sscanf(ptr, "'descr': '%c%c%d'%n", &endian, &type_code, &type_bytes, &consumed) == 3) { + int consumed = 0; + // %n is not counted in sscanf's return value and is only reached + // if the trailing quote matches, so a descr missing its closing + // quote returns 3 with consumed still 0. Requiring consumed > 0 + // avoids advancing ptr by an indeterminate amount. + if (std::sscanf(ptr, "'descr': '%c%c%d'%n", &endian, &type_code, &type_bytes, &consumed) == 3 && consumed > 0) { if (endian != '<' && endian != '|') { return false; }