Fix reading multi-sector files whose sectors are stored verbatim

mpyq.py

@@ -224,7 +224,8 @@ def decompress(data):
                 # File consists of many sectors. They all need to be
                 # decompressed separately and united.
                 sector_size = 512 << self.header['sector_size_shift']
-                sectors = block_entry.size // sector_size + 1
+                sectors = (block_entry.size + sector_size - 1) // sector_size
+                sectors = sectors or 1
                 if block_entry.flags & MPQ_FILE_SECTOR_CRC:
                     crc = True
                     sectors += 1
@@ -236,8 +237,14 @@ def decompress(data):
                 sector_bytes_left = block_entry.size
                 for i in range(len(positions) - (2 if crc else 1)):
                     sector = file_data[positions[i]:positions[i+1]]
+                    # A sector is only compressed if its stored size is smaller
+                    # than the amount of plain data it holds, which is a full
+                    # sector except for the last one. Comparing against every
+                    # remaining byte instead would treat verbatim-stored
+                    # sectors as compressed in any multi-sector file.
+                    plain_size = min(sector_size, sector_bytes_left)
                     if (block_entry.flags & MPQ_FILE_COMPRESS and
-                        (force_decompress or sector_bytes_left > len(sector))):
+                        (force_decompress or plain_size > len(sector))):
                         sector = decompress(sector)
 
                     sector_bytes_left -= len(sector)

test/mpqbuilder.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+"""
+A minimal MPQ writer, used by the tests to build archives that exercise
+specific storage layouts.
+
+Only what the tests need: format version 1, unencrypted, files split into
+sectors with a sector offset table. Each sector is deflated and kept
+compressed only when that actually saves space, which is what Storm does --
+so an incompressible file ends up with its sectors stored verbatim even
+though MPQ_FILE_COMPRESS is set on the block entry.
+"""
+
+from __future__ import division
+
+import struct
+import zlib
+
+MPQ_FILE_COMPRESS = 0x00000200
+MPQ_FILE_EXISTS = 0x80000000
+
+HASH_TABLE_OFFSET = 0
+HASH_NAME_A = 1
+HASH_NAME_B = 2
+HASH_FILE_KEY = 3
+
+HEADER_SIZE = 32
+HASH_ENTRY_FREE = 0xFFFFFFFF
+
+
+def _build_crypt_table():
+    table = [0] * 0x500
+    seed = 0x00100001
+    for i in range(0x100):
+        index = i
+        for _ in range(5):
+            seed = (seed * 125 + 3) % 0x2AAAAB
+            temp1 = (seed & 0xFFFF) << 16
+            seed = (seed * 125 + 3) % 0x2AAAAB
+            temp2 = seed & 0xFFFF
+            table[index] = temp1 | temp2
+            index += 0x100
+    return table
+
+
+_CRYPT_TABLE = _build_crypt_table()
+
+
+def _hash(string, hash_type):
+    seed1 = 0x7FED7FED
+    seed2 = 0xEEEEEEEE
+    for ch in string.upper():
+        value = ord(ch)
+        seed1 = _CRYPT_TABLE[(hash_type << 8) + value] ^ ((seed1 + seed2) & 0xFFFFFFFF)
+        seed2 = (value + seed1 + seed2 + (seed2 << 5) + 3) & 0xFFFFFFFF
+    return seed1
+
+
+def _encrypt(values, key):
+    seed = 0xEEEEEEEE
+    out = []
+    for value in values:
+        seed = (seed + _CRYPT_TABLE[0x400 + (key & 0xFF)]) & 0xFFFFFFFF
+        out.append(value ^ ((key + seed) & 0xFFFFFFFF))
+        key = ((((~key) & 0xFFFFFFFF) << 0x15) + 0x11111111 | (key >> 0x0B)) & 0xFFFFFFFF
+        seed = (value + seed + (seed << 5) + 3) & 0xFFFFFFFF
+    return struct.pack('<%dI' % len(out), *out)
+
+
+def _pack_sectors(data, sector_size):
+    """Split into sectors, deflating each one only when it pays off."""
+    sectors = [data[i:i + sector_size]
+               for i in range(0, len(data), sector_size)]
+    if not sectors:
+        sectors = [b'']
+
+    stored = []
+    for sector in sectors:
+        squeezed = b'\x02' + zlib.compress(sector, 9)
+        stored.append(squeezed if len(squeezed) < len(sector) else sector)
+
+    positions = [4 * (len(stored) + 1)]
+    for sector in stored:
+        positions.append(positions[-1] + len(sector))
+
+    table = struct.pack('<%dI' % len(positions), *positions)
+    return table + b''.join(stored)
+
+
+def build_archive(files, sector_size_shift=3, hash_table_size=16):
+    """Build an MPQ archive.
+
+    `files` is a list of (archive name, contents) pairs. Returns the raw
+    archive bytes.
+    """
+    sector_size = 512 << sector_size_shift
+
+    block_table = []
+    payload = []
+    offset = HEADER_SIZE
+    for _name, data in files:
+        stored = _pack_sectors(data, sector_size)
+        block_table.append((offset, len(stored), len(data),
+                            MPQ_FILE_EXISTS | MPQ_FILE_COMPRESS))
+        payload.append(stored)
+        offset += len(stored)
+
+    hash_table = [[HASH_ENTRY_FREE] * 4 for _ in range(hash_table_size)]
+    for block_index, (name, _data) in enumerate(files):
+        start = _hash(name, HASH_TABLE_OFFSET) & (hash_table_size - 1)
+        for probe in range(hash_table_size):
+            slot = (start + probe) % hash_table_size
+            if hash_table[slot][3] == HASH_ENTRY_FREE:
+                hash_table[slot] = [_hash(name, HASH_NAME_A),
+                                    _hash(name, HASH_NAME_B),
+                                    0,
+                                    block_index]
+                break
+        else:
+            raise ValueError("hash table is full")
+
+    hash_values = []
+    for entry in hash_table:
+        hash_values.extend(entry)
+    block_values = []
+    for entry in block_table:
+        block_values.extend(entry)
+
+    hash_data = _encrypt(hash_values, _hash("(hash table)", HASH_FILE_KEY))
+    block_data = _encrypt(block_values, _hash("(block table)", HASH_FILE_KEY))
+
+    body = b''.join(payload)
+    hash_offset = HEADER_SIZE + len(body)
+    block_offset = hash_offset + len(hash_data)
+
+    header = struct.pack('<4s2I2H4I',
+                         b'MPQ\x1a',
+                         HEADER_SIZE,
+                         block_offset + len(block_data),
+                         0,
+                         sector_size_shift,
+                         hash_offset,
+                         block_offset,
+                         hash_table_size,
+                         len(block_table))
+
+    return header + body + hash_data + block_data

test/test_multi_sector.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+"""
+Regression tests for reading files whose sectors are stored verbatim.
+
+MPQ_FILE_COMPRESS on a block entry means "this file has a sector offset
+table", not "every sector is compressed". Storm deflates each sector on its
+own and keeps the compressed form only when it is smaller, so an
+incompressible file has MPQ_FILE_COMPRESS set while every one of its sectors
+is stored as-is.
+
+Deciding whether a sector is compressed therefore has to compare its stored
+size against the plain size of that sector -- a full sector except for the
+last one. Comparing against all the bytes still owed by the file instead
+holds for every sector but the last, so verbatim sectors in a multi-sector
+file get handed to the decompressor and it raises on the first byte.
+"""
+
+import hashlib
+import os
+import tempfile
+import unittest
+
+from mpyq import MPQArchive
+
+from .mpqbuilder import build_archive
+
+SECTOR_SIZE = 4096
+
+
+def incompressible(size):
+    """Deterministic bytes that deflate larger than they started."""
+    out = bytearray()
+    digest = hashlib.sha256(b'mpyq-multi-sector').digest()
+    while len(out) < size:
+        digest = hashlib.sha256(digest).digest()
+        out.extend(digest)
+    return bytes(out[:size])
+
+
+def compressible(size):
+    return (b'the quick brown fox ' * (size // 20 + 1))[:size]
+
+
+class TestMultiSectorStoredSectors(unittest.TestCase):
+
+    # Each case is (name, contents). Sizes are chosen to span several
+    # sectors, including one that is an exact multiple of the sector size.
+    CASES = [
+        ('stored.bin', incompressible(3 * SECTOR_SIZE + 123)),
+        ('deflated.bin', compressible(3 * SECTOR_SIZE + 123)),
+        ('mixed.bin', compressible(SECTOR_SIZE) + incompressible(2 * SECTOR_SIZE)),
+        ('exact.bin', incompressible(2 * SECTOR_SIZE)),
+        ('single.bin', incompressible(64)),
+        ('empty.bin', b''),
+    ]
+
+    def setUp(self):
+        listfile = b'\r\n'.join(name.encode('ascii') for name, _ in self.CASES)
+        contents = self.CASES + [('(listfile)', listfile)]
+        handle, self.path = tempfile.mkstemp(suffix='.mpq')
+        with os.fdopen(handle, 'wb') as archive:
+            archive.write(build_archive(contents))
+        self.archive = MPQArchive(self.path)
+
+    def tearDown(self):
+        self.archive.close()
+        self.archive = None
+        os.unlink(self.path)
+
+    def test_sectors_roundtrip(self):
+        for name, expected in self.CASES:
+            actual = self.archive.read_file(name)
+            if not expected:
+                # An empty file has no sectors to read back.
+                self.assertFalse(actual)
+                continue
+            self.assertEqual(len(actual), len(expected),
+                             "%s: wrong length" % name)
+            self.assertEqual(actual, expected, "%s: wrong contents" % name)
+
+    def test_incompressible_file_is_stored_verbatim(self):
+        # Guards the premise of the test: if this file were compressed, it
+        # would not exercise the verbatim path at all.
+        name, expected = self.CASES[0]
+        entry = self.archive.get_hash_table_entry(name)
+        block = self.archive.block_table[entry.block_table_index]
+        self.assertTrue(block.flags & 0x00000200)  # MPQ_FILE_COMPRESS
+        self.assertGreater(block.archived_size, len(expected))
+
+    def test_exact_multiple_of_sector_size(self):
+        name, expected = self.CASES[3]
+        self.assertEqual(len(expected) % SECTOR_SIZE, 0)
+        self.assertEqual(self.archive.read_file(name), expected)
+
+
+if __name__ == '__main__':
+    unittest.main()