fix: validate catalog reads and bundle signatures, isolate master writes

src/nnnotes/addressables.py

@@ -88,37 +88,26 @@ def remote_path(internal_id: str) -> str | None:
 def parse(data: bytes) -> list[dict]:
     """Every location of a binary catalog: {offset, primary_key, internal_id, dependencies (location offsets)}."""
     data = catalog_bytes(data)
-    def u32(offset):
-        return struct.unpack_from("<I", data, offset)[0]
-
-    def array(offset):
-        if offset == NONE:
-            return []
-        return struct.unpack_from(f"<{u32(offset - 4) // 4}I", data, offset)
-
-    def text(offset):
-        if offset == NONE:
-            return ""
-        if offset & 0x40000000:                     # a path of parts, stored last part first
-            parts = []
-            while offset != NONE:
-                part, offset = struct.unpack_from("<II", data, offset & 0x3FFFFFFF)
-                parts.append(text(part))
-            return "/".join(reversed(parts))
-        pos = offset & 0x3FFFFFFF
-        return data[pos:pos + u32(pos - 4)].decode("utf-16-le" if offset & 0x80000000 else "ascii")
-
-    magic, version, keys = struct.unpack_from("<III", data)
+    if len(data) < 12:
+        raise ValueError("unsupported catalog format")
+    buf = _Buffer(data)
+    magic, version, keys = buf.unpack("<III", 0)
     if magic != CATALOG_MAGIC or version != CATALOG_VERSION:
         raise ValueError("unsupported catalog format")
+    n = buf.u32(keys - 4)
+    if n % 8:
+        raise ValueError(f"catalog key table: byte length {n} is not a multiple of 8")
+    buf.check(keys, n)
     locations = set()
-    for pos in range(keys, keys + u32(keys - 4), 8):
-        locations.update(array(u32(pos + 4)))
+    for pos in range(keys, keys + n, 8):
+        locations.update(buf.array(buf.u32(pos + 4)))
     result = []
     for pos in sorted(locations):
-        primary, internal, _, deps = struct.unpack_from("<4I", data, pos)
-        result.append({"offset": pos, "primary_key": text(primary),
-                       "internal_id": text(internal), "dependencies": array(deps)})
+        buf.check(pos, LOCATION.size)
+        primary, internal, _, deps = buf.unpack("<4I", pos)
+        result.append({"offset": pos, "primary_key": buf.string(primary, "/") or "",
+                       "internal_id": buf.string(internal, "/") or "",
+                       "dependencies": [] if deps == NONE else tuple(buf.array(deps))})
     return result
 
 
@@ -149,8 +138,16 @@ def __init__(self, data: bytes):
         self._plain: dict[int, str] = {}
         self._types: dict[int, tuple] = {}
 
+    def check(self, offset: int, size: int) -> None:
+        if offset < 0 or size < 0 or offset > len(self.data) - size:
+            raise ValueError(f"catalog data at {offset}: past the end of the catalog")
+
+    def unpack(self, fmt: str, offset: int) -> tuple:
+        self.check(offset, struct.calcsize(fmt))
+        return struct.unpack_from(fmt, self.data, offset)
+
     def u32(self, offset: int) -> int:
-        return struct.unpack_from("<I", self.data, offset)[0]
+        return self.unpack("<I", offset)[0]
 
     def array(self, offset: int) -> list[int]:
         """A u32 array (its byte length is the u32 before it); [] for null."""
@@ -159,13 +156,15 @@ def array(self, offset: int) -> list[int]:
         n = self.u32(offset - 4)
         if n % 4:
             raise ValueError(f"catalog array at {offset}: byte length {n} is not a multiple of 4")
-        return list(struct.unpack_from(f"<{n // 4}I", self.data, offset))
+        return list(self.unpack(f"<{n // 4}I", offset))
 
     def plain(self, offset: int) -> str:
         s = self._plain.get(offset)
         if s is None:
             pos = offset & OFFSET_MASK
-            raw = self.data[pos:pos + self.u32(pos - 4)]
+            size = self.u32(pos - 4)
+            self.check(pos, size)
+            raw = self.data[pos:pos + size]
             s = self._plain[offset] = raw.decode("utf-16-le" if offset & UNICODE else "ascii")
         return s
 
@@ -183,7 +182,7 @@ def string(self, offset: int, sep: str) -> str | None:
             if pos in seen:
                 raise ValueError(f"catalog string at {offset & OFFSET_MASK}: the part chain loops")
             seen.add(pos)
-            part, link = struct.unpack_from("<II", self.data, pos)
+            part, link = self.unpack("<II", pos)
             if part != NONE and part & DYNAMIC:
                 raise ValueError(f"catalog string at {pos}: a nested part list")
             parts.append("" if part == NONE else self.plain(part))
@@ -195,7 +194,7 @@ def type_name(self, offset: int) -> tuple[str | None, str | None] | None:
             return None
         t = self._types.get(offset)
         if t is None:
-            assembly, cls = struct.unpack_from("<II", self.data, offset)
+            assembly, cls = self.unpack("<II", offset)
             t = self._types[offset] = (self.string(assembly, "."), self.string(cls, "."))
         return t
 
@@ -204,28 +203,28 @@ def object_init(self, offset: int) -> dict | None:
         TypeSerializer.Data."""
         if offset == NONE:
             return None
-        ident, typ, data = struct.unpack_from("<III", self.data, offset)
+        ident, typ, data = self.unpack("<III", offset)
         t = self.type_name(typ) or (None, None)
         return {"id": self.string(ident, ""), "assembly": t[0], "type": t[1], "data": self.string(data, "")}
 
     def key(self, offset: int) -> tuple[str | None, object]:
         """A key object (ObjectTypeData {type, object}) -> (type name, value). System.String keys are an
         ObjectToStringRemap {u32 string, u16 separator}; System.Int32 keys a 4-byte integer; other types None."""
-        typ, obj = struct.unpack_from("<II", self.data, offset)
+        typ, obj = self.unpack("<II", offset)
         cls = (self.type_name(typ) or (None, None))[1]
         if cls == "System.String":
-            sid, sep = struct.unpack_from("<IH", self.data, obj)
+            sid, sep = self.unpack("<IH", obj)
             return cls, self.string(sid, chr(sep) if sep else "")
         if cls == "System.Int32":
-            return cls, struct.unpack_from("<i", self.data, obj)[0]
+            return cls, self.unpack("<i", obj)[0]
         return cls, None
 
     def extra(self, offset: int) -> dict | None:
         """A location's extra data (ObjectTypeData {type, object}): {"type": class name} plus, for
         AssetBundleRequestOptions, the decoded options; None for null."""
         if offset == NONE:
             return None
-        typ, obj = struct.unpack_from("<II", self.data, offset)
+        typ, obj = self.unpack("<II", offset)
         cls = (self.type_name(typ) or (None, None))[1]
         out: dict = {"type": cls}
         if cls == REQUEST_OPTIONS and obj != NONE:
@@ -236,13 +235,15 @@ def request_options(self, offset: int) -> dict:
         """AssetBundleRequestOptionsSerializationAdapter.SerializedData {u32 hash, u32 bundleName, u32 crc,
         u32 bundleSize, u32 common}: hash a Hash128 (16 raw bytes, hex in stored order), bundleName read with '_',
         common a SerializedData.Common {i16 timeout, u8 redirectLimit, u8 retryCount, i32 flags}."""
-        hash_id, name_id, crc, size, common = struct.unpack_from("<5I", self.data, offset)
+        hash_id, name_id, crc, size, common = self.unpack("<5I", offset)
+        if hash_id != NONE:
+            self.check(hash_id, 16)
         out = {"hash": None if hash_id == NONE else self.data[hash_id:hash_id + 16].hex(),
                "bundleName": self.string(name_id, "_"), "crc": crc, "bundleSize": size}
         if common == NONE:
             out.update(timeout=None, redirectLimit=None, retryCount=None, flags=None)
             return out
-        timeout, redirects, retries, flags = struct.unpack_from("<hBBi", self.data, common)
+        timeout, redirects, retries, flags = self.unpack("<hBBi", common)
         out.update(timeout=timeout, redirectLimit=redirects, retryCount=retries, flags=flags)
         for name, bit in FLAG_BITS:
             out[name] = (flags & bit) // bit if name == "assetLoadMode" else bool(flags & bit)
@@ -276,6 +277,7 @@ def parse_keys(data: bytes) -> list[dict]:
     n = buf.u32(table - 4)
     if n % 8:
         raise ValueError(f"catalog key table: byte length {n} is not a multiple of 8")
+    buf.check(table, n)
     for pos in range(table, table + n, 8):
         key_obj, locations = struct.unpack_from("<II", data, pos)
         cls, value = buf.key(key_obj) if key_obj != NONE else (None, None)
@@ -297,7 +299,11 @@ def parse_locations(data: bytes) -> list[dict]:
     buf = _Buffer(data)
     offsets: set[int] = set()
     table = header["keysOffset"]
-    for pos in range(table, table + buf.u32(table - 4), 8):
+    n = buf.u32(table - 4)
+    if n % 8:
+        raise ValueError(f"catalog key table: byte length {n} is not a multiple of 8")
+    buf.check(table, n)
+    for pos in range(table, table + n, 8):
         offsets.update(buf.array(buf.u32(pos + 4)))
     out = []
     for pos in sorted(offsets):

src/nnnotes/catalog.py

@@ -16,7 +16,7 @@
 from pathlib import Path
 
 from .apkset import ApkSet
-from .addressables import REMOTE_PREFIX, BundleKey, decrypt, parse, parse_locations, remote_path
+from .addressables import UNITYFS, REMOTE_PREFIX, BundleKey, decrypt, parse, parse_locations, remote_path
 from .cache import write_atomic as _write_atomic
 from .config import ConfigError, apk_missing, check_catalog_version
 
@@ -59,11 +59,26 @@ def file_name(internal_id: str) -> str:
 
 
 def _unityfs(data: bytes, name: str, key: BundleKey | None) -> bytes:
-    if data[:7] == b"UnityFS":
+    if data.startswith(UNITYFS):
         return data
     if key is None:
         raise RuntimeError(f"bundle {name} is encrypted and no bundle key was given")
-    return decrypt(data, name, key)
+    decoded = decrypt(data, name, key)
+    if not decoded.startswith(UNITYFS):
+        raise ValueError(f"bundle {name} is not UnityFS after decryption (wrong key or damaged file)")
+    return decoded
+
+
+def _cached_bundle(path: Path) -> Path | None:
+    """A decrypted bundle cache hit requires the complete UnityFS signature."""
+    if path.is_file():
+        try:
+            with path.open("rb") as stream:
+                if stream.read(len(UNITYFS)) == UNITYFS:
+                    return path
+        except FileNotFoundError:
+            pass
+    return None
 
 
 class Catalog:
@@ -281,7 +296,7 @@ def resolve(self, key: str) -> list[Bundle]:
     def cached(self, b: Bundle) -> Path | None:
         """The file fetch(b) returns when the bundle is in the cache already, else None."""
         dst = (self.cache_dir if b.remote else self.local_cache_dir()) / "bundles" / b.name
-        return dst if dst.is_file() and dst.stat().st_size > 0 else None
+        return _cached_bundle(dst)
 
     def cached_raw(self, e: dict) -> Path | None:
         """The file fetch_raw(e) returns when it is in the cache already, else None."""
@@ -293,8 +308,9 @@ def fetch(self, b: Bundle) -> Path:
         """Local path to the decrypted bundle (CDN download or APK read)."""
         dst = (self.cache_dir if b.remote else self.local_cache_dir()) / "bundles" / b.name
         dst.parent.mkdir(parents=True, exist_ok=True)
-        if dst.exists() and dst.stat().st_size > 0:
-            return dst
+        hit = _cached_bundle(dst)
+        if hit is not None:
+            return hit
         if b.remote:
             url = self._url(b.internal_id)
             key = self._setting("bundle_key", b.name)   # before the download: a missing key fails first
@@ -306,7 +322,7 @@ def fetch(self, b: Bundle) -> Path:
             with ApkSet(self.apk) as z:
                 self.check_apk(z)
                 data = z.read(APK_AA_DIR + rel)
-            key = self._setting("bundle_key", b.name) if data[:7] != b"UnityFS" else None
+            key = self._setting("bundle_key", b.name) if not data.startswith(UNITYFS) else None
         _write_atomic(dst, _unityfs(data, b.name, key))
         return dst
 
@@ -354,11 +370,11 @@ def apk_bundle(self, name_contains: str) -> Path:
                 raise KeyError(f"{name_contains!r}: {len(names)} APK bundles match")
             name = names[0].rsplit("/", 1)[1]
             dst = self.local_cache_dir() / "bundles" / name
-            if not (dst.exists() and dst.stat().st_size > 0):
+            if _cached_bundle(dst) is None:
                 dst.parent.mkdir(parents=True, exist_ok=True)
                 self.check_apk(z)
                 data = z.read(names[0])
-                key = self._setting("bundle_key", name) if data[:7] != b"UnityFS" else None
+                key = self._setting("bundle_key", name) if not data.startswith(UNITYFS) else None
                 _write_atomic(dst, _unityfs(data, name, key))
         return dst
 

src/nnnotes/master.py

@@ -17,7 +17,6 @@
 import gzip
 import hashlib
 import json
-import os
 import re
 import time
 import urllib.error
@@ -29,6 +28,7 @@
 import numpy as np
 
 from . import jsonio
+from .store import write_file
 
 PREFIX = 64                     # bytes before the ciphertext
 BLOCK = 32                      # 256-bit block
@@ -246,8 +246,7 @@ def download(cdn: str, version: str, out_dir: Path, workers: int = 16, *, get=No
             seen.add(name)
     else:
         # Preserve the original partial-download receipt even if a worker raises.
-        from .cache import write_atomic
-        write_atomic(out_dir / "MasterManifest.json", manifest_raw)
+        write_file(out_dir / "MasterManifest.json", manifest_raw)
 
     def one(f: dict):
         name, sha = f["name"], (f.get("hash") or "").lower()
@@ -262,16 +261,13 @@ def one(f: dict):
             return {"file": name, "error": "size or hash missing/different from the manifest"}
         if sha and hashlib.sha256(data).hexdigest() != sha:
             return {"file": name, "error": "sha256 differs from the manifest"}
-        tmp = dst.with_name(f"{name}.{os.getpid()}.part")
-        tmp.write_bytes(data)
-        os.replace(tmp, dst)
+        write_file(dst, data)
         return "downloaded"
 
     with ThreadPoolExecutor(max(1, workers)) as ex:
         results = list(ex.map(one, files))
     if not strict or not any(isinstance(r, dict) for r in results):
-        from .cache import write_atomic
-        write_atomic(out_dir / "MasterManifest.json", manifest_raw)
+        write_file(out_dir / "MasterManifest.json", manifest_raw)
     return {"version": manifest.get("version", version), "files": len(files),
             "downloaded": results.count("downloaded"), "kept": results.count("kept"),
             "failed": [r for r in results if isinstance(r, dict)], "out": str(out_dir)}