reject non-positive extents in load_npy and load_mat

test/correctness/image_io.cpp

@@ -268,6 +268,76 @@ void test_mat_header() {
     }
 }
 
+// A negative extent in the header makes size_in_bytes() exceed the
+// allocation, so the payload read runs off the end of the buffer.
+void test_negative_extents() {
+    const auto write_file = [](const std::string &filename, const std::vector<uint8_t> &bytes) {
+        std::ofstream fs(filename.c_str(), std::ofstream::binary);
+        fs.write((const char *)bytes.data(), bytes.size());
+        fs.close();
+    };
+    const auto append32 = [](std::vector<uint8_t> &bytes, uint32_t value) {
+        for (int i = 0; i < 4; i++) {
+            bytes.push_back((uint8_t)(value >> (8 * i)));
+        }
+    };
+
+    {
+        // .npy with a shape of (2, -3)
+        std::string dict = "{'descr': '<i4', 'fortran_order': False, 'shape': (2, -3), }";
+        while ((6 + 2 + 2 + dict.size()) % 64 != 0) {
+            dict += ' ';
+        }
+        std::vector<uint8_t> bytes = {0x93, 'N', 'U', 'M', 'P', 'Y', 1, 0};
+        bytes.push_back((uint8_t)(dict.size() & 0xff));
+        bytes.push_back((uint8_t)((dict.size() >> 8) & 0xff));
+        bytes.insert(bytes.end(), dict.begin(), dict.end());
+        bytes.resize(bytes.size() + 16, 0);
+
+        std::ostringstream o;
+        o << Internal::get_test_tmp_dir() << "test_negative_extents.npy";
+        std::string filename = o.str();
+        write_file(filename, bytes);
+
+        Buffer<> buf;
+        if (Tools::load<Buffer<>, Tools::Internal::CheckReturn>(filename, &buf)) {
+            std::cout << "Loaded " << filename << " despite a negative extent\n";
+            exit(1);
+        }
+    }
+
+    {
+        // .mat with extents of (4, -4)
+        std::vector<uint8_t> bytes(128, 0);
+        append32(bytes, Tools::Internal::miMATRIX);
+        append32(bytes, 64);
+        append32(bytes, Tools::Internal::miUINT32);
+        append32(bytes, 8);
+        append32(bytes, Tools::Internal::miUINT8);
+        append32(bytes, 1);
+        append32(bytes, Tools::Internal::miINT32);
+        append32(bytes, 8);
+        append32(bytes, 4);
+        append32(bytes, (uint32_t)-4);
+        append32(bytes, Tools::Internal::miINT8);
+        append32(bytes, 0);
+        append32(bytes, Tools::Internal::miUINT8);
+        append32(bytes, 16);
+        bytes.resize(bytes.size() + 16, 0);
+
+        std::ostringstream o;
+        o << Internal::get_test_tmp_dir() << "test_negative_extents.mat";
+        std::string filename = o.str();
+        write_file(filename, bytes);
+
+        Buffer<> buf;
+        if (Tools::load<Buffer<>, Tools::Internal::CheckReturn>(filename, &buf)) {
+            std::cout << "Loaded " << filename << " despite a negative extent\n";
+            exit(1);
+        }
+    }
+}
+
 int main(int argc, char **argv) {
     do_test<int8_t>();
     do_test<int16_t>();
@@ -283,6 +353,7 @@ int main(int argc, char **argv) {
 #endif
     do_test<double>();
     test_mat_header();
+    test_negative_extents();
     printf("Success!\n");
     return 0;
 }

tools/halide_image_io.h

@@ -1339,6 +1339,12 @@ bool load_npy(const std::string &filename, ImageType *im) {
         return false;
     }
 
+    for (const int extent : h.extents) {
+        if (!check(extent > 0, "Bad extent in .npy file")) {
+            return false;
+        }
+    }
+
     halide_type_t im_type((halide_type_code_t)0, 0);
     for (const auto &d : npy_dtypes) {
         if (h.type_code == d.second.type_code && h.type_bytes == d.second.type_bytes) {
@@ -1783,6 +1789,12 @@ bool load_mat(const std::string &filename, ImageType *im) {
         }
     }
 
+    for (const int extent : extents) {
+        if (!check(extent > 0, "Bad extent in .mat file\n")) {
+            return false;
+        }
+    }
+
     // Skip over the name
     uint32_t name_header[2];
     if (!check(f.read_array(name_header), "Could not read .mat header\n")) {