unity_formats/
tpk.rs

1//! Type tree packages (TPK, "TPK*" files): the layouts of Unity's built-in classes for every Unity version, which a
2//! player build strips from its files. Format versions 1 and 2, type tree data, uncompressed or LZ4 (LZMA and
3//! Brotli packages are refused: re-pack them uncompressed).
4
5use primase_foundations::compress::{Decompressor, LZ4_MAX_RATIO, Lz4};
6use std::collections::BTreeMap;
7use std::sync::Arc;
8
9use crate::reader::Reader;
10use crate::typetree::{Node, TypeTree};
11use crate::version::UnityVersion;
12use crate::{Error, Result, invalid};
13
14const MAGIC: u32 = 0x2A4B_5054;
15const HAS_EDITOR_ROOT: u8 = 64;
16const HAS_RELEASE_ROOT: u8 = 128;
17
18#[derive(Clone, Debug)]
19struct Class {
20    name: u16,
21    flags: u8,
22    release_root: Option<u16>,
23}
24
25#[derive(Clone, Debug)]
26struct PackedNode {
27    type_name: u16,
28    name: u16,
29    byte_size: i32,
30    version: i16,
31    meta_flag: u32,
32    children: Vec<u16>,
33}
34
35/// A parsed type tree package.
36#[derive(Clone, Debug)]
37pub struct TypeTreePackage {
38    /// Per class id, the class from each listed version on (`None`: absent from that version).
39    classes: BTreeMap<i32, Vec<(UnityVersion, Option<Class>)>>,
40    nodes: Vec<PackedNode>,
41    strings: Vec<Arc<str>>,
42}
43
44#[cfg(test)]
45impl TypeTreePackage {
46    /// A package without classes, for tests that need a package but decode no object.
47    pub(crate) fn empty() -> Self {
48        TypeTreePackage {
49            classes: BTreeMap::new(),
50            nodes: Vec::new(),
51            strings: Vec::new(),
52        }
53    }
54}
55
56fn varint_string(r: &mut Reader) -> Result<Arc<str>> {
57    let (mut len, mut shift) = (0usize, 0u32);
58    loop {
59        let b = r.u8()?;
60        if shift > 28 {
61            return Err(invalid("TPK: string length varint too long"));
62        }
63        len |= usize::from(b & 0x7F) << shift;
64        shift = shift.saturating_add(7);
65        if b & 0x80 == 0 {
66            break;
67        }
68    }
69    let s = std::str::from_utf8(r.take(len)?).map_err(|_| invalid("TPK: string not UTF-8"))?;
70    Ok(Arc::from(s))
71}
72
73/// Skip `n` entries of `size` bytes each.
74fn skip(r: &mut Reader, n: usize, size: usize, what: &str) -> Result<()> {
75    let bytes = n.checked_mul(size).ok_or_else(|| {
76        invalid(format!(
77            "TPK: {n} {what} entries of {size} bytes at offset {}",
78            r.pos
79        ))
80    })?;
81    r.take(bytes).map(|_| ())
82}
83
84fn count(r: &mut Reader, what: &str) -> Result<usize> {
85    let n = r.i32()?;
86    usize::try_from(n).map_err(|_| invalid(format!("TPK: negative {what} count {n}")))
87}
88
89/// The most nodes one class's type tree may expand to (the trees seen have a few thousand).
90const MAX_TREE_NODES: usize = 1 << 20;
91
92impl TypeTreePackage {
93    /// Parse a type tree package (TPK format 1 or 2, uncompressed or LZ4). Every index in it is checked, the size an
94    /// LZ4 block claims is checked against what LZ4 can inflate to before anything is allocated for it, and a package
95    /// that is LZMA- or Brotli-compressed is refused as [`crate::Error::Unsupported`].
96    pub fn parse(bytes: &[u8]) -> Result<Self> {
97        let mut r = Reader::new(bytes, false);
98        if r.u32()? != MAGIC {
99            return Err(invalid("not a TPK file"));
100        }
101        let format = r.u8()?;
102        if !(1..=2).contains(&format) {
103            return Err(Error::Unsupported(format!("TPK format version {format}")));
104        }
105        let compression = r.u8()?;
106        let data_type = r.u8()?;
107        let _reserved = r.u8()?;
108        let _reserved2 = r.u32()?;
109        let compressed = r.u32()? as usize;
110        let uncompressed = r.u32()? as usize;
111        if data_type != 0 {
112            return Err(Error::Unsupported(format!(
113                "TPK data type {data_type} (not type trees)"
114            )));
115        }
116        let packed = r.take(compressed)?;
117        let blob: Vec<u8> = match compression {
118            0 => packed.to_vec(),
119            1 => {
120                // the header's size is the file's word: what the packed bytes cannot inflate to is not allocated
121                if uncompressed > packed.len().saturating_mul(LZ4_MAX_RATIO) {
122                    return Err(invalid(format!(
123                        "TPK: {uncompressed} bytes claimed of {} bytes of LZ4, which inflate to at most {}",
124                        packed.len(),
125                        packed.len().saturating_mul(LZ4_MAX_RATIO)
126                    )));
127                }
128                Lz4.decompress(packed, uncompressed)
129                    .map_err(|e| invalid(format!("TPK LZ4: {e}")))?
130            }
131            2 => {
132                return Err(Error::Unsupported(
133                    "LZMA-compressed TPK: use an uncompressed or LZ4 package".into(),
134                ));
135            }
136            3 => {
137                return Err(Error::Unsupported(
138                    "Brotli-compressed TPK: use an uncompressed or LZ4 package".into(),
139                ));
140            }
141            c => return Err(invalid(format!("TPK compression type {c}"))),
142        };
143        if blob.len() != uncompressed {
144            return Err(invalid(format!(
145                "TPK: {} bytes, header says {uncompressed}",
146                blob.len()
147            )));
148        }
149        Self::parse_blob(&blob, format)
150    }
151
152    fn parse_blob(blob: &[u8], format: u8) -> Result<Self> {
153        let mut r = Reader::new(blob, false);
154        let _created = r.i64()?;
155        let versions = count(&mut r, "version")?;
156        skip(&mut r, versions, 8, "version")?;
157        let mut classes = BTreeMap::new();
158        for _ in 0..count(&mut r, "class")? {
159            let id = r.i32()?;
160            let mut entries = Vec::new();
161            for _ in 0..count(&mut r, "class version")? {
162                let version = UnityVersion(r.u64()?);
163                let class = if r.u8()? != 0 {
164                    let name = r.u16()?;
165                    let _base = r.u16()?;
166                    let flags = r.u8()?;
167                    if flags & HAS_EDITOR_ROOT != 0 {
168                        r.u16()?;
169                    }
170                    let release_root = if flags & HAS_RELEASE_ROOT != 0 {
171                        Some(r.u16()?)
172                    } else {
173                        None
174                    };
175                    Some(Class {
176                        name,
177                        flags,
178                        release_root,
179                    })
180                } else {
181                    None
182                };
183                entries.push((version, class));
184            }
185            classes.insert(id, entries);
186        }
187        // the common string table (offsets into Unity's built-in strings): not needed for stripped files
188        match format {
189            1 => {
190                let n = count(&mut r, "common string version")?;
191                skip(&mut r, n, 9, "common string version")?;
192                let n = count(&mut r, "common string index")?;
193                skip(&mut r, n, 2, "common string index")?;
194            }
195            _ => {
196                for _ in 0..count(&mut r, "common string version")? {
197                    r.u64()?;
198                    let n = count(&mut r, "common string entry")?;
199                    skip(&mut r, n, 4, "common string entry")?;
200                }
201            }
202        }
203        let mut nodes = Vec::new();
204        for _ in 0..count(&mut r, "node")? {
205            let type_name = r.u16()?;
206            let name = r.u16()?;
207            let byte_size = r.i32()?;
208            let version = r.i16()?;
209            let _type_flags = r.u8()?;
210            let meta_flag = r.u32()?;
211            let n = r.u16()?;
212            let children = (0..n).map(|_| r.u16()).collect::<Result<Vec<_>>>()?;
213            nodes.push(PackedNode {
214                type_name,
215                name,
216                byte_size,
217                version,
218                meta_flag,
219                children,
220            });
221        }
222        let strings = (0..count(&mut r, "string")?)
223            .map(|_| varint_string(&mut r))
224            .collect::<Result<Vec<_>>>()?;
225        if r.remaining() != 0 {
226            return Err(invalid(format!(
227                "TPK: {} bytes after the string buffer",
228                r.remaining()
229            )));
230        }
231        let pkg = TypeTreePackage {
232            classes,
233            nodes,
234            strings,
235        };
236        pkg.check()?;
237        Ok(pkg)
238    }
239
240    /// Every index in range, so lookups need no checks later.
241    fn check(&self) -> Result<()> {
242        let (ns, nn) = (self.strings.len(), self.nodes.len());
243        for n in &self.nodes {
244            if usize::from(n.type_name) >= ns || usize::from(n.name) >= ns {
245                return Err(invalid("TPK: node string index out of range"));
246            }
247            if n.children.iter().any(|&c| usize::from(c) >= nn) {
248                return Err(invalid("TPK: child node index out of range"));
249            }
250        }
251        for c in self.classes.values().flatten().filter_map(|(_, c)| c.as_ref()) {
252            if usize::from(c.name) >= ns || c.release_root.is_some_and(|r| usize::from(r) >= nn) {
253                return Err(invalid("TPK: class index out of range"));
254            }
255        }
256        Ok(())
257    }
258
259    fn class(&self, class_id: i32, version: UnityVersion) -> Option<&Class> {
260        self.classes
261            .get(&class_id)?
262            .iter()
263            .take_while(|(v, _)| *v <= version)
264            .last()
265            .and_then(|(_, c)| c.as_ref())
266    }
267
268    /// The class's name in that Unity version.
269    pub fn class_name(&self, class_id: i32, version: UnityVersion) -> Option<&str> {
270        self.class(class_id, version)
271            .and_then(|c| self.strings.get(usize::from(c.name)))
272            .map(|s| &**s)
273    }
274
275    /// The class's release (player) type tree in that Unity version.
276    pub fn type_tree(&self, class_id: i32, version: UnityVersion) -> Result<TypeTree> {
277        let class = self.class(class_id, version).ok_or_else(|| {
278            invalid(format!(
279                "class {class_id} not in the type tree package for {version}"
280            ))
281        })?;
282        let root = class.release_root.ok_or_else(|| {
283            invalid(format!(
284                "class {class_id} ({}) has no release type tree in {version} (flags {:#x})",
285                self.strings.get(usize::from(class.name)).map_or("?", |s| &**s),
286                class.flags
287            ))
288        })?;
289        let mut nodes: Vec<Node> = Vec::new();
290        // depth first, children in order; each node's children by their index in `nodes`
291        fn add(pkg: &TypeTreePackage, id: u16, depth: u32, nodes: &mut Vec<Node>) -> Result<usize> {
292            if depth > 64 {
293                return Err(invalid("TPK: type tree deeper than 64"));
294            }
295            // nodes shared by several parents are copied under each: a package whose nodes name each other twice over
296            // would copy 2^depth of them
297            if nodes.len() >= MAX_TREE_NODES {
298                return Err(invalid(format!(
299                    "TPK: a type tree of more than {MAX_TREE_NODES} nodes"
300                )));
301            }
302            let p = pkg
303                .nodes
304                .get(usize::from(id))
305                .ok_or_else(|| invalid(format!("TPK: node {id} of {}", pkg.nodes.len())))?;
306            let string = |i: u16| {
307                pkg.strings
308                    .get(usize::from(i))
309                    .cloned()
310                    .ok_or_else(|| invalid(format!("TPK: string {i} of {}", pkg.strings.len())))
311            };
312            let at = nodes.len();
313            nodes.push(Node {
314                ty: string(p.type_name)?,
315                name: string(p.name)?,
316                byte_size: p.byte_size,
317                version: i32::from(p.version),
318                meta_flag: p.meta_flag,
319                children: Vec::new(),
320            });
321            for &c in &p.children {
322                let ci = add(pkg, c, depth.saturating_add(1), nodes)?;
323                nodes
324                    .get_mut(at)
325                    .ok_or_else(|| invalid("TPK: a type tree node went missing"))?
326                    .children
327                    .push(ci);
328            }
329            Ok(at)
330        }
331        add(self, root, 0, &mut nodes)?;
332        TypeTree::new(nodes)
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use super::*;
339
340    fn package(compression: u8, packed: &[u8], expected: u32) -> Vec<u8> {
341        let mut bytes = MAGIC.to_le_bytes().to_vec();
342        bytes.extend([2, compression, 0, 0]);
343        bytes.extend(0u32.to_le_bytes());
344        bytes.extend((packed.len() as u32).to_le_bytes());
345        bytes.extend(expected.to_le_bytes());
346        bytes.extend(packed);
347        bytes
348    }
349
350    /// An LZ4 block of one literal run (no matches), encoded by hand: the token's high nibble 15, then the rest of the
351    /// length in bytes of 255 and a final byte.
352    fn literal_lz4(blob: &[u8]) -> Vec<u8> {
353        let mut out = vec![0xf0];
354        let mut rest = blob.len() - 15;
355        while rest >= 255 {
356            out.push(255);
357            rest -= 255;
358        }
359        out.push(rest as u8);
360        out.extend(blob);
361        out
362    }
363
364    /// A format-2 package encoded by hand: one listed version, class 1 named by string 0 ("Base") from version 0
365    /// with release root node 0, no common strings, one node (type and name string 0, size -1, version 1), one string.
366    fn one_class_blob() -> Vec<u8> {
367        let mut b = Vec::new();
368        b.extend(0x0102_0304_0506_0708i64.to_le_bytes()); // created
369        b.extend(1i32.to_le_bytes()); // versions
370        b.extend(0u64.to_le_bytes());
371        b.extend(1i32.to_le_bytes()); // classes
372        b.extend(1i32.to_le_bytes()); // class id
373        b.extend(1i32.to_le_bytes()); // its versions
374        b.extend(0u64.to_le_bytes()); // from version 0
375        b.push(1); // present
376        b.extend(0u16.to_le_bytes()); // name: string 0
377        b.extend(0u16.to_le_bytes()); // base
378        b.push(HAS_RELEASE_ROOT);
379        b.extend(0u16.to_le_bytes()); // release root: node 0
380        b.extend(0i32.to_le_bytes()); // common string versions
381        b.extend(1i32.to_le_bytes()); // nodes
382        b.extend(0u16.to_le_bytes()); // type name: string 0
383        b.extend(0u16.to_le_bytes()); // name: string 0
384        b.extend((-1i32).to_le_bytes()); // byte size
385        b.extend(1i16.to_le_bytes()); // version
386        b.push(0); // type flags
387        b.extend(0u32.to_le_bytes()); // meta flag
388        b.extend(0u16.to_le_bytes()); // children
389        b.extend(1i32.to_le_bytes()); // strings
390        b.push(4);
391        b.extend(b"Base");
392        b
393    }
394
395    /// The stored and the LZ4 package decode to the same classes, nodes and strings, the ones encoded; wrong sizes
396    /// and a cut block are refused.
397    #[test]
398    fn literal_lz4_package_matches_uncompressed_and_rejects_bad_sizes() {
399        let blob = one_class_blob();
400        let n = blob.len() as u32;
401        let packed = literal_lz4(&blob);
402        let raw = TypeTreePackage::parse(&package(0, &blob, n)).unwrap();
403        let compressed = TypeTreePackage::parse(&package(1, &packed, n)).unwrap();
404        for p in [&raw, &compressed] {
405            assert_eq!(p.class_name(1, UnityVersion(1)), Some("Base"));
406            assert_eq!(p.class_name(2, UnityVersion(1)), None);
407            assert_eq!(p.classes[&1][0].1.as_ref().unwrap().release_root, Some(0));
408            assert_eq!(&*p.strings, &[std::sync::Arc::<str>::from("Base")]);
409            let node = &p.nodes[0];
410            assert_eq!(
411                (node.type_name, node.name, node.byte_size, node.version),
412                (0, 0, -1, 1)
413            );
414            assert!(p.type_tree(1, UnityVersion(1)).is_ok());
415        }
416        for expected in [0, n - 1, n + 1] {
417            assert!(TypeTreePackage::parse(&package(1, &packed, expected)).is_err());
418        }
419        assert!(TypeTreePackage::parse(&package(1, &packed[..packed.len() - 1], n)).is_err());
420        assert!(TypeTreePackage::parse(&package(1, &[], 0)).is_err());
421        // a size the packed bytes cannot inflate to is refused before anything is allocated for it
422        let e = TypeTreePackage::parse(&package(1, &packed, u32::MAX))
423            .unwrap_err()
424            .to_string();
425        assert!(e.contains("inflate to at most"), "{e}");
426    }
427
428    /// Nodes that name each other twice over are copied under each parent: thirty levels would be 2^30 nodes. The copy
429    /// stops at a count.
430    #[test]
431    fn a_type_tree_that_expands_past_a_million_nodes_is_refused() {
432        let levels = 30u16;
433        let nodes = (0..levels)
434            .map(|i| PackedNode {
435                type_name: 0,
436                name: 0,
437                byte_size: -1,
438                version: 1,
439                meta_flag: 0,
440                children: if i + 1 < levels {
441                    vec![i + 1, i + 1]
442                } else {
443                    vec![]
444                },
445            })
446            .collect();
447        let class = Class {
448            name: 0,
449            flags: HAS_RELEASE_ROOT,
450            release_root: Some(0),
451        };
452        let pkg = TypeTreePackage {
453            classes: BTreeMap::from([(1, vec![(UnityVersion(0), Some(class))])]),
454            nodes,
455            strings: vec![std::sync::Arc::from("Base")],
456        };
457        let e = pkg.type_tree(1, UnityVersion(1)).unwrap_err().to_string();
458        assert!(e.contains("more than 1048576 nodes"), "{e}");
459    }
460}
unity_formats/tpk.rs

unity_formats/
tpk.rs
