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::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    pub fn new() -> DexDedup {
68        DexDedup::default()
69    }
70
71    /// Claim `entry`'s content. Claim every entry of the sweep before
72    /// searching any of them, so a leader knows whether anyone follows.
73    pub fn claim(&self, entry: &Entry) -> Claim {
74        self.claim_rules(entry, &[])
75    }
76
77    /// [`DexDedup::claim`] for an entry searched for the rules `rules` (indexes
78    /// into the sweep's rule set): it shares rows only with entries searched
79    /// for the same rules.
80    pub fn claim_rules(&self, entry: &Entry, rules: &[usize]) -> Claim {
81        let mut groups = self.groups.lock().unwrap();
82        let group = groups.entry((entry.crc32, entry.uncomp_size, rules.to_vec())).or_default();
83        let lead = group.members.fetch_add(1, Ordering::Relaxed) == 0;
84        Claim { group: Arc::clone(group), lead, shared: Arc::clone(&self.shared) }
85    }
86
87    /// Whether a claim follows another entry (and is cheap to search).
88    pub fn follows(claim: &Claim) -> bool {
89        !claim.lead
90    }
91
92    /// Entries answered from another APK's search so far.
93    pub fn shared(&self) -> usize {
94        self.shared.load(Ordering::Relaxed)
95    }
96}
97
98/// Publishes `None` if the leader unwinds before publishing, so followers
99/// never wait for it forever.
100struct Publish<'a>(&'a Group);
101
102impl Publish<'_> {
103    fn set(self, value: Option<Arc<Shared>>) {
104        let mut r = self.0.result.lock().unwrap();
105        *r = Some(value);
106        self.0.ready.notify_all();
107        drop(r);
108        std::mem::forget(self);
109    }
110}
111
112impl Drop for Publish<'_> {
113    fn drop(&mut self) {
114        let mut r = self.0.result.lock().unwrap_or_else(|e| e.into_inner());
115        if r.is_none() {
116            *r = Some(None);
117        }
118        self.0.ready.notify_all();
119    }
120}
121
122impl Engine {
123    /// The rows an earlier search of `query` on this APK left in L3: what
124    /// [`Engine::find_refs_with`] replays, so the one-APK and the sweep paths
125    /// share results.
126    pub fn cached_refs(&self, target: &SearchTarget, query: &Query, aggregate: bool) -> Option<Vec<RefBatch>> {
127        let (disk, key) = self.rendered_tier(target)?;
128        // A sweep takes no condition resolved over the program, so no notes.
129        tiers::decode_batches(disk.get_rendered_view(key, &tiers::findrefs_key(query, sweep_opts(aggregate), Shape::Results))?)
130            .map(|r| r.batches)
131    }
132
133    /// Keep the batches of a complete search of `query` on this APK in L3
134    /// (not when one failed, see [`Engine::cached_refs`]).
135    pub fn store_refs(&self, query: &Query, aggregate: bool, batches: &[RefBatch]) {
136        if !self.cache.as_ref().and_then(|c| c.disk.as_deref()).is_some_and(|d| d.uses(Tier::Rendered)) {
137            return;
138        }
139        let Some(apk) = self.open().ok() else { return };
140        let Some(key) = self.apk_key(&apk) else { return };
141        let mut recorder = tiers::ResultRecorder::new(tiers::L3_MAX_RESULT);
142        for b in batches {
143            recorder.record(b);
144        }
145        if let Some(payload) = recorder.finish(&[]) {
146            let query = tiers::findrefs_key(query, sweep_opts(aggregate), Shape::Results);
147            let job = WriteJob { key, apk: Arc::from(apk.path()), kind: Job::Rendered { query, payload } };
148            self.cache.as_ref().expect("an L3 tier implies a cache").put_async(job);
149        }
150    }
151
152    fn rendered_tier<'a>(&'a self, target: &'a SearchTarget) -> Option<(&'a asc_cache::DiskCache, &'a asc_cache::ApkKey)> {
153        let disk = self.cache.as_ref()?.disk.as_deref().filter(|d| d.uses(Tier::Rendered))?;
154        Some((disk, target.key.as_ref()?))
155    }
156
157    /// [`Engine::search_target_entry`] for an entry claimed in a
158    /// [`DexDedup`]: a leading entry is searched and its rows offered to the
159    /// entries that follow it; a following one reuses them when its bytes
160    /// hash the same.
161    pub fn search_target_entry_dedup(
162        &self,
163        target: &SearchTarget,
164        entry: &Entry,
165        query: &CompiledQuery,
166        aggregate: bool,
167        cancel: &AtomicBool,
168        claim: &Claim,
169    ) -> RefBatch {
170        let search = || vec![self.search_target_entry(target, entry, query, aggregate, cancel)];
171        self.dedup_search(target, entry, cancel, claim, search).pop().expect("one batch per query")
172    }
173
174    /// [`Engine::search_target_entry_rules`] for an entry claimed with
175    /// [`DexDedup::claim_rules`] for these `rules`.
176    pub fn search_target_entry_rules_dedup(
177        &self,
178        target: &SearchTarget,
179        entry: &Entry,
180        rules: &[&CompiledQuery],
181        aggregate: bool,
182        cancel: &AtomicBool,
183        claim: &Claim,
184    ) -> Vec<RefBatch> {
185        let search = || self.search_target_entry_rules(target, entry, rules, aggregate, cancel);
186        self.dedup_search(target, entry, cancel, claim, search)
187    }
188
189    fn dedup_search(
190        &self,
191        target: &SearchTarget,
192        entry: &Entry,
193        cancel: &AtomicBool,
194        claim: &Claim,
195        search: impl FnOnce() -> Vec<RefBatch>,
196    ) -> Vec<RefBatch> {
197        let group = &claim.group;
198        if claim.lead {
199            let publish = Publish(group);
200            let batches = search();
201            if group.members.load(Ordering::Relaxed) > 1 {
202                // The search left the bytes in the memory cache, if there is one.
203                let ident = content_id(target, entry, || {
204                    self.entry_bytes(&target.apk, target.key.as_ref(), entry).ok().map(|b| Sha256::digest(&**b).into())
205                });
206                publish.set(ident.map(|ident| Arc::new(Shared { ident, batches: batches.clone() })));
207            } else {
208                publish.set(None);
209            }
210            return batches;
211        }
212        if let Some(shared) = wait_shared(group, cancel) {
213            // Hashed once per APK version: later queries of the sweep's
214            // session neither read nor inflate a verified follower.
215            let ident = content_id(target, entry, || {
216                target.apk.read_entry_cancellable(entry, Some(cancel)).ok().map(|b| Sha256::digest(&*b).into())
217            });
218            if ident == Some(shared.ident) {
219                claim.shared.fetch_add(1, Ordering::Relaxed);
220                return shared.batches.iter().map(|b| renamed(b.clone(), &entry.name)).collect();
221            }
222        }
223        search()
224    }
225}
226
227fn sweep_opts(aggregate: bool) -> FindOptions {
228    FindOptions { aggregate, context: RefContext::None }
229}
230
231/// The leader's result, once it is published.
232fn wait_shared(group: &Group, cancel: &AtomicBool) -> Option<Arc<Shared>> {
233    let mut r = group.result.lock().unwrap();
234    while r.is_none() {
235        if cancel.load(Ordering::Relaxed) {
236            return None;
237        }
238        // The leader is skipped if the sweep is cancelled: poll the flag.
239        r = group.ready.wait_timeout(r, Duration::from_millis(50)).unwrap().0;
240    }
241    r.clone().flatten()
242}
243
244/// An entry of one version of an APK file.
245type ContentKey = (PathBuf, u64, Option<SystemTime>, String, u32, u64);
246
247/// Most content ids kept; past it the memo starts over.
248const MAX_CONTENT_IDS: usize = 1 << 16;
249
250/// `entry`'s SHA-256, from `hash` the first time it is asked for this
251/// version of the APK file.
252fn content_id(target: &SearchTarget, entry: &Entry, hash: impl FnOnce() -> Option<Ident>) -> Option<Ident> {
253    static IDS: OnceLock<Mutex<HashMap<ContentKey, Ident>>> = OnceLock::new();
254    let (size, modified) = target.apk.stamp().ok()?;
255    let key = (target.apk.path().to_path_buf(), size, modified, entry.name.clone(), entry.crc32, entry.uncomp_size);
256    let ids = IDS.get_or_init(Mutex::default);
257    if let Some(id) = ids.lock().unwrap().get(&key) {
258        return Some(*id);
259    }
260    let id = hash()?;
261    let mut ids = ids.lock().unwrap();
262    if ids.len() >= MAX_CONTENT_IDS {
263        ids.clear();
264    }
265    ids.insert(key, id);
266    Some(id)
267}
268
269/// `batch` as if `name` had produced it: the leader's entry name prefixes
270/// the rows' DEX names (`classes2.dex`, `classes2.dex!classes1.dex`).
271fn renamed(mut batch: RefBatch, name: &str) -> RefBatch {
272    let from = std::mem::replace(&mut batch.entry_name, name.to_string());
273    if from != name {
274        for r in &mut batch.results {
275            if let Some(rest) = r.dex_name.strip_prefix(&from) {
276                r.dex_name = format!("{name}{rest}");
277            }
278        }
279        if let Some(e) = &mut batch.error {
280            *e = e.replace(&format!("{from}!"), &format!("{name}!"));
281        }
282    }
283    batch.source = "shared";
284    batch
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290    use crate::RefResult;
291
292    #[test]
293    fn renamed_moves_rows_and_container_errors_to_the_follower() {
294        let row = |dex: &str| RefResult {
295            dex_name: dex.into(),
296            method_idx: 0,
297            method: "La;->m".into(),
298            signature: "La;->m()V".into(),
299            matched: vec![],
300            relation: None,
301            sites: vec![],
302            context: None,
303        };
304        let batch = RefBatch {
305            entry_name: "classes2.dex".into(),
306            results: vec![row("classes2.dex"), row("classes2.dex!classes1.dex")],
307            text: None,
308            inflate_us: 0,
309            search_us: 0,
310            source: "inflate",
311            error: Some("classes2.dex!classes2.dex: bad".into()),
312        };
313        let b = renamed(batch, "classes5.dex");
314        assert_eq!(b.entry_name, "classes5.dex");
315        assert_eq!(b.results.iter().map(|r| r.dex_name.as_str()).collect::<Vec<_>>(), ["classes5.dex", "classes5.dex!classes1.dex"]);
316        assert_eq!(b.error.as_deref(), Some("classes5.dex!classes2.dex: bad"));
317        assert_eq!(b.source, "shared");
318    }
319
320    #[test]
321    fn a_cancelled_or_unwound_leader_releases_followers() {
322        let group = Group::default();
323        drop(Publish(&group));
324        assert!(wait_shared(&group, &AtomicBool::new(false)).is_none());
325        let pending = Group::default();
326        assert!(wait_shared(&pending, &AtomicBool::new(true)).is_none());
327    }
328}
asc_engine/sweep.rs

asc_engine/
sweep.rs
