unity-formats 0.2.1

//! Type tree packages (TPK, "TPK*" files): the layouts of Unity's built-in classes for every Unity version, which a
//! player build strips from its files. Format versions 1 and 2, type tree data, uncompressed or LZ4 (LZMA and
//! Brotli packages are refused: re-pack them uncompressed).

use primase_foundations::compress::{Decompressor, LZ4_MAX_RATIO, Lz4};
use std::collections::BTreeMap;
use std::sync::Arc;

use crate::reader::Reader;
use crate::typetree::{Node, TypeTree};
use crate::version::UnityVersion;
use crate::{Error, Result, invalid};

const MAGIC: u32 = 0x2A4B_5054;
const HAS_EDITOR_ROOT: u8 = 64;
const HAS_RELEASE_ROOT: u8 = 128;

#[derive(Clone, Debug)]
struct Class {
    name: u16,
    flags: u8,
    release_root: Option<u16>,
}

#[derive(Clone, Debug)]
struct PackedNode {
    type_name: u16,
    name: u16,
    byte_size: i32,
    version: i16,
    meta_flag: u32,
    children: Vec<u16>,
}

/// One Unity version's common string table: the names every embedded type tree of that version shares, as
/// `(offset into the engine's table, index into the package's strings)` sorted by offset.
#[derive(Clone, Debug)]
struct CommonTable {
    from: UnityVersion,
    entries: Vec<(u32, u16)>,
}

/// A parsed type tree package.
#[derive(Clone, Debug)]
pub struct TypeTreePackage {
    /// Per class id, the class from each listed version on (`None`: absent from that version).
    classes: BTreeMap<i32, Vec<(UnityVersion, Option<Class>)>>,
    nodes: Vec<PackedNode>,
    strings: Vec<Arc<str>>,
    /// The common string table from each listed version on, in version order.
    common: Vec<CommonTable>,
}

/// The common string table of one Unity version, from [`TypeTreePackage::common_strings`]: what an embedded type
/// tree's string offset with the top bit set names.
#[derive(Clone, Copy, Debug)]
pub struct CommonStrings<'a> {
    package: &'a TypeTreePackage,
    entries: &'a [(u32, u16)],
}

impl CommonStrings<'_> {
    /// The string at `offset` of the engine's table (`0` is `AABB`, ...); `None` for an offset the table has no
    /// string at.
    pub fn get(&self, offset: u32) -> Option<&str> {
        let at = self.entries.binary_search_by_key(&offset, |&(o, _)| o).ok()?;
        let &(_, string) = self.entries.get(at)?;
        self.package.strings.get(usize::from(string)).map(|s| &**s)
    }

    /// How many strings the table holds.
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Whether the table is empty.
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// The strings with their offsets, in offset order.
    pub fn iter(&self) -> impl Iterator<Item = (u32, &str)> + '_ {
        self.entries.iter().filter_map(|&(offset, string)| {
            self.package
                .strings
                .get(usize::from(string))
                .map(|s| (offset, &**s))
        })
    }
}

#[cfg(test)]
impl TypeTreePackage {
    /// A package without classes, for tests that need a package but decode no object.
    pub(crate) fn empty() -> Self {
        TypeTreePackage {
            classes: BTreeMap::new(),
            nodes: Vec::new(),
            strings: Vec::new(),
            common: Vec::new(),
        }
    }
}

fn varint_string(r: &mut Reader) -> Result<Arc<str>> {
    let (mut len, mut shift) = (0usize, 0u32);
    loop {
        let b = r.u8()?;
        if shift > 28 {
            return Err(invalid("TPK: string length varint too long"));
        }
        len |= usize::from(b & 0x7F) << shift;
        shift = shift.saturating_add(7);
        if b & 0x80 == 0 {
            break;
        }
    }
    let s = std::str::from_utf8(r.take(len)?).map_err(|_| invalid("TPK: string not UTF-8"))?;
    Ok(Arc::from(s))
}

/// Skip `n` entries of `size` bytes each.
fn skip(r: &mut Reader, n: usize, size: usize, what: &str) -> Result<()> {
    let bytes = n.checked_mul(size).ok_or_else(|| {
        invalid(format!(
            "TPK: {n} {what} entries of {size} bytes at offset {}",
            r.pos
        ))
    })?;
    r.take(bytes).map(|_| ())
}

fn count(r: &mut Reader, what: &str) -> Result<usize> {
    let n = r.i32()?;
    usize::try_from(n).map_err(|_| invalid(format!("TPK: negative {what} count {n}")))
}

/// The most nodes one class's type tree may expand to (the trees seen have a few thousand).
const MAX_TREE_NODES: usize = 1 << 20;

impl TypeTreePackage {
    /// Parse a type tree package (TPK format 1 or 2, uncompressed or LZ4). Every index in it is checked, the size an
    /// LZ4 block claims is checked against what LZ4 can inflate to before anything is allocated for it, and a package
    /// that is LZMA- or Brotli-compressed is refused as [`crate::Error::Unsupported`].
    pub fn parse(bytes: &[u8]) -> Result<Self> {
        let mut r = Reader::new(bytes, false);
        if r.u32()? != MAGIC {
            return Err(invalid("not a TPK file"));
        }
        let format = r.u8()?;
        if !(1..=2).contains(&format) {
            return Err(Error::Unsupported(format!("TPK format version {format}")));
        }
        let compression = r.u8()?;
        let data_type = r.u8()?;
        let _reserved = r.u8()?;
        let _reserved2 = r.u32()?;
        let compressed = r.u32()? as usize;
        let uncompressed = r.u32()? as usize;
        if data_type != 0 {
            return Err(Error::Unsupported(format!(
                "TPK data type {data_type} (not type trees)"
            )));
        }
        let packed = r.take(compressed)?;
        let blob: Vec<u8> = match compression {
            0 => packed.to_vec(),
            1 => {
                // the header's size is the file's word: what the packed bytes cannot inflate to is not allocated
                if uncompressed > packed.len().saturating_mul(LZ4_MAX_RATIO) {
                    return Err(invalid(format!(
                        "TPK: {uncompressed} bytes claimed of {} bytes of LZ4, which inflate to at most {}",
                        packed.len(),
                        packed.len().saturating_mul(LZ4_MAX_RATIO)
                    )));
                }
                Lz4.decompress(packed, uncompressed)
                    .map_err(|e| invalid(format!("TPK LZ4: {e}")))?
            }
            2 => {
                return Err(Error::Unsupported(
                    "LZMA-compressed TPK: use an uncompressed or LZ4 package".into(),
                ));
            }
            3 => {
                return Err(Error::Unsupported(
                    "Brotli-compressed TPK: use an uncompressed or LZ4 package".into(),
                ));
            }
            c => return Err(invalid(format!("TPK compression type {c}"))),
        };
        if blob.len() != uncompressed {
            return Err(invalid(format!(
                "TPK: {} bytes, header says {uncompressed}",
                blob.len()
            )));
        }
        Self::parse_blob(&blob, format)
    }

    fn parse_blob(blob: &[u8], format: u8) -> Result<Self> {
        let mut r = Reader::new(blob, false);
        let _created = r.i64()?;
        let versions = count(&mut r, "version")?;
        skip(&mut r, versions, 8, "version")?;
        let mut classes = BTreeMap::new();
        for _ in 0..count(&mut r, "class")? {
            let id = r.i32()?;
            let mut entries = Vec::new();
            for _ in 0..count(&mut r, "class version")? {
                let version = UnityVersion(r.u64()?);
                let class = if r.u8()? != 0 {
                    let name = r.u16()?;
                    let _base = r.u16()?;
                    let flags = r.u8()?;
                    if flags & HAS_EDITOR_ROOT != 0 {
                        r.u16()?;
                    }
                    let release_root = if flags & HAS_RELEASE_ROOT != 0 {
                        Some(r.u16()?)
                    } else {
                        None
                    };
                    Some(Class {
                        name,
                        flags,
                        release_root,
                    })
                } else {
                    None
                };
                entries.push((version, class));
            }
            classes.insert(id, entries);
        }
        // the common string table: what the offsets with the top bit set in an embedded type tree stand for. Format 1
        // lists, per version, how many of a shared index list's first strings the version has (the strings follow each
        // other, a NUL after each); format 2 lists, per version, each string's offset and index.
        let mut common_raw: Vec<(UnityVersion, Vec<(u16, u16)>)> = Vec::new();
        let mut v1_counts: Vec<(UnityVersion, u8)> = Vec::new();
        let mut v1_indices: Vec<u16> = Vec::new();
        match format {
            1 => {
                for _ in 0..count(&mut r, "common string version")? {
                    let version = UnityVersion(r.u64()?);
                    v1_counts.push((version, r.u8()?));
                }
                for _ in 0..count(&mut r, "common string index")? {
                    v1_indices.push(r.u16()?);
                }
            }
            _ => {
                for _ in 0..count(&mut r, "common string version")? {
                    let version = UnityVersion(r.u64()?);
                    let n = count(&mut r, "common string entry")?;
                    let mut entries = Vec::new();
                    for _ in 0..n {
                        let offset = r.u16()?;
                        entries.push((offset, r.u16()?));
                    }
                    common_raw.push((version, entries));
                }
            }
        }
        let mut nodes = Vec::new();
        for _ in 0..count(&mut r, "node")? {
            let type_name = r.u16()?;
            let name = r.u16()?;
            let byte_size = r.i32()?;
            let version = r.i16()?;
            let _type_flags = r.u8()?;
            let meta_flag = r.u32()?;
            let n = r.u16()?;
            let children = (0..n).map(|_| r.u16()).collect::<Result<Vec<_>>>()?;
            nodes.push(PackedNode {
                type_name,
                name,
                byte_size,
                version,
                meta_flag,
                children,
            });
        }
        let strings = (0..count(&mut r, "string")?)
            .map(|_| varint_string(&mut r))
            .collect::<Result<Vec<_>>>()?;
        if r.remaining() != 0 {
            return Err(invalid(format!(
                "TPK: {} bytes after the string buffer",
                r.remaining()
            )));
        }
        for (version, n) in v1_counts {
            let mut offset = 0u16;
            let mut entries = Vec::new();
            for &index in v1_indices.iter().take(usize::from(n)) {
                entries.push((offset, index));
                let len = strings
                    .get(usize::from(index))
                    .ok_or_else(|| invalid("TPK: common string index out of range"))?
                    .len();
                offset = u16::try_from(len)
                    .ok()
                    .and_then(|l| offset.checked_add(l))
                    .and_then(|o| o.checked_add(1))
                    .ok_or_else(|| invalid("TPK: common string table longer than 64 KiB"))?;
            }
            common_raw.push((version, entries));
        }
        let common = common_raw
            .into_iter()
            .map(|(from, entries)| {
                let mut entries: Vec<(u32, u16)> =
                    entries.into_iter().map(|(o, s)| (u32::from(o), s)).collect();
                entries.sort_unstable();
                if entries
                    .windows(2)
                    .any(|w| w.first().map(|e| e.0) == w.get(1).map(|e| e.0))
                {
                    return Err(invalid("TPK: two common strings at one offset"));
                }
                Ok(CommonTable { from, entries })
            })
            .collect::<Result<Vec<_>>>()?;
        let pkg = TypeTreePackage {
            classes,
            nodes,
            strings,
            common,
        };
        pkg.check()?;
        Ok(pkg)
    }

    /// Every index in range, so lookups need no checks later.
    fn check(&self) -> Result<()> {
        let (ns, nn) = (self.strings.len(), self.nodes.len());
        for n in &self.nodes {
            if usize::from(n.type_name) >= ns || usize::from(n.name) >= ns {
                return Err(invalid("TPK: node string index out of range"));
            }
            if n.children.iter().any(|&c| usize::from(c) >= nn) {
                return Err(invalid("TPK: child node index out of range"));
            }
        }
        for t in &self.common {
            if t.entries.iter().any(|&(_, s)| usize::from(s) >= ns) {
                return Err(invalid("TPK: common string index out of range"));
            }
        }
        for c in self.classes.values().flatten().filter_map(|(_, c)| c.as_ref()) {
            if usize::from(c.name) >= ns || c.release_root.is_some_and(|r| usize::from(r) >= nn) {
                return Err(invalid("TPK: class index out of range"));
            }
        }
        Ok(())
    }

    fn class(&self, class_id: i32, version: UnityVersion) -> Option<&Class> {
        self.classes
            .get(&class_id)?
            .iter()
            .take_while(|(v, _)| *v <= version)
            .last()
            .and_then(|(_, c)| c.as_ref())
    }

    /// The common string table of a Unity version: the table of the newest listed version that is not newer than
    /// `version`. `None` for a package that lists none up to that version.
    pub fn common_strings(&self, version: UnityVersion) -> Option<CommonStrings<'_>> {
        let table = self.common.iter().take_while(|t| t.from <= version).last()?;
        Some(CommonStrings {
            package: self,
            entries: &table.entries,
        })
    }

    /// The class's name in that Unity version.
    pub fn class_name(&self, class_id: i32, version: UnityVersion) -> Option<&str> {
        self.class(class_id, version)
            .and_then(|c| self.strings.get(usize::from(c.name)))
            .map(|s| &**s)
    }

    /// The class's release (player) type tree in that Unity version.
    pub fn type_tree(&self, class_id: i32, version: UnityVersion) -> Result<TypeTree> {
        let class = self.class(class_id, version).ok_or_else(|| {
            invalid(format!(
                "class {class_id} not in the type tree package for {version}"
            ))
        })?;
        let root = class.release_root.ok_or_else(|| {
            invalid(format!(
                "class {class_id} ({}) has no release type tree in {version} (flags {:#x})",
                self.strings.get(usize::from(class.name)).map_or("?", |s| &**s),
                class.flags
            ))
        })?;
        let mut nodes: Vec<Node> = Vec::new();
        // depth first, children in order; each node's children by their index in `nodes`
        fn add(pkg: &TypeTreePackage, id: u16, depth: u32, nodes: &mut Vec<Node>) -> Result<usize> {
            if depth > 64 {
                return Err(invalid("TPK: type tree deeper than 64"));
            }
            // nodes shared by several parents are copied under each: a package whose nodes name each other twice over
            // would copy 2^depth of them
            if nodes.len() >= MAX_TREE_NODES {
                return Err(invalid(format!(
                    "TPK: a type tree of more than {MAX_TREE_NODES} nodes"
                )));
            }
            let p = pkg
                .nodes
                .get(usize::from(id))
                .ok_or_else(|| invalid(format!("TPK: node {id} of {}", pkg.nodes.len())))?;
            let string = |i: u16| {
                pkg.strings
                    .get(usize::from(i))
                    .cloned()
                    .ok_or_else(|| invalid(format!("TPK: string {i} of {}", pkg.strings.len())))
            };
            let at = nodes.len();
            nodes.push(Node {
                ty: string(p.type_name)?,
                name: string(p.name)?,
                byte_size: p.byte_size,
                version: i32::from(p.version),
                meta_flag: p.meta_flag,
                children: Vec::new(),
            });
            for &c in &p.children {
                let ci = add(pkg, c, depth.saturating_add(1), nodes)?;
                nodes
                    .get_mut(at)
                    .ok_or_else(|| invalid("TPK: a type tree node went missing"))?
                    .children
                    .push(ci);
            }
            Ok(at)
        }
        add(self, root, 0, &mut nodes)?;
        TypeTree::new(nodes)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn package(compression: u8, packed: &[u8], expected: u32) -> Vec<u8> {
        let mut bytes = MAGIC.to_le_bytes().to_vec();
        bytes.extend([2, compression, 0, 0]);
        bytes.extend(0u32.to_le_bytes());
        bytes.extend((packed.len() as u32).to_le_bytes());
        bytes.extend(expected.to_le_bytes());
        bytes.extend(packed);
        bytes
    }

    /// An LZ4 block of one literal run (no matches), encoded by hand: the token's high nibble 15, then the rest of the
    /// length in bytes of 255 and a final byte.
    fn literal_lz4(blob: &[u8]) -> Vec<u8> {
        let mut out = vec![0xf0];
        let mut rest = blob.len() - 15;
        while rest >= 255 {
            out.push(255);
            rest -= 255;
        }
        out.push(rest as u8);
        out.extend(blob);
        out
    }

    /// A format-2 package encoded by hand: one listed version, class 1 named by string 0 ("Base") from version 0
    /// with release root node 0, no common strings, one node (type and name string 0, size -1, version 1), one string.
    fn one_class_blob() -> Vec<u8> {
        let mut b = Vec::new();
        b.extend(0x0102_0304_0506_0708i64.to_le_bytes()); // created
        b.extend(1i32.to_le_bytes()); // versions
        b.extend(0u64.to_le_bytes());
        b.extend(1i32.to_le_bytes()); // classes
        b.extend(1i32.to_le_bytes()); // class id
        b.extend(1i32.to_le_bytes()); // its versions
        b.extend(0u64.to_le_bytes()); // from version 0
        b.push(1); // present
        b.extend(0u16.to_le_bytes()); // name: string 0
        b.extend(0u16.to_le_bytes()); // base
        b.push(HAS_RELEASE_ROOT);
        b.extend(0u16.to_le_bytes()); // release root: node 0
        b.extend(0i32.to_le_bytes()); // common string versions
        b.extend(1i32.to_le_bytes()); // nodes
        b.extend(0u16.to_le_bytes()); // type name: string 0
        b.extend(0u16.to_le_bytes()); // name: string 0
        b.extend((-1i32).to_le_bytes()); // byte size
        b.extend(1i16.to_le_bytes()); // version
        b.push(0); // type flags
        b.extend(0u32.to_le_bytes()); // meta flag
        b.extend(0u16.to_le_bytes()); // children
        b.extend(1i32.to_le_bytes()); // strings
        b.push(4);
        b.extend(b"Base");
        b
    }

    /// The stored and the LZ4 package decode to the same classes, nodes and strings, the ones encoded; wrong sizes
    /// and a cut block are refused.
    #[test]
    fn literal_lz4_package_matches_uncompressed_and_rejects_bad_sizes() {
        let blob = one_class_blob();
        let n = blob.len() as u32;
        let packed = literal_lz4(&blob);
        let raw = TypeTreePackage::parse(&package(0, &blob, n)).unwrap();
        let compressed = TypeTreePackage::parse(&package(1, &packed, n)).unwrap();
        for p in [&raw, &compressed] {
            assert_eq!(p.class_name(1, UnityVersion(1)), Some("Base"));
            assert_eq!(p.class_name(2, UnityVersion(1)), None);
            assert_eq!(p.classes[&1][0].1.as_ref().unwrap().release_root, Some(0));
            assert_eq!(&*p.strings, &[std::sync::Arc::<str>::from("Base")]);
            let node = &p.nodes[0];
            assert_eq!(
                (node.type_name, node.name, node.byte_size, node.version),
                (0, 0, -1, 1)
            );
            assert!(p.type_tree(1, UnityVersion(1)).is_ok());
        }
        for expected in [0, n - 1, n + 1] {
            assert!(TypeTreePackage::parse(&package(1, &packed, expected)).is_err());
        }
        assert!(TypeTreePackage::parse(&package(1, &packed[..packed.len() - 1], n)).is_err());
        assert!(TypeTreePackage::parse(&package(1, &[], 0)).is_err());
        // a size the packed bytes cannot inflate to is refused before anything is allocated for it
        let e = TypeTreePackage::parse(&package(1, &packed, u32::MAX))
            .unwrap_err()
            .to_string();
        assert!(e.contains("inflate to at most"), "{e}");
    }

    /// Nodes that name each other twice over are copied under each parent: thirty levels would be 2^30 nodes. The copy
    /// stops at a count.
    #[test]
    fn a_type_tree_that_expands_past_a_million_nodes_is_refused() {
        let levels = 30u16;
        let nodes = (0..levels)
            .map(|i| PackedNode {
                type_name: 0,
                name: 0,
                byte_size: -1,
                version: 1,
                meta_flag: 0,
                children: if i + 1 < levels {
                    vec![i + 1, i + 1]
                } else {
                    vec![]
                },
            })
            .collect();
        let class = Class {
            name: 0,
            flags: HAS_RELEASE_ROOT,
            release_root: Some(0),
        };
        let pkg = TypeTreePackage {
            classes: BTreeMap::from([(1, vec![(UnityVersion(0), Some(class))])]),
            nodes,
            strings: vec![std::sync::Arc::from("Base")],
            common: Vec::new(),
        };
        let e = pkg.type_tree(1, UnityVersion(1)).unwrap_err().to_string();
        assert!(e.contains("more than 1048576 nodes"), "{e}");
    }
}