From c5c648a6a33372b90e9fbf0c0a7bc4e75bdd238c Mon Sep 17 00:00:00 2001
From: wiredfool <eric-github@soroos.net>
Date: Mon, 29 Dec 2014 17:10:27 -0800
Subject: [PATCH] Limit total text chunk size to 64k

---
 PIL/PngImagePlugin.py  | 23 +++++++++++++++++++++--
 Tests/check_png_dos.py | 30 +++++++++++++++++++++++++++---
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/PIL/PngImagePlugin.py b/PIL/PngImagePlugin.py
index 8bfdc2ac7..598395c0a 100644
--- a/PIL/PngImagePlugin.py
+++ b/PIL/PngImagePlugin.py
@@ -72,9 +72,15 @@ _MODES = {
 
 _simple_palette = re.compile(b'^\xff+\x00\xff*$')
 
+# Maximum decompressed size for a iTXt or zTXt chunk.
+# Eliminates decompression bombs where compressed chunks can expand 1000x
+MAX_TEXT_CHUNK = ImageFile.SAFEBLOCK
+# Set the maximum total text chunk size.
+MAX_TEXT_MEMORY = 64 * MAX_TEXT_CHUNK
+
 def _safe_zlib_decompress(s):
     dobj = zlib.decompressobj()
-    plaintext = dobj.decompress(s, ImageFile.SAFEBLOCK)
+    plaintext = dobj.decompress(s, MAX_TEXT_CHUNK)
     if dobj.unconsumed_tail:
         raise ValueError("Decompressed Data Too Large")
     return plaintext
@@ -234,6 +240,14 @@ class PngStream(ChunkStream):
         self.im_tile = None
         self.im_palette = None
 
+        self.text_memory = 0
+
+    def check_text_memory(self, chunklen):
+        self.text_memory += chunklen
+        if self.text_memory > MAX_TEXT_MEMORY:
+            raise ValueError("Too much memory used in text chunks: %s>MAX_TEXT_MEMORY" %
+                             self.text_memory)
+
     def chunk_iCCP(self, pos, length):
 
         # ICC profile
@@ -346,6 +360,8 @@ class PngStream(ChunkStream):
                 v = v.decode('latin-1', 'replace')
 
             self.im_info[k] = self.im_text[k] = v
+            self.check_text_memory(len(v))
+
         return s
 
     def chunk_zTXt(self, pos, length):
@@ -375,6 +391,8 @@ class PngStream(ChunkStream):
                 v = v.decode('latin-1', 'replace')
 
             self.im_info[k] = self.im_text[k] = v
+            self.check_text_memory(len(v))
+
         return s
 
     def chunk_iTXt(self, pos, length):
@@ -410,7 +428,8 @@ class PngStream(ChunkStream):
                 return s
 
         self.im_info[k] = self.im_text[k] = iTXt(v, lang, tk)
-
+        self.check_text_memory(len(v))
+            
         return s
 
 
diff --git a/Tests/check_png_dos.py b/Tests/check_png_dos.py
index 4e9b76537..8f974d293 100644
--- a/Tests/check_png_dos.py
+++ b/Tests/check_png_dos.py
@@ -1,13 +1,12 @@
 from helper import unittest, PillowTestCase
 import sys
-from PIL import Image
+from PIL import Image, PngImagePlugin
 from io import BytesIO
+import zlib
 
 test_file = "Tests/images/png_decompression_dos.png"
 
-@unittest.skipIf(sys.platform.startswith('win32'), "requires Unix or MacOS")
 class TestPngDos(PillowTestCase):
-
     def test_dos_text(self):
 
         try:
@@ -20,5 +19,30 @@ class TestPngDos(PillowTestCase):
         for s in im.text.values():
             self.assert_(len(s) < 1024*1024, "Text chunk larger than 1M")
 
+    def test_dos_total_memory(self):
+        im = Image.new('L',(1,1))
+        compressed_data = zlib.compress('a'*1024*1023)
+
+        info = PngImagePlugin.PngInfo()
+
+        for x in range(64):
+            info.add_text('t%s'%x, compressed_data, 1)
+            info.add_itxt('i%s'%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:
+            self.assert_("Too much memory" in msg)
+            return
+
+        total_len = 0
+        for txt in im2.text.values():
+            total_len += len(txt)
+        self.assert_(total_len < 64*1024*1024)
+        
 if __name__ == '__main__':
     unittest.main()