mirror of
https://github.com/python-pillow/Pillow.git
synced 2024-09-20 19:08:58 +03:00
66 lines
1.6 KiB
Python
66 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
import zlib
|
|
from io import BytesIO
|
|
|
|
from PIL import Image, ImageFile, PngImagePlugin
|
|
|
|
TEST_FILE = "Tests/images/png_decompression_dos.png"
|
|
|
|
|
|
def test_ignore_dos_text() -> None:
|
|
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
|
|
|
try:
|
|
im = Image.open(TEST_FILE)
|
|
im.load()
|
|
finally:
|
|
ImageFile.LOAD_TRUNCATED_IMAGES = False
|
|
|
|
assert isinstance(im, PngImagePlugin.PngImageFile)
|
|
for s in im.text.values():
|
|
assert len(s) < 1024 * 1024, "Text chunk larger than 1M"
|
|
|
|
for s in im.info.values():
|
|
assert len(s) < 1024 * 1024, "Text chunk larger than 1M"
|
|
|
|
|
|
def test_dos_text() -> None:
|
|
try:
|
|
im = Image.open(TEST_FILE)
|
|
im.load()
|
|
except ValueError as msg:
|
|
assert msg, "Decompressed Data Too Large"
|
|
return
|
|
|
|
assert isinstance(im, PngImagePlugin.PngImageFile)
|
|
for s in im.text.values():
|
|
assert len(s) < 1024 * 1024, "Text chunk larger than 1M"
|
|
|
|
|
|
def test_dos_total_memory() -> None:
|
|
im = Image.new("L", (1, 1))
|
|
compressed_data = zlib.compress(b"a" * 1024 * 1023)
|
|
|
|
info = PngImagePlugin.PngInfo()
|
|
|
|
for x in range(64):
|
|
info.add_text(f"t{x}", compressed_data, zip=True)
|
|
info.add_itxt(f"i{x}", compressed_data, zip=True)
|
|
|
|
b = BytesIO()
|
|
im.save(b, "PNG", pnginfo=info)
|
|
b.seek(0)
|
|
|
|
try:
|
|
im2 = Image.open(b)
|
|
except ValueError as msg:
|
|
assert "Too much memory" in str(msg)
|
|
return
|
|
|
|
total_len = 0
|
|
assert isinstance(im2, PngImagePlugin.PngImageFile)
|
|
for txt in im2.text.values():
|
|
total_len += len(txt)
|
|
assert total_len < 64 * 1024 * 1024, "Total text chunks greater than 64M"
|