From 22e9bee4ef225c0edbb9323f94c26cee0c623497 Mon Sep 17 00:00:00 2001
From: Eric Soroos <eric-github@soroos.net>
Date: Sun, 7 Mar 2021 19:04:25 +0100
Subject: [PATCH] Fix DOS in PSDImagePlugin -- CVE-2021-28675

Conflict:NA
Reference:https://github.com/python-pillow/Pillow/commit/22e9bee4ef225c0edbb9323f94c26cee0c623497
---
 Tests/test_decompression_bomb.py |  1 +
 Tests/test_file_apng.py          |  2 +-
 Tests/test_file_blp.py           |  1 +
 Tests/test_file_tiff.py          |  6 ++++--
 src/PIL/ImageFile.py             | 14 ++++++++++++--
 src/PIL/PsdImagePlugin.py        | 32 +++++++++++++++++++++-----------
 6 files changed, 40 insertions(+), 16 deletions(-)

diff --git a/Tests/test_decompression_bomb.py b/Tests/test_decompression_bomb.py
index 7671cdc..f96a15a 100644
--- a/Tests/test_decompression_bomb.py
+++ b/Tests/test_decompression_bomb.py
@@ -52,6 +52,7 @@ class TestDecompressionBomb:
             with Image.open(TEST_FILE):
                 pass
 
+    @pytest.mark.xfail(reason="different exception")
     def test_exception_ico(self):
         with pytest.raises(Image.DecompressionBombError):
             Image.open("Tests/images/decompression_bomb.ico")
diff --git a/Tests/test_file_apng.py b/Tests/test_file_apng.py
index 97e2a15..8348da4 100644
--- a/Tests/test_file_apng.py
+++ b/Tests/test_file_apng.py
@@ -312,7 +312,7 @@ def test_apng_syntax_errors():
             exception = e
         assert exception is None
 
-    with pytest.raises(SyntaxError):
+    with pytest.raises(OSError):
         with Image.open("Tests/images/apng/syntax_num_frames_high.png") as im:
             im.seek(im.n_frames - 1)
             im.load()
diff --git a/Tests/test_file_blp.py b/Tests/test_file_blp.py
index 94c469c..1510614 100644
--- a/Tests/test_file_blp.py
+++ b/Tests/test_file_blp.py
@@ -1,4 +1,5 @@
 from PIL import Image
+import pytest
 
 from .helper import assert_image_equal
 
diff --git a/Tests/test_file_tiff.py b/Tests/test_file_tiff.py
index bb1bbda..1500ac8 100644
--- a/Tests/test_file_tiff.py
+++ b/Tests/test_file_tiff.py
@@ -612,8 +612,10 @@ class TestFileTiff:
     )
     def test_string_dimension(self):
         # Assert that an error is raised if one of the dimensions is a string
-        with pytest.raises(ValueError):
-            Image.open("Tests/images/string_dimension.tiff")
+        with pytest.raises(OSError):
+            with Image.open("Tests/images/string_dimension.tiff") as im:
+                im.load()
+
 
 
 @pytest.mark.skipif(not is_win32(), reason="Windows only")
diff --git a/src/PIL/ImageFile.py b/src/PIL/ImageFile.py
index f2a55cb..468314b 100644
--- a/src/PIL/ImageFile.py
+++ b/src/PIL/ImageFile.py
@@ -555,12 +555,18 @@ def _safe_read(fp, size):
 
     :param fp: File handle.  Must implement a <b>read</b> method.
     :param size: Number of bytes to read.
-    :returns: A string containing up to <i>size</i> bytes of data.
+    :returns: A string containing <i>size</i> bytes of data.
+
+    Raises an OSError if the file is truncated and the read can not be completed
+
     """
     if size <= 0:
         return b""
     if size <= SAFEBLOCK:
-        return fp.read(size)
+        data = fp.read(size)
+        if len(data) < size:
+            raise OSError("Truncated File Read")
+        return data
     data = []
     while size > 0:
         block = fp.read(min(size, SAFEBLOCK))
@@ -568,9 +574,13 @@ def _safe_read(fp, size):
             break
         data.append(block)
         size -= len(block)
+    if sum(len(d) for d in data) < size:
+        raise OSError("Truncated File Read")
     return b"".join(data)
 
 
+
+
 class PyCodecState:
     def __init__(self):
         self.xsize = 0
diff --git a/src/PIL/PsdImagePlugin.py b/src/PIL/PsdImagePlugin.py
index d3799ed..96de58f 100644
--- a/src/PIL/PsdImagePlugin.py
+++ b/src/PIL/PsdImagePlugin.py
@@ -119,7 +119,8 @@ class PsdImageFile(ImageFile.ImageFile):
             end = self.fp.tell() + size
             size = i32(read(4))
             if size:
-                self.layers = _layerinfo(self.fp)
+                _layer_data = io.BytesIO(ImageFile._safe_read(self.fp, size))
+                self.layers = _layerinfo(_layer_data, size)
             self.fp.seek(end)
         self.n_frames = len(self.layers)
         self.is_animated = self.n_frames > 1
@@ -170,12 +171,20 @@ class PsdImageFile(ImageFile.ImageFile):
         finally:
             self.__fp = None
 
-
-def _layerinfo(file):
+def _layerinfo(fp, ct_bytes):
     # read layerinfo block
     layers = []
-    read = file.read
-    for i in range(abs(i16(read(2)))):
+
+    def read(size):
+        return ImageFile._safe_read(fp, size)
+
+    ct = i16(read(2))
+
+    # sanity check
+    if ct_bytes < (abs(ct) * 20):
+        raise SyntaxError("Layer block too short for number of layers requested")
+
+    for i in range(abs(ct)):
 
         # bounding box
         y0 = i32(read(4))
@@ -186,7 +195,8 @@ def _layerinfo(file):
         # image info
         info = []
         mode = []
-        types = list(range(i16(read(2))))
+        ct_types = i16(read(2))
+        types = list(range(ct_types))
         if len(types) > 4:
             continue
 
@@ -219,16 +229,16 @@ def _layerinfo(file):
         size = i32(read(4))  # length of the extra data field
         combined = 0
         if size:
-            data_end = file.tell() + size
+            data_end = fp.tell() + size
 
             length = i32(read(4))
             if length:
-                file.seek(length - 16, io.SEEK_CUR)
+                fp.seek(length - 16, io.SEEK_CUR)
             combined += length + 4
 
             length = i32(read(4))
             if length:
-                file.seek(length, io.SEEK_CUR)
+                fp.seek(length, io.SEEK_CUR)
             combined += length + 4
 
             length = i8(read(1))
@@ -238,7 +248,7 @@ def _layerinfo(file):
                 name = read(length).decode("latin-1", "replace")
             combined += length + 1
 
-            file.seek(data_end)
+            fp.seek(data_end)
         layers.append((name, mode, (x0, y0, x1, y1)))
 
     # get tiles
@@ -246,7 +256,7 @@ def _layerinfo(file):
     for name, mode, bbox in layers:
         tile = []
         for m in mode:
-            t = _maketile(file, m, bbox, 1)
+            t = _maketile(fp, m, bbox, 1)
             if t:
                 tile.extend(t)
         layers[i] = name, mode, bbox, tile
-- 
2.23.0

