1use crate::compression::{self, CompressionType};
8use crate::error::{BinaryError, Result};
9use crate::typetree::{TypeTree, TypeTreeNode, TypeTreeRegistry};
10use crate::unity_version::{UnityVersion, UnityVersionType};
11use std::collections::HashMap;
12use std::io::{Cursor, Read};
13use std::path::Path;
14use std::sync::{Arc, RwLock};
15
16type TypeTreeCache = Arc<RwLock<HashMap<(i32, u64), Arc<TypeTree>>>>;
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19#[repr(i8)]
20enum TpkCompressionType {
21 None = 0,
22 Lz4 = 1,
23 Lzma = 2,
24 Brotli = 3,
25}
26
27impl TryFrom<i8> for TpkCompressionType {
28 type Error = BinaryError;
29
30 fn try_from(value: i8) -> Result<Self> {
31 match value {
32 0 => Ok(Self::None),
33 1 => Ok(Self::Lz4),
34 2 => Ok(Self::Lzma),
35 3 => Ok(Self::Brotli),
36 other => Err(BinaryError::invalid_data(format!(
37 "Invalid TPK compression type: {}",
38 other
39 ))),
40 }
41 }
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45#[repr(i8)]
46enum TpkDataType {
47 TypeTreeInformation = 0,
48 Collection = 1,
49 FileSystem = 2,
50 Json = 3,
51 ReferenceAssemblies = 4,
52 EngineAssets = 5,
53}
54
55impl TryFrom<i8> for TpkDataType {
56 type Error = BinaryError;
57
58 fn try_from(value: i8) -> Result<Self> {
59 match value {
60 0 => Ok(Self::TypeTreeInformation),
61 1 => Ok(Self::Collection),
62 2 => Ok(Self::FileSystem),
63 3 => Ok(Self::Json),
64 4 => Ok(Self::ReferenceAssemblies),
65 5 => Ok(Self::EngineAssets),
66 other => Err(BinaryError::invalid_data(format!(
67 "Invalid TPK data type: {}",
68 other
69 ))),
70 }
71 }
72}
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
75#[repr(u8)]
76enum TpkUnityClassFlags {
77 HasEditorRootNode = 64,
78 HasReleaseRootNode = 128,
79}
80
81#[derive(Debug, Clone)]
82struct TpkFileHeader {
83 compression: TpkCompressionType,
84 data_type: TpkDataType,
85 compressed_size: u32,
86 uncompressed_size: u32,
87}
88
89#[derive(Debug, Clone)]
90struct TpkUnityClass {
91 #[allow(dead_code)]
92 name: u16,
93 #[allow(dead_code)]
94 base: u16,
95 #[allow(dead_code)]
96 flags: u8,
97 #[allow(dead_code)]
98 editor_root_node: Option<u16>,
99 release_root_node: Option<u16>,
100}
101
102#[derive(Debug, Clone)]
103struct TpkClassInformation {
104 #[allow(dead_code)]
105 id: i32,
106 classes: Vec<(u64, Option<TpkUnityClass>)>,
107}
108
109#[derive(Debug, Clone)]
110struct TpkUnityNode {
111 type_name: u16,
112 name: u16,
113 byte_size: i32,
114 version: i16,
115 type_flags: i8,
116 meta_flag: u32,
117 sub_nodes: Vec<u16>,
118}
119
120#[derive(Debug, Clone)]
121struct TpkTypeTreeBlob {
122 #[allow(dead_code)]
123 creation_time: i64,
124 #[allow(dead_code)]
125 versions: Vec<u64>,
126 class_information: HashMap<i32, TpkClassInformation>,
127 nodes: Vec<TpkUnityNode>,
128 strings: Vec<String>,
129}
130
131#[derive(Debug)]
132struct TpkReader<'a> {
133 cur: Cursor<&'a [u8]>,
134}
135
136impl<'a> TpkReader<'a> {
137 fn new(data: &'a [u8]) -> Self {
138 Self {
139 cur: Cursor::new(data),
140 }
141 }
142
143 fn read_exact<const N: usize>(&mut self) -> Result<[u8; N]> {
144 let mut buf = [0u8; N];
145 self.cur
146 .read_exact(&mut buf)
147 .map_err(|e| BinaryError::generic(format!("TPK read failed: {}", e)))?;
148 Ok(buf)
149 }
150
151 fn read_u8(&mut self) -> Result<u8> {
152 Ok(self.read_exact::<1>()?[0])
153 }
154
155 fn read_i8(&mut self) -> Result<i8> {
156 Ok(self.read_u8()? as i8)
157 }
158
159 fn read_u16_le(&mut self) -> Result<u16> {
160 Ok(u16::from_le_bytes(self.read_exact::<2>()?))
161 }
162
163 fn read_i16_le(&mut self) -> Result<i16> {
164 Ok(i16::from_le_bytes(self.read_exact::<2>()?))
165 }
166
167 fn read_u32_le(&mut self) -> Result<u32> {
168 Ok(u32::from_le_bytes(self.read_exact::<4>()?))
169 }
170
171 fn read_i32_le(&mut self) -> Result<i32> {
172 Ok(i32::from_le_bytes(self.read_exact::<4>()?))
173 }
174
175 fn read_i64_le(&mut self) -> Result<i64> {
176 Ok(i64::from_le_bytes(self.read_exact::<8>()?))
177 }
178
179 fn read_u64_le(&mut self) -> Result<u64> {
180 Ok(u64::from_le_bytes(self.read_exact::<8>()?))
181 }
182
183 fn read_bytes(&mut self, n: usize) -> Result<Vec<u8>> {
184 let mut buf = vec![0u8; n];
185 self.cur
186 .read_exact(&mut buf)
187 .map_err(|e| BinaryError::generic(format!("TPK read failed: {}", e)))?;
188 Ok(buf)
189 }
190
191 fn read_varint_len(&mut self) -> Result<usize> {
192 let mut shift = 0u32;
193 let mut len: u64 = 0;
194 loop {
195 let b = self.read_u8()?;
196 len |= ((b & 0x7F) as u64) << shift;
197 if (b & 0x80) == 0 {
198 break;
199 }
200 shift = shift.saturating_add(7);
201 if shift > 63 {
202 return Err(BinaryError::invalid_data(
203 "TPK varint too large".to_string(),
204 ));
205 }
206 }
207 Ok(len as usize)
208 }
209
210 fn read_string(&mut self) -> Result<String> {
211 let len = self.read_varint_len()?;
212 let bytes = self.read_bytes(len)?;
213 String::from_utf8(bytes)
214 .map_err(|e| BinaryError::invalid_data(format!("TPK invalid utf8: {}", e)))
215 }
216}
217
218fn unity_version_to_u64(v: &UnityVersion) -> u64 {
219 let type_byte: u8 = match v.version_type {
220 UnityVersionType::A => 0,
221 UnityVersionType::B => 1,
222 UnityVersionType::C => 2,
223 UnityVersionType::F => 3,
224 UnityVersionType::P => 4,
225 UnityVersionType::X => 5,
226 UnityVersionType::U => 255,
227 };
228 ((v.major as u64) << 48)
229 | ((v.minor as u64) << 32)
230 | ((v.build as u64) << 16)
231 | ((type_byte as u64) << 8)
232 | (v.type_number as u64)
233}
234
235fn select_versioned_class(
236 version: u64,
237 classes: &[(u64, Option<TpkUnityClass>)],
238) -> Option<&TpkUnityClass> {
239 let mut ret: Option<&TpkUnityClass> = None;
240 for (v, item) in classes {
241 if version >= *v {
242 if let Some(c) = item.as_ref() {
243 ret = Some(c);
244 }
245 } else {
246 break;
247 }
248 }
249 ret
250}
251
252fn build_tree_from_blob(blob: &TpkTypeTreeBlob, class: &TpkUnityClass) -> Result<TypeTree> {
253 let root_id = class
254 .release_root_node
255 .ok_or_else(|| BinaryError::invalid_data("TPK class has no ReleaseRootNode".to_string()))?
256 as usize;
257
258 fn build_node(
259 blob: &TpkTypeTreeBlob,
260 node_id: usize,
261 level: i32,
262 next_index: &mut i32,
263 ) -> Result<TypeTreeNode> {
264 let node = blob.nodes.get(node_id).ok_or_else(|| {
265 BinaryError::invalid_data(format!("TPK node out of range: {}", node_id))
266 })?;
267 let type_name = blob
268 .strings
269 .get(node.type_name as usize)
270 .ok_or_else(|| {
271 BinaryError::invalid_data("TPK type string index out of range".to_string())
272 })?
273 .clone();
274 let name = blob
275 .strings
276 .get(node.name as usize)
277 .ok_or_else(|| {
278 BinaryError::invalid_data("TPK name string index out of range".to_string())
279 })?
280 .clone();
281
282 let mut out = TypeTreeNode::new();
283 out.type_name = type_name;
284 out.name = name;
285 out.byte_size = node.byte_size;
286 out.index = *next_index;
287 out.version = node.version as i32;
288 out.type_flags = node.type_flags as i32;
289 out.meta_flags = node.meta_flag as i32;
290 out.level = level;
291
292 *next_index = next_index.saturating_add(1);
293 out.children = node
294 .sub_nodes
295 .iter()
296 .map(|id| build_node(blob, *id as usize, level + 1, next_index))
297 .collect::<Result<Vec<_>>>()?;
298 Ok(out)
299 }
300
301 let mut next_index: i32 = 0;
302 let root = build_node(blob, root_id, 0, &mut next_index)?;
303 let mut tree = TypeTree::new();
304 tree.add_node(root);
305 Ok(tree)
306}
307
308fn parse_tpk_header(reader: &mut TpkReader<'_>) -> Result<TpkFileHeader> {
309 let magic = reader.read_u32_le()?;
310 const TPK_MAGIC: u32 = 0x2A4B5054;
311 if magic != TPK_MAGIC {
312 return Err(BinaryError::invalid_data(
313 "Invalid TPK magic bytes".to_string(),
314 ));
315 }
316
317 let version_number = reader.read_i8()?;
318 if version_number != 1 {
319 return Err(BinaryError::invalid_data(format!(
320 "Invalid TPK version number: {}",
321 version_number
322 )));
323 }
324
325 let compression = TpkCompressionType::try_from(reader.read_i8()?)?;
326 let data_type = TpkDataType::try_from(reader.read_i8()?)?;
327 let _unused_b = reader.read_i8()?;
328 let _unused_u32 = reader.read_u32_le()?;
329 let compressed_size = reader.read_u32_le()?;
330 let uncompressed_size = reader.read_u32_le()?;
331
332 Ok(TpkFileHeader {
333 compression,
334 data_type,
335 compressed_size,
336 uncompressed_size,
337 })
338}
339
340fn decompress_tpk_payload(header: &TpkFileHeader, compressed: &[u8]) -> Result<Vec<u8>> {
341 let (ctype, expected) = match header.compression {
342 TpkCompressionType::None => (CompressionType::None, compressed.len()),
343 TpkCompressionType::Lz4 => (CompressionType::Lz4, header.uncompressed_size as usize),
344 TpkCompressionType::Lzma => (CompressionType::Lzma, header.uncompressed_size as usize),
345 TpkCompressionType::Brotli => (CompressionType::Brotli, header.uncompressed_size as usize),
346 };
347 if ctype == CompressionType::None {
348 return Ok(compressed.to_vec());
349 }
350 compression::decompress(compressed, ctype, expected)
351}
352
353fn parse_tpk_typetree_blob(data: &[u8]) -> Result<TpkTypeTreeBlob> {
354 let mut r = TpkReader::new(data);
355 let creation_time = r.read_i64_le()?;
356 let version_count = r.read_i32_le()?;
357 if version_count < 0 {
358 return Err(BinaryError::invalid_data(
359 "Negative TPK version count".to_string(),
360 ));
361 }
362 let mut versions: Vec<u64> = Vec::with_capacity(version_count as usize);
363 for _ in 0..version_count {
364 versions.push(r.read_u64_le()?);
365 }
366
367 let class_count = r.read_i32_le()?;
368 if class_count < 0 {
369 return Err(BinaryError::invalid_data(
370 "Negative TPK class count".to_string(),
371 ));
372 }
373 let mut class_information: HashMap<i32, TpkClassInformation> = HashMap::new();
374 for _ in 0..class_count {
375 let id = r.read_i32_le()?;
376 let count = r.read_i32_le()?;
377 if count < 0 {
378 return Err(BinaryError::invalid_data(
379 "Negative TPK class version count".to_string(),
380 ));
381 }
382 let mut classes: Vec<(u64, Option<TpkUnityClass>)> = Vec::with_capacity(count as usize);
383 for _ in 0..count {
384 let version = r.read_u64_le()?;
385 let present = r.read_u8()?;
386 let class = if present != 0 {
387 let name = r.read_u16_le()?;
388 let base = r.read_u16_le()?;
389 let flags = r.read_u8()?;
390 let mut editor_root_node: Option<u16> = None;
391 let mut release_root_node: Option<u16> = None;
392 if (flags & TpkUnityClassFlags::HasEditorRootNode as u8) != 0 {
393 editor_root_node = Some(r.read_u16_le()?);
394 }
395 if (flags & TpkUnityClassFlags::HasReleaseRootNode as u8) != 0 {
396 release_root_node = Some(r.read_u16_le()?);
397 }
398 Some(TpkUnityClass {
399 name,
400 base,
401 flags,
402 editor_root_node,
403 release_root_node,
404 })
405 } else {
406 None
407 };
408 classes.push((version, class));
409 }
410 class_information.insert(id, TpkClassInformation { id, classes });
411 }
412
413 let common_version_count = r.read_i32_le()?;
415 if common_version_count < 0 {
416 return Err(BinaryError::invalid_data(
417 "Negative TPK common string version count".to_string(),
418 ));
419 }
420 for _ in 0..common_version_count {
421 let _ver = r.read_u64_le()?;
422 let _count = r.read_u8()?;
423 }
424 let indices_count = r.read_i32_le()?;
425 if indices_count < 0 {
426 return Err(BinaryError::invalid_data(
427 "Negative TPK common string indices count".to_string(),
428 ));
429 }
430 for _ in 0..indices_count {
431 let _idx = r.read_u16_le()?;
432 }
433
434 let node_count = r.read_i32_le()?;
436 if node_count < 0 {
437 return Err(BinaryError::invalid_data(
438 "Negative TPK node count".to_string(),
439 ));
440 }
441 let mut nodes: Vec<TpkUnityNode> = Vec::with_capacity(node_count as usize);
442 for _ in 0..node_count {
443 let type_name = r.read_u16_le()?;
444 let name = r.read_u16_le()?;
445 let byte_size = r.read_i32_le()?;
446 let version = r.read_i16_le()?;
447 let type_flags = r.read_i8()?;
448 let meta_flag = r.read_u32_le()?;
449 let count = r.read_u16_le()? as usize;
450 let mut sub_nodes: Vec<u16> = Vec::with_capacity(count);
451 for _ in 0..count {
452 sub_nodes.push(r.read_u16_le()?);
453 }
454 nodes.push(TpkUnityNode {
455 type_name,
456 name,
457 byte_size,
458 version,
459 type_flags,
460 meta_flag,
461 sub_nodes,
462 });
463 }
464
465 let string_count = r.read_i32_le()?;
467 if string_count < 0 {
468 return Err(BinaryError::invalid_data(
469 "Negative TPK string count".to_string(),
470 ));
471 }
472 let mut strings: Vec<String> = Vec::with_capacity(string_count as usize);
473 for _ in 0..string_count {
474 strings.push(r.read_string()?);
475 }
476
477 Ok(TpkTypeTreeBlob {
478 creation_time,
479 versions,
480 class_information,
481 nodes,
482 strings,
483 })
484}
485
486#[derive(Debug, Clone)]
488pub struct TpkTypeTreeRegistry {
489 blob: Arc<TpkTypeTreeBlob>,
490 cache: TypeTreeCache,
491}
492
493impl TpkTypeTreeRegistry {
494 pub fn from_bytes(data: &[u8]) -> Result<Self> {
495 let mut r = TpkReader::new(data);
496 let header = parse_tpk_header(&mut r)?;
497 if header.data_type != TpkDataType::TypeTreeInformation {
498 return Err(BinaryError::unsupported(format!(
499 "Unsupported TPK data type: {:?}",
500 header.data_type
501 )));
502 }
503 let compressed = r.read_bytes(header.compressed_size as usize)?;
504 if compressed.len() != header.compressed_size as usize {
505 return Err(BinaryError::invalid_data(
506 "Invalid TPK compressed size".to_string(),
507 ));
508 }
509 let decompressed = decompress_tpk_payload(&header, &compressed)?;
510 let blob = parse_tpk_typetree_blob(&decompressed)?;
511 Ok(Self {
512 blob: Arc::new(blob),
513 cache: Arc::new(RwLock::new(HashMap::new())),
514 })
515 }
516
517 pub fn from_path(path: impl AsRef<Path>) -> Result<Self> {
518 let data = std::fs::read(path.as_ref()).map_err(|e| {
519 BinaryError::generic(format!(
520 "Failed to read TPK file {:?}: {}",
521 path.as_ref(),
522 e
523 ))
524 })?;
525 Self::from_bytes(&data)
526 }
527}
528
529impl TypeTreeRegistry for TpkTypeTreeRegistry {
530 fn resolve(&self, unity_version: &str, class_id: i32) -> Option<Arc<TypeTree>> {
531 let Ok(v) = UnityVersion::parse_version(unity_version) else {
532 return None;
533 };
534 let encoded = unity_version_to_u64(&v);
535
536 if let Ok(cache) = self.cache.read()
537 && let Some(found) = cache.get(&(class_id, encoded))
538 {
539 return Some(found.clone());
540 }
541
542 let ci = self.blob.class_information.get(&class_id)?;
543 let class = select_versioned_class(encoded, &ci.classes)?;
544 let built = build_tree_from_blob(&self.blob, class).ok()?;
545 let built = Arc::new(built);
546
547 match self.cache.write() {
548 Ok(mut cache) => {
549 cache.insert((class_id, encoded), built.clone());
550 }
551 Err(e) => {
552 let mut cache = e.into_inner();
553 cache.insert((class_id, encoded), built.clone());
554 }
555 }
556
557 Some(built)
558 }
559}
560
561#[cfg(test)]
562mod tests {
563 use super::*;
564 use crate::reader::{BinaryReader, ByteOrder};
565 use crate::typetree::{TypeTreeParseOptions, TypeTreeSerializer};
566 use unity_asset_core::UnityValue;
567
568 fn write_varint(mut n: usize, out: &mut Vec<u8>) {
569 loop {
570 let mut b = (n & 0x7F) as u8;
571 n >>= 7;
572 if n != 0 {
573 b |= 0x80;
574 }
575 out.push(b);
576 if n == 0 {
577 break;
578 }
579 }
580 }
581
582 fn write_tpk_string(s: &str, out: &mut Vec<u8>) {
583 write_varint(s.len(), out);
584 out.extend_from_slice(s.as_bytes());
585 }
586
587 fn build_minimal_tpk() -> Vec<u8> {
588 let mut blob: Vec<u8> = Vec::new();
590 blob.extend_from_slice(&0i64.to_le_bytes()); blob.extend_from_slice(&1i32.to_le_bytes()); let v = UnityVersion::parse_version("2020.3.0f1").unwrap();
594 let v_u64 = unity_version_to_u64(&v);
595 blob.extend_from_slice(&v_u64.to_le_bytes()); blob.extend_from_slice(&1i32.to_le_bytes()); blob.extend_from_slice(&(28i32).to_le_bytes()); blob.extend_from_slice(&1i32.to_le_bytes()); blob.extend_from_slice(&v_u64.to_le_bytes()); blob.push(1u8); blob.extend_from_slice(&(0u16).to_le_bytes()); blob.extend_from_slice(&(0u16).to_le_bytes()); blob.push(TpkUnityClassFlags::HasReleaseRootNode as u8); blob.extend_from_slice(&(0u16).to_le_bytes()); blob.extend_from_slice(&0i32.to_le_bytes());
611 blob.extend_from_slice(&0i32.to_le_bytes());
612
613 blob.extend_from_slice(&2i32.to_le_bytes());
615 blob.extend_from_slice(&(0u16).to_le_bytes()); blob.extend_from_slice(&(1u16).to_le_bytes()); blob.extend_from_slice(&(-1i32).to_le_bytes()); blob.extend_from_slice(&(1i16).to_le_bytes()); blob.push(0i8 as u8); blob.extend_from_slice(&(0u32).to_le_bytes()); blob.extend_from_slice(&(1u16).to_le_bytes()); blob.extend_from_slice(&(1u16).to_le_bytes()); blob.extend_from_slice(&(2u16).to_le_bytes()); blob.extend_from_slice(&(3u16).to_le_bytes()); blob.extend_from_slice(&(-1i32).to_le_bytes()); blob.extend_from_slice(&(1i16).to_le_bytes()); blob.push(0i8 as u8); blob.extend_from_slice(&(0u32).to_le_bytes()); blob.extend_from_slice(&(0u16).to_le_bytes()); blob.extend_from_slice(&4i32.to_le_bytes());
635 write_tpk_string("RootType", &mut blob); write_tpk_string("Base", &mut blob); write_tpk_string("string", &mut blob); write_tpk_string("m_Name", &mut blob); let mut out: Vec<u8> = Vec::new();
641 out.extend_from_slice(&0x2A4B5054u32.to_le_bytes()); out.push(1u8); out.push(TpkCompressionType::None as i8 as u8); out.push(TpkDataType::TypeTreeInformation as i8 as u8); out.push(0u8); out.extend_from_slice(&0u32.to_le_bytes()); out.extend_from_slice(&(blob.len() as u32).to_le_bytes()); out.extend_from_slice(&(blob.len() as u32).to_le_bytes()); out.extend_from_slice(&blob);
651 out
652 }
653
654 #[test]
655 fn tpk_registry_resolves_typetree_and_parses_name() {
656 let tpk = build_minimal_tpk();
657 let registry = TpkTypeTreeRegistry::from_bytes(&tpk).unwrap();
658 let tree = registry.resolve("2020.3.0f1", 28).unwrap();
659
660 let mut bytes: Vec<u8> = Vec::new();
661 bytes.extend_from_slice(&(3i32).to_le_bytes());
662 bytes.extend_from_slice(b"foo");
663 bytes.push(0); let mut reader = BinaryReader::new(&bytes, ByteOrder::Little);
666 let serializer = TypeTreeSerializer::new(tree.as_ref());
667 let out = serializer
668 .parse_object_prefix_detailed(&mut reader, TypeTreeParseOptions::default(), 1)
669 .unwrap();
670 assert_eq!(
671 out.properties.get("m_Name").and_then(|v| v.as_str()),
672 Some("foo")
673 );
674 assert_eq!(reader.remaining(), 0);
675 assert_eq!(out.warnings.len(), 0);
676 assert_eq!(out.properties.len(), 1);
677 assert!(matches!(
678 out.properties.get("m_Name"),
679 Some(UnityValue::String(_))
680 ));
681 }
682}