Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 33 additions & 0 deletions test/correctness/image_io.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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': '<f8"; // no closing quote after the type
const size_t header_len = 54;
dict.resize(header_len, ' ');

std::ofstream fs(filename.c_str(), std::ofstream::binary);
const char magic[8] = {'\x93', 'N', 'U', 'M', 'P', 'Y', '\x01', '\x00'};
fs.write(magic, 8);
const char len_le[2] = {(char)(header_len & 0xff), (char)((header_len >> 8) & 0xff)};
fs.write(len_le, 2);
fs.write(dict.data(), dict.size());
const std::vector<char> payload(64, 0);
fs.write(payload.data(), payload.size());
fs.close();

Buffer<> im;
if (Tools::load<Buffer<>>(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
Expand Down Expand Up @@ -352,6 +384,7 @@ int main(int argc, char **argv) {
do_test<double>();
test_mat_header();
test_read_big_endian_row_channel_offset();
test_malformed_npy_header();
#ifndef HALIDE_NO_PNG
test_png_unsupported_bit_depth();
#endif
Expand Down
8 changes: 6 additions & 2 deletions tools/halide_image_io.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
Loading