-
-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
check_j2k_leaks.py
executable file
·44 lines (35 loc) · 1.38 KB
/
check_j2k_leaks.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
from .helper import unittest, PillowTestCase
import sys
from PIL import Image
from io import BytesIO
# Limits for testing the leak
mem_limit = 1024 * 1048576
stack_size = 8 * 1048576
iterations = int((mem_limit / stack_size) * 2)
codecs = dir(Image.core)
test_file = "Tests/images/rgb_trns_ycbc.jp2"
@unittest.skipIf(sys.platform.startswith("win32"), "requires Unix or macOS")
class TestJpegLeaks(PillowTestCase):
def setUp(self):
if "jpeg2k_encoder" not in codecs or "jpeg2k_decoder" not in codecs:
self.skipTest("JPEG 2000 support not available")
def test_leak_load(self):
from resource import setrlimit, RLIMIT_AS, RLIMIT_STACK
setrlimit(RLIMIT_STACK, (stack_size, stack_size))
setrlimit(RLIMIT_AS, (mem_limit, mem_limit))
for _ in range(iterations):
with Image.open(test_file) as im:
im.load()
def test_leak_save(self):
from resource import setrlimit, RLIMIT_AS, RLIMIT_STACK
setrlimit(RLIMIT_STACK, (stack_size, stack_size))
setrlimit(RLIMIT_AS, (mem_limit, mem_limit))
for _ in range(iterations):
with Image.open(test_file) as im:
im.load()
test_output = BytesIO()
im.save(test_output, "JPEG2000")
test_output.seek(0)
test_output.read()
if __name__ == "__main__":
unittest.main()