Pillow/Tests/check_png_dos.py
Andrew Murray e19a1496c2
Use monkeypatch (#8707)
Co-authored-by: Andrew Murray <radarhere@users.noreply.github.com>
2025-01-27 12:17:51 +02:00

65 lines
1.7 KiB
Python

from __future__ import annotations
import zlib
from io import BytesIO
import pytest
from PIL import Image, ImageFile, PngImagePlugin
TEST_FILE = "Tests/images/png_decompression_dos.png"
def test_ignore_dos_text(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
with Image.open(TEST_FILE) as im:
im.load()
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"