Read compiled shaders from Larian v13 archives

dm/gfx/tools/shader_split.py

@@ -616,6 +616,116 @@ def linkdata_shaders(path, rel, rng, out):
             off, n, stored, comp = struct.unpack_from('<QQQQ', idx, i * 32)
             if n: kt_shaders(d[off:off + (stored if comp else n)], f'/{rel}#{i}', out)
 
+LSPK = b'LSPK'
+LSPK_LIMIT = 1 << 30   # bound each file list, archive member and decoded .bshd before allocation
+
+def lspk_block(data, size, method, label):
+    """One stored, zlib or LZ4 block with an exact decoded size, never an unbounded inflate."""
+    import lz4.block
+    if size > LSPK_LIMIT: red(f'{label}: LSPK decoded size {size:#x} exceeds the byte budget')
+    try:
+        if method == 0: out = data
+        elif method == 1:
+            stream = zlib.decompressobj()
+            out = stream.decompress(data, size + 1)
+            if not stream.eof or stream.unused_data or stream.unconsumed_tail:
+                red(f'{label}: LSPK zlib stream is truncated, oversized or has trailing bytes')
+        elif method == 2:
+            # A raw LZ4 block expands by less than 255 bytes per input byte;
+            # reject an impossible declared size before the decoder allocates it.
+            if size > 255 * len(data): red(f'{label}: LSPK LZ4 decoded size cannot fit its compressed block')
+            out = lz4.block.decompress(data, uncompressed_size=size)
+        else: red(f'{label}: LSPK compression {method} is not read')
+    except (ValueError, lz4.block.LZ4BlockError, zlib.error) as e: red(f'{label}: invalid LSPK compressed block ({e})')
+    if len(out) != size: red(f'{label}: LSPK block decoded {len(out)} bytes, expected {size}')
+    return out
+
+def lspk_entries(path):
+    """Larian v13: footer {u32 header size, LSPK}, 32-byte header, LZ4 list of 280-byte FileEntry10 records.
+    Layout: Norbyte/lslib PackageFormat.cs LSPKHeader13/FileEntry10; .bshd wrapper reported in issue #1679."""
+    with open(path, 'rb') as f:
+        f.seek(0, 2); length = f.tell()
+        if length < 8: return None
+        f.seek(length - 8); hsize, magic = struct.unpack('<I4s', f.read(8))
+        if magic != LSPK:
+            f.seek(0); head = f.read(8)
+            if head[:4] == LSPK: red(f'{path}: LSPK version {int.from_bytes(head[4:], "little")}: only footer-based version 13 is read')
+            return None
+        if hsize != 40 or length < hsize: red(f'{path}: invalid LSPK header size {hsize}')
+        f.seek(length - hsize); version, off, size, parts, flags, _ = struct.unpack('<IIIHBB', f.read(16))
+        if version != 13: red(f'{path}: LSPK version {version}: only 13 is read')
+        if flags & 4: red(f'{path}: solid LSPK archives are not read')
+        if not parts: red(f'{path}: LSPK has no archive parts')
+        if size < 4 or off + size > length - hsize: red(f'{path}: LSPK file list lies outside the data area')
+        if size > LSPK_LIMIT: red(f'{path}: LSPK file list exceeds the byte budget')
+        f.seek(off); count = struct.unpack('<I', f.read(4))[0]
+        table = lspk_block(f.read(size - 4), count * 280, 2, f'{path}: file list')
+    entries = []
+    for name, at, stored, decoded, part, method, _crc in struct.iter_unpack('<256s6I', table):
+        try: name = name.split(b'\0', 1)[0].decode('utf-8')
+        except UnicodeDecodeError: red(f'{path}: invalid UTF-8 name in LSPK file list')
+        if part >= parts: red(f'{path}#{name}: LSPK part {part} exceeds the part count {parts}')
+        if part == 0 and at + stored > off: red(f'{path}#{name}: LSPK entry overlaps its file list or header')
+        entries.append((name, at, stored, decoded, part, method))
+    return parts, entries
+
+def lspk_part(path, part):
+    stem, ext = os.path.splitext(path)
+    return f'{stem}_{part}{ext}' if part else path
+
+def lspk_jobs(files, parts=64):
+    """Split only .bshd entries into jobs; archive parts and shaderless archives are not scanned whole."""
+    jobs, owned, archives = [], set(), []
+    for path, rel in files:
+        if not path.lower().endswith('.pak'): continue
+        try: info = lspk_entries(path)
+        except (Exception, SystemExit):
+            # Preserve scan_file's unreadable accounting and filename diagnostics.
+            jobs.append((path, rel, (LSPK, None))); owned.add(path); continue
+        if info is None: continue
+        count, entries = info
+        owned.update(lspk_part(path, i) for i in range(count))
+        archives.append((path, rel, [e for e in entries if e[0].lower().endswith('.bshd')]))
+    step = max(1, sum(e[2] for _, _, entries in archives for e in entries) // max(1, parts))
+    for path, rel, entries in archives:
+        group, size = [], 0
+        for e in entries:
+            group.append(e); size += e[2]
+            if size >= step: jobs.append((path, rel, (LSPK, tuple(group)))); group, size = [], 0
+        if group: jobs.append((path, rel, (LSPK, tuple(group))))
+    return jobs, owned
+
+def lspk_shaders(path, rel, entries, out):
+    """Archive-member compression, then a .bshd's big-endian size/flag and zlib-wrapped GLSLC program."""
+    import contextlib
+    if entries is None:
+        info = lspk_entries(path)   # reproduce a malformed list's diagnostic inside scan_file
+        if info is None: red(f'/{rel}: LSPK header disappeared')
+        entries = [e for e in info[1] if e[0].lower().endswith('.bshd')]
+    with contextlib.ExitStack() as stack:
+        handles = {}
+        for name, off, stored, decoded, part, flags in entries:
+            label = f'/{rel}#{name}'
+            if stored > LSPK_LIMIT: red(f'{label}: LSPK stored size exceeds the byte budget')
+            src = lspk_part(path, part)
+            try:
+                if src not in handles: handles[src] = stack.enter_context(open(src, 'rb'))
+                f = handles[src]; f.seek(0, 2)
+                if off + stored > f.tell(): red(f'{label}: LSPK entry lies outside archive part {part}')
+                f.seek(off); raw = f.read(stored)
+            except OSError as e: red(f'{label}: cannot read LSPK part {part} ({e})')
+            how = flags & 15
+            data = lspk_block(raw, decoded if how else stored, how, label)
+            if len(data) < 4: red(f'{label}: truncated LSPK .bshd header')
+            size = int.from_bytes(data[:4], 'big')
+            if not size & 0x80000000: red(f'{label}: LSPK .bshd is not zlib-wrapped')
+            data = lspk_block(data[4:], size & 0x7fffffff, 1, label)
+            found = glslc_shaders(data, label, need=False)
+            if sorted(st for _, st, _, _ in found) == ['ps', 'vs']:
+                found = [(f'{label}#GLSLC#var', st, ctl, code) for _, st, ctl, code in found]
+            out.extend(found)
+
+
 KPKA = b'KPKA'
 KPKA_BLOCKS = 0x20         # header feature: a block table follows the entries
 KPKA_IN_BLOCKS = 0x1000000   # entry attribute: the entry's data is in the block table's blocks
@@ -3222,6 +3332,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] == LSPK:
+        lspk_shaders(path, rel, args[2][1], 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
@@ -3455,8 +3568,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]
+    ls_jobs, lspks = lspk_jobs(files)   # Divinity: Original Sin 2's Larian archives: only the compiled .bshd shaders
+    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 + ls_jobs
+    files = [(p, r) for p, r in files if not linkdata_jobs(p, r) and p not in paks and p not in lspks 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]
     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)

dm/tools/tests/test_split_lspk.py

@@ -0,0 +1,190 @@
+"""Larian v13 archives deliver real synthetic GLSLC programs to the shader split (#1679)."""
+import contextlib
+import io
+import os
+import struct
+import subprocess
+import sys
+import tempfile
+import unittest
+import zlib
+from pathlib import Path
+from unittest import mock
+
+import lz4.block
+
+REPO = Path(__file__).resolve().parents[3]
+sys.path.insert(0, str(REPO / 'dm/tools/fuzz'))
+import fuzz_shader_split as fz
+ss = fz.ss
+import romfs_image
+
+
+def bshd(data):
+    return struct.pack('>I', 0x80000000 | len(data)) + zlib.compress(data)
+
+
+def archive(files, *, flags=0):
+    """Independent v13 fixture: (name, decoded bytes, compression flags, part)."""
+    count = max((part for _, _, _, part in files), default=0) + 1
+    parts = [bytearray() for _ in range(count)]
+    table = bytearray()
+    for name, raw, method, part in files:
+        packed = zlib.compress(raw) if method & 15 == 1 else lz4.block.compress(raw, store_size=False) if method & 15 == 2 else raw
+        table += struct.pack('<256s6I', name.encode(), len(parts[part]), len(packed), len(raw) if method & 15 else 0, part, method, zlib.crc32(packed))
+        parts[part] += packed
+    offset = len(parts[0])
+    listing = struct.pack('<I', len(files)) + lz4.block.compress(bytes(table), store_size=False)
+    header = struct.pack('<IIIHBB16s', 13, offset, len(listing), count, flags, 0, bytes(16))
+    parts[0] += listing + header + struct.pack('<I4s', 40, b'LSPK')
+    return parts
+
+
+class LspkShaders(unittest.TestCase):
+    def setUp(self):
+        self.tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(self.tmp.cleanup)
+        self.root = Path(self.tmp.name)
+        self.program = fz.glslc([fz.VS, fz.PS])
+
+    def write(self, parts, name='Materials.pak'):
+        path = self.root / name
+        for i, data in enumerate(parts):
+            dest = path if not i else path.with_name(f'{path.stem}_{i}{path.suffix}')
+            dest.write_bytes(data)
+        return str(path)
+
+    def scan(self, path):
+        jobs, owned = ss.lspk_jobs([(path, 'Data/Materials.pak')], parts=1)
+        self.assertEqual(len(jobs), 1)
+        with contextlib.redirect_stdout(io.StringIO()) as log:
+            result = ss.scan_file(jobs[0])
+        return result, log.getvalue(), owned
+
+    def test_each_compression_and_external_part_keeps_real_stages_and_pair(self):
+        for flags in (0, 1, 2, 0x21, 0x42):
+            for part in (0, 1):
+                with self.subTest(flags=flags, part=part):
+                    path = self.write(archive([('Shaders/main.bshd', bshd(self.program), flags, part)]))
+                    (_, keys, pairs, extra), log, owned = self.scan(path)
+                    self.assertEqual({v[1]: v[2:] for v in keys.values()}, {'vs': fz.VS, 'ps': fz.PS})
+                    by_stage = {v[1]: k for k, v in keys.items()}
+                    self.assertEqual(pairs, {(by_stage['vs'], by_stage['ps'])})
+                    self.assertEqual((extra, log), ({}, ''))
+                    self.assertIn(path, owned)
+                    if part: self.assertIn(str(self.root / 'Materials_1.pak'), owned)
+
+    def test_nonshader_members_are_not_decoded(self):
+        # Compression 15 is unsupported, but an unrelated file must not hide shaders.
+        path = self.write(archive([('readme.txt', b'not compressed', 15, 0), ('shader.BSHD', bshd(self.program), 0, 0)]), 'Materials.PAK')
+        (_, keys, pairs, extra), log, _ = self.scan(path)
+        self.assertEqual((len(keys), len(pairs), extra, log), (2, 1, {}, ''))
+
+    def test_shaderless_archive_and_empty_file_list_are_consumed(self):
+        for files in ([], [('names.bshd', bshd(b'reflection names only'), 0, 0)], [('readme', b'plain', 0, 0)]):
+            with self.subTest(files=len(files)):
+                path = self.write(archive(files))
+                jobs, owned = ss.lspk_jobs([(path, 'Data/Materials.pak')])
+                self.assertEqual(owned, {path})
+                for job in jobs:
+                    self.assertEqual(ss.scan_file(job)[1:], ({}, set(), {}))
+
+    def test_malformed_header_and_file_list_are_named_unreadable(self):
+        good = archive([('shader.bshd', bshd(self.program), 0, 0)])[0]
+        cases = {}
+        for label, offset, value in [('header size', -8, 4), ('version', -40, 15), ('list offset', -36, len(good)), ('list length', -32, 3)]:
+            bad = bytearray(good); struct.pack_into('<I', bad, len(bad) + offset, value); cases[label] = bad
+        bad = bytearray(good); lo = struct.unpack_from('<I', bad, len(bad) - 36)[0]
+        struct.pack_into('<I', bad, lo, 0xffffffff); cases['count budget'] = bad
+        bad = bytearray(good); struct.pack_into('<I', bad, lo, 1000000); cases['impossible LZ4 expansion'] = bad
+        bad = bytearray(good); bad[lo + 4] = 0; cases['corrupt list'] = bad
+        cases['solid archive'] = archive([('shader.bshd', bshd(self.program), 0, 0)], flags=4)[0]
+        for label, data in cases.items():
+            with self.subTest(case=label):
+                path = self.write([data])
+                (_, keys, pairs, extra), log, _ = self.scan(path)
+                self.assertEqual((keys, pairs), ({}, set()))
+                self.assertEqual(extra[0], 'unreadable')
+                self.assertIn('LSPK', log)
+                self.assertNotIn('Traceback', log)
+
+    def test_bad_inner_streams_do_not_publish_partial_shaders(self):
+        packed = bshd(self.program)
+        cases = [packed[:-1], packed + b'trailing', b'\0\0\0\0' + packed[4:],
+                 struct.pack('>I', 0x80000000 | (len(self.program) - 1)) + packed[4:],
+                 struct.pack('>I', 0xffffffff) + packed[4:]]
+        for bad in cases:
+            with self.subTest(size=len(bad), header=bad[:4].hex()):
+                path = self.write(archive([('good.bshd', packed, 0, 0), ('bad.bshd', bad, 0, 0)]))
+                (_, keys, pairs, extra), log, _ = self.scan(path)
+                self.assertEqual((keys, pairs), ({}, set()))
+                self.assertEqual(extra[0], 'unreadable')
+                self.assertIn('bad.bshd', log)
+
+    def test_missing_or_truncated_part_is_unreadable(self):
+        for missing in (True, False):
+            with self.subTest(missing=missing):
+                path = self.write(archive([('shader.bshd', bshd(self.program), 2, 1)]))
+                part = self.root / 'Materials_1.pak'
+                if missing: part.unlink()
+                else: part.write_bytes(part.read_bytes()[:-1])
+                (_, keys, pairs, extra), log, _ = self.scan(path)
+                self.assertEqual((keys, pairs), ({}, set()))
+                self.assertEqual(extra[0], 'unreadable')
+                self.assertIn('shader.bshd', log)
+
+    def test_bad_member_ranges_sizes_and_compression_are_unreadable(self):
+        for field, value in [(1, 0xffffffff), (2, 0xffffffff), (3, 1), (4, 2), (5, 3)]:
+            with self.subTest(field=field):
+                data = archive([('shader.bshd', bshd(self.program), 1, 0)])[0]
+                off = struct.unpack_from('<I', data, len(data) - 36)[0]
+                table = bytearray(lz4.block.decompress(data[off + 4:-40], uncompressed_size=280))
+                struct.pack_into('<I', table, 256 + (field - 1) * 4, value)
+                listing = struct.pack('<I', 1) + lz4.block.compress(bytes(table), store_size=False)
+                header = bytearray(data[-40:]); struct.pack_into('<I', header, 8, len(listing))
+                path = self.write([data[:off] + listing + header])
+                (_, keys, pairs, extra), log, _ = self.scan(path)
+                self.assertEqual((keys, pairs), ({}, set()))
+                self.assertEqual(extra[0], 'unreadable')
+                self.assertIn('LSPK', log)
+
+    def test_compute_shader_remains_unpaired(self):
+        ctl, code = fz.VS
+        code = ss.CS_MAGIC + code[4:]
+        path = self.write(archive([('compute.bshd', bshd(fz.glslc([(ctl, code)])), 0, 0)]))
+        (_, keys, pairs, extra), log, _ = self.scan(path)
+        self.assertEqual([v[1:] for v in keys.values()], [('cs', ctl, code)])
+        self.assertEqual((pairs, extra, log), (set(), {}, ''))
+
+    def test_budget_is_checked_before_decompressing(self):
+        path = self.write(archive([('shader.bshd', bshd(self.program), 1, 0)]))
+        # A small injected resource budget makes a valid member over-budget.
+        with mock.patch.object(ss, 'LSPK_LIMIT', 1024):
+            (_, keys, pairs, extra), log, _ = self.scan(path)
+        self.assertEqual((keys, pairs), ({}, set()))
+        self.assertEqual(extra[0], 'unreadable')
+        self.assertIn('budget', log)
+
+    def test_cli_extracts_archived_program_without_scanning_part_as_loose_file(self):
+        rom = self.root / 'RomFS'; rom.mkdir()
+        self.write(archive([('shader.bshd', bshd(self.program), 2, 1)]))
+        for p in self.root.glob('Materials*.pak'): p.rename(rom / p.name)
+        main = self.root / 'main'; main.write_bytes(b'NSO0' + bytes(0xfc))
+        image = self.root / 'RomFS.img'; romfs_image.pack(rom, image)
+        results = []
+        for i, source in enumerate((rom, image)):
+            dest = self.root / f'out{i}'; dest.mkdir()
+            outputs = [dest / n for n in ('pkg', 'map.tsv', 'facts.tsv', 'pairs.tsv', 'pso.tsv')]
+            proc = subprocess.run([sys.executable, REPO / 'dm/gfx/tools/shader_split.py', source, main, *outputs],
+                                  capture_output=True, text=True, env=dict(os.environ, DM_GAME='0100C1F0051B6000', J='2'))
+            self.assertEqual(proc.returncode, 0, proc.stdout + proc.stderr)
+            self.assertIn('shaders: 2 (', proc.stdout)
+            self.assertNotIn('[RED]', proc.stdout + proc.stderr)
+            rows = [r for r in outputs[1].read_text().splitlines() if r and not r.startswith('#')][1:]
+            self.assertEqual(len(rows), 2)
+            self.assertTrue(all('shader.bshd' in row for row in rows))
+            results.append({str(p.relative_to(dest)): p.read_bytes() for p in dest.rglob('*') if p.is_file()})
+        self.assertEqual(results[0], results[1])
+
+
+if __name__ == '__main__': unittest.main()