asc_engine/
sweep.rs

1//! Engine support for sweeps over many APKs, whose DEX entries the caller
2//! schedules ([`Engine::search_target`]): L3 results per APK, and identical
3//! DEX entries across the APKs of one sweep searched once.
4//!
5//! App variants, re-signed copies and split builds often carry byte-identical
6//! `classes*.dex`. While planning, every entry claims its content key, the
7//! CRC-32 and size the central directory already records. The first claimant
8//! of a key leads and searches; the others follow and reuse its rows, renamed
9//! to their own entry. The CRC-32, the size and the DEX header's checksum and
10//! signature are all whatever the APK's author wrote, so an APK could claim
11//! another one's DEX to hide its own code from an audit: a follower reuses
12//! rows only when the SHA-256 of its inflated bytes equals the leader's, and
13//! searches by itself otherwise. That costs a read (an inflate for a deflated
14//! entry) but still saves the parse, the index build and the scan.
15//!
16//! A rule set claims with the rules it still has to search (the others came
17//! from L3), so only entries searched for the same rules share rows.
18
19use crate::{CompiledQuery, Engine, FindOptions, Job, Query, RefBatch, RefContext, SearchTarget, Shape, WriteJob, tiers};
20use asc_apk::Entry;
21use asc_cache::{LockExt, Tier};
22use sha2::{Digest, Sha256};
23use std::collections::HashMap;
24use std::path::PathBuf;
25use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
26use std::sync::{Arc, Condvar, Mutex, OnceLock};
27use std::time::{Duration, SystemTime};
28
29/// SHA-256 of an entry's inflated bytes.
30type Ident = [u8; 32];
31
32/// What claimants share rows by: CRC-32, inflated size, and the rules
33/// searched (empty for a single query).
34type GroupKey = (u32, u64, Vec<usize>);
35
36/// Content keys claimed by the entries of one sweep.
37#[derive(Default)]
38pub struct DexDedup {
39    groups: Mutex<HashMap<GroupKey, Arc<Group>>>,
40    shared: Arc<AtomicUsize>,
41}
42
43/// One entry's part in a [`DexDedup`].
44#[derive(Clone)]
45pub struct Claim {
46    group: Arc<Group>,
47    lead: bool,
48    /// The sweep's count of entries answered by their leader.
49    shared: Arc<AtomicUsize>,
50}
51
52#[derive(Default)]
53struct Group {
54    members: AtomicUsize,
55    /// Set once by the leader: `None` when it has nothing to share.
56    result: Mutex<Option<Option<Arc<Shared>>>>,
57    ready: Condvar,
58}
59
60struct Shared {
61    ident: Ident,
62    /// One batch per query searched.
63    batches: Vec<RefBatch>,
64}
65
66impl DexDedup {
67    /// A dedup with no claims yet, for one sweep.
68    pub fn new() -> DexDedup {
69        DexDedup::default()
70    }
71
72    /// Claim `entry`'s content. Claim every entry of the sweep before
73    /// searching any of them, so a leader knows whether anyone follows.
74    pub fn claim(&self, entry: &Entry) -> Claim {
75        self.claim_rules(entry, &[])
76    }
77
78    /// [`DexDedup::claim`] for an entry searched for the rules `rules` (indexes
79    /// into the sweep's rule set): it shares rows only with entries searched
80    /// for the same rules.
81    pub fn claim_rules(&self, entry: &Entry, rules: &[usize]) -> Claim {
82        let mut groups = self.groups.lock_or_default();
83        let group = groups.entry((entry.crc32, entry.uncomp_size, rules.to_vec())).or_default();
84        let lead = group.members.fetch_add(1, Ordering::Relaxed) == 0;
85        Claim { group: Arc::clone(group), lead, shared: Arc::clone(&self.shared) }
86    }
87
88    /// Whether a claim follows another entry (and is cheap to search).
89    pub fn follows(claim: &Claim) -> bool {
90        !claim.lead
91    }
92
93    /// Entries answered from another APK's search so far.
94    pub fn shared(&self) -> usize {
95        self.shared.load(Ordering::Relaxed)
96    }
97}
98
99/// Publishes `None` if the leader unwinds before publishing, so followers
100/// never wait for it forever.
101struct Publish<'a>(&'a Group);
102
103impl Publish<'_> {
104    fn set(self, value: Option<Arc<Shared>>) {
105        let mut r = self.0.result.lock().expect("shared result mutex poisoned");
106        *r = Some(value);
107        self.0.ready.notify_all();
108        drop(r);
109        std::mem::forget(self);
110    }
111}
112
113impl Drop for Publish<'_> {
114    fn drop(&mut self) {
115        let mut r = self.0.result.lock_or_default();
116        if r.is_none() {
117            *r = Some(None);
118        }
119        self.0.ready.notify_all();
120    }
121}
122
123impl Engine {
124    /// The rows an earlier search of `query` on this APK left in L3: what
125    /// [`Engine::find_refs_with`] replays, so the one-APK and the sweep paths
126    /// share results.
127    pub fn cached_refs(&self, target: &SearchTarget, query: &Query, aggregate: bool) -> Option<Vec<RefBatch>> {
128        let (disk, key) = self.rendered_tier(target)?;
129        // A sweep takes no condition resolved over the program, so no notes.
130        tiers::decode_batches(disk.get_rendered_view(key, &tiers::findrefs_key(query, sweep_opts(aggregate), Shape::Results))?)
131            .map(|r| r.batches)
132    }
133
134    /// Keep the batches of a complete search of `query` on this APK in L3
135    /// (not when one failed, see [`Engine::cached_refs`]).
136    ///
137    /// # Panics
138    ///
139    /// Never: it returns early unless a disk cache with L3 is in use.
140    pub fn store_refs(&self, query: &Query, aggregate: bool, batches: &[RefBatch]) {
141        if !self.cache.as_ref().and_then(|c| c.disk.as_deref()).is_some_and(|d| d.uses(Tier::Rendered)) {
142            return;
143        }
144        let Some(apk) = self.open().ok() else { return };
145        let Some(key) = self.apk_key(&apk) else { return };
146        let mut recorder = tiers::ResultRecorder::new(tiers::L3_MAX_RESULT);
147        for b in batches {
148            recorder.record(b);
149        }
150        if let Some(payload) = recorder.finish(&[]) {
151            let query = tiers::findrefs_key(query, sweep_opts(aggregate), Shape::Results);
152            let job = WriteJob { key, apk: Arc::from(apk.path()), kind: Job::Rendered { query, payload } };
153            self.cache.as_ref().expect("an L3 tier implies a cache").put_async(job);
154        }
155    }
156
157    fn rendered_tier<'a>(&'a self, target: &'a SearchTarget) -> Option<(&'a asc_cache::DiskCache, &'a asc_cache::ApkKey)> {
158        let disk = self.cache.as_ref()?.disk.as_deref().filter(|d| d.uses(Tier::Rendered))?;
159        Some((disk, target.key.as_ref()?))
160    }
161
162    /// [`Engine::search_target_entry`] for an entry claimed in a
163    /// [`DexDedup`]: a leading entry is searched and its rows offered to the
164    /// entries that follow it; a following one reuses them when its bytes
165    /// hash the same.
166    ///
167    /// # Panics
168    ///
169    /// Never: a search of one query yields exactly one batch.
170    pub fn search_target_entry_dedup(
171        &self,
172        target: &SearchTarget,
173        entry: &Entry,
174        query: &CompiledQuery,
175        aggregate: bool,
176        cancel: &AtomicBool,
177        claim: &Claim,
178    ) -> RefBatch {
179        let search = || vec![self.search_target_entry(target, entry, query, aggregate, cancel)];
180        self.dedup_search(target, entry, cancel, claim, search).pop().expect("one batch per query")
181    }
182
183    /// [`Engine::search_target_entry_rules`] for an entry claimed with
184    /// [`DexDedup::claim_rules`] for these `rules`.
185    pub fn search_target_entry_rules_dedup(
186        &self,
187        target: &SearchTarget,
188        entry: &Entry,
189        rules: &[&CompiledQuery],
190        aggregate: bool,
191        cancel: &AtomicBool,
192        claim: &Claim,
193    ) -> Vec<RefBatch> {
194        let search = || self.search_target_entry_rules(target, entry, rules, aggregate, cancel);
195        self.dedup_search(target, entry, cancel, claim, search)
196    }
197
198    fn dedup_search(
199        &self,
200        target: &SearchTarget,
201        entry: &Entry,
202        cancel: &AtomicBool,
203        claim: &Claim,
204        search: impl FnOnce() -> Vec<RefBatch>,
205    ) -> Vec<RefBatch> {
206        let group = &claim.group;
207        if claim.lead {
208            let publish = Publish(group);
209            let batches = search();
210            if group.members.load(Ordering::Relaxed) > 1 {
211                // The search left the bytes in the memory cache, if there is one.
212                let ident = content_id(target, entry, || {
213                    self.entry_bytes(&target.apk, target.key.as_ref(), entry).ok().map(|b| Sha256::digest(&**b).into())
214                });
215                publish.set(ident.map(|ident| Arc::new(Shared { ident, batches: batches.clone() })));
216            } else {
217                publish.set(None);
218            }
219            return batches;
220        }
221        if let Some(shared) = wait_shared(group, cancel) {
222            // Hashed once per APK version: later queries of the sweep's
223            // session neither read nor inflate a verified follower.
224            let ident = content_id(target, entry, || {
225                target.apk.read_entry_cancellable(entry, Some(cancel)).ok().map(|b| Sha256::digest(&*b).into())
226            });
227            if ident == Some(shared.ident) {
228                claim.shared.fetch_add(1, Ordering::Relaxed);
229                return shared.batches.iter().map(|b| renamed(b.clone(), &entry.name)).collect();
230            }
231        }
232        search()
233    }
234}
235
236fn sweep_opts(aggregate: bool) -> FindOptions {
237    FindOptions { aggregate, context: RefContext::None }
238}
239
240/// The leader's result, once it is published.
241fn wait_shared(group: &Group, cancel: &AtomicBool) -> Option<Arc<Shared>> {
242    let mut r = group.result.lock().expect("shared result mutex poisoned");
243    while r.is_none() {
244        if cancel.load(Ordering::Relaxed) {
245            return None;
246        }
247        // The leader is skipped if the sweep is cancelled: poll the flag.
248        r = group.ready.wait_timeout(r, Duration::from_millis(50)).expect("shared result mutex poisoned").0;
249    }
250    r.clone().flatten()
251}
252
253/// An entry of one version of an APK file.
254type ContentKey = (PathBuf, u64, Option<SystemTime>, String, u32, u64);
255
256/// Most content ids kept; past it the memo starts over.
257const MAX_CONTENT_IDS: usize = 1 << 16;
258
259/// `entry`'s SHA-256, from `hash` the first time it is asked for this
260/// version of the APK file.
261fn content_id(target: &SearchTarget, entry: &Entry, hash: impl FnOnce() -> Option<Ident>) -> Option<Ident> {
262    static IDS: OnceLock<Mutex<HashMap<ContentKey, Ident>>> = OnceLock::new();
263    let (size, modified) = target.apk.stamp().ok()?;
264    let key = (target.apk.path().to_path_buf(), size, modified, entry.name.clone(), entry.crc32, entry.uncomp_size);
265    let ids = IDS.get_or_init(Mutex::default);
266    if let Some(id) = ids.lock_or_default().get(&key) {
267        return Some(*id);
268    }
269    let id = hash()?;
270    let mut ids = ids.lock_or_default();
271    if ids.len() >= MAX_CONTENT_IDS {
272        ids.clear();
273    }
274    ids.insert(key, id);
275    Some(id)
276}
277
278/// `batch` as if `name` had produced it: the leader's entry name prefixes
279/// the rows' DEX names (`classes2.dex`, `classes2.dex!classes1.dex`).
280fn renamed(mut batch: RefBatch, name: &str) -> RefBatch {
281    let from = std::mem::replace(&mut batch.entry_name, name.to_string());
282    if from != name {
283        for r in &mut batch.results {
284            if let Some(rest) = r.dex_name.strip_prefix(&from) {
285                r.dex_name = format!("{name}{rest}");
286            }
287        }
288        if let Some(e) = &mut batch.error {
289            *e = e.replace(&format!("{from}!"), &format!("{name}!"));
290        }
291    }
292    batch.source = "shared";
293    batch
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299    use crate::RefResult;
300
301    #[test]
302    fn renamed_moves_rows_and_container_errors_to_the_follower() {
303        let row = |dex: &str| RefResult {
304            dex_name: dex.into(),
305            method_idx: 0,
306            method: "La;->m".into(),
307            signature: "La;->m()V".into(),
308            matched: vec![],
309            relation: None,
310            sites: vec![],
311            context: None,
312        };
313        let batch = RefBatch {
314            entry_name: "classes2.dex".into(),
315            results: vec![row("classes2.dex"), row("classes2.dex!classes1.dex")],
316            text: None,
317            inflate_us: 0,
318            search_us: 0,
319            source: "inflate",
320            error: Some("classes2.dex!classes2.dex: bad".into()),
321        };
322        let b = renamed(batch, "classes5.dex");
323        assert_eq!(b.entry_name, "classes5.dex");
324        assert_eq!(b.results.iter().map(|r| r.dex_name.as_str()).collect::<Vec<_>>(), ["classes5.dex", "classes5.dex!classes1.dex"]);
325        assert_eq!(b.error.as_deref(), Some("classes5.dex!classes2.dex: bad"));
326        assert_eq!(b.source, "shared");
327    }
328
329    #[test]
330    fn a_cancelled_or_unwound_leader_releases_followers() {
331        let group = Group::default();
332        drop(Publish(&group));
333        assert!(wait_shared(&group, &AtomicBool::new(false)).is_none());
334        let pending = Group::default();
335        assert!(wait_shared(&pending, &AtomicBool::new(true)).is_none());
336    }
337}
asc_engine/sweep.rs

asc_engine/
sweep.rs
