Catch truncated DDS files

PIL/DdsImagePlugin.py

@@ -215,8 +215,11 @@ class DdsImageFile(ImageFile.ImageFile):
         magic, header_size = struct.unpack("<II", self.fp.read(8))
         if header_size != 124:
             raise IOError("Unsupported header size %r" % (header_size))
-        header = BytesIO(self.fp.read(header_size - 4))
-
+        header_bytes = self.fp.read(header_size - 4)
+        if len(header_bytes) != 120:
+            raise IOError("Incomplete header: %s bytes" % len(header_bytes))
+        header = BytesIO(header_bytes)
+        
         flags, height, width = struct.unpack("<3I", header.read(12))
         self.size = (width, height)
         self.mode = "RGBA"
@@ -244,11 +247,16 @@ class DdsImageFile(ImageFile.ImageFile):
             raise NotImplementedError("Unimplemented pixel format %r" %
                                       (fourcc))
 
-        decoded_data = codec(self.fp, self.width, self.height)
-
-        self.fp.close()
+        try:
+            decoded_data = codec(self.fp, self.width, self.height)
+        except struct.error:
+            raise IOError("Truncated DDS file")
+        finally:
+            self.fp.close()
+            
         self.fp = BytesIO(decoded_data)
 
+
     def load_seek(self, pos):
         pass
 

Tests/test_file_dds.py

@@ -1,5 +1,6 @@
-from helper import unittest, PillowTestCase
+from io import BytesIO
 
+from helper import unittest, PillowTestCase
 from PIL import Image, DdsImagePlugin
 
 TEST_FILE_DXT1 = "Tests/images/dxt1-rgb-4bbp-noalpha_MipMaps-1.dds"
@@ -44,7 +45,6 @@ class TestFileDds(PillowTestCase):
         # I'm currently seeing about a 3 for the epsilon. 
         self.assert_image_similar(target, im, 5)
 
-
     def test_sanity_dxt3(self):
         """Check DXT3 images are not supported"""
         self.assertRaises(NotImplementedError,
@@ -72,6 +72,26 @@ class TestFileDds(PillowTestCase):
         # Assert
         self.assertFalse(output)
 
+    def test_short_header(self):
+        """ Check a short header"""
+        with open(TEST_FILE_DXT5, 'rb') as f:
+            img_file = f.read()
+
+        def short_header():
+            im = Image.open(BytesIO(img_file[:119]))
+
+        self.assertRaises(IOError, short_header)
+
+    def test_short_file(self):
+        """ Check that the appropriate error is thrown for a short file"""
+        
+        with open(TEST_FILE_DXT5, 'rb') as f:
+            img_file = f.read()
+
+        def short_file():
+            im = Image.open(BytesIO(img_file[:-100]))
+
+        self.assertRaises(IOError, short_file)
 
 if __name__ == '__main__':
     unittest.main()