From 6b9b11445739e3ad1d3ddd7f050fff0b64807d8f Mon Sep 17 00:00:00 2001 From: Eli Kobrin Date: Wed, 21 Jun 2023 17:54:18 +0300 Subject: [PATCH 1/2] Add data length check for png. --- test/assets/toosmall_png/heapbof.png | Bin 0 -> 7 bytes test/test_image.py | 3 +++ torchvision/csrc/io/image/cpu/decode_png.cpp | 1 + 3 files changed, 4 insertions(+) create mode 100644 test/assets/toosmall_png/heapbof.png diff --git a/test/assets/toosmall_png/heapbof.png b/test/assets/toosmall_png/heapbof.png new file mode 100644 index 0000000000000000000000000000000000000000..e720d1833423d20f7df5a5bab5411956ed01a879 GIT binary patch literal 7 OcmeAS@N;KiU;qFJNC9sE literal 0 HcmV?d00001 diff --git a/test/test_image.py b/test/test_image.py index 4c210ea7eef..0409453f894 100644 --- a/test/test_image.py +++ b/test/test_image.py @@ -32,6 +32,7 @@ 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") +TOOSMALL_PNG = os.path.join(IMAGE_ROOT, "toosmall_png") IS_WINDOWS = sys.platform in ("win32", "cygwin") PILLOW_VERSION = tuple(int(x) for x in PILLOW_VERSION.split(".")) @@ -193,6 +194,8 @@ def test_decode_png_errors(): 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"))) + with pytest.raises(RuntimeError, match="Content is too small for png"): + decode_png(read_file(os.path.join(TOOSMALL_PNG, "toosmall.png"))) @pytest.mark.parametrize( diff --git a/torchvision/csrc/io/image/cpu/decode_png.cpp b/torchvision/csrc/io/image/cpu/decode_png.cpp index b1ceaf1badd..d27eafe45a7 100644 --- a/torchvision/csrc/io/image/cpu/decode_png.cpp +++ b/torchvision/csrc/io/image/cpu/decode_png.cpp @@ -49,6 +49,7 @@ torch::Tensor decode_png( png_destroy_read_struct(&png_ptr, &info_ptr, nullptr); TORCH_CHECK(false, "Internal error."); } + TORCH_CHECK(datap_len >= 8, "Content is too small for png!") auto is_png = !png_sig_cmp(datap, 0, 8); TORCH_CHECK(is_png, "Content is not png!") From 4dd5ac4fe8494d06f30ac34eb2023425b6bbb6b5 Mon Sep 17 00:00:00 2001 From: Eli Kobrin Date: Wed, 21 Jun 2023 17:57:20 +0300 Subject: [PATCH 2/2] Fix test. --- test/test_image.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_image.py b/test/test_image.py index 0409453f894..b08dc2026d4 100644 --- a/test/test_image.py +++ b/test/test_image.py @@ -195,7 +195,7 @@ def test_decode_png_errors(): with pytest.raises(RuntimeError, match="Out of bound read in decode_png"): decode_png(read_file(os.path.join(DAMAGED_PNG, "sigsegv.png"))) with pytest.raises(RuntimeError, match="Content is too small for png"): - decode_png(read_file(os.path.join(TOOSMALL_PNG, "toosmall.png"))) + decode_png(read_file(os.path.join(TOOSMALL_PNG, "heapbof.png"))) @pytest.mark.parametrize(