Skip to main content

asc_engine/
hierarchy.rs

1//! Class hierarchy across every DEX of an APK: transitive subclasses /
2//! implementers, supertypes, and which of them define (override) a method.
3//!
4//! The graph is built from the `class_def_item`s of every DEX (superclass and
5//! interface list per class, interned descriptors, edges sorted both ways,
6//! the class annotations of each class) and kept on the engine for later
7//! queries until the APK changes or [`Engine::release_memory_cache`] drops
8//! it. With a disk cache it is also persisted to L2 ([`GRAPH_KIND`]), so a
9//! later process maps it back instead of reading every class definition. A type the APK only references
10//! (a framework class) is a node without a definition: its own supertypes are
11//! unknown.
12
13use crate::resources::{FileIdent, file_ident};
14use crate::{DexView, Engine, EngineError, Job, NOT_CANCELLED, Result, WriteJob, cancelled};
15use asc_cache::Tier;
16use asc_dex::Dex;
17use asc_findrefs::ClassRelation;
18use std::collections::{HashMap, HashSet, VecDeque};
19use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
20use std::sync::{Arc, Mutex};
21use std::time::Instant;
22
23mod packed;
24
25use packed::{Builder, Packed};
26
27/// No `superclass_idx` (`java.lang.Object` itself).
28const NO_INDEX: u32 = u32::MAX;
29
30/// L2 kind of the persisted class graph, stored under [`GRAPH_NAME`].
31pub(crate) const GRAPH_KIND: &str = "classgraph";
32/// The graph spans the whole APK (or program): one per APK key.
33pub(crate) const GRAPH_NAME: &str = "apk";
34/// Layout version of [`ClassGraph::bytes`]; bump on any change.
35pub(crate) const GRAPH_FORMAT: u32 = 1;
36
37/// Which way to walk the hierarchy.
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum Lineage {
40    /// Classes extending / implementing the target, transitively.
41    Subclasses,
42    /// The target's superclasses and interfaces, transitively.
43    Supers,
44}
45
46impl Lineage {
47    pub fn as_str(self) -> &'static str {
48        match self {
49            Lineage::Subclasses => "subclasses",
50            Lineage::Supers => "supers",
51        }
52    }
53
54    /// The word before a class reached through `rel`.
55    fn word(self, rel: ClassRelation) -> &'static str {
56        match (self, rel) {
57            (Lineage::Subclasses, ClassRelation::Extends) => "subclass",
58            (Lineage::Subclasses, ClassRelation::Implements) => "implementer",
59            (Lineage::Supers, rel) => rel.verb(),
60        }
61    }
62}
63
64/// A hierarchy query.
65#[derive(Debug, Clone)]
66pub struct HierarchyQuery {
67    /// Dalvik descriptor (`Lcom/poc/Base;`).
68    pub class: String,
69    pub direction: Lineage,
70    /// Levels to walk (1 = direct subclasses / supertypes).
71    pub depth: usize,
72    /// Stop discovering classes beyond this many (the root included).
73    pub max_nodes: usize,
74    /// Mark the classes that define this method: a name (every overload) or
75    /// exact `name(params)ret`.
76    pub method: Option<String>,
77}
78
79/// One line of the tree, in depth-first order.
80#[derive(Debug, Clone, PartialEq, Eq)]
81pub struct HierarchyRow {
82    /// 0 for the root.
83    pub depth: usize,
84    /// Index of the parent row (`None` for the root).
85    pub parent: Option<usize>,
86    pub class: String,
87    /// How the row relates to its parent row (`extends` / `implements`, from
88    /// the subclass's side).
89    pub relation: Option<ClassRelation>,
90    /// The DEX defining the class; `None` when the APK only references it.
91    pub dex: Option<String>,
92    /// Shown again: expanded at another row.
93    pub repeated: bool,
94    /// Signatures (`name(params)ret`) of the queried method this class defines.
95    pub defines: Vec<String>,
96}
97
98/// Result of [`Engine::hierarchy`].
99#[derive(Debug, Clone)]
100pub struct HierarchyTree {
101    pub direction: Lineage,
102    pub depth: usize,
103    pub rows: Vec<HierarchyRow>,
104    /// Distinct classes discovered (the root included).
105    pub classes: usize,
106    /// More classes exist than `max_nodes` allowed.
107    pub truncated: bool,
108    /// The queried method, as given.
109    pub method: Option<String>,
110    /// DEX entries that could not be loaded or read.
111    pub errors: Vec<String>,
112}
113
114impl HierarchyTree {
115    /// One line per row: indentation and depth, how it relates to its parent
116    /// (`subclass` / `implementer` going down, `extends` / `implements` going
117    /// up), the class, its DEX or `[not in APK]`, `[defines m(..)V]` and
118    /// `(repeated)` flags.
119    pub fn lines(&self) -> Vec<String> {
120        self.rows
121            .iter()
122            .map(|r| {
123                let mut s = "  ".repeat(r.depth);
124                match r.relation {
125                    Some(rel) if r.depth > 0 => s.push_str(&format!("{} {} {}", r.depth, self.direction.word(rel), r.class)),
126                    _ => s.push_str(&format!("0 {}", r.class)),
127                }
128                match &r.dex {
129                    Some(d) => s.push_str(&format!(" {d}")),
130                    None => s.push_str(" [not in APK]"),
131                }
132                for d in &r.defines {
133                    s.push_str(&format!(" [defines {d}]"));
134                }
135                if r.repeated {
136                    s.push_str(" (repeated)");
137                }
138                s
139            })
140            .collect()
141    }
142
143    /// `Lcls;->m(..)V dex` of every definition of the queried method below
144    /// (or above) the root, one per line, prefixed `override` / `inherited`.
145    pub fn definition_lines(&self) -> Vec<String> {
146        let word = match self.direction {
147            Lineage::Subclasses => "override",
148            Lineage::Supers => "inherited",
149        };
150        let mut seen = HashSet::new();
151        let mut out = Vec::new();
152        for r in self.rows.iter().filter(|r| r.depth > 0 && seen.insert(&r.class)) {
153            for d in &r.defines {
154                out.push(format!("{word} {}->{d} {}", r.class, r.dex.as_deref().unwrap_or("-")));
155            }
156        }
157        out
158    }
159
160    pub fn footer(&self) -> String {
161        let mut s =
162            format!("# {} rows, {} distinct classes ({}, depth {})", self.rows.len(), self.classes, self.direction.as_str(), self.depth);
163        if let Some(m) = &self.method {
164            let n = self.definition_lines().len();
165            s.push_str(&format!("; {n} define {m}"));
166        }
167        if self.truncated {
168            s.push_str("; truncated (raise max_nodes or lower depth)");
169        }
170        s
171    }
172}
173
174/// Every class the APK defines or references through a class definition,
175/// with the edges between them and the class annotations of each.
176pub(crate) struct ClassGraph(Packed);
177
178/// One class definition of a DEX, by descriptor.
179struct DefinedClass {
180    class: String,
181    superclass: Option<String>,
182    interfaces: Vec<String>,
183    /// Its index in the DEX's `class_defs`.
184    def_idx: u32,
185    /// Descriptors of its class annotations (any visibility).
186    annotations: Vec<String>,
187}
188
189fn dex_classes(v: &DexView) -> Result<Vec<DefinedClass>> {
190    let dex = Dex::with_name(&v.data, &v.name)?;
191    let mut out = Vec::with_capacity(dex.class_def_count() as usize);
192    for i in 0..dex.class_def_count() {
193        let def = dex.class_def(i)?;
194        let sup = (def.superclass_idx != NO_INDEX).then(|| dex.type_descriptor(def.superclass_idx)).transpose()?;
195        let ifaces = if def.interfaces_off == 0 {
196            Vec::new()
197        } else {
198            dex.type_list(def.interfaces_off)?.into_iter().map(|t| dex.type_descriptor(t as u32)).collect::<std::result::Result<_, _>>()?
199        };
200        let annotations = class_annotations(&dex, def.annotations_off)?;
201        out.push(DefinedClass { class: dex.type_descriptor(def.class_idx)?, superclass: sup, interfaces: ifaces, def_idx: i, annotations });
202    }
203    Ok(out)
204}
205
206/// Descriptors of the class annotations of the `annotations_directory_item`
207/// at `off` (0: none). Only the directory's first word is read: the member
208/// annotations it lists are left alone.
209fn class_annotations(dex: &Dex<'_>, off: u32) -> Result<Vec<String>> {
210    if off == 0 {
211        return Ok(Vec::new());
212    }
213    let buf = dex.buf();
214    let set_off = buf
215        .get(off as usize..off as usize + 4)
216        .map(|b| u32::from_le_bytes(b.try_into().expect("4 bytes")))
217        .ok_or(asc_dex::DexError::BadRange { what: "annotations_directory_item", offset: off as usize })?;
218    if set_off == 0 {
219        return Ok(Vec::new());
220    }
221    let mut out = Vec::new();
222    for item_off in asc_dex::annotations::parse_annotation_set(buf, set_off as usize)? {
223        let (item, _) = asc_dex::encoded::parse_annotation_item(buf, item_off as usize)?;
224        out.push(dex.type_descriptor(item.annotation.type_idx)?);
225    }
226    Ok(out)
227}
228
229impl ClassGraph {
230    /// Heap held by the graph, roughly: none once mapped from L2.
231    pub(crate) fn heap_bytes(&self) -> usize {
232        self.0.heap_bytes()
233    }
234
235    /// `None` once `cancel` is raised.
236    fn build(views: &[DexView], threads: usize, errors: &mut Vec<String>, cancel: &AtomicBool) -> Option<ClassGraph> {
237        let next = AtomicUsize::new(0);
238        let parsed: Mutex<Vec<(usize, Vec<DefinedClass>)>> = Mutex::new(Vec::new());
239        let errs = Mutex::new(Vec::new());
240        std::thread::scope(|s| {
241            for _ in 0..threads.min(views.len()).max(1) {
242                s.spawn(|| {
243                    loop {
244                        if cancel.load(Ordering::Relaxed) {
245                            break;
246                        }
247                        let i = next.fetch_add(1, Ordering::Relaxed);
248                        let Some(v) = views.get(i) else { break };
249                        match dex_classes(v) {
250                            Ok(c) => parsed.lock().unwrap().push((i, c)),
251                            Err(e) => errs.lock().unwrap().push(format!("{}: {e}", v.name)),
252                        }
253                    }
254                });
255            }
256        });
257        if cancel.load(Ordering::Relaxed) {
258            return None;
259        }
260        errors.extend(errs.into_inner().unwrap());
261        let mut parsed = parsed.into_inner().unwrap();
262        // Views in entry order, so "first definition" means the first DEX.
263        parsed.sort_by_key(|(i, _)| *i);
264        let mut b = Builder::default();
265        for (view, classes) in parsed {
266            for DefinedClass { class, superclass, interfaces, def_idx, annotations } in classes {
267                let c = b.intern(&class);
268                if !b.define(c, view as u32, def_idx) {
269                    continue; // a duplicate definition: the first one wins
270                }
271                for a in annotations {
272                    b.annotate(c, a);
273                }
274                let parents =
275                    superclass.iter().map(|s| (s, ClassRelation::Extends)).chain(interfaces.iter().map(|i| (i, ClassRelation::Implements)));
276                for (p, rel) in parents {
277                    let p = b.intern(p);
278                    b.edge(c, p, rel);
279                }
280            }
281        }
282        Some(ClassGraph(b.finish(views.iter().map(|v| v.name.as_str()))))
283    }
284
285    /// The graph as persisted in L2 (read back by [`ClassGraph::from_mapped`]).
286    pub(crate) fn bytes(&self) -> &[u8] {
287        self.0.bytes()
288    }
289
290    /// A graph [`ClassGraph::bytes`] stored, used in place; `None` for bytes
291    /// that are not one.
292    fn from_mapped(d: asc_cache::Derived) -> Option<ClassGraph> {
293        Packed::from_mapped(d).map(ClassGraph)
294    }
295
296    fn id(&self, name: &str) -> Option<u32> {
297        self.0.id(name)
298    }
299
300    fn name(&self, id: u32) -> &str {
301        self.0.name(id)
302    }
303
304    /// Neighbours of `id` one level up or down.
305    fn neighbours(&self, id: u32, dir: Lineage) -> impl Iterator<Item = (u32, ClassRelation)> + '_ {
306        self.0.edges_of(id, dir == Lineage::Subclasses).map(move |(c, p, rel)| match dir {
307            Lineage::Supers => (p, rel),
308            Lineage::Subclasses => (c, rel),
309        })
310    }
311
312    /// Every class reachable from `class` going `dir` (not `class` itself),
313    /// nearest first; at most `limit`.
314    pub(crate) fn reachable(&self, class: &str, dir: Lineage, limit: usize) -> Vec<String> {
315        let Some(root) = self.id(class) else { return Vec::new() };
316        let mut seen = HashSet::from([root]);
317        let mut queue = VecDeque::from([root]);
318        let mut out = Vec::new();
319        while let Some(c) = queue.pop_front() {
320            for (n, _) in self.neighbours(c, dir) {
321                if out.len() >= limit {
322                    return out;
323                }
324                if seen.insert(n) {
325                    out.push(self.name(n).to_string());
326                    queue.push_back(n);
327                }
328            }
329        }
330        out
331    }
332
333    fn dex_of(&self, id: u32) -> Option<&str> {
334        let (view, _) = self.0.defined(id);
335        (view != 0).then(|| self.0.dex_name(view as usize - 1))
336    }
337
338    /// Descriptors of the classes this APK defines whose descriptor matches
339    /// `re` (for `in:{type:...}`).
340    pub(crate) fn defined_classes_matching(&self, re: &regex::bytes::Regex) -> Vec<String> {
341        (0..self.0.classes() as u32)
342            .filter(|&id| self.0.defined(id).0 != 0)
343            .map(|id| self.name(id))
344            .filter(|name| re.is_match(name.as_bytes()))
345            .map(String::from)
346            .collect()
347    }
348
349    /// Descriptors of the classes carrying the class annotation `annotation`
350    /// (for `in:{annotation:...}`), sorted.
351    pub(crate) fn classes_annotated(&self, annotation: &str) -> Vec<String> {
352        let mut out: Vec<String> = self.0.annotated(annotation).map(|id| self.name(id).to_string()).collect();
353        out.sort_unstable();
354        out
355    }
356
357    /// `name(params)ret` of the methods `class` defines whose name is `name`
358    /// (and proto `proto` when given).
359    pub(crate) fn defines(&self, views: &[DexView], class: &str, name: &str, proto: Option<&str>) -> Result<Vec<String>> {
360        let Some(id) = self.id(class) else { return Ok(Vec::new()) };
361        let (view, def_idx) = self.0.defined(id);
362        if view == 0 {
363            return Ok(Vec::new());
364        }
365        let v = &views[view as usize - 1];
366        let dex = Dex::with_name(&v.data, &v.name)?;
367        let def = dex.class_def(def_idx)?;
368        let data = dex.class_data(def.class_data_off)?;
369        let mut out = Vec::new();
370        for m in data.direct_methods.iter().chain(&data.virtual_methods) {
371            let mid = dex.method_id(m.method_idx)?;
372            if dex.string_bytes(mid.name_idx)? != name.as_bytes() {
373                continue;
374            }
375            let p = asc_smali::proto_descriptor(&dex, mid.proto_idx as u32)?;
376            if proto.is_none_or(|want| want == p) {
377                out.push(format!("{name}{p}"));
378            }
379        }
380        Ok(out)
381    }
382
383    /// Where the program first defines `class`: `(view index, class_def
384    /// index)`; `None` when it does not.
385    pub(crate) fn definition(&self, class: &str) -> Option<(usize, u32)> {
386        let (view, def_idx) = self.0.defined(self.id(class)?);
387        (view != 0).then(|| (view as usize - 1, def_idx))
388    }
389
390    /// The classes of `class`'s subtree (itself included) `defines` says
391    /// define the method: those a call naming it on `class` may run. The
392    /// walk stops after `max_classes` classes.
393    pub(crate) fn defining(&self, class: &str, max_classes: usize, mut defines: impl FnMut(&str) -> Result<bool>) -> Result<Defining> {
394        let Some(root) = self.id(class) else { return Ok(Defining { count: 0, complete: true }) };
395        let mut seen = HashSet::from([root]);
396        let mut queue = VecDeque::from([root]);
397        let mut count = usize::from(defines(class)?);
398        while let Some(c) = queue.pop_front() {
399            for (n, _) in self.neighbours(c, Lineage::Subclasses) {
400                if seen.len() >= max_classes {
401                    return Ok(Defining { count, complete: false });
402                }
403                if seen.insert(n) {
404                    count += usize::from(defines(self.name(n))?);
405                    queue.push_back(n);
406                }
407            }
408        }
409        Ok(Defining { count, complete: true })
410    }
411
412    /// `class` and the subclasses that inherit the method from it: a subtree
413    /// stops at a class `defines` says defines it itself. Nearest first, at
414    /// most `limit`.
415    pub(crate) fn inheritors(&self, class: &str, limit: usize, mut defines: impl FnMut(&str) -> Result<bool>) -> Result<Vec<String>> {
416        let Some(root) = self.id(class) else { return Ok(Vec::new()) };
417        let mut seen = HashSet::from([root]);
418        let mut queue = VecDeque::from([root]);
419        let mut out = vec![class.to_string()];
420        while let Some(c) = queue.pop_front() {
421            for (n, _) in self.neighbours(c, Lineage::Subclasses) {
422                if out.len() >= limit {
423                    return Ok(out);
424                }
425                if seen.insert(n) && !defines(self.name(n))? {
426                    out.push(self.name(n).to_string());
427                    queue.push_back(n);
428                }
429            }
430        }
431        Ok(out)
432    }
433}
434
435/// The classes defining a method in a subtree, as [`ClassGraph::defining`]
436/// counted them.
437#[derive(Debug, Clone, Copy)]
438pub(crate) struct Defining {
439    pub count: usize,
440    /// The whole subtree was walked: `count` is exact, not a lower bound.
441    pub complete: bool,
442}
443
444/// The cached graph and the APK file it was built from.
445pub(crate) type CachedGraph = (Option<FileIdent>, Arc<ClassGraph>);
446
447/// `name` or `name(params)ret` split into name and normalised proto.
448pub(crate) fn split_method(method: &str) -> Result<(String, Option<String>)> {
449    match method.find('(') {
450        Some(_) => {
451            let sig = asc_smali::normalize_method_signature(method)?;
452            let p = sig.find('(').expect("normalised signature has a paren");
453            Ok((sig[..p].to_string(), Some(sig[p..].to_string())))
454        }
455        None => Ok((method.to_string(), None)),
456    }
457}
458
459impl Engine {
460    /// The class graph of the APK, built over `views` unless one built from
461    /// the same APK file is kept, in memory or in L2. A cancelled build is
462    /// not kept; one that met unreadable DEX files is not persisted (a later
463    /// process should report them too).
464    pub(crate) fn class_graph(&self, views: &[DexView], errors: &mut Vec<String>, cancel: &AtomicBool) -> Result<Arc<ClassGraph>> {
465        let apk = self.open()?;
466        let ident = file_ident(&apk);
467        let mut slot = self.class_graph.lock().unwrap();
468        if let Some((id, g)) = slot.as_ref()
469            && ident.is_some()
470            && *id == ident
471        {
472            return Ok(g.clone());
473        }
474        let disk = self.cache.as_ref().and_then(|c| c.disk.as_deref()).filter(|d| d.uses(Tier::Derived));
475        let key = disk.and_then(|_| self.apk_key(&apk));
476        if let (Some(disk), Some(key)) = (disk, &key) {
477            let stored = disk.get_derived(key, GRAPH_NAME, GRAPH_KIND, GRAPH_FORMAT).and_then(ClassGraph::from_mapped);
478            // The key covers the archives; the DEX names pin the view order
479            // the definitions index.
480            if let Some(g) = stored.filter(|g| g.0.dex_names().eq(views.iter().map(|v| v.name.as_str()))) {
481                self.stats.derived_hits.fetch_add(1, Ordering::Relaxed);
482                let g = Arc::new(g);
483                *slot = Some((ident, g.clone()));
484                return Ok(g);
485            }
486        }
487        let t0 = Instant::now();
488        let errors_before = errors.len();
489        let g = Arc::new(ClassGraph::build(views, self.threads(), errors, cancel).ok_or_else(cancelled)?);
490        let cost = t0.elapsed();
491        if let (Some(disk), Some(key)) = (disk, key)
492            && errors.len() == errors_before
493            && disk.admits_derived(cost)
494        {
495            let job = WriteJob { key, apk: Arc::from(apk.path()), kind: Job::Graph { graph: g.clone(), cost } };
496            self.cache.as_ref().expect("a disk implies a cache").put_async(job);
497        }
498        *slot = Some((ident, g.clone()));
499        Ok(g)
500    }
501
502    /// Transitive subclasses / implementers (or supertypes) of a class, as a
503    /// tree, with the classes defining `method` marked.
504    pub fn hierarchy(&self, q: &HierarchyQuery) -> Result<HierarchyTree> {
505        self.hierarchy_cancellable(q, &NOT_CANCELLED)
506    }
507
508    /// [`Engine::hierarchy`] that stops, with `ApkError::Cancelled`, once
509    /// `cancel` is raised.
510    pub fn hierarchy_cancellable(&self, q: &HierarchyQuery, cancel: &AtomicBool) -> Result<HierarchyTree> {
511        let method = q.method.as_deref().map(split_method).transpose()?;
512        let (views, mut errors) = self.load_views(None, cancel)?;
513        let g = self.class_graph(&views, &mut errors, cancel)?;
514        let Some(root) = g.id(&q.class) else {
515            return Err(EngineError::Query(format!("no class {} in this APK", q.class)));
516        };
517        // Breadth-first: each class's shallowest depth and its children.
518        let max_nodes = q.max_nodes.max(1);
519        let mut depth_of: HashMap<u32, usize> = HashMap::from([(root, 0)]);
520        let mut children: HashMap<u32, Vec<(u32, ClassRelation)>> = HashMap::new();
521        let mut frontier = vec![root];
522        let mut truncated = false;
523        for level in 0..q.depth {
524            let mut following = Vec::new();
525            for &c in &frontier {
526                let mut kept = Vec::new();
527                for (n, rel) in g.neighbours(c, q.direction) {
528                    if depth_of.contains_key(&n) {
529                        kept.push((n, rel));
530                    } else if depth_of.len() < max_nodes {
531                        depth_of.insert(n, level + 1);
532                        following.push(n);
533                        kept.push((n, rel));
534                    } else {
535                        truncated = true;
536                    }
537                }
538                children.insert(c, kept);
539            }
540            if following.is_empty() {
541                break;
542            }
543            frontier = following;
544        }
545        // Classes at the last level may still have neighbours.
546        truncated |= frontier.iter().any(|&c| !children.contains_key(&c) && g.neighbours(c, q.direction).next().is_some());
547
548        let mut defines: HashMap<u32, Vec<String>> = HashMap::new();
549        if let Some((name, proto)) = &method {
550            for &c in depth_of.keys() {
551                if cancel.load(Ordering::Relaxed) {
552                    return Err(cancelled());
553                }
554                match g.defines(&views, g.name(c), name, proto.as_deref()) {
555                    Ok(d) if !d.is_empty() => {
556                        defines.insert(c, d);
557                    }
558                    Ok(_) => {}
559                    Err(e) => errors.push(format!("{}: {e}", g.name(c))),
560                }
561            }
562        }
563
564        // Depth-first rendering, expanding each class once at its BFS depth.
565        let mut rows = Vec::new();
566        let mut expanded = HashSet::new();
567        let mut stack: Vec<(u32, usize, Option<usize>, Option<ClassRelation>)> = vec![(root, 0, None, None)];
568        while let Some((c, depth, parent, relation)) = stack.pop() {
569            let expand_here = depth_of.get(&c) == Some(&depth) && !expanded.contains(&c);
570            let has_children = children.get(&c).is_some_and(|n| !n.is_empty());
571            rows.push(HierarchyRow {
572                depth,
573                parent,
574                class: g.name(c).to_string(),
575                relation,
576                dex: g.dex_of(c).map(String::from),
577                repeated: !expand_here && has_children,
578                defines: defines.get(&c).cloned().unwrap_or_default(),
579            });
580            if expand_here {
581                expanded.insert(c);
582                let me = rows.len() - 1;
583                let mut kids: Vec<_> = children.get(&c).cloned().unwrap_or_default();
584                kids.sort_by(|a, b| g.name(b.0).cmp(g.name(a.0)));
585                for (n, rel) in kids {
586                    stack.push((n, depth + 1, Some(me), Some(rel)));
587                }
588            }
589        }
590        errors.sort();
591        errors.dedup();
592        Ok(HierarchyTree {
593            direction: q.direction,
594            depth: q.depth,
595            rows,
596            classes: depth_of.len(),
597            truncated,
598            method: q.method.clone(),
599            errors,
600        })
601    }
602}
603
604#[cfg(test)]
605mod tests {
606    use super::*;
607    use std::collections::BTreeSet;
608
609    /// `B extends A implements I`, `C extends B` (in the second DEX), `D`
610    /// with a non-ASCII name; `B` and `C` carry `@M`; `I` is only referenced.
611    fn sample() -> ClassGraph {
612        let mut b = Builder::default();
613        let [a, bb, c, i, d] = ["LA;", "LB;", "LC;", "LI;", "Lé;"].map(|n| b.intern(n));
614        assert!(b.define(a, 0, 0) && b.define(bb, 0, 1) && b.define(c, 1, 0) && b.define(d, 0, 2));
615        assert!(!b.define(c, 0, 9), "the first definition wins");
616        b.edge(bb, a, ClassRelation::Extends);
617        b.edge(bb, i, ClassRelation::Implements);
618        b.edge(c, bb, ClassRelation::Extends);
619        b.edge(c, bb, ClassRelation::Extends);
620        b.annotate(c, "LM;".into());
621        b.annotate(bb, "LM;".into());
622        b.annotate(bb, "LN;".into());
623        ClassGraph(b.finish(["classes.dex", "lib.jar/classes.dex"].into_iter()))
624    }
625
626    fn check(g: &ClassGraph) {
627        assert_eq!(g.reachable("LA;", Lineage::Subclasses, 10), ["LB;", "LC;"]);
628        assert_eq!(g.reachable("LC;", Lineage::Supers, 10), ["LB;", "LA;", "LI;"]);
629        assert_eq!(g.reachable("LC;", Lineage::Supers, 2), ["LB;", "LA;"]);
630        assert!(g.reachable("LZ;", Lineage::Supers, 10).is_empty());
631        assert_eq!(g.classes_annotated("LM;"), ["LB;", "LC;"]);
632        assert_eq!(g.classes_annotated("LN;"), ["LB;"]);
633        assert!(g.classes_annotated("LL;").is_empty() && g.classes_annotated("LZ;").is_empty());
634        assert_eq!(g.dex_of(g.id("LC;").unwrap()), Some("lib.jar/classes.dex"));
635        assert_eq!(g.dex_of(g.id("LI;").unwrap()), None);
636        assert_eq!(g.name(g.id("Lé;").unwrap()), "Lé;");
637        let re = regex::bytes::Regex::new("^L[A-C];$").unwrap();
638        assert_eq!(
639            g.defined_classes_matching(&re).into_iter().collect::<BTreeSet<_>>(),
640            BTreeSet::from(["LA;".into(), "LB;".into(), "LC;".into()])
641        );
642        assert_eq!(g.0.dex_names().collect::<Vec<_>>(), ["classes.dex", "lib.jar/classes.dex"]);
643    }
644
645    /// An L2 directory to hand bytes back as the disk cache would.
646    struct Store {
647        _dir: tempfile::TempDir,
648        disk: asc_cache::DiskCache,
649        key: asc_cache::ApkKey,
650    }
651
652    impl Store {
653        fn new() -> Store {
654            let dir = tempfile::tempdir().unwrap();
655            let disk = asc_cache::DiskCache::open(dir.path().join("cache")).with_capacity(1 << 30);
656            let apk = dir.path().join("a.apk");
657            std::fs::write(&apk, asc_fixture::make_apk(&[("classes.dex", &asc_fixture::make_rich_dex())], true)).unwrap();
658            let key = asc_cache::ApkKey::for_apk(&asc_apk::Apk::open(&apk).unwrap()).unwrap();
659            Store { _dir: dir, disk, key }
660        }
661
662        fn mapped(&self, bytes: &[u8]) -> Option<ClassGraph> {
663            self.disk.put_derived(&self.key, GRAPH_NAME, GRAPH_KIND, GRAPH_FORMAT, std::time::Duration::ZERO, bytes).unwrap();
664            ClassGraph::from_mapped(self.disk.get_derived(&self.key, GRAPH_NAME, GRAPH_KIND, GRAPH_FORMAT).unwrap())
665        }
666    }
667
668    #[test]
669    fn graphs_answer_alike_built_and_mapped() {
670        let store = Store::new();
671        let g = sample();
672        check(&g);
673        let back = store.mapped(g.bytes()).unwrap();
674        check(&back);
675        assert_eq!(back.heap_bytes(), 0);
676        assert!(g.heap_bytes() >= g.bytes().len());
677        // No classes at all.
678        let empty = ClassGraph(Builder::default().finish(std::iter::empty()));
679        assert!(store.mapped(empty.bytes()).unwrap().reachable("LA;", Lineage::Supers, 1).is_empty());
680        // Enough classes to grow the table, all still found.
681        let mut b = Builder::default();
682        let names: Vec<String> = (0..5000).map(|i| format!("Lc{i};")).collect();
683        let ids: Vec<u32> = names.iter().map(|n| b.intern(n)).collect();
684        assert_eq!(names.iter().map(|n| b.intern(n)).collect::<Vec<_>>(), ids);
685        let g = store.mapped(ClassGraph(b.finish(std::iter::empty())).bytes()).unwrap();
686        assert!(names.iter().zip(&ids).all(|(n, &id)| g.id(n) == Some(id) && g.name(id) == n));
687    }
688
689    #[test]
690    fn damaged_bytes_never_index_out_of_bounds() {
691        let store = Store::new();
692        let bytes = sample().bytes().to_vec();
693        for n in 0..bytes.len() {
694            assert!(store.mapped(&bytes[..n]).is_none(), "truncated to {n}");
695        }
696        // Any single-byte change is rejected or still answers without
697        // panicking.
698        let any = regex::bytes::Regex::new(".").unwrap();
699        for at in 0..bytes.len() {
700            let mut b = bytes.clone();
701            b[at] ^= 0xa5;
702            if let Some(g) = store.mapped(&b) {
703                for name in ["LA;", "LB;", "LC;", "LI;", "Lé;"] {
704                    g.reachable(name, Lineage::Subclasses, 10);
705                    g.reachable(name, Lineage::Supers, 10);
706                    g.id(name).map(|id| g.dex_of(id));
707                }
708                g.classes_annotated("LM;");
709                g.defined_classes_matching(&any);
710            }
711        }
712    }
713}