bpo-5950: Support reading zips with comments in zipimport (#9548)

ZackerySpytz · warsaw · commit 5a5ce064b3ba · 2018-09-25T15:15:47.000-04:00
* bpo-5950: Support reading zips with comments in zipimport
diff --git a/Doc/library/zipimport.rst b/Doc/library/zipimport.rst
@@ -28,7 +28,8 @@ Any files may be present in the ZIP archive, but only files :file:`.py` and
 corresponding :file:`.pyc` file, meaning that if a ZIP archive
 doesn't contain :file:`.pyc` files, importing may be rather slow.
 
-ZIP archives with an archive comment are currently not supported.
+.. versionchanged:: 3.8
+   Previously, ZIP archives with an archive comment were not supported.
 
 .. seealso::
 
diff --git a/Lib/test/test_zipimport.py b/Lib/test/test_zipimport.py
@@ -116,6 +116,9 @@ def makeZip(self, files, zipName=TEMP_ZIP, **kw):
                 zinfo = ZipInfo(name, time.localtime(mtime))
                 zinfo.compress_type = self.compression
                 z.writestr(zinfo, data)
+            comment = kw.get("comment", None)
+            if comment is not None:
+                z.comment = comment
 
         stuff = kw.get("stuff", None)
         if stuff is not None:
@@ -665,6 +668,18 @@ def testBytesPath(self):
         with self.assertRaises(TypeError):
             zipimport.zipimporter(memoryview(os.fsencode(filename)))
 
+    def testComment(self):
+        files = {TESTMOD + ".py": (NOW, test_src)}
+        self.doTest(".py", files, TESTMOD, comment=b"comment")
+
+    def testBeginningCruftAndComment(self):
+        files = {TESTMOD + ".py": (NOW, test_src)}
+        self.doTest(".py", files, TESTMOD, stuff=b"cruft" * 64, comment=b"hi")
+
+    def testLargestPossibleComment(self):
+        files = {TESTMOD + ".py": (NOW, test_src)}
+        self.doTest(".py", files, TESTMOD, comment=b"c" * ((1 << 16) - 1))
+
 
 @support.requires_zlib
 class CompressedZipImportTestCase(UncompressedZipImportTestCase):
diff --git a/Lib/zipimport.py b/Lib/zipimport.py
@@ -38,6 +38,9 @@ class ZipImportError(ImportError):
 
 _module_type = type(sys)
 
+END_CENTRAL_DIR_SIZE = 22
+STRING_END_ARCHIVE = b'PK\x05\x06'
+MAX_COMMENT_LEN = (1 << 16) - 1
 
 class zipimporter:
     """zipimporter(archivepath) -> zipimporter object
@@ -354,16 +357,39 @@ def _read_directory(archive):
 
     with fp:
         try:
-            fp.seek(-22, 2)
+            fp.seek(-END_CENTRAL_DIR_SIZE, 2)
             header_position = fp.tell()
-            buffer = fp.read(22)
+            buffer = fp.read(END_CENTRAL_DIR_SIZE)
         except OSError:
             raise ZipImportError(f"can't read Zip file: {archive!r}", path=archive)
-        if len(buffer) != 22:
+        if len(buffer) != END_CENTRAL_DIR_SIZE:
             raise ZipImportError(f"can't read Zip file: {archive!r}", path=archive)
-        if buffer[:4] != b'PK\x05\x06':
+        if buffer[:4] != STRING_END_ARCHIVE:
             # Bad: End of Central Dir signature
-            raise ZipImportError(f'not a Zip file: {archive!r}', path=archive)
+            # Check if there's a comment.
+            try:
+                fp.seek(0, 2)
+                file_size = fp.tell()
+            except OSError:
+                raise ZipImportError(f"can't read Zip file: {archive!r}",
+                                     path=archive)
+            max_comment_start = max(file_size - MAX_COMMENT_LEN -
+                                    END_CENTRAL_DIR_SIZE, 0)
+            try:
+                fp.seek(max_comment_start)
+                data = fp.read()
+            except OSError:
+                raise ZipImportError(f"can't read Zip file: {archive!r}",
+                                     path=archive)
+            pos = data.rfind(STRING_END_ARCHIVE)
+            if pos < 0:
+                raise ZipImportError(f'not a Zip file: {archive!r}',
+                                     path=archive)
+            buffer = data[pos:pos+END_CENTRAL_DIR_SIZE]
+            if len(buffer) != END_CENTRAL_DIR_SIZE:
+                raise ZipImportError(f"corrupt Zip file: {archive!r}",
+                                     path=archive)
+            header_position = file_size - len(data) + pos
 
         header_size = _unpack_uint32(buffer[12:16])
         header_offset = _unpack_uint32(buffer[16:20])
diff --git a/Misc/NEWS.d/next/Library/2018-09-24-14-21-58.bpo-5950.xH0ekQ.rst b/Misc/NEWS.d/next/Library/2018-09-24-14-21-58.bpo-5950.xH0ekQ.rst
@@ -0,0 +1 @@
+Support reading zip files with archive comments in :mod:`zipimport`.
diff --git a/Python/importlib_zipimport.h b/Python/importlib_zipimport.h

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Support reading zip files with archive comments in :mod:`zipimport`.