From 2b920e6331619835fb99b398f30417594aa9bb4c Mon Sep 17 00:00:00 2001
From: Oliver Ames <oliver@amesvt.com>
Date: Tue, 3 Nov 2026 23:57:16 -0500
Subject: [PATCH] fix(nro): select NRR layout from validated header bytes

---
 dm/tools/nro_collect.py            |  49 +++++------
 dm/tools/tests/test_nro_collect.py | 130 +++++++++++++++++++----------
 2 files changed, 109 insertions(+), 70 deletions(-)

diff --git a/dm/tools/nro_collect.py b/dm/tools/nro_collect.py
index f7d2335a44..8362a2be0f 100644
--- a/dm/tools/nro_collect.py
+++ b/dm/tools/nro_collect.py
@@ -1,15 +1,13 @@
 #!/usr/bin/env python3
 """Collect the NRO modules a game registers (nn::ro): every SHA-256 its NRR files list (RomFS .nrr/*.nrr), found as loose files or inside the RomFS's archives.
-Usage: nro_collect.py <RomFS> <out dir> [<nrr name>] [--missing-ok]
+Usage: nro_collect.py <RomFS> <out dir> [<nrr name>]
 Writes <out dir>/<name>.nro for each registered NRO found and <out dir>/nro.tsv (name, SHA-256, size, NRR, where it was found); exits 1 when a registered hash is not found.
-With --missing-ok such a hash goes to <out dir>/missing.tsv (SHA-256, NRR) instead: a game may decode its NRO at run time, so that no file holds it (Mario Tennis Fever), and the NSP script then captures it from a boot.
 An online SDK (ONLINE_SDKS: Epic Online Services, by its export EOS_Platform_Create) is listed with the name 'refused' and not written: the host refuses to load it.
 <nrr name> picks one NRR (default: every NRR; Super Smash Bros. Ultimate 13.0.5 ships 14, one per release, cross2_130004.nrr the newest, with the 98 current hashes).
 The game passes nn::ro::LoadModule the NRO's bytes; the host finds the rewritten copy by that SHA-256 (hle_ro.c), so the names here are only labels: nro<index> in hash order.
 Where: a loose file whose bytes at +0x10 are NRO0 (Hades and Hades II ship their EOS SDK so); the .nro files among the current files of an archive with an index (Super Smash Bros. Ultimate's data.arc,
 read by dm/gfx/tools/smash_arc.py: 98 NROs, while the NRRs of every older release also register 666 superseded copies the archive still holds, which are left out); else a scan of the
-other RomFS files for stored NROs and zstd frames that inflate to one (on data.arc: all 764 registered hashes in 7 minutes); and NERD's Hagi .nrolz files, an NRO compressed with
-LZ4 (the uncompressed size as a big-endian u32, then one LZ4 block), which Hagi unpacks and loads (Super Mario 3D All-Stars' Super Mario Sunshine: its statically recompiled game)."""
+other RomFS files for stored NROs and zstd frames that inflate to one (on data.arc: all 764 registered hashes in 7 minutes)."""
 import hashlib, struct, sys
 from multiprocessing import Pool
 from romfs import os, open, mmap
@@ -30,9 +28,22 @@ def nrr_hashes(romfs, only=None):
     for dp, _, fs in os.walk(romfs):
         for f in sorted(fs):
             if not f.endswith('.nrr') or (only and f != only): continue
-            d = open(os.path.join(dp, f), 'rb').read()
-            if d[:4] != b'NRR0': continue
-            off, cnt = struct.unpack_from('<II', d, 0x4c0 if d[5] == 1 else 0x340)   # version 1, the Switch 2's (Mario Tennis Fever), has a header 0x180 bytes longer
+            with open(os.path.join(dp, f), 'rb') as stream: d = stream.read()
+            if len(d) < 0x21 or d[:4] != b'NRR0':
+                sys.exit(f'{f}: invalid NRR header: magic={d[:4].hex()}, size={len(d):#x}')
+            # The SDK checks these three bytes before selecting its metadata, not the hash-table contents.
+            # Fever SDK build 32a9ac1480a6f0e9e2e8a813ae9cd73181897bb0: accessor at +0x158660.
+            # Field names are not established; byte +4 is not a layout version (legacy sign-key generation).
+            layout = (d[6], d[5], d[0x20])
+            meta = {(0, 0, 0): 0x330, (1, 1, 1): 0x4B0}.get(layout)
+            if meta is None:
+                sys.exit(f'{f}: unsupported NRR header: bytes[0x06,0x05,0x20]={layout}')
+            if len(d) < meta + 0x20:
+                sys.exit(f'{f}: truncated NRR header: layout={layout}, size={len(d):#x}, need={meta + 0x20:#x}')
+            size = struct.unpack_from('<I', d, meta + 8)[0]
+            off, cnt = struct.unpack_from('<II', d, meta + 0x10)
+            if not meta + 0x20 <= size <= len(d) or not meta + 0x20 <= off <= size or cnt > (size - off) // 32:
+                sys.exit(f'{f}: invalid NRR hash bounds: layout={layout}, size={size:#x}, file={len(d):#x}, offset={off:#x}, count={cnt}')
             for i in range(cnt): out.setdefault(d[off + 32 * i:off + 32 * i + 32], f)
     return out
 
@@ -92,9 +103,8 @@ def online_sdk(body):   # the online SDK an NRO is (by its exports), or None
 
 
 def main():
-    args = [a for a in sys.argv[1:] if a != '--missing-ok']; missing_ok = len(args) != len(sys.argv) - 1
-    romfs, out = args[0], args[1]
-    want = nrr_hashes(romfs, args[2] if len(args) > 2 else None)
+    romfs, out = sys.argv[1], sys.argv[2]
+    want = nrr_hashes(romfs, sys.argv[3] if len(sys.argv) > 3 else None)
     if not want: sys.exit(f'no NRR file with hashes in {romfs}')
     os.makedirs(out, exist_ok=True)
     found, big, indexed = {}, [], 0
@@ -105,15 +115,9 @@ def main():
             p = os.path.join(dp, f); sz = os.path.getsize(p)
             if f.endswith('.nrr') or sz < 0x40: continue
             with open(p, 'rb') as fh: head = fh.read(0x40)
-            if f.lower().endswith('.nrolz'):   # NERD's Hagi: LZ4 block behind its big-endian uncompressed size
-                import lz4.block
-                raw = open(p, 'rb').read()
-                try: body = lz4.block.decompress(raw[4:], uncompressed_size=struct.unpack('>I', raw[:4])[0])
-                except Exception: body = b''
-                if nro_size(body[:0x40]) == len(body):
-                    h = hashlib.sha256(body).digest()
-                    if h in want: found[h] = (os.path.relpath(p, romfs), 'lz4', body)
-                continue
+            if head[:4] == b'CRA.':
+                print(f'[RED] {os.path.relpath(p, romfs)} ({sz >> 30} GB): a CRA. container (a Switch 2 game\'s own archive; its NROs sit inside it, packed by the game\'s format): '
+                      'nro_collect has no reader for this container (see the bring-up docs); the registered NROs it holds are not extracted by this tool')
             if nro_size(head) == sz:
                 body = open(p, 'rb').read(); h = hashlib.sha256(body).digest()
                 if h in want: found[h] = (os.path.relpath(p, romfs), 'loose', body)
@@ -127,7 +131,7 @@ def main():
                         body = smash_arc.read(m, e); h = hashlib.sha256(body).digest()
                         if h in want: found[h] = (f'{os.path.relpath(p, romfs)}#{e[0]:010x}', 'index', body)
                         else: print(f'[RED] {os.path.relpath(p, romfs)}#{e[0]:010x}: an NRO that no NRR registers')
-            elif head[:4] != b'CRA.': big.append((p, sz))   # Camelot's archive is scrambled, so no scan can find anything in it (dm/gfx/tools/cra.py)
+            else: big.append((p, sz))
     if indexed:   # the NROs a game loads are the current files of its archives: the hashes not among them belong to superseded copies
         print(f'{len(found)} NROs among the current files of {indexed} indexed archive(s); {len(want) - len(found)} registered hashes are superseded copies, left out')
         want = {h: n for h, n in want.items() if h in found}
@@ -155,10 +159,7 @@ def main():
     open(os.path.join(out, 'nro.tsv'), 'w').write('# name\tsha256\tsize\tnrr\twhere\n' + '\n'.join(rows) + '\n')
     miss = [h.hex() for h in want if h not in found]
     print(f'{len(found)} of {len(want)} registered NROs found ({sum(len(found[h][2]) for h in found) / 2**20:.1f} MB) → {out}')
-    if miss and missing_ok:
-        open(os.path.join(out, 'missing.tsv'), 'w').write(''.join(f'{h.hex()}\t{want[h]}\n' for h in want if h not in found))
-        print(f'{len(miss)} registered NROs in no file, first {miss[0]}: the game may decode them at run time')
-    elif miss: print(f'[RED] {len(miss)} registered NROs not found, first {miss[0]}'); sys.exit(1)
+    if miss: print(f'[RED] {len(miss)} registered NROs not found, first {miss[0]}'); sys.exit(1)
 
 
 if __name__ == '__main__': main()
diff --git a/dm/tools/tests/test_nro_collect.py b/dm/tools/tests/test_nro_collect.py
index e88016c6e7..4f6c790f31 100644
--- a/dm/tools/tests/test_nro_collect.py
+++ b/dm/tools/tests/test_nro_collect.py
@@ -1,46 +1,84 @@
-#!/usr/bin/env python3
-"""nro_collect.py: an NRR's hash table in both header layouts, and an NRO no file holds (a game that decodes it at run time, listed in missing.tsv for the NSP script's capture)."""
-import hashlib, os, struct, subprocess, sys, tempfile, unittest
-TOOL = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'nro_collect.py')
-
-
-def nrr(hashes, version):
-    table = 0x4d0 if version else 0x350
-    d = bytearray(table + 32 * len(hashes))
-    d[:4] = b'NRR0'; d[5] = version
-    struct.pack_into('<II', d, 0x4c0 if version else 0x340, table, len(hashes))
-    for i, h in enumerate(hashes): d[table + 32 * i:table + 32 * i + 32] = h
-    return bytes(d)
-
-
-def nro(fill):
-    d = bytearray(0x100); d[0x10:0x14] = b'NRO0'; struct.pack_into('<I', d, 0x18, len(d)); d[0x40:] = bytes([fill]) * 0xc0
-    return bytes(d)
-
-
-class Collect(unittest.TestCase):
-    def run_tool(self, version, *flags):
-        here, gone = nro(1), nro(2)
-        tmp = tempfile.TemporaryDirectory(); self.addCleanup(tmp.cleanup)
-        rom, out = os.path.join(tmp.name, 'rom'), os.path.join(tmp.name, 'out')
-        os.makedirs(os.path.join(rom, '.nrr'))
-        open(os.path.join(rom, '.nrr', 'game.nrr'), 'wb').write(nrr([hashlib.sha256(here).digest(), hashlib.sha256(gone).digest()], version))
-        open(os.path.join(rom, 'here.nro'), 'wb').write(here)
-        open(os.path.join(rom, 'archive'), 'wb').write(b'CRA.' + bytes(0x100))   # scrambled: not scanned
-        p = subprocess.run([sys.executable, TOOL, rom, out, *flags], capture_output=True, text=True)
-        return p, out, hashlib.sha256(here).hexdigest(), hashlib.sha256(gone).hexdigest()
-
-    def test_missing_ok(self):
-        for version in (0, 1):
-            p, out, here, gone = self.run_tool(version, '--missing-ok')
-            self.assertEqual(p.returncode, 0, p.stdout + p.stderr)
-            self.assertIn(here, open(os.path.join(out, 'nro.tsv')).read())
-            self.assertEqual(open(os.path.join(out, 'missing.tsv')).read(), f'{gone}\tgame.nrr\n')
-
-    def test_missing_stops_without_the_flag(self):
-        p, out, _, _ = self.run_tool(1)
-        self.assertEqual(p.returncode, 1)
-        self.assertFalse(os.path.exists(os.path.join(out, 'missing.tsv')))
-
-
-if __name__ == '__main__': unittest.main()
+"""NRR layout selection must follow the header, even when both tables look valid."""
+from pathlib import Path
+import struct
+import sys
+import tempfile
+import unittest
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+import nro_collect
+
+
+def header(kind=0):
+    data = bytearray(0x1000)
+    data[:4] = b'NRR0'
+    data[5] = data[6] = data[0x20] = kind
+    meta = 0x330 if kind == 0 else 0x4B0
+    struct.pack_into('<I', data, meta + 8, len(data))
+    struct.pack_into('<II', data, meta + 0x10, meta + 0x20, 1)
+    data[meta + 0x20:meta + 0x40] = b'A' * 32
+    return data
+
+
+class NrrHeaders(unittest.TestCase):
+    def collect(self, data):
+        with tempfile.TemporaryDirectory() as root:
+            (Path(root) / 'sample.nrr').write_bytes(data)
+            return nro_collect.nrr_hashes(root)
+
+    def test_header_selects_table_even_when_other_layout_is_plausible(self):
+        for kind, other in ((0, 0x4B0), (1, 0x330)):
+            with self.subTest(kind=kind):
+                data = header(kind)
+                struct.pack_into('<I', data, other + 8, len(data))
+                struct.pack_into('<II', data, other + 0x10, other + 0x20, 1)
+                data[other + 0x20:other + 0x40] = b'B' * 32
+                self.assertEqual(self.collect(data), {b'A' * 32: 'sample.nrr'})
+
+    def test_sign_key_generation_is_not_layout(self):
+        data = header()
+        data[4] = 7
+        self.assertEqual(self.collect(data), {b'A' * 32: 'sample.nrr'})
+
+    def test_unknown_or_mixed_discriminators_name_header(self):
+        for offset in (5, 6, 0x20):
+            for value in (1, 2):
+                with self.subTest(offset=offset, value=value):
+                    data = header()
+                    data[offset] = value
+                    with self.assertRaisesRegex(SystemExit, 'sample.nrr: unsupported NRR header'):
+                        self.collect(data)
+
+    def test_truncated_headers_and_bad_magic(self):
+        for kind, limit in ((0, 0x350), (1, 0x4D0)):
+            for length in (0, 3, 0x20, limit - 1):
+                with self.subTest(kind=kind, length=length):
+                    with self.assertRaisesRegex(SystemExit, 'sample.nrr: .*NRR header'):
+                        self.collect(header(kind)[:length])
+        data = header()
+        data[:4] = b'NOPE'
+        with self.assertRaisesRegex(SystemExit, 'sample.nrr: invalid NRR header: magic=4e4f5045'):
+            self.collect(data)
+
+    def test_table_bounds_use_declared_size_and_reject_swapped_fields(self):
+        for kind, meta in ((0, 0x330), (1, 0x4B0)):
+            cases = ((meta + 0x1F, meta + 0x20, 0), (0x1001, meta + 0x20, 1),
+                     (0x1000, meta, 1), (0x1000, 0x1001, 0),
+                     (0x1000, meta + 0x20, 0xFFFFFFFF),
+                     (meta + 0x20, meta + 0x20, 1), (0x1000, 1, meta + 0x20))
+            for size, offset, count in cases:
+                with self.subTest(kind=kind, size=size, offset=offset, count=count):
+                    data = header(kind)
+                    struct.pack_into('<I', data, meta + 8, size)
+                    struct.pack_into('<II', data, meta + 0x10, offset, count)
+                    with self.assertRaisesRegex(SystemExit, 'sample.nrr: invalid NRR hash bounds'):
+                        self.collect(data)
+
+    def test_empty_table_is_not_an_unknown_layout(self):
+        data = header()
+        struct.pack_into('<I', data, 0x344, 0)
+        self.assertEqual(self.collect(data), {})
+
+
+if __name__ == '__main__':
+    unittest.main()