packing/shader_split: read Phyre shaders from split KIT archives

dm/gfx/tests/test_pg_shader.py

@@ -47,6 +47,15 @@ def test_truncated_or_unrelated_data_is_not_a_shader(self):
                 split.scan_blob(bad, '/unrelated', found)
                 self.assertEqual(found, [])
 
+    def test_minimal_binding_table_is_not_a_phyre_record(self):
+        data, ctl, code = object_data('ps')
+        data = struct.pack('<I', 4) + data[0x88:]
+        found = []
+        split.scan_blob(archive(data), '/minimal.dat', found)
+        self.assertEqual(len(found), 1)
+        self.assertEqual([(s, c, k) for _, s, c, k in found], [('ps', ctl, code)])
+        self.assertIn('#PGSO', found[0][0])
+
     def test_control_and_code_lengths_must_agree(self):
         data = bytearray(object_data('vs')[0])
         struct.pack_into('<I', data, 0x8c + 0x100 + 4 + 0x704, 0x80)

dm/gfx/tests/test_vbf.py

@@ -1,11 +1,14 @@
-"""VBF block decoding and unaligned Phyre NVN programs, without game files."""
+"""VBF/KIT block decoding and unaligned Phyre NVN programs, without game files."""
 from pathlib import Path
+import hashlib
 import os
+import random
 import struct
 import sys
 import tempfile
 import unittest
 import zlib
+import lz4.block
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'tools'))
 os.environ.setdefault('DM_GAME', '010015100B514000')
@@ -51,6 +54,121 @@ def test_effect_and_loose_programs_use_same_key(self):
         bad = bytearray(effect); struct.pack_into('<I', bad, 13 + 0x100 + 0x704, 0x104)
         self.assertEqual([x[1] for x in split.phyre_shaders(bad, 'bad')], ['ps'])
         self.assertEqual(split.phyre_shaders(b'RYHP' + split.CTL_MAGIC, 'short'), [])
+        literal = bytearray(vs); literal[0x80:0x84] = split.CTL_MAGIC
+        self.assertEqual([x[1] for x in split.phyre_shaders(b'RYHP' + literal, 'literal')], ['vs'])
+
+    def test_headerless_compute_is_bounded_by_the_effect(self):
+        cs = bytearray(program('vs')); cs[:0x100] = bytes(0x100)
+        struct.pack_into('<3I', cs, 0x100 + 0x748, 8, 4, 1)
+        effect = b'RYHP' + bytes(9) + cs + program('ps')
+        out = []; split.scan_blob(effect, 'effect', out)
+        self.assertEqual([x[1] for x in out], ['cs', 'ps'])
+        self.assertIsNone(split.check('effect', *out[0][1:]))
+        self.assertEqual(split.phyre_shaders(cs, 'unframed'), [])
+        for at, value in ((0x748, 0), (0x704, 0x100000), (0x20, 0x100000)):
+            bad = bytearray(effect); struct.pack_into('<I', bad, 13 + 0x100 + at, value)
+            self.assertEqual([x[1] for x in split.phyre_shaders(bad, 'bad')], ['ps'])
+        bad = bytearray(effect); bad[13 + 4:13 + 8] = split.CTL_MAGIC
+        self.assertEqual([x[1] for x in split.phyre_shaders(bad, 'overlap')], ['ps'])
+
+    def test_effect_keeps_the_parent_archive_budget(self):
+        effect = b'RYHP' + bytes(9) + program('ps')
+        with self.assertRaisesRegex(SystemExit, 'archive members add up to more than'):
+            split.scan_blob(effect, '/member', [], left=[1, 1])
+
+
+class Kit(unittest.TestCase):
+    def archive(self, tmp, payloads):
+        """Two shards, mixed stored/compressed members, and a padded index."""
+        path = Path(tmp) / 'Data.kit'; header = bytearray(0x25 + 17)
+        header[:4] = split.KIT
+        struct.pack_into('<QI', header, 0x10, 17, len(payloads))
+        shards = [bytearray(), bytearray()]
+        for i, (payload, compressed) in enumerate(payloads):
+            data = bytearray()
+            if compressed:
+                for at in range(0, len(payload), 65536):
+                    block = lz4.block.compress(payload[at:at + 65536], store_size=False, dict=payload[max(0, at - 65536):at])
+                    data += struct.pack('<I', len(block)) + block
+            else: data += payload
+            part = i % 2
+            header += struct.pack('<IHIHQQ20sII', 9 if compressed else 1, part, len(shards[part]), 0, len(data), len(payload), hashlib.sha1(payload).digest(), 0, 0)
+            shards[part] += data
+        path.write_bytes(header)
+        for i, shard in enumerate(shards): Path(f'{path}.{i}').write_bytes(shard)
+        return path
+
+    def test_chained_dictionary_raw_members_and_scanner(self):
+        first = random.Random(7).randbytes(65536)
+        payload = first + first[-16000:]
+        effect = b'RYHP' + bytes(9) + program('vs') + program('ps')
+        with tempfile.TemporaryDirectory() as tmp:
+            path = self.archive(tmp, [(payload, True), (effect, False)])
+            entries = split.kit_entries(str(path))
+            self.assertEqual([split.kit_read(str(path), e) for e in entries], [payload, effect])
+            stored = Path(f'{path}.0').read_bytes(); at = 4 + struct.unpack_from('<I', stored)[0]
+            n = struct.unpack_from('<I', stored, at)[0]
+            with self.assertRaises(lz4.block.LZ4BlockError): lz4.block.decompress(stored[at + 4:at + 4 + n], uncompressed_size=16000)
+            files = [(str(p), p.name) for p in Path(tmp).iterdir()]
+            unrelated = str(path) + '.notes'; files.append((unrelated, 'Data.kit.notes'))
+            jobs, paths = split.kit_jobs(files, parts=1)
+            self.assertEqual(paths, {str(path), str(path) + '.0', str(path) + '.1'})
+            self.assertNotIn(unrelated, paths)
+            self.assertEqual(len(jobs), 1)
+            _, found, _, extra = split.scan_file(jobs[0])
+            self.assertEqual(extra, {})
+            self.assertEqual(set(found), {split.skey(c, k) for _, _, c, k in split.phyre_shaders(effect, 'loose')})
+
+    def test_raw_reads_decoded_size(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            path = self.archive(tmp, [(b'raw member', False)])
+            d = bytearray(path.read_bytes()); struct.pack_into('<Q', d, 0x25 + 17 + 12, 9999); path.write_bytes(d)
+            self.assertEqual(split.kit_read(str(path), split.kit_entries(str(path))[0]), b'raw member')
+
+    def test_shared_effect_keeps_game_rows_with_or_without_dlc(self):
+        effect = b'RYHP' + bytes(9) + program('vs') + program('ps')
+        with tempfile.TemporaryDirectory() as tmp:
+            path = self.archive(tmp, [(effect, True)])
+            jobs, _ = split.kit_jobs([(str(path), 'Data.kit')], parts=1)
+            game = split.scan_file(jobs[0])
+            dlc = split.scan_file((str(path), 'aoc/1/Data.kit', jobs[0][2]))
+            self.assertEqual(len(game[1]), 2)
+            self.assertTrue(all(v[0].startswith('/aoc/1/') for v in dlc[1].values()))
+            expected = split.merge_keys([game])
+            for results in ([dlc, game], [game, dlc]):
+                keys, pairs, per_src = split.merge_keys(results)
+                self.assertEqual((keys, pairs), expected[:2])
+                self.assertEqual(+per_src, +expected[2])
+
+    def test_index_and_member_bounds(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            path = self.archive(tmp, [(b'member', True)])
+            original = path.read_bytes(); base = 0x25 + 17
+            variants = [original[:20], original[:-1]]
+            for at, fmt, value in ((0x10, '<Q', len(original)), (0x18, '<I', 100), (base + 10, '<H', 1), (base + 6, '<I', 99999), (base + 12, '<Q', 99999)):
+                bad = bytearray(original); struct.pack_into(fmt, bad, at, value); variants.append(bad)
+            for bad in variants:
+                with self.subTest(index=bytes(bad).hex()):
+                    path.write_bytes(bad)
+                    with self.assertRaises(SystemExit): split.kit_entries(str(path))
+            path.write_bytes(original); Path(f'{path}.0').unlink()
+            with self.assertRaises(SystemExit): split.kit_entries(str(path))
+
+    def test_lz4_bounds_lengths_and_hash(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            path = self.archive(tmp, [(b'payload' * 100, True)])
+            entry = split.kit_entries(str(path))[0]; shard = Path(f'{path}.0'); original = shard.read_bytes()
+            cases = [original[:2], original[:-1], struct.pack('<I', 0), struct.pack('<I', len(original)), original + b'trailing']
+            for data in cases:
+                with self.subTest(data=data[:16]):
+                    shard.write_bytes(data); e = (*entry[:3], len(data), *entry[4:])
+                    with self.assertRaises(SystemExit): split.kit_read(str(path), e)
+            shard.write_bytes(original)
+            for size in (entry[4] - 1, entry[4] + 1):
+                with self.assertRaises(SystemExit): split.kit_read(str(path), (*entry[:4], size, *entry[5:]))
+            with self.assertRaises(SystemExit): split.kit_read(str(path), (*entry[:-1], bytes(20)))
+            shard.write_bytes(original[:-1])
+            with self.assertRaises(SystemExit): split.kit_read(str(path), entry)
 
 
 if __name__ == '__main__': unittest.main()

dm/gfx/tools/shader_split.py

@@ -2169,18 +2169,81 @@ def sumo_shaders(d, label):
         out.append(sumo_program(d, h, size, co, ko, cn, kn, f'{label}#NVNSH@{h:#x}', STAGES[stage]))
     return out
 
+KIT = b'kit!'
+KIT_CHUNK = 0x10000
+
+def kit_entries(path):
+    """Unravel's KIT index: a 0x25-byte header and an offset to 56-byte records, naming a .kit.N shard, stored/decoded sizes and the decoded SHA1. Flag 8 means chained LZ4 blocks."""
+    with open(path, 'rb') as f: d = f.read()
+    if len(d) < 0x25 or d[:4] != KIT: red(f'{path}: incomplete KIT header')
+    base = 0x25 + struct.unpack_from('<Q', d, 0x10)[0]
+    count = struct.unpack_from('<I', d, 0x18)[0]
+    if base > len(d) or count > (len(d) - base) // 56: red(f'{path}: KIT directory is outside the index')
+    entries, sizes = [], {}
+    for i in range(count):
+        flags, part, off, high, stored, size, sha, _, _ = struct.unpack_from('<IHIHQQ20sII', d, base + 56 * i)
+        if high: red(f'{path}: KIT entry {i} has an unsupported high offset word')
+        shard = f'{path}.{part}'
+        if part not in sizes:
+            if not os.path.isfile(shard): red(f'{path}: KIT entry {i} has no shard {part}')
+            sizes[part] = os.path.getsize(shard)
+        n = stored if flags & 8 else size
+        if off > sizes[part] or n > sizes[part] - off: red(f'{path}: KIT entry {i} is outside shard {part}')
+        entries.append((i, part, off, n, size, bool(flags & 8), sha))
+    return entries
+
+def kit_read(path, entry):
+    import lz4.block
+    i, part, off, n, size, compressed, sha = entry
+    label = f'{path}#{i}'
+    with open(f'{path}.{part}', 'rb') as f: f.seek(off); d = f.read(n)
+    if len(d) != n: red(f'{label}: truncated KIT member')
+    if compressed:
+        out, pos = bytearray(), 0
+        while pos + 4 <= len(d) and len(out) < size:
+            stored = struct.unpack_from('<I', d, pos)[0]; pos += 4
+            want = min(KIT_CHUNK, size - len(out))
+            if not 0 < stored <= len(d) - pos: red(f'{label}: LZ4 block is outside the KIT member')
+            try: block = lz4.block.decompress(d[pos:pos + stored], uncompressed_size=want, dict=bytes(out[-KIT_CHUNK:]))
+            except lz4.block.LZ4BlockError as e: red(f'{label}: LZ4 block at {pos - 4:#x}: {e}')
+            if len(block) != want: red(f'{label}: LZ4 block makes {len(block)} of {want} bytes')
+            out += block; pos += stored
+        if pos != len(d) or len(out) != size: red(f'{label}: incomplete KIT LZ4 stream ({len(out)} of {size} bytes, {pos} of {len(d)} stored)')
+        d = bytes(out)
+    if hashlib.sha1(d).digest() != sha: red(f'{label}: KIT SHA1 mismatch')
+    return d
+
+def kit_jobs(files, parts=64):
+    """Read KIT shards only through their index, in ranges of members. Hash-only names cannot select shaders, so every member is decoded and scanned."""
+    kits = []
+    for p, r in files:
+        if not p.lower().endswith('.kit'): continue
+        with open(p, 'rb') as f: head = f.read(4)
+        if head == KIT: kits.append((p, r, kit_entries(p)))
+    jobs, paths = [], set()
+    step = sum(e[3] for _, _, es in kits for e in es) / parts
+    for p, r, es in kits:
+        paths.add(p); paths.update(f'{p}.{e[1]}' for e in es)
+        cur, size = [], 0
+        for e in es:
+            cur.append(e); size += e[3]
+            if size >= step: jobs.append((p, r, (KIT, tuple(cur)))); cur, size = [], 0
+        if cur: jobs.append((p, r, (KIT, tuple(cur))))
+    return jobs, paths
+
 def phyre_shaders(d, label):
-    """Phyre's NVN code immediately precedes its control section, unaligned in RYHP effects; loose blobs also prefix their two lengths (FFX/X-2)."""
+    """Phyre's NVN code immediately precedes its control section, unaligned in RYHP effects; loose blobs also prefix their two lengths (FFX/X-2). Unravel TWO also has headerless compute code."""
     out = []
     for m in re.finditer(re.escape(CTL_MAGIC), d):
         c = m.start()
         if c + 0x708 > len(d): continue
         cn = struct.unpack_from('<I', d, c + 0x20)[0]
         kn = struct.unpack_from('<I', d, c + 0x704)[0]
         k = c - kn
-        if k < 0 or kn < 0x80 or cn < 0x760 or c + cn > len(d) or d[k:k + 4] != GFX_MAGIC: continue
+        if k < 0 or kn < 0x80 or cn < 0x760 or c + cn > len(d): continue
         ctl, code = d[c:c + cn], d[k:c]
-        st = sph_stage(code)
+        st = glslc_stage(ctl, code)
+        if st == 'cs' and code[:4] != CS_MAGIC and (d[:4] != b'RYHP' or CTL_MAGIC in code): continue
         if st and not check(label, st, ctl, code): out.append((f'{label}#PHYRE@{k:#x}', st, ctl, code))
     return out
 
@@ -3038,6 +3101,8 @@ def scan_blob(d, label, out, depth=0, left=None):
     if d[:4] in PG_STAGES and len(d) >= 0x18 and d[(o := struct.unpack_from('<I', d, 0xc)[0]):o + 4] == CTL_MAGIC: out.extend(pg_shader(d, label)); return
     if shaders := pg_sized_shader(d, label): out.extend(shaders); return
     if (shaders := viper_shader(d, label)) is not None: out.extend(shaders); return   # Star Fox's shader binaries
+    if d[:4] == b'RYHP' or d[8:12] == GFX_MAGIC:
+        if shaders := phyre_shaders(d, label): out.extend(shaders); return
     if d[:4] == b'VIPR':   # Star Fox's resource containers: two of them wrap the engine's NDSR upscaler networks (ounce-nvn2-upper_left_ndsr.bin), 40 shaders each, as a file inside
         for m in re.finditer(b'NDSR', d):
             if m.start() + 0x20 <= len(d) and d[m.start() + 4:m.start() + 6] == b'\x00\x01': out.extend(ndsr_shaders(d[m.start():], f'{label}#NDSR@{m.start():#x}'))
@@ -3239,6 +3304,9 @@ def scan_file_in(args):
     if ext in ('.data', '.patch', '.debug', '.nxpc') and os.path.exists(stem + '.dict'): return rel, {}, set(), {}   # read through its Next Level Games index (nlg_shaders)
     with open(path, 'rb') as f: head = f.read(16)
     found = []; trees = {}
+    if len(args) > 2 and args[2][0] == KIT:
+        for e in args[2][1]: scan_blob(kit_read(path, e), f'/{rel}#{e[0]}', found)
+        return rel, uniq_keys(found), var_pairs(found), trees
     if len(args) > 2 and args[2][0] == b'ARH1':
         arh_shaders(path, rel, args[2][1], found)
         return rel, uniq_keys(found), var_pairs(found), trees
@@ -3471,8 +3539,9 @@ def main():
     fb_jobs, fbs = frostbite_jobs(files)   # EA SPORTS FC 27's Frostbite bundles, Oodle Leviathan decoded by the game's own decoder in the Mac host
     pd_jobs, pdats = pdat_jobs(files)   # Kirby Air Riders' PDAT archives: the PMAP records in byte ranges
     rp_jobs, rpfs = rpf_jobs(files, main_v)   # Red Dead Redemption's RPF6 archives, their tables decrypted with the key in main (rpf6.py)
-    big += pak_jobs + tr_jobs + ue_jobs + bd_jobs + arc_jobs + pz_jobs + ck_jobs + d3_jobs + ca_jobs + cr_jobs + zp_jobs + ah_jobs + fb_jobs + pd_jobs + rp_jobs
-    files = [(p, r) for p, r in files if not linkdata_jobs(p, r) and p not in paks and p not in trp and p not in uep and p not in bdts and p not in arcs and p not in pkzs and p not in cpks and p not in d3s and p not in cascs and p not in cras and p not in zips and p not in arhs and p not in fbs and p not in pdats and p not in rpfs]
+    kt_jobs, kits = kit_jobs(files)   # Unravel TWO's KIT shards: chained LZ4 resources and Phyre effects
+    big += pak_jobs + tr_jobs + ue_jobs + bd_jobs + arc_jobs + pz_jobs + ck_jobs + d3_jobs + ca_jobs + cr_jobs + zp_jobs + ah_jobs + fb_jobs + pd_jobs + rp_jobs + kt_jobs
+    files = [(p, r) for p, r in files if not linkdata_jobs(p, r) and p not in paks and p not in trp and p not in uep and p not in bdts and p not in arcs and p not in pkzs and p not in cpks and p not in d3s and p not in cascs and p not in cras and p not in zips and p not in arhs and p not in fbs and p not in pdats and p not in rpfs and p not in kits]
     with cf.ProcessPoolExecutor(J, initializer=_pool_init, initargs=(common, {}, zdicts)) as ex:
         n = len(big) + len(files)
         results += pooled(ex, big, n) + pooled(ex, files, n, len(big), chunk=16)

docs/agents/symptoms.md

@@ -933,6 +933,7 @@ One line per symptom in the [subsystem references](../reference/): what you see
 - the split finds only a small part of a game's shaders, with no `[RED]` line, and its RomFS holds files that start `NARC` (`.narc`, `.nptl`, `.ntex` on Mario Tennis Aces). → [Packed shader banks in loose NARC files](../reference/shaders.md#packed-shader-banks-in-loose-narc-files)
 - the split finds no shaders in `main` or RomFS, while reporting control sections in `main` with no code area. The game has `.sp4` resources and loose `.nvnbin` files rather than BNSH archives or whole GLSLC outputs. → [Sumo Ninty_NX shaders and MuxFDMR streams](../reference/shaders.md#sumo-ninty_nx-shaders-and-muxfdmr-streams)
 - the shader split finds no game programs despite loose `.fx.phyre` effects and multi-GB `.vbf` archives. → [Phyre NVN effects in Virtuos VBF archives](../reference/shaders.md#phyre-nvn-effects-in-virtuos-vbf-archives)
+- the shader split finds no shaders in `main` or RomFS, which holds a `.kit` index and numbered `.kit.N` parts. Reading the effects alone still misses their compute programs. → [Phyre effects in split KIT archives](../reference/shaders.md#phyre-effects-in-split-kit-archives)
 - the split finds no shaders (`[RED] the shader split found no shaders in main or the RomFS`), the RomFS holds only `.cpk` files, and `cpk.py` refuses them (`not a CPK`); the run prints `[RED] shader … is not in the shader map` for every program. → [Diablo III's console packages](../reference/shaders.md#diablo-iiis-console-packages)
 - the split finds no shaders (`[RED] the shader split found no shaders in main or the RomFS`) and names control sections in `main` with `no code area: no data record points at one`; the RomFS is a `data` folder of archives named `00-00000001.data` … `17-00000002.data` beside `.psv` manifests and `build_key.txt`; the run prints `[RED] shader … is not in the shader map` for every program and draws black. → [Blizzard's CASC storage and SHAD containers](../reference/shaders.md#blizzards-casc-storage-and-shad-containers)
 - the shader split stops on one archive with a Python traceback (an `IndexError` or `struct.error` naming no file), or a reader's `[RED]` line about a format it does not know, and nothing is written; the NSP script loses the build at the split. → [A file no reader reads](../reference/shaders.md#a-file-no-reader-reads)