Fix out-of-bounds read in decode_png (#6456)

* Fix out-of-bounds read in decode_png

* Fix lint

* Use size_t instead of int to fix sign-compare check

* Add a unit-test to decode currupted png file

* Fix linter

* Update torchvision/csrc/io/image/cpu/decode_png.cpp

* Replace png_error with TORCH_CHECK
Co-authored-by: Vasilis Vryniotis <datumbox@users.noreply.github.com>

* Update error message

* Update error message
Co-authored-by: vfdev <vfdev.5@gmail.com>

* Fix linter
Co-authored-by: Vasilis Vryniotis <datumbox@users.noreply.github.com>
Co-authored-by: vfdev <vfdev.5@gmail.com>

test/test_image.py

@@ -29,6 +29,7 @@ IMAGE_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
 FAKEDATA_DIR = os.path.join(IMAGE_ROOT, "fakedata")
 IMAGE_DIR = os.path.join(FAKEDATA_DIR, "imagefolder")
 DAMAGED_JPEG = os.path.join(IMAGE_ROOT, "damaged_jpeg")
+DAMAGED_PNG = os.path.join(IMAGE_ROOT, "damaged_png")
 ENCODE_JPEG = os.path.join(IMAGE_ROOT, "encode_jpeg")
 INTERLACED_PNG = os.path.join(IMAGE_ROOT, "interlaced_png")
 IS_WINDOWS = sys.platform in ("win32", "cygwin")
@@ -190,6 +191,8 @@ def test_decode_png_errors():
         decode_png(torch.empty((), dtype=torch.uint8))
     with pytest.raises(RuntimeError, match="Content is not png"):
         decode_png(torch.randint(3, 5, (300,), dtype=torch.uint8))
+    with pytest.raises(RuntimeError, match="Out of bound read in decode_png"):
+        decode_png(read_file(os.path.join(DAMAGED_PNG, "sigsegv.png")))
 
 
 @pytest.mark.parametrize(

torchvision/csrc/io/image/cpu/decode_png.cpp

@@ -41,7 +41,9 @@ torch::Tensor decode_png(
     TORCH_CHECK(info_ptr, "libpng info structure allocation failed!")
   }
 
-  auto datap = data.accessor<unsigned char, 1>().data();
+  auto accessor = data.accessor<unsigned char, 1>();
+  auto datap = accessor.data();
+  auto datap_len = accessor.size(0);
 
   if (setjmp(png_jmpbuf(png_ptr)) != 0) {
     png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
@@ -52,15 +54,22 @@ torch::Tensor decode_png(
 
   struct Reader {
     png_const_bytep ptr;
+    png_size_t count;
   } reader;
   reader.ptr = png_const_bytep(datap) + 8;
-
-  auto read_callback =
-      [](png_structp png_ptr, png_bytep output, png_size_t bytes) {
-        auto reader = static_cast<Reader*>(png_get_io_ptr(png_ptr));
-        std::copy(reader->ptr, reader->ptr + bytes, output);
-        reader->ptr += bytes;
-      };
+  reader.count = datap_len - 8;
+
+  auto read_callback = [](png_structp png_ptr,
+                          png_bytep output,
+                          png_size_t bytes) {
+    auto reader = static_cast<Reader*>(png_get_io_ptr(png_ptr));
+    TORCH_CHECK(
+        reader->count >= bytes,
+        "Out of bound read in decode_png. Probably, the input image is corrupted");
+    std::copy(reader->ptr, reader->ptr + bytes, output);
+    reader->ptr += bytes;
+    reader->count -= bytes;
+  };
   png_set_sig_bytes(png_ptr, 8);
   png_set_read_fn(png_ptr, &reader, read_callback);
   png_read_info(png_ptr, info_ptr);