use crate::hierarchy::{ClassGraph, Lineage};
use crate::{DexView, Engine, EngineError, NOT_CANCELLED, Result, cancelled, names};
use asc_dex::Dex;
use asc_findrefs::{FindRefs, Query};
use std::collections::{BTreeSet, HashMap, HashSet};
use std::sync::Mutex;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Direction {
Callers,
Callees,
}
impl Direction {
pub fn as_str(self) -> &'static str {
match self {
Direction::Callers => "callers",
Direction::Callees => "callees",
}
}
fn arrow(self) -> &'static str {
match self {
Direction::Callers => "<-",
Direction::Callees => "->",
}
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum CallRoots {
Method {
class: String,
method: String,
},
Matching(Query),
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct CallQuery {
pub roots: CallRoots,
pub direction: Direction,
pub depth: usize,
pub max_nodes: usize,
pub virtual_calls: bool,
pub from: Option<Query>,
pub unless: Vec<Query>,
}
impl CallQuery {
#[must_use]
pub fn depth(self, depth: usize) -> Self {
CallQuery { depth, ..self }
}
#[must_use]
pub fn max_nodes(self, max_nodes: usize) -> Self {
CallQuery { max_nodes, ..self }
}
#[must_use]
pub fn virtual_calls(self, virtual_calls: bool) -> Self {
CallQuery { virtual_calls, ..self }
}
#[must_use]
pub fn from(self, from: Option<Query>) -> Self {
CallQuery { from, ..self }
}
#[must_use]
pub fn unless(self, unless: Vec<Query>) -> Self {
CallQuery { unless, ..self }
}
}
impl CallQuery {
pub fn method(class: impl Into<String>, method: impl Into<String>, direction: Direction) -> CallQuery {
CallQuery::new(CallRoots::Method { class: class.into(), method: method.into() }, direction)
}
pub fn new(roots: CallRoots, direction: Direction) -> CallQuery {
CallQuery { roots, direction, depth: 2, max_nodes: 500, virtual_calls: false, from: None, unless: Vec::new() }
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct CallRow {
pub depth: usize,
pub parent: Option<usize>,
pub method: String,
pub dex: Vec<String>,
pub repeated: bool,
pub via: Option<String>,
pub from: bool,
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct CallTree {
pub direction: Direction,
pub depth: usize,
pub rows: Vec<CallRow>,
pub methods: usize,
pub truncated: bool,
pub pruned: usize,
pub reached: Option<usize>,
pub errors: Vec<String>,
pub notes: Vec<String>,
}
impl CallTree {
pub fn lines(&self) -> Vec<String> {
self.rows
.iter()
.map(|r| {
let mut s = " ".repeat(r.depth);
if r.depth == 0 {
s.push_str(&format!("0 {}", r.method));
} else {
s.push_str(&format!("{} {} {}", r.depth, self.direction.arrow(), r.method));
}
if r.dex.is_empty() {
s.push_str(" [no body]");
}
let mut archives: Vec<&str> = r.dex.iter().filter_map(|d| crate::archive_of(d)).collect();
archives.sort_unstable();
archives.dedup();
if !archives.is_empty() {
s.push_str(&format!(" [{}]", archives.join(", ")));
}
if let Some(v) = &r.via {
match self.direction {
Direction::Callers => s.push_str(&format!(" (via {v})")),
Direction::Callees => s.push_str(&format!(" (override of {v})")),
}
}
if r.from {
s.push_str(" [from]");
}
if r.repeated {
s.push_str(" (repeated)");
}
s
})
.collect()
}
pub fn footer(&self) -> String {
let mut s =
format!("# {} rows, {} distinct methods ({}, depth {})", self.rows.len(), self.methods, self.direction.as_str(), self.depth);
if let Some(n) = self.reached {
s.push_str(&format!("; {n} from methods reached"));
}
if self.pruned > 0 {
s.push_str(&format!("; {} methods left out by unless", self.pruned));
}
if self.truncated {
s.push_str("; truncated (raise max_nodes or lower depth)");
}
for n in &self.notes {
s.push_str(&format!("\n# {n}"));
}
s
}
}
struct Signature<'s> {
class: &'s str,
name: &'s str,
proto: &'s str,
}
fn split_signature(sig: &str) -> Option<Signature<'_>> {
let (class, rest) = sig.split_once("->")?;
let p = rest.find('(')?;
Some(Signature { class, name: &rest[..p], proto: &rest[p..] })
}
fn method_ids(dex: &Dex<'_>, class: &str, name: &str, proto: Option<&str>) -> Result<Vec<u32>> {
let Some(class_idx) = dex.find_type_idx(&asc_dex::mutf8::encode(class))? else { return Ok(Vec::new()) };
let Some(name_idx) = dex.find_string_idx(&asc_dex::mutf8::encode(name))? else { return Ok(Vec::new()) };
let target = (class_idx, name_idx);
let (mut lo, mut hi) = (0u32, dex.method_count());
while lo < hi {
let mid = lo + (hi - lo) / 2;
let id = dex.method_id(mid)?;
if ((id.class_idx as u32), id.name_idx) < target {
lo = mid + 1;
} else {
hi = mid;
}
}
let mut out = Vec::new();
for i in lo..dex.method_count() {
let id = dex.method_id(i)?;
if ((id.class_idx as u32), id.name_idx) != target {
break;
}
if proto.is_none_or(|p| asc_smali::proto_descriptor(dex, id.proto_idx as u32).is_ok_and(|d| names::proto_matches(&d, p))) {
out.push(i);
}
}
Ok(out)
}
const MAX_OVERRIDE_SCAN: usize = 4096;
const MAX_SUPERS: usize = 256;
#[derive(Default)]
struct Found {
next: BTreeSet<String>,
bodies: BTreeSet<String>,
}
impl Engine {
pub fn call_graph(&self, q: &CallQuery) -> Result<CallTree> {
self.call_graph_cancellable(q, &NOT_CANCELLED)
}
pub fn call_graph_cancellable(&self, q: &CallQuery, cancel: &AtomicBool) -> Result<CallTree> {
if q.from.is_some() && q.direction == Direction::Callees {
return Err(EngineError::Query("`from` keeps paths of a callers walk: walk the callers of the target instead".into()));
}
let t0 = std::time::Instant::now();
let matching_roots = match &q.roots {
CallRoots::Matching(c) => Some(c),
CallRoots::Method { .. } => None,
};
let conditions: Vec<Query> = matching_roots.into_iter().chain(&q.from).chain(&q.unless).cloned().collect();
if let Some(engine) = self.transient(!conditions.is_empty())? {
return engine.call_graph_cancellable(q, cancel);
}
let matching = if conditions.is_empty() { None } else { Some(self.matching_methods(&conditions, cancel)?) };
let mut notes = matching.as_ref().map(|m| m.notes.clone()).unwrap_or_default();
let mut sets = matching.map(|m| m.sets).unwrap_or_default().into_iter();
let matched_roots = matching_roots.and_then(|_| sets.next());
let from = q.from.as_ref().and_then(|_| sets.next());
let unless: HashSet<String> = sets.flatten().collect();
let _query = self.prebuild.query();
let (views, mut errors) = self.load_views(None, cancel)?;
let stop = |views: &[DexView], e: EngineError| {
self.recharge_views(views);
e
};
let graph = if q.virtual_calls { Some(self.class_graph(&views, &mut errors, cancel).map_err(|e| stop(&views, e))?) } else { None };
let load_us = t0.elapsed().as_micros();
let mut roots = match (&q.roots, matched_roots) {
(CallRoots::Method { class, method }, _) => self.method_roots(&views, class, method, &mut notes)?,
(CallRoots::Matching(_), Some(set)) if !set.is_empty() => set,
(CallRoots::Matching(_), _) => return Err(EngineError::Query("no method's code matches the root condition".into())),
};
let mut pruned: HashSet<String> = roots.iter().filter(|r| unless.contains(*r)).cloned().collect();
roots.retain(|r| !pruned.contains(r));
let is_from = |sig: &str| from.as_ref().is_some_and(|f| f.contains(sig));
let max_nodes = q.max_nodes.max(roots.len());
let mut depth_of: HashMap<String, usize> = roots.iter().map(|r| (r.clone(), 0)).collect();
let mut next: HashMap<String, Vec<String>> = HashMap::new();
let mut bodies: HashMap<String, Vec<String>> = HashMap::new();
let mut via: HashMap<(String, String), String> = HashMap::new();
let mut frontier: Vec<String> = roots.iter().cloned().collect();
let mut truncated = false;
for level in 0..=q.depth {
if frontier.is_empty() {
break;
}
let expand = level < q.depth;
let aliases = match (&graph, q.direction) {
(Some(g), Direction::Callers) if expand => super_aliases(g, &frontier),
_ => HashMap::new(),
};
let mut lookup = frontier.clone();
lookup.extend(aliases.values().flatten().map(|(alias, _)| alias.clone()));
lookup.sort();
lookup.dedup();
let found = self.expand(&views, &lookup, q.direction, expand, &mut errors, cancel);
if cancel.load(Ordering::Relaxed) {
return Err(stop(&views, cancelled()));
}
let mut following = Vec::new();
for sig in &frontier {
let f = found.get(sig);
bodies.insert(sig.clone(), f.map(|f| f.bodies.iter().cloned().collect()).unwrap_or_default());
if !expand || is_from(sig) {
continue;
}
let mut neighbours: Vec<(String, Option<String>)> =
f.map(|f| f.next.iter()).into_iter().flatten().map(|n| (n.clone(), None)).collect();
for (alias, class) in aliases.get(sig).into_iter().flatten() {
for n in found.get(alias).map(|f| f.next.iter()).into_iter().flatten() {
neighbours.push((n.clone(), Some(class.clone())));
}
}
if let (Some(g), Direction::Callees) = (&graph, q.direction) {
let direct: Vec<String> = neighbours.iter().map(|(n, _)| n.clone()).collect();
for n in direct {
if cancel.load(Ordering::Relaxed) {
return Err(stop(&views, cancelled()));
}
let (overrides, cut) = overrides_of(g, &views, &n, &mut errors);
truncated |= cut;
neighbours.extend(overrides.into_iter().map(|o| (o, split_signature(&n).map(|p| p.class.to_string()))));
}
}
let mut seen = HashSet::new();
neighbours.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.is_some().cmp(&b.1.is_some())));
neighbours.retain(|(n, _)| seen.insert(n.clone()));
let mut kept = Vec::new();
for (n, through) in &neighbours {
if unless.contains(n) {
pruned.insert(n.clone());
continue;
}
if let Some(c) = through {
via.insert((sig.clone(), n.clone()), c.clone());
}
if depth_of.contains_key(n) {
kept.push(n.clone());
} else if depth_of.len() < max_nodes {
depth_of.insert(n.clone(), level + 1);
following.push(n.clone());
kept.push(n.clone());
} else {
truncated = true;
}
}
next.insert(sig.clone(), kept);
}
frontier = following;
}
self.recharge_views(&views);
let mut methods = depth_of.len();
let reached = from.as_ref().map(|from| {
let keep = on_paths_to(&next, depth_of.keys().filter(|m| from.contains(*m)));
for kids in next.values_mut() {
kids.retain(|k| keep.contains(k));
}
roots.retain(|r| keep.contains(r));
methods = keep.len();
depth_of.keys().filter(|m| from.contains(*m)).count()
});
if self.debug {
let graphs: Vec<_> = views.iter().filter_map(|v| v.state.call_graph()).collect();
self.log(format!(
"[APK] callgraph load={load_us} us total={} us dex={} graphs={} edges={} graph_bytes={}",
t0.elapsed().as_micros(),
views.len(),
graphs.len(),
graphs.iter().map(|g| g.edge_count()).sum::<usize>(),
graphs.iter().map(|g| g.heap_bytes()).sum::<usize>()
));
}
let mut rows: Vec<CallRow> = Vec::new();
let mut expanded = HashSet::new();
let mut stack: Vec<(String, usize, Option<usize>)> = roots.iter().rev().map(|r| (r.clone(), 0, None)).collect();
while let Some((sig, depth, parent)) = stack.pop() {
let expand_here = depth_of.get(&sig) == Some(&depth) && !expanded.contains(&sig);
let has_children = next.get(&sig).is_some_and(|n| !n.is_empty());
rows.push(CallRow {
depth,
parent,
method: sig.clone(),
dex: bodies.get(&sig).cloned().unwrap_or_default(),
repeated: !expand_here && has_children,
via: parent.and_then(|p| via.get(&(rows[p].method.clone(), sig.clone())).cloned()),
from: is_from(&sig),
});
if expand_here {
expanded.insert(sig.clone());
let me = rows.len() - 1;
for n in next.get(&sig).into_iter().flatten().rev() {
stack.push((n.clone(), depth + 1, Some(me)));
}
}
}
errors.sort();
errors.dedup();
Ok(CallTree { direction: q.direction, depth: q.depth, rows, methods, truncated, pruned: pruned.len(), reached, errors, notes })
}
fn method_roots(&self, views: &[DexView], class: &str, method: &str, notes: &mut Vec<String>) -> Result<BTreeSet<String>> {
let split = |spec: &str| -> Result<(String, Option<String>)> {
let Some(p) = spec.find('(') else { return Ok((spec.to_string(), None)) };
match asc_smali::normalize_method_signature(spec) {
Ok(sig) => {
let p2 = sig.find('(').unwrap_or(p);
Ok((sig[..p2].to_string(), Some(sig[p2..].to_string())))
}
Err(e) => names::partial_proto(spec).map(|(n, proto)| (n, Some(proto))).ok_or_else(|| e.into()),
}
};
let collect = |cls: &str, name: &str, proto: Option<&str>| -> Result<BTreeSet<String>> {
let mut roots = BTreeSet::new();
for v in views {
let dex = Dex::with_name(&v.data, &v.name)?;
for m in method_ids(&dex, cls, name, proto)? {
roots.insert(asc_smali::method_signature(&dex, m)?);
}
}
Ok(roots)
};
let (name, proto) = split(method)?;
let roots = collect(class, &name, proto.as_deref())?;
if !roots.is_empty() {
return Ok(roots);
}
let classes = std::iter::once(class.to_string()).chain(names::denormalize_class_candidates(class));
for cls in classes {
let methods = std::iter::once(method.to_string()).chain(names::denormalize_method_candidates(Some(&cls), method));
for m in methods {
if cls == class && m == method {
continue;
}
let Ok((name, proto)) = split(&m) else { continue };
let roots = collect(&cls, &name, proto.as_deref())?;
if !roots.is_empty() {
notes.push(format!("{class}->{method} read as the Java/Kotlin spelling of {cls}->{m}"));
return Ok(roots);
}
}
}
let hint = self.member_not_found_hint(class, Some(&name), true);
Err(EngineError::Query(format!("no method {class}->{method} in this APK{hint}")))
}
fn expand(
&self,
views: &[DexView],
frontier: &[String],
dir: crate::Direction,
expand: bool,
errors: &mut Vec<String>,
cancel: &AtomicBool,
) -> HashMap<String, Found> {
let next_view = AtomicUsize::new(0);
let merged: Mutex<HashMap<String, Found>> = Mutex::new(HashMap::new());
let errs = Mutex::new(Vec::new());
let workers = self.threads().min(views.len()).max(1);
std::thread::scope(|s| {
for _ in 0..workers {
s.spawn(|| {
loop {
if cancel.load(Ordering::Relaxed) {
break;
}
let i = next_view.fetch_add(1, Ordering::Relaxed);
let Some(v) = views.get(i) else { break };
match view_neighbours(v, frontier, dir, expand) {
Ok(local) => {
let mut m = merged.lock().expect("merged callers mutex poisoned");
for (sig, f) in local {
let e = m.entry(sig).or_default();
e.next.extend(f.next);
e.bodies.extend(f.bodies);
}
}
Err(e) => errs.lock().expect("errors mutex poisoned").push(format!("{}: {e}", v.name)),
}
}
});
}
});
errors.extend(errs.into_inner().expect("errors mutex poisoned"));
merged.into_inner().expect("merged callers mutex poisoned")
}
}
fn on_paths_to<'s>(next: &HashMap<String, Vec<String>>, ends: impl Iterator<Item = &'s String>) -> HashSet<String> {
let mut parents: HashMap<&str, Vec<&str>> = HashMap::new();
for (p, kids) in next {
for k in kids {
parents.entry(k).or_default().push(p);
}
}
let mut keep: HashSet<String> = ends.cloned().collect();
let mut todo: Vec<String> = keep.iter().cloned().collect();
while let Some(m) = todo.pop() {
for &p in parents.get(m.as_str()).into_iter().flatten() {
if keep.insert(p.to_string()) {
todo.push(p.to_string());
}
}
}
keep
}
fn super_aliases(g: &ClassGraph, frontier: &[String]) -> HashMap<String, Vec<(String, String)>> {
let mut out = HashMap::new();
for sig in frontier {
let Some(p) = split_signature(sig) else { continue };
if p.name.starts_with('<') {
continue;
}
let aliases: Vec<(String, String)> =
g.reachable(p.class, Lineage::Supers, MAX_SUPERS).into_iter().map(|s| (format!("{s}->{}{}", p.name, p.proto), s)).collect();
if !aliases.is_empty() {
out.insert(sig.clone(), aliases);
}
}
out
}
fn overrides_of(g: &ClassGraph, views: &[DexView], sig: &str, errors: &mut Vec<String>) -> (Vec<String>, bool) {
let Some(p) = split_signature(sig) else { return (Vec::new(), false) };
if p.name.starts_with('<') {
return (Vec::new(), false);
}
let subs = g.reachable(p.class, Lineage::Subclasses, MAX_OVERRIDE_SCAN);
let cut = subs.len() >= MAX_OVERRIDE_SCAN;
let mut out = Vec::new();
for s in subs {
match g.defines(views, &s, p.name, Some(p.proto)) {
Ok(d) if !d.is_empty() => out.push(format!("{s}->{}{}", p.name, p.proto)),
Ok(_) => {}
Err(e) => errors.push(format!("{s}: {e}")),
}
}
(out, cut)
}
fn view_neighbours(v: &DexView, frontier: &[String], dir: Direction, expand: bool) -> Result<HashMap<String, Found>> {
let dex = Dex::with_name(&v.data, &v.name)?;
let fr = FindRefs::with_state(&dex, &v.state);
let mut out: HashMap<String, Found> = HashMap::new();
let mut graph = None;
for sig in frontier {
let Some(parts) = split_signature(sig) else { continue };
for m in method_ids(&dex, parts.class, parts.name, Some(parts.proto))? {
let g = match graph {
Some(g) => g,
None => *graph.insert(fr.call_graph()?),
};
let f = out.entry(sig.clone()).or_default();
if g.has_body(m) {
f.bodies.insert(v.name.clone());
}
if !expand {
continue;
}
let ids = match dir {
Direction::Callers => g.callers(m),
Direction::Callees => g.callees(m),
};
for &n in ids {
f.next.insert(asc_smali::method_signature(&dex, n as u32)?);
}
}
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn signatures_split_at_the_arrow_and_paren() {
let s = split_signature("Lcom/a;->run(ILjava/lang/String;)V").unwrap();
assert_eq!((s.class, s.name, s.proto), ("Lcom/a;", "run", "(ILjava/lang/String;)V"));
assert!(split_signature("Lcom/a;->run").is_none());
assert!(split_signature("run()V").is_none());
}
}