unity_asset_binary/typetree/
registry.rs

1//! External TypeTree registry (UnityPy TPK-like fallback).
2//!
3//! Unity assets can be built with stripped TypeTrees (`enableTypeTree = false`). In those cases,
4//! consumers may still want a best-effort parser by supplying an external registry of TypeTrees.
5//!
6//! This module provides an injectable registry abstraction and a simple JSON-backed implementation.
7
8use crate::typetree::TypeTree;
9use crate::{error::BinaryError, error::Result};
10use serde::Deserialize;
11use std::collections::HashMap;
12use std::io::Read;
13use std::path::Path;
14use std::sync::Arc;
15
16pub trait TypeTreeRegistry: Send + Sync + std::fmt::Debug {
17    fn resolve(&self, unity_version: &str, class_id: i32) -> Option<Arc<TypeTree>>;
18}
19
20/// A registry that resolves by trying multiple registries in order (first match wins).
21#[derive(Debug, Default, Clone)]
22pub struct CompositeTypeTreeRegistry {
23    registries: Vec<Arc<dyn TypeTreeRegistry>>,
24}
25
26impl CompositeTypeTreeRegistry {
27    pub fn new(registries: Vec<Arc<dyn TypeTreeRegistry>>) -> Self {
28        Self { registries }
29    }
30
31    pub fn push(&mut self, registry: Arc<dyn TypeTreeRegistry>) {
32        self.registries.push(registry);
33    }
34
35    pub fn extend(&mut self, registries: impl IntoIterator<Item = Arc<dyn TypeTreeRegistry>>) {
36        self.registries.extend(registries);
37    }
38
39    pub fn is_empty(&self) -> bool {
40        self.registries.is_empty()
41    }
42}
43
44impl TypeTreeRegistry for CompositeTypeTreeRegistry {
45    fn resolve(&self, unity_version: &str, class_id: i32) -> Option<Arc<TypeTree>> {
46        for r in &self.registries {
47            if let Some(t) = r.resolve(unity_version, class_id) {
48                return Some(t);
49            }
50        }
51        None
52    }
53}
54
55#[derive(Debug, Clone)]
56enum VersionSelector {
57    Any,
58    Exact(String),
59    Prefix(String),
60}
61
62#[derive(Debug, Clone)]
63struct RegistryEntry {
64    selector: VersionSelector,
65    tree: Arc<TypeTree>,
66}
67
68/// A simple in-memory registry keyed by Unity class ID.
69#[derive(Debug, Default, Clone)]
70pub struct InMemoryTypeTreeRegistry {
71    by_class_id: HashMap<i32, Vec<RegistryEntry>>,
72}
73
74impl InMemoryTypeTreeRegistry {
75    pub fn insert_any(&mut self, class_id: i32, tree: TypeTree) {
76        self.insert_internal(class_id, VersionSelector::Any, tree);
77    }
78
79    pub fn insert_exact(&mut self, unity_version: String, class_id: i32, tree: TypeTree) {
80        self.insert_internal(class_id, VersionSelector::Exact(unity_version), tree);
81    }
82
83    pub fn insert_prefix(&mut self, unity_version_prefix: String, class_id: i32, tree: TypeTree) {
84        self.insert_internal(
85            class_id,
86            VersionSelector::Prefix(unity_version_prefix),
87            tree,
88        );
89    }
90
91    fn insert_internal(&mut self, class_id: i32, selector: VersionSelector, tree: TypeTree) {
92        self.by_class_id
93            .entry(class_id)
94            .or_default()
95            .push(RegistryEntry {
96                selector,
97                tree: Arc::new(tree),
98            });
99    }
100}
101
102impl TypeTreeRegistry for InMemoryTypeTreeRegistry {
103    fn resolve(&self, unity_version: &str, class_id: i32) -> Option<Arc<TypeTree>> {
104        let entries = self.by_class_id.get(&class_id)?;
105
106        // 1) exact match
107        for e in entries {
108            if matches!(&e.selector, VersionSelector::Exact(v) if v == unity_version) {
109                return Some(e.tree.clone());
110            }
111        }
112
113        // 2) best (longest) prefix match
114        let mut best: Option<(&RegistryEntry, usize)> = None;
115        for e in entries {
116            let VersionSelector::Prefix(prefix) = &e.selector else {
117                continue;
118            };
119            if unity_version.starts_with(prefix) {
120                let len = prefix.len();
121                match best {
122                    Some((_prev, prev_len)) if prev_len >= len => {}
123                    _ => best = Some((e, len)),
124                }
125            }
126        }
127        if let Some((e, _)) = best {
128            return Some(e.tree.clone());
129        }
130
131        // 3) any
132        for e in entries {
133            if matches!(e.selector, VersionSelector::Any) {
134                return Some(e.tree.clone());
135            }
136        }
137
138        None
139    }
140}
141
142#[derive(Debug, Deserialize)]
143struct JsonRegistryFile {
144    schema: u32,
145    entries: Vec<JsonRegistryEntry>,
146}
147
148#[derive(Debug, Deserialize)]
149struct JsonRegistryEntry {
150    #[serde(default)]
151    unity_version: Option<String>,
152    class_id: i32,
153    type_tree: TypeTree,
154}
155
156/// JSON-backed TypeTree registry.
157///
158/// Format:
159/// ```json
160/// { "schema": 1, "entries": [ { "unity_version": "2020.3.*", "class_id": 28, "type_tree": { ... } } ] }
161/// ```
162#[derive(Debug, Default, Clone)]
163pub struct JsonTypeTreeRegistry {
164    inner: InMemoryTypeTreeRegistry,
165}
166
167impl JsonTypeTreeRegistry {
168    pub fn from_reader(mut reader: impl Read) -> Result<Self> {
169        let mut buf = String::new();
170        reader
171            .read_to_string(&mut buf)
172            .map_err(|e| BinaryError::generic(format!("Failed to read registry JSON: {}", e)))?;
173        let parsed: JsonRegistryFile = serde_json::from_str(&buf)
174            .map_err(|e| BinaryError::invalid_data(format!("Invalid registry JSON: {}", e)))?;
175        if parsed.schema != 1 {
176            return Err(BinaryError::invalid_data(format!(
177                "Unsupported registry schema: {}",
178                parsed.schema
179            )));
180        }
181
182        let mut inner = InMemoryTypeTreeRegistry::default();
183        for e in parsed.entries {
184            match e.unity_version {
185                None => inner.insert_any(e.class_id, e.type_tree),
186                Some(v) => {
187                    if v.is_empty() {
188                        inner.insert_any(e.class_id, e.type_tree);
189                    } else if let Some(prefix) = v.strip_suffix('*') {
190                        inner.insert_prefix(prefix.to_string(), e.class_id, e.type_tree);
191                    } else {
192                        inner.insert_exact(v, e.class_id, e.type_tree);
193                    }
194                }
195            }
196        }
197
198        Ok(Self { inner })
199    }
200
201    pub fn from_path(path: impl AsRef<Path>) -> Result<Self> {
202        let mut f = std::fs::File::open(path.as_ref()).map_err(|e| {
203            BinaryError::generic(format!(
204                "Failed to open registry JSON {:?}: {}",
205                path.as_ref(),
206                e
207            ))
208        })?;
209        Self::from_reader(&mut f)
210    }
211}
212
213impl TypeTreeRegistry for JsonTypeTreeRegistry {
214    fn resolve(&self, unity_version: &str, class_id: i32) -> Option<Arc<TypeTree>> {
215        self.inner.resolve(unity_version, class_id)
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    fn dummy_tree(tag: u32) -> TypeTree {
224        let mut t = TypeTree::new();
225        t.version = tag;
226        t.platform = tag;
227        t
228    }
229
230    #[test]
231    fn in_memory_registry_version_precedence() {
232        let class_id = 28;
233
234        let mut reg = InMemoryTypeTreeRegistry::default();
235        reg.insert_any(class_id, dummy_tree(1));
236        reg.insert_prefix("2020.3.".to_string(), class_id, dummy_tree(2));
237        reg.insert_exact("2020.3.48f1".to_string(), class_id, dummy_tree(3));
238
239        let exact = reg.resolve("2020.3.48f1", class_id).unwrap();
240        assert_eq!(exact.version, 3);
241
242        let prefix = reg.resolve("2020.3.9f1", class_id).unwrap();
243        assert_eq!(prefix.version, 2);
244
245        let any = reg.resolve("2019.4.40f1", class_id).unwrap();
246        assert_eq!(any.version, 1);
247    }
248
249    #[test]
250    fn in_memory_registry_longest_prefix_wins() {
251        let class_id = 28;
252
253        let mut reg = InMemoryTypeTreeRegistry::default();
254        reg.insert_prefix("2020.".to_string(), class_id, dummy_tree(1));
255        reg.insert_prefix("2020.3.".to_string(), class_id, dummy_tree(2));
256
257        let t = reg.resolve("2020.3.48f1", class_id).unwrap();
258        assert_eq!(t.version, 2);
259    }
260
261    #[test]
262    fn composite_registry_first_match_wins() {
263        let class_id = 28;
264
265        let mut a = InMemoryTypeTreeRegistry::default();
266        a.insert_any(class_id, dummy_tree(1));
267        let mut b = InMemoryTypeTreeRegistry::default();
268        b.insert_any(class_id, dummy_tree(2));
269
270        let composite_ab = CompositeTypeTreeRegistry::new(vec![Arc::new(a), Arc::new(b)]);
271        let t = composite_ab.resolve("2020.3.48f1", class_id).unwrap();
272        assert_eq!(t.version, 1);
273
274        let mut a2 = InMemoryTypeTreeRegistry::default();
275        a2.insert_any(class_id, dummy_tree(1));
276        let mut b2 = InMemoryTypeTreeRegistry::default();
277        b2.insert_any(class_id, dummy_tree(2));
278
279        let composite_ba = CompositeTypeTreeRegistry::new(vec![Arc::new(b2), Arc::new(a2)]);
280        let t = composite_ba.resolve("2020.3.48f1", class_id).unwrap();
281        assert_eq!(t.version, 2);
282    }
283
284    #[test]
285    fn json_registry_supports_wildcard_and_exact() {
286        let json = r#"
287        {
288          "schema": 1,
289          "entries": [
290            { "unity_version": "2020.3.*", "class_id": 28, "type_tree": { "nodes": [], "string_buffer": [], "version": 2, "platform": 2, "has_type_dependencies": false } },
291            { "unity_version": "2020.3.48f1", "class_id": 28, "type_tree": { "nodes": [], "string_buffer": [], "version": 3, "platform": 3, "has_type_dependencies": false } },
292            { "class_id": 28, "type_tree": { "nodes": [], "string_buffer": [], "version": 1, "platform": 1, "has_type_dependencies": false } }
293          ]
294        }
295        "#;
296
297        let reg = JsonTypeTreeRegistry::from_reader(json.as_bytes()).unwrap();
298
299        let exact = reg.resolve("2020.3.48f1", 28).unwrap();
300        assert_eq!(exact.version, 3);
301
302        let prefix = reg.resolve("2020.3.9f1", 28).unwrap();
303        assert_eq!(prefix.version, 2);
304
305        let any = reg.resolve("2019.4.40f1", 28).unwrap();
306        assert_eq!(any.version, 1);
307    }
308}
unity_asset_binary/typetree/registry.rs

unity_asset_binary/typetree/
registry.rs
