diff --git a/src/asc_client/apk/dex_reader.rs b/src/asc_client/apk/dex_reader.rs index 4d13fea..016dad5 100644 --- a/src/asc_client/apk/dex_reader.rs +++ b/src/asc_client/apk/dex_reader.rs @@ -117,6 +117,226 @@ fn _decode_name(raw_name: &[u8]) -> String { } } +/// Smallest DEX prefix that contains every class descriptor string. +/// +/// R8/D8 emit `class_defs` before the data section, so the id tables pin down +/// each descriptor's `string_data_off`; only those strings (plus a NUL) need to +/// be decompressed to list class names. +/// Smallest DEX prefix that can answer "does this DEX define `target`?". +/// +/// The answer needs only the id tables plus the descriptor string; code items, +/// debug info and class_data are irrelevant. A class lookup therefore inflates +/// ~10-30% of each DEX it rejects instead of the whole file. +pub(crate) fn class_defines_prefix_len(buf: &[u8], target: &[u8]) -> Result { + const SLACK: usize = 1024; + if buf.len() < 0x70 { + return Ok(0x70); + } + if buf[..3] != *b"dex" { + // Not a DEX prefix we can reason about; ask for everything. + return Ok(buf.len().saturating_add(SLACK)); + } + let string_ids_size = _STRUCT_I::unpack_from(buf, 0x38) as usize; + let string_ids_off = _STRUCT_I::unpack_from(buf, 0x3C) as usize; + let type_ids_size = _STRUCT_I::unpack_from(buf, 0x40) as usize; + let type_ids_off = _STRUCT_I::unpack_from(buf, 0x44) as usize; + let class_defs_size = _STRUCT_I::unpack_from(buf, 0x60) as usize; + let class_defs_off = _STRUCT_I::unpack_from(buf, 0x64) as usize; + + let tables_end = string_ids_off + .saturating_add(string_ids_size * 4) + .max(type_ids_off.saturating_add(type_ids_size * 4)) + .max(class_defs_off.saturating_add(class_defs_size * 32)); + if buf.len() < tables_end { + return Ok(tables_end); + } + + let Ok(type_idx) = _find_type_idx(buf, target) else { + return Ok(buf.len().saturating_add(SLACK)); + }; + if type_idx < 0 { + // Descriptor absent: the id tables alone already give the answer. + return Ok(tables_end); + } + let string_idx = _STRUCT_I::unpack_from(buf, type_ids_off + (type_idx as usize) * 4) as usize; + if string_idx >= string_ids_size { + return Ok(tables_end); + } + let string_off = _STRUCT_I::unpack_from(buf, string_ids_off + (string_idx << 2)) as usize; + match _read_string_data_bytes(buf, string_off) { + Ok(raw) => { + let start = raw.as_ptr() as usize - buf.as_ptr() as usize; + Ok((start + raw.len() + 1).max(tables_end)) + } + Err(_) => Ok(buf.len().saturating_add(SLACK)), + } +} + +pub(crate) fn class_names_prefix_len(buf: &[u8]) -> Result { + const SLACK: usize = 4096; + if buf.len() < 0x70 { + return Ok(0x70); + } + if buf[..3] != *b"dex" { + return Err("invalid DEX header".to_owned()); + } + let string_ids_size = _STRUCT_I::unpack_from(buf, 0x38) as usize; + let string_ids_off = _STRUCT_I::unpack_from(buf, 0x3C) as usize; + let type_ids_size = _STRUCT_I::unpack_from(buf, 0x40) as usize; + let type_ids_off = _STRUCT_I::unpack_from(buf, 0x44) as usize; + let class_defs_size = _STRUCT_I::unpack_from(buf, 0x60) as usize; + let class_defs_off = _STRUCT_I::unpack_from(buf, 0x64) as usize; + + // Probe one string_data offset before requesting the complete id tables. + // A late string pool is a hint that prefix decoding will be expensive. Ask + // for the full entry early so the loader's existing fallback can restart + // after only a header-sized probe, rather than redoing all of the id tables. + // Physical string order is not guaranteed by DEX: this is only a choice of + // decoder, never an assumption used to omit classes or validate an offset. + if string_ids_size != 0 { + let first_id_end = string_ids_off + .checked_add(4) + .ok_or_else(|| "bad string_ids range".to_owned())?; + if buf.len() < first_id_end { + return Ok(first_id_end); + } + let first_data_off = _STRUCT_I::unpack_from(buf, string_ids_off) as u64; + let file_size = _STRUCT_I::unpack_from(buf, 0x20) as u64; + if first_data_off < file_size && first_data_off * 5 > file_size * 3 { + return Ok(usize::MAX); + } + } + + let Some(id_end) = class_defs_off.checked_add(class_defs_size * 32) else { + return Ok(buf.len()); + }; + let tables_end = id_end + .max(string_ids_off.saturating_add(string_ids_size * 4)) + .max(type_ids_off.saturating_add(type_ids_size * 4)); + if buf.len() < tables_end { + return Ok(tables_end); + } + + let descriptor_off = |class_index: usize| -> Option { + let class_off = class_defs_off + class_index * 32; + let type_idx = _STRUCT_I::unpack_from(buf, class_off) as usize; + if type_idx >= type_ids_size { + return None; + } + let string_idx = _STRUCT_I::unpack_from(buf, type_ids_off + (type_idx << 2)) as usize; + if string_idx >= string_ids_size { + return None; + } + Some(_STRUCT_I::unpack_from(buf, string_ids_off + (string_idx << 2)) as usize) + }; + + let mut max_start = tables_end; + for class_index in 0..class_defs_size { + if let Some(offset) = descriptor_off(class_index) { + max_start = max_start.max(offset); + } + } + if buf.len() < max_start.saturating_add(SLACK) { + return Ok(max_start.saturating_add(SLACK)); + } + + let mut max_end = max_start.saturating_add(1); + for class_index in 0..class_defs_size { + let Some(offset) = descriptor_off(class_index) else { + continue; + }; + match _read_string_data_bytes(buf, offset) { + Ok(raw) => { + let start = raw.as_ptr() as usize - buf.as_ptr() as usize; + max_end = max_end.max(start + raw.len() + 1); + } + // The terminator is past the current prefix; ask for more. + Err(_) => return Ok(buf.len().saturating_add(SLACK)), + } + } + Ok(max_end) +} + +pub(crate) fn write_dex_class_names( + buf: &[u8], + prefix: Option<&str>, + prefix_bytes: Option<&[u8]>, + out: &mut Vec, +) -> Result { + if buf.len() < 0x70 || buf[..3] != *b"dex" { + return Err("invalid DEX header".to_owned()); + } + + let string_ids_size = _STRUCT_I::unpack_from(buf, 0x38) as usize; + let string_ids_off = _STRUCT_I::unpack_from(buf, 0x3C) as usize; + let type_ids_size = _STRUCT_I::unpack_from(buf, 0x40) as usize; + let type_ids_off = _STRUCT_I::unpack_from(buf, 0x44) as usize; + let class_defs_size = _STRUCT_I::unpack_from(buf, 0x60) as usize; + let class_defs_off = _STRUCT_I::unpack_from(buf, 0x64) as usize; + + if class_defs_size == 0 { + return Ok(0); + } + if string_ids_off + string_ids_size * 4 > buf.len() { + return Err("bad string_ids range".to_owned()); + } + if type_ids_off + type_ids_size * 4 > buf.len() { + return Err("bad type_ids range".to_owned()); + } + if class_defs_off + class_defs_size * 32 > buf.len() { + return Err("bad class_defs range".to_owned()); + } + + let start_len = out.len(); + let mut count = 0usize; + let class_defs_end = class_defs_off + class_defs_size * 32; + for class_def_off in (class_defs_off..class_defs_end).step_by(32) { + let type_idx = _STRUCT_I::unpack_from(buf, class_def_off) as usize; + if type_idx >= type_ids_size { + return Err("bad class_def->type_idx".to_owned()); + } + let string_idx = _STRUCT_I::unpack_from(buf, type_ids_off + (type_idx << 2)) as usize; + if string_idx >= string_ids_size { + return Err("bad type_id->string_idx".to_owned()); + } + let string_off = _STRUCT_I::unpack_from(buf, string_ids_off + (string_idx << 2)) as usize; + if string_off >= buf.len() { + return Err("bad string_data_off".to_owned()); + } + let raw_name = _read_string_data_bytes(buf, string_off) + .map_err(|_| "bad string_data_off".to_owned())?; + + // ASCII prefixes can be rejected before decoding every class name. + if let Some(prefix_bytes) = prefix_bytes { + if !raw_name.starts_with(prefix_bytes) { + continue; + } + } + let name = std::str::from_utf8(raw_name); + match name { + Ok(name) => { + if prefix.is_some() && prefix_bytes.is_none() && !name.starts_with(prefix.unwrap()) { + continue; + } + out.reserve(name.len() + 1); + out.extend_from_slice(name.as_bytes()); + } + Err(_) => { + let name = _decode_name(raw_name); + if prefix.is_some_and(|prefix| !name.starts_with(prefix)) { + continue; + } + out.reserve(name.len() + 1); + out.extend_from_slice(name.as_bytes()); + } + } + out.push(b'\n'); + count += 1; + } + let _ = start_len; + Ok(count) +} + pub(crate) fn list_dex_classes( buf: &[u8], prefix: Option<&str>, diff --git a/src/asc_client/apk/loader.rs b/src/asc_client/apk/loader.rs index e68d459..876540c 100644 --- a/src/asc_client/apk/loader.rs +++ b/src/asc_client/apk/loader.rs @@ -198,6 +198,126 @@ fn _validate_inflated_size(entry: &DexEntry, data: Vec) -> Result, S Ok(data) } +/// Inflate only as much of an entry as `need` asks for. +/// +/// `listclass` only reads the DEX id tables plus the class descriptor strings, +/// which live in the first ~20-35% of a D8/R8 DEX; the remaining code/debug +/// sections never need to leave the deflate stream. `need` is re-evaluated on +/// each round and returns the prefix length required by the current data. +pub(crate) fn _inflate_dex_prefix_into( + buf: &[u8], + entry: &DexEntry, + out: &mut Vec, + decompressor: &mut Decompressor, + need: F, +) -> Result<(), String> +where + F: Fn(&[u8]) -> Result, +{ + let comp_view = _compressed_data(buf, entry)?; + out.clear(); + match entry.comp_method { + 0 => { + out.extend_from_slice(comp_view); + return Ok(()); + } + 8 => {} + method => return Err(format!("unsupported compression method: {method}")), + } + + let total = entry.uncomp_size as usize; + // Low-level streaming API: `flate2`'s Read adapter copies through an + // intermediate buffer, while `decompress` writes straight into `out`. + let mut decoder = flate2::Decompress::new(false); + let mut produced = 0usize; + let mut step = 1usize << 20; + let mut round = 0usize; + loop { + let target = need(&out[..produced])?.min(total); + if produced >= target { + out.truncate(produced); + return Ok(()); + } + // Some APKs keep class descriptors near the end of the DEX. Streaming + // (flate2) decodes slower than the one-shot libdeflate path, so switch + // back when the saving would be small. + if round > 0 && target * 5 > total * 3 { + return _inflate_dex_into_with(buf, entry, out, decompressor); + } + round += 1; + + let want = (target - produced).min(step); + let consumed = decoder.total_in() as usize; + out.resize(produced + want, 0); + decoder + .decompress( + &comp_view[consumed..], + &mut out[produced..produced + want], + flate2::FlushDecompress::None, + ) + .map_err(|error| format!("flate2 cannot decompress {}: {error}", entry.name))?; + produced = decoder.total_out() as usize; + if produced == 0 || (produced < target && consumed == decoder.total_in() as usize) { + out.truncate(produced); + return Err(format!( + "size mismatch: expect {}, got {}", + entry.uncomp_size, + produced + )); + } + step = (step * 2).min(8 << 20); + } +} + +/// Inflate one entry into a caller-owned buffer so a worker can reuse the same +/// pages across entries. Fresh multi-MB allocations must be zero-filled on +/// first touch, which roughly halves libdeflate's throughput; reusing a warm +/// buffer measured ~915 MB/s against ~460 MB/s for a fresh `Vec`. +pub(crate) fn _inflate_dex_into( + buf: &[u8], + entry: &DexEntry, + out: &mut Vec, +) -> Result<(), String> { + _inflate_dex_into_with(buf, entry, out, &mut Decompressor::new()) +} + +/// Same as [`_inflate_dex_into`] with a reusable decompressor: libdeflate +/// builds its decode tables per instance, which measured ~7% of inflate time. +pub(crate) fn _inflate_dex_into_with( + buf: &[u8], + entry: &DexEntry, + out: &mut Vec, + decompressor: &mut Decompressor, +) -> Result<(), String> { + let comp_view = _compressed_data(buf, entry)?; + out.clear(); + match entry.comp_method { + 0 => out.extend_from_slice(comp_view), + 8 => { + let size = entry.uncomp_size as usize; + out.reserve(size); + // SAFETY: `reserve` guarantees capacity for `size` bytes and every + // byte is written by the decompressor before it is observed. + unsafe { out.set_len(size) }; + let actual_size = decompressor + .deflate_decompress(comp_view, out) + .map_err(|error| { + format!("libdeflate cannot decompress {}: {error}", entry.name) + })?; + out.truncate(actual_size); + } + method => return Err(format!("unsupported compression method: {method}")), + } + if entry.uncomp_size != 0 && out.len() != entry.uncomp_size as usize { + return Err(format!( + "size mismatch: expect {}, got {}", + entry.uncomp_size, + out.len() + )); + } + Ok(()) +} + pub(crate) fn _inflate_dex( buf: &[u8], entry: &DexEntry, diff --git a/src/asc_client/apk/worker.rs b/src/asc_client/apk/worker.rs index 4e2768d..08fe14a 100644 --- a/src/asc_client/apk/worker.rs +++ b/src/asc_client/apk/worker.rs @@ -7,8 +7,10 @@ use std::time::Instant; use memmap2::{Mmap, MmapOptions}; -use crate::asc_client::apk::loader::{_inflate_dex, DexEntry}; -use crate::asc_client::dex_container::into_logical_dex_buffers; +use crate::asc_client::apk::loader::{_inflate_dex_into_with, DexEntry}; +use crate::asc_client::dex_container::{ + dex041_logical_offsets, into_logical_dex_buffers, +}; use crate::asc_core::findrefs::findrefs_manager::FindRefInput; static _WORKER_APK_MM: OnceLock)>>> = OnceLock::new(); @@ -45,15 +47,49 @@ pub(crate) fn _findrefs_worker_from_buf( // Python worker imports AscHandler only after receiving a findrefs task. use crate::asc_client::asc_handler::AscHandler; + let mut inflate_buf = Vec::new(); + let mut decompressor = libdeflater::Decompressor::new(); + _findrefs_worker_from_buf_into( + apk_buf, + entry, + find_type, + find, + aggregate, + &mut inflate_buf, + &mut decompressor, + ) +} + +/// Same as [`_findrefs_worker_from_buf`] but lets a worker thread keep one +/// warm inflate buffer across DEX entries. +pub(crate) fn _findrefs_worker_from_buf_into( + apk_buf: &[u8], + entry: &DexEntry, + find_type: &str, + find: &FindRefInput, + aggregate: bool, + inflate_buf: &mut Vec, + decompressor: &mut libdeflater::Decompressor, +) -> Result<(String, Vec, f64, f64, u32), String> { + use crate::asc_client::asc_handler::AscHandler; + let t0 = Instant::now(); - let data = _inflate_dex(apk_buf, entry, None)?.unwrap_or_default(); + _inflate_dex_into_with(apk_buf, entry, inflate_buf, decompressor)?; let t1 = Instant::now(); let mut lines = Vec::new(); let handler = AscHandler::new(false); // A physical classes*.dex entry may contain several logical DEX files in // the DEX 041 container format; each needs an independent reference scan. - for (dex_name, dex_buf) in into_logical_dex_buffers(&entry.name, data) { - lines.extend(handler.findrefs_vec(&dex_name, dex_buf, find_type, find, aggregate)); + if dex041_logical_offsets(inflate_buf).len() <= 1 { + let data = std::mem::take(inflate_buf); + let dex = crate::asc_core::utils::tinydex::DEX::from_vec(data, &entry.name); + lines.extend(handler._findrefs(&entry.name, &dex, find_type, find, aggregate)); + *inflate_buf = dex.into_vec(); + } else { + let data = std::mem::take(inflate_buf); + for (dex_name, dex_buf) in into_logical_dex_buffers(&entry.name, data) { + lines.extend(handler.findrefs_vec(&dex_name, dex_buf, find_type, find, aggregate)); + } } let t2 = Instant::now(); Ok(( diff --git a/src/asc_client/apk_handler.rs b/src/asc_client/apk_handler.rs index 7c95fe7..40b086a 100644 --- a/src/asc_client/apk_handler.rs +++ b/src/asc_client/apk_handler.rs @@ -3,7 +3,7 @@ use std::fs::File; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Condvar, Mutex}; use std::thread; use std::time::Instant; @@ -11,20 +11,54 @@ use cesu8::to_java_cesu8; use memmap2::{Mmap, MmapOptions}; use crate::asc_client::apk::dex_reader::{ - _dex_defines_class as _defines_class, list_dex_classes as _list_classes, + _dex_defines_class as _defines_class, class_defines_prefix_len, class_names_prefix_len, + list_dex_classes, write_dex_class_names as _write_classes, }; use crate::asc_client::apk::loader::{ - _inflate_dex as _load_dex, _parse_cd_dex_entries as _parse_entries, DexEntry, + _inflate_dex as _load_dex, _inflate_dex_into, _inflate_dex_into_with, + _inflate_dex_prefix_into, _parse_cd_dex_entries as _parse_entries, DexEntry, +}; +use crate::asc_client::dex_container::{ + dex041_logical_offsets, into_logical_dex_buffers, iter_logical_dex_buffers, }; -use crate::asc_client::dex_container::{into_logical_dex_buffers, iter_logical_dex_buffers}; use crate::asc_core::findrefs::findrefs_manager::FindRefInput; use crate::asc_core::findrefs::scan::method_id_scan::render_dex_methods; +type DexChunkResult = Result<(Vec>, usize), String>; + +/// Shared state for the ordered streaming pipeline. +struct StreamState { + slots: Vec>, + /// Next entry index the consumer still has to emit. Producers may not run + /// further than `window` ahead of it, which bounds retained buffers. + next_emit: usize, + live: usize, + stop: bool, +} + +/// Decrements the live-worker count even if a worker unwinds, so the consumer +/// never blocks forever on a dead producer. +struct LiveGuard<'a> { + state: &'a Mutex, + ready: &'a Condvar, +} + +impl Drop for LiveGuard<'_> { + fn drop(&mut self) { + if let Ok(mut guard) = self.state.lock() { + guard.live -= 1; + } + self.ready.notify_all(); + } +} + fn _inflate_and_hit( mm: &[u8], entry: &DexEntry, target_bytes: &[u8], stop_event: Option<&AtomicBool>, + scratch: &mut Vec, + decompressor: &mut libdeflater::Decompressor, log: &dyn Fn(&str), ) -> (bool, Option, Option>) { let tid = format!("{:?}", thread::current().id()); @@ -41,10 +75,23 @@ fn _inflate_and_hit( } let t0 = Instant::now(); - let data = match _load_dex(mm, entry, stop_event) { - Ok(Some(data)) => data, - Ok(None) | Err(_) => return (false, None, None), - }; + // Membership only needs the DEX id tables plus one descriptor string, so + // most rejected entries decompress a small prefix instead of the whole file. + let prefix_need = |data: &[u8]| class_defines_prefix_len(data, target_bytes); + if _inflate_dex_prefix_into(mm, entry, scratch, decompressor, prefix_need).is_err() { + return (false, None, None); + } + if stop_event.is_some_and(|stop_event| stop_event.load(Ordering::SeqCst)) { + return (false, None, None); + } + if !matches!(iter_logical_dex_buffers(name, scratch).first(), Some((_, dex_buf)) if _defines_class(dex_buf, target_bytes).unwrap_or(false)) + { + return (false, None, None); + } + // Confirmed: inflate this entry in full, then locate the class inside it. + if _inflate_dex_into_with(mm, entry, scratch, decompressor).is_err() { + return (false, None, None); + } let t1 = Instant::now(); if stop_event.is_some_and(|stop_event| stop_event.load(Ordering::SeqCst)) { return (false, None, None); @@ -53,7 +100,7 @@ fn _inflate_and_hit( let mut hit = false; let mut hit_name = None; let mut hit_data = None; - for (dex_name, dex_buf) in iter_logical_dex_buffers(name, &data) { + for (dex_name, dex_buf) in into_logical_dex_buffers(name, std::mem::take(scratch)) { if stop_event.is_some_and(|stop_event| stop_event.load(Ordering::SeqCst)) { return (false, None, None); } @@ -140,8 +187,17 @@ impl ApkHandler { let (hit_name, hit_data) = if entries.len() == 1 && self.max_workers > 0 { // A thread pool cannot parallelize one item. Skipping it also lets // Deflate use the faster non-cancellable one-shot path. - let (ok, hit_name, hit_data) = - _inflate_and_hit(&mm, &entries[0], &target_bytes, None, &|msg| self._log(msg)); + let mut scratch = Vec::new(); + let mut decompressor = libdeflater::Decompressor::new(); + let (ok, hit_name, hit_data) = _inflate_and_hit( + &mm, + &entries[0], + &target_bytes, + None, + &mut scratch, + &mut decompressor, + &|msg| self._log(msg), + ); if ok { (hit_name, hit_data) } else { @@ -167,7 +223,11 @@ impl ApkHandler { let stop_event = Arc::clone(&stop_event); let target_bytes = target_bytes.clone(); let debug = self.debug; - handles.push(thread::spawn(move || loop { + handles.push(thread::spawn(move || { + // Reused across entries so decoded pages stay warm. + let mut scratch = Vec::new(); + let mut decompressor = libdeflater::Decompressor::new(); + loop { if stop_event.load(Ordering::SeqCst) { return; } @@ -185,6 +245,8 @@ impl ApkHandler { &entries[index], &target_bytes, Some(&stop_event), + &mut scratch, + &mut decompressor, &|msg| { if debug { println!("{msg}"); @@ -200,6 +262,7 @@ impl ApkHandler { } return; } + } })); } for handle in handles { @@ -221,6 +284,8 @@ impl ApkHandler { hit_name.zip(hit_data) } + /// List class descriptors as individual strings, preserving the row API. + /// CLI consumers can use `for_each_class_chunk` to avoid collecting rows. pub fn list_classes(&self, prefix: Option<&str>) -> Result, String> { if self.max_workers == 0 { return Err("Worker count must be greater than zero".to_owned()); @@ -242,12 +307,25 @@ impl ApkHandler { .as_deref() .filter(|prefix| prefix.is_ascii()) .map(str::as_bytes); - - self._list_dex_items("class", |buf| { - let names = _list_classes(buf, prefix.as_deref(), prefix_bytes)?; - let count = names.len(); - Ok((names, count)) - }) + let mut rows = Vec::new(); + self._stream_dex_items( + "class", + |buf| { + // Keep record boundaries instead of splitting serialized CLI + // lines; the historical API does not validate descriptor syntax. + let names = list_dex_classes(buf, prefix.as_deref(), prefix_bytes)?; + let count = names.len(); + Ok((names.into_iter().map(String::into_bytes).collect(), count)) + }, + Some(&class_names_prefix_len), + |bytes| { + let name = String::from_utf8(bytes) + .map_err(|error| format!("invalid UTF-8 class name: {error}"))?; + rows.push(name); + Ok(()) + }, + )?; + Ok(rows) } /// List defined method signatures in DEX declaration order. @@ -268,6 +346,19 @@ impl ApkHandler { class: Option<&str>, filter: Option<&str>, ) -> Result, String> { + self.list_method_bytes(class, filter).map(|chunks| { + chunks + .into_iter() + .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) + .collect() + }) + } + + pub(crate) fn list_method_bytes( + &self, + class: Option<&str>, + filter: Option<&str>, + ) -> Result>, String> { let class = class .map(str::trim) .map(|name| { @@ -284,72 +375,261 @@ impl ApkHandler { let filter = filter.map(str::to_lowercase); self._list_dex_items("method", |buf| { let (text, count) = render_dex_methods(buf, class.as_deref(), filter.as_deref())?; - Ok((vec![text], count)) + Ok((vec![text.into_bytes()], count)) }) } - fn _list_dex_items(&self, kind: &str, list: F) -> Result, String> + /// Streams per-DEX output in central-directory order through `emit`. + /// + /// Entries are inflated and rendered in parallel, but only one window of + /// rendered buffers is retained at a time: a full list of every method in a + /// 400 MB APK is hundreds of MB of text, and holding it all until the end + /// dominated peak memory. Output order is unchanged (entry order), so the + /// emitted bytes are identical. + fn _stream_dex_items( + &self, + kind: &str, + list: F, + need: Option<&(dyn Fn(&[u8]) -> Result + Sync)>, + mut emit: E, + ) -> Result where - F: Fn(&[u8]) -> Result<(Vec, usize), String> + Sync, + F: Fn(&[u8]) -> Result<(Vec>, usize), String> + Sync, + E: FnMut(Vec) -> Result<(), String>, { if self.max_workers == 0 { return Err("Worker count must be greater than zero".to_owned()); } + let t_start = Instant::now(); let mm = self._open_apk()?; let entries = _parse_entries(&mm); if entries.is_empty() { - return Ok(Vec::new()); + return Ok(0); } let workers = entries.len().min(self.max_workers); + let mut total_count; + + // Bounded reorder pipeline: workers inflate and render out of order, + // the consumer emits in central-directory order and the producer side + // stalls once `window` rendered buffers are waiting. No barriers, so + // parallelism is preserved while retained text stays bounded. + let window = workers.max(1); + let state = Mutex::new(StreamState { + slots: vec![None; entries.len()], + next_emit: 0, + live: workers, + stop: false, + }); + let ready = Condvar::new(); let next = AtomicUsize::new(0); - let results = Mutex::new(vec![None; entries.len()]); - thread::scope(|scope| { + + total_count = thread::scope(|scope| -> Result { for _ in 0..workers { scope.spawn(|| { + let _live = LiveGuard { state: &state, ready: &ready }; + // Reused across entries: keeps the decompressed pages warm. + let mut inflate_buf = Vec::new(); + let mut decompressor = libdeflater::Decompressor::new(); loop { let index = next.fetch_add(1, Ordering::Relaxed); let Some(entry) = entries.get(index) else { break; }; - let result = _load_dex(&mm, entry, None).and_then(|data| { - let data = data.ok_or_else(|| "inflate stopped".to_owned())?; + let mut inflate = |buf: &mut Vec| match need { + Some(need) => { + _inflate_dex_prefix_into(&mm, entry, buf, &mut decompressor, need) + } + None => { + _inflate_dex_into_with(&mm, entry, buf, &mut decompressor) + } + }; + let result = inflate(&mut inflate_buf).and_then(|()| { let mut items = Vec::new(); let mut count = 0; - for (_, dex_buf) in into_logical_dex_buffers(&entry.name, data) { - let (dex_items, dex_count) = list(&dex_buf)?; + let mut offsets = dex041_logical_offsets(&inflate_buf); + if offsets.len() <= 1 { + let (dex_items, dex_count) = list(&inflate_buf)?; items.extend(dex_items); count += dex_count; + } else { + // DEX 041 containers need every logical DEX in + // full; the prefix path cannot be used for them. + offsets.clear(); + _inflate_dex_into_with( + &mm, + entry, + &mut inflate_buf, + &mut decompressor, + )?; + for (_, dex_buf) in + iter_logical_dex_buffers(&entry.name, &inflate_buf) + { + let (dex_items, dex_count) = list(&dex_buf)?; + items.extend(dex_items); + count += dex_count; + } } Ok((items, count)) }); - results.lock().unwrap()[index] = Some(result); + let mut guard = state.lock().unwrap(); + // The consumer is waiting for the lowest unfinished + // index and never blocks producers at that index, so + // this wait can always be released. + while index >= guard.next_emit + window && !guard.stop { + guard = ready.wait(guard).unwrap(); + } + if guard.stop { + return; + } + guard.slots[index] = Some(result); + drop(guard); + ready.notify_all(); } }); } - }); - let mut items = Vec::new(); - let mut total_count = 0; - for (entry, result) in entries.iter().zip(results.into_inner().unwrap()) { - let (dex_items, count) = - result.ok_or_else(|| format!("list{kind} worker did not return a result"))??; - if self.debug { - self._log(&format!("[APK] '{}' {kind}_count={count}", entry.name)); + let mut emitted = 0usize; + for index in 0..entries.len() { + let result = { + let mut guard = state.lock().unwrap(); + while guard.slots[index].is_none() { + if guard.live == 0 { + break; + } + guard = ready.wait(guard).unwrap(); + } + match guard.slots[index].take() { + Some(result) => { + guard.next_emit = index + 1; + ready.notify_all(); + result + } + None => Err(format!( + "list{kind} worker did not return a result for '{}'", + entries[index].name + )), + } + }; + let (dex_items, count) = match result { + Ok(value) => value, + Err(error) => { + // Release blocked producers before unwinding. + state.lock().unwrap().stop = true; + ready.notify_all(); + return Err(error); + } + }; + if self.debug { + self._log(&format!( + "[APK] '{}' {kind}_count={count}", + entries[index].name + )); + } + for item in dex_items { + if let Err(error) = emit(item) { + state.lock().unwrap().stop = true; + ready.notify_all(); + return Err(error); + } + } + emitted += count; } - items.extend(dex_items); - total_count += count; - } + Ok(emitted) + })?; + if self.debug { self._log(&format!( "[APK] list{kind} total={:.2} us count={total_count} workers={workers}", t_start.elapsed().as_secs_f64() * 1_000_000.0, )); } + Ok(total_count) + } + + fn _list_dex_items(&self, kind: &str, list: F) -> Result>, String> + where + F: Fn(&[u8]) -> Result<(Vec>, usize), String> + Sync, + { + let mut items = Vec::new(); + self._stream_dex_items(kind, list, None, |chunk| { + items.push(chunk); + Ok(()) + })?; Ok(items) } + /// Stream class names (one per line) in DEX declaration order. + pub fn for_each_class_chunk(&self, prefix: Option<&str>, mut emit: E) -> Result + where + E: FnMut(&[u8]) -> Result<(), String>, + { + let prefix = prefix + .map(str::trim) + .map(|prefix| { + if prefix.is_empty() { + return Err("Class prefix cannot be empty".to_owned()); + } + Ok(if prefix.starts_with('L') { + prefix.to_owned() + } else { + format!("L{}", prefix.replace('.', "/")) + }) + }) + .transpose()?; + let prefix_bytes = prefix + .as_deref() + .filter(|prefix| prefix.is_ascii()) + .map(str::as_bytes); + + self._stream_dex_items( + "class", + |buf| { + let mut names = Vec::new(); + let count = _write_classes(buf, prefix.as_deref(), prefix_bytes, &mut names)?; + Ok((vec![names], count)) + }, + Some(&class_names_prefix_len), + |chunk| emit(&chunk), + ) + } + + /// Stream rendered method lines in DEX declaration order. + pub fn for_each_method_chunk( + &self, + class: Option<&str>, + filter: Option<&str>, + mut emit: E, + ) -> Result + where + E: FnMut(&[u8]) -> Result<(), String>, + { + let class = class + .map(str::trim) + .map(|name| { + if name.is_empty() { + return Err("Class name cannot be empty".to_owned()); + } + Ok(if name.starts_with('L') && name.ends_with(';') { + name.to_owned() + } else { + format!("L{};", name.replace('.', "/")) + }) + }) + .transpose()?; + let filter = filter.map(str::to_lowercase); + self._stream_dex_items( + "method", + |buf| { + let (text, count) = render_dex_methods(buf, class.as_deref(), filter.as_deref())?; + Ok((vec![text.into_bytes()], count)) + }, + None, + |chunk| emit(&chunk), + ) + } + pub fn for_each_findrefs( &self, find_type: &str, @@ -382,6 +662,9 @@ impl ApkHandler { let find = Arc::clone(&find); handles.push(thread::spawn(move || { let mut results = Vec::new(); + // One warm decompression buffer per worker, reused per entry. + let mut inflate_buf = Vec::new(); + let mut decompressor = libdeflater::Decompressor::new(); loop { let index = { let mut next = work_queue.lock().unwrap(); @@ -392,13 +675,17 @@ impl ApkHandler { *next += 1; index }; - if let Ok(result) = crate::asc_client::apk::worker::_findrefs_worker_from_buf( - &apk_buf, - &entries[index], - &find_type, - &find, - true, - ) { + if let Ok(result) = + crate::asc_client::apk::worker::_findrefs_worker_from_buf_into( + &apk_buf, + &entries[index], + &find_type, + &find, + true, + &mut inflate_buf, + &mut decompressor, + ) + { results.push(result); } } @@ -407,6 +694,8 @@ impl ApkHandler { } let mut final_results = Vec::new(); + let mut inflate_buf = Vec::new(); + let mut decompressor = libdeflater::Decompressor::new(); loop { let index = { let mut next = work_queue.lock().unwrap(); @@ -417,13 +706,17 @@ impl ApkHandler { *next += 1; index }; - if let Ok(result) = crate::asc_client::apk::worker::_findrefs_worker_from_buf( - &apk_buf, - &entries[index], - find_type, - &find, - true, - ) { + if let Ok(result) = + crate::asc_client::apk::worker::_findrefs_worker_from_buf_into( + &apk_buf, + &entries[index], + find_type, + &find, + true, + &mut inflate_buf, + &mut decompressor, + ) + { final_results.push(result); } } diff --git a/src/asc_client/asc_handler.rs b/src/asc_client/asc_handler.rs index 0e4eb48..de26c6d 100644 --- a/src/asc_client/asc_handler.rs +++ b/src/asc_client/asc_handler.rs @@ -88,7 +88,7 @@ impl AscHandler { self._findrefs(dex_name, &dex, find_type, find, aggregate) } - fn _findrefs( + pub(crate) fn _findrefs( &self, dex_name: &str, dex: &DEX, @@ -97,12 +97,18 @@ impl AscHandler { aggregate: bool, ) -> Vec { let mut ref_manager = FindRefManager::new(dex, self.debug); + let profile = std::env::var_os("RASC_PROFILE").is_some(); + let tf0 = profile.then(std::time::Instant::now); let output = ref_manager.find_ref(find, true); + let tf1 = profile.then(std::time::Instant::now); let Some(out) = output else { return Vec::new(); }; + if profile { + eprintln!("[PROF] find_ref={:.3}ms", (tf1.unwrap() - tf0.unwrap()).as_secs_f64() * 1e3); + } let matched_idxs: Vec = out.mark_idx.iter().map(|&x| x as usize).collect(); let mids: &Vec>> = match find_type { diff --git a/src/asc_core/findrefs/findrefs_manager.rs b/src/asc_core/findrefs/findrefs_manager.rs index 5f85c00..a4d8c1e 100644 --- a/src/asc_core/findrefs/findrefs_manager.rs +++ b/src/asc_core/findrefs/findrefs_manager.rs @@ -141,20 +141,35 @@ impl<'a> FindRefManager<'a> { pub fn find_ref(&mut self, find: &FindRefInput, mark: bool) -> Option { let mut has_idx = false; + let profile = std::env::var_os("RASC_PROFILE").is_some(); + let ph0 = profile.then(std::time::Instant::now); // Phase 0: build all required locators first (drop mutable refs before Phase 1) if find.string.is_some() || find.type_.is_some() || find.method.is_some() || find.field.is_some() { let _ = self._get_str_locator(true); } if find.type_.is_some() || find.method.is_some() || find.field.is_some() { - let _ = self._get_type_locator(true); + // Only type queries or fuzzy declaring-class filters use the type + // adjacency map. Exact/unscoped member queries read the id tables. + let fuzzy_method_class = find.method.as_ref().is_some_and(|query| { + query.class.as_ref().is_some_and(|(name, precise)| !precise && !name.is_empty()) + }); + let fuzzy_field_class = find.field.as_ref().is_some_and(|query| { + query.class.as_ref().is_some_and(|(name, precise)| !precise && !name.is_empty()) + }); + let _ = self._get_type_locator( + find.type_.is_some() || fuzzy_method_class || fuzzy_field_class, + ); } if find.method.is_some() { - let _ = self._get_method_locator(true); + let _ = self._get_method_locator(false); } if find.field.is_some() { - let _ = self._get_field_locator(true); + // A find_ref call is a one-shot query over this DEX. Resolve its + // indices directly instead of building maps for every field. + let _ = self._get_field_locator(false); } + let ph1 = profile.then(std::time::Instant::now); // Phase 1: locate all indices (immutable borrows only) let string_result: Option>; let type_result: Option>; @@ -179,7 +194,7 @@ impl<'a> FindRefManager<'a> { let str_loc = self.str_locator.as_ref().unwrap(); let type_loc = self.type_locator.as_ref().unwrap(); let method_loc = self.method_locator.as_ref().unwrap(); - let r = method_loc.locate(m, str_loc, type_loc); + let r = method_loc.locate_unindexed(m, str_loc, type_loc); if !r.is_empty() { has_idx = true; } method_result = Some(r); } else { method_result = None; } @@ -188,11 +203,19 @@ impl<'a> FindRefManager<'a> { let str_loc = self.str_locator.as_ref().unwrap(); let type_loc = self.type_locator.as_ref().unwrap(); let field_loc = self.field_locator.as_ref().unwrap(); - let r = field_loc.locate(f, str_loc, type_loc); + let r = field_loc.locate_unindexed(f, str_loc, type_loc); if !r.is_empty() { has_idx = true; } field_result = Some(r); } else { field_result = None; } + let ph2 = profile.then(std::time::Instant::now); + if profile { + eprintln!( + "[PROF] build_locators={:.3}ms locate_idxs={:.3}ms", + (ph1.unwrap() - ph0.unwrap()).as_secs_f64() * 1e3, + (ph2.unwrap() - ph1.unwrap()).as_secs_f64() * 1e3 + ); + } if find.is_empty() { return None; } @@ -207,7 +230,9 @@ impl<'a> FindRefManager<'a> { } // Phase 2: build scan map and run scanner + let tp0 = profile.then(std::time::Instant::now); let insn_locator = self._ensure_insn_locator(); + let tp1 = profile.then(std::time::Instant::now); let mut scanner = CodeItemScanner::new(insn_locator); let mut scan_map: HashMap>>> = HashMap::new(); @@ -229,6 +254,14 @@ impl<'a> FindRefManager<'a> { } scanner.scan(&mut scan_map, mark, insn_locator); + let tp2 = profile.then(std::time::Instant::now); + if profile { + eprintln!( + "[PROF] insn_parse={:.3}ms scan={:.3}ms", + (tp1.unwrap() - tp0.unwrap()).as_secs_f64() * 1e3, + (tp2.unwrap() - tp1.unwrap()).as_secs_f64() * 1e3 + ); + } let mark_idx = std::mem::take(&mut scanner.mark_idx); let get_or_default = |key: &str| -> Vec>> { diff --git a/src/asc_core/findrefs/locator/field_locator.rs b/src/asc_core/findrefs/locator/field_locator.rs index b26216c..c8142cb 100644 --- a/src/asc_core/findrefs/locator/field_locator.rs +++ b/src/asc_core/findrefs/locator/field_locator.rs @@ -136,6 +136,81 @@ impl<'a> FieldLocator<'a> { ret } + /// Locate a one-shot field query without materializing two maps of sets. + /// Each field row already stores its owner and name indices; selecting rows + /// after resolving the query is cheaper than indexing every field first. + /// The indexed `locate` API remains available for callers reusing a locator. + pub fn locate_unindexed( + &self, + find: &FieldFind, + str_locator: &StringLocator, + type_locator: &TypeLocator, + ) -> HashSet { + let class = find.class.as_ref().filter(|(name, _)| !name.is_empty()); + let field = find.field.as_deref().filter(|name| !name.is_empty()); + if class.is_none() && field.is_none() { + return HashSet::new(); + } + + let precise_class = class.is_some_and(|(_, precise)| *precise); + let class_indices = class.map(|(name, precise)| { + if *precise { + let index = self._find_type_idx_precisely(name); + if index < 0 { + HashSet::new() + } else { + HashSet::from([index as u32]) + } + } else { + type_locator.locate(name, str_locator) + } + }); + if class_indices.as_ref().is_some_and(HashSet::is_empty) { + return HashSet::new(); + } + + // Preserve both established match paths: exact-class queries compare + // decoded member names; unscoped/fuzzy-class substring queries resolve + // matching string indices through StringLocator first. + let name_indices = field + .filter(|_| !find.field_precise && !precise_class) + .map(|name| str_locator.locate(name)); + if name_indices.as_ref().is_some_and(HashSet::is_empty) { + return HashSet::new(); + } + + let (offset, count) = self.base.dex.header.fields; + let mut found = HashSet::new(); + for index in 0..count { + let (owner, _, name_index) = + _STRUCT_HHI::unpack_from(self.base.buf, offset as usize + index as usize * 8); + if class_indices + .as_ref() + .is_some_and(|classes| !classes.contains(&(owner as u32))) + { + continue; + } + if let Some(names) = &name_indices { + if !names.contains(&name_index) { + continue; + } + } else if let Some(wanted) = field { + let member = self.base.dex.get_field(index); + let name = member.name(); + let matches = if find.field_precise { + name == wanted + } else { + name.contains(wanted) + }; + if !matches { + continue; + } + } + found.insert(index); + } + found + } + // find struct: {"class" : ["None|clz", precise], "field" : "None|fuzzy_field"} // the two values cannot both be None // if class not None, input clz must be precise dalvik format value or fuzzy clz diff --git a/src/asc_core/findrefs/locator/method_locator.rs b/src/asc_core/findrefs/locator/method_locator.rs index 0cfe80f..724807e 100644 --- a/src/asc_core/findrefs/locator/method_locator.rs +++ b/src/asc_core/findrefs/locator/method_locator.rs @@ -5,7 +5,7 @@ use crate::asc_core::findrefs::locator::string_locator::StringLocator; use crate::asc_core::findrefs::locator::type_locator::TypeLocator; use crate::asc_core::utils::structs::{_STRUCT_HHI, _STRUCT_I}; use crate::asc_core::utils::tinydex::DEX; -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::time::Instant; #[derive(Clone)] @@ -33,11 +33,18 @@ impl MethodFind { } } +/// Groups `method_id_item` rows two ways (by owning class, by name string) using +/// CSR arrays instead of `HashMap>`. A large DEX holds +/// hundreds of thousands of methods spread over tens of thousands of distinct +/// names, so the map-of-sets variant allocated one hash set per name and paid a +/// hash per insert. Counting sort is two linear passes over the method table. pub struct MethodLocator<'a> { pub base: BaseLocator<'a>, pub parsed: bool, - pub clz_maps: HashMap>, // {type_idx : {method_idx, ...}} - pub method_maps: HashMap>, // {name_idx : {method_idx, ...}} + clz_offsets: Vec, // len = type_ids_size + 1 + clz_mids: Vec, // method idxs grouped by class_idx + name_offsets: Vec, // len = string_ids_size + 1 + name_mids: Vec, // method idxs grouped by name_idx } impl<'a> MethodLocator<'a> { @@ -45,8 +52,10 @@ impl<'a> MethodLocator<'a> { MethodLocator { base: BaseLocator::new(dex), parsed: false, - clz_maps: HashMap::new(), - method_maps: HashMap::new(), + clz_offsets: Vec::new(), + clz_mids: Vec::new(), + name_offsets: Vec::new(), + name_mids: Vec::new(), } } @@ -58,23 +67,73 @@ impl<'a> MethodLocator<'a> { _debug_log("MethodLocator", stage, t_start, count, self.base.debug); } - // table build is very fast, dont worry about performance + fn clz_slice(&self, class_idx: u32) -> &[u32] { + let index = class_idx as usize; + match (self.clz_offsets.get(index), self.clz_offsets.get(index + 1)) { + (Some(&start), Some(&end)) => &self.clz_mids[start as usize..end as usize], + _ => &[], + } + } + + fn name_slice(&self, name_idx: u32) -> &[u32] { + let index = name_idx as usize; + match (self.name_offsets.get(index), self.name_offsets.get(index + 1)) { + (Some(&start), Some(&end)) => &self.name_mids[start as usize..end as usize], + _ => &[], + } + } + pub fn _build_map(&mut self) { if self.parsed { return; } let t_start = if self.base.debug { Some(Instant::now()) } else { None }; let (method_ids_off, method_ids_size) = self.base.dex.header.methods; - let mut method_ids_off = method_ids_off as usize; - let clz_maps = &mut self.clz_maps; - let method_maps = &mut self.method_maps; + let (_, type_ids_size) = self.base.dex.header.types; + let (_, string_ids_size) = self.base.dex.header.strings; let buf = self.base.buf; + let mut method_ids_off = method_ids_off as usize; + + self.clz_offsets = vec![0u32; type_ids_size as usize + 1]; + self.name_offsets = vec![0u32; string_ids_size as usize + 1]; + + // Bounded counters: malformed indices fall outside the table and are + // dropped rather than trusted, matching the previous map behaviour. + for _ in 0..method_ids_size { + let (class_idx, _proto_idx, name_idx) = _STRUCT_HHI::unpack_from(buf, method_ids_off); + method_ids_off += 8; + if let Some(slot) = self.clz_offsets.get_mut(class_idx as usize + 1) { + *slot += 1; + } + if let Some(slot) = self.name_offsets.get_mut(name_idx as usize + 1) { + *slot += 1; + } + } + for index in 1..self.clz_offsets.len() { + self.clz_offsets[index] += self.clz_offsets[index - 1]; + } + for index in 1..self.name_offsets.len() { + self.name_offsets[index] += self.name_offsets[index - 1]; + } + self.clz_mids = vec![0u32; self.clz_offsets.last().copied().unwrap_or(0) as usize]; + self.name_mids = vec![0u32; self.name_offsets.last().copied().unwrap_or(0) as usize]; + let mut clz_cursor = self.clz_offsets.clone(); + let mut name_cursor = self.name_offsets.clone(); + let mut method_ids_off = self.base.dex.header.methods.0 as usize; for method_idx in 0..method_ids_size { - let (class_idx, _type_idx, name_idx) = _STRUCT_HHI::unpack_from(buf, method_ids_off); + let (class_idx, _proto_idx, name_idx) = _STRUCT_HHI::unpack_from(buf, method_ids_off); method_ids_off += 8; - clz_maps.entry(class_idx as u32).or_insert_with(HashSet::new).insert(method_idx); - method_maps.entry(name_idx).or_insert_with(HashSet::new).insert(method_idx); + let clz_index = class_idx as usize; + if let Some(cursor) = clz_cursor.get_mut(clz_index) { + self.clz_mids[*cursor as usize] = method_idx; + *cursor += 1; + } + let name_index = name_idx as usize; + if let Some(cursor) = name_cursor.get_mut(name_index) { + self.name_mids[*cursor as usize] = method_idx; + *cursor += 1; + } } self.parsed = true; self._debug_log("build_map", t_start, Some(method_ids_size as usize)); @@ -106,11 +165,8 @@ impl<'a> MethodLocator<'a> { pub fn _collect_clz_mids(&self, type_idxs: &HashSet) -> HashSet { let mut ret = HashSet::new(); - let clz_maps = &self.clz_maps; for &type_idx in type_idxs { - if let Some(mids) = clz_maps.get(&type_idx) { - ret.extend(mids); - } + ret.extend(self.clz_slice(type_idx)); } ret } @@ -136,6 +192,81 @@ impl<'a> MethodLocator<'a> { ret } + /// Locate a one-shot method query without materializing two maps of sets. + /// Each method row already stores its owner and name indices; selecting rows + /// after resolving the query is cheaper than indexing every method first. + /// The indexed `locate` API remains available for callers reusing a locator. + pub fn locate_unindexed( + &self, + find: &MethodFind, + str_locator: &StringLocator, + type_locator: &TypeLocator, + ) -> HashSet { + let class = find.class.as_ref().filter(|(name, _)| !name.is_empty()); + let method = find.method.as_deref().filter(|name| !name.is_empty()); + if class.is_none() && method.is_none() { + return HashSet::new(); + } + + let precise_class = class.is_some_and(|(_, precise)| *precise); + let class_indices = class.map(|(name, precise)| { + if *precise { + let index = self._find_type_idx_precisely(name); + if index < 0 { + HashSet::new() + } else { + HashSet::from([index as u32]) + } + } else { + type_locator.locate(name, str_locator) + } + }); + if class_indices.as_ref().is_some_and(HashSet::is_empty) { + return HashSet::new(); + } + + // Preserve both established match paths: exact-class queries compare + // decoded member names; unscoped/fuzzy-class substring queries resolve + // matching string indices through StringLocator first. + let name_indices = method + .filter(|_| !find.method_precise && !precise_class) + .map(|name| str_locator.locate(name)); + if name_indices.as_ref().is_some_and(HashSet::is_empty) { + return HashSet::new(); + } + + let (offset, count) = self.base.dex.header.methods; + let mut found = HashSet::new(); + for index in 0..count { + let (owner, _, name_index) = + _STRUCT_HHI::unpack_from(self.base.buf, offset as usize + index as usize * 8); + if class_indices + .as_ref() + .is_some_and(|classes| !classes.contains(&(owner as u32))) + { + continue; + } + if let Some(names) = &name_indices { + if !names.contains(&name_index) { + continue; + } + } else if let Some(wanted) = method { + let member = self.base.dex.get_method(index); + let name = member.name(); + let matches = if find.method_precise { + name == wanted + } else { + name.contains(wanted) + }; + if !matches { + continue; + } + } + found.insert(index); + } + found + } + // find struct: {"class" : ["None|clz", precise], "method" : "None|fuzzy_method"} // the two values cannot both be None // if class not None, input clz must be precise dalvik format value or fuzzy clz @@ -173,12 +304,9 @@ impl<'a> MethodLocator<'a> { return ret; } let name_idxs = str_locator.locate(method.unwrap()); - let method_maps = &self.method_maps; let mut ret = HashSet::new(); for name_idx in name_idxs { - if let Some(mids) = method_maps.get(&name_idx) { - ret.extend(mids); - } + ret.extend(self.name_slice(name_idx)); } self._debug_log("locate", t_start, Some(ret.len())); return ret; @@ -191,26 +319,7 @@ impl<'a> MethodLocator<'a> { self._debug_log("locate", t_start, Some(0)); return HashSet::new(); } - let clz_mids = self.clz_maps.get(&(type_idx as u32)); - match clz_mids { - None => { - self._debug_log("locate", t_start, Some(0)); - return HashSet::new(); - } - Some(mids) => { - let clz_mids_set: HashSet = mids.iter().cloned().collect(); - if method.is_none() { - self._debug_log("locate", t_start, Some(clz_mids_set.len())); - return clz_mids_set; - } - let ret = - self._match_clz_mids(&clz_mids_set, method.unwrap(), method_precise); - self._debug_log("locate", t_start, Some(ret.len())); - return ret; - } - } - } else { - let clz_mids = self._collect_clz_mids(&type_locator.locate(clz_str, str_locator)); + let clz_mids: HashSet = self.clz_slice(type_idx as u32).iter().copied().collect(); if clz_mids.is_empty() { self._debug_log("locate", t_start, Some(0)); return HashSet::new(); @@ -219,23 +328,33 @@ impl<'a> MethodLocator<'a> { self._debug_log("locate", t_start, Some(clz_mids.len())); return clz_mids; } - if method_precise { - return self._match_clz_mids(&clz_mids, method.unwrap(), true); - } - // fuzzy class + fuzzy method - let name_idxs = str_locator.locate(method.unwrap()); - let mut ret = HashSet::new(); - for name_idx in name_idxs { - if let Some(mids) = self.method_maps.get(&name_idx) { - for mid in mids { - if clz_mids.contains(mid) { - ret.insert(*mid); - } - } - } - } + let ret = self._match_clz_mids(&clz_mids, method.unwrap(), method_precise); self._debug_log("locate", t_start, Some(ret.len())); return ret; } + let clz_mids = self._collect_clz_mids(&type_locator.locate(clz_str, str_locator)); + if clz_mids.is_empty() { + self._debug_log("locate", t_start, Some(0)); + return HashSet::new(); + } + if method.is_none() { + self._debug_log("locate", t_start, Some(clz_mids.len())); + return clz_mids; + } + if method_precise { + return self._match_clz_mids(&clz_mids, method.unwrap(), true); + } + // fuzzy class + fuzzy method + let name_idxs = str_locator.locate(method.unwrap()); + let mut ret = HashSet::new(); + for name_idx in name_idxs { + for mid in self.name_slice(name_idx) { + if clz_mids.contains(mid) { + ret.insert(*mid); + } + } + } + self._debug_log("locate", t_start, Some(ret.len())); + ret } } diff --git a/src/asc_core/findrefs/locator/type_locator.rs b/src/asc_core/findrefs/locator/type_locator.rs index 34901d6..7daf668 100644 --- a/src/asc_core/findrefs/locator/type_locator.rs +++ b/src/asc_core/findrefs/locator/type_locator.rs @@ -4,13 +4,17 @@ use crate::asc_core::findrefs::locator::base_locator::{_debug_log, BaseLocator}; use crate::asc_core::findrefs::locator::string_locator::StringLocator; use crate::asc_core::utils::structs::_STRUCT_I; use crate::asc_core::utils::tinydex::DEX; -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::time::Instant; +/// Maps string index -> type idxs holding that descriptor using CSR arrays. +/// `type_ids` is already ordered by string index in a valid DEX, so a single +/// counting pass replaces one hash set per distinct descriptor. pub struct TypeLocator<'a> { pub base: BaseLocator<'a>, pub parsed: bool, - pub type_maps: HashMap>, // {str_idx : {type_idx, ...}} + str_offsets: Vec, // len = string_ids_size + 1 + type_idxs: Vec, // type idxs grouped by descriptor string index } impl<'a> TypeLocator<'a> { @@ -18,7 +22,8 @@ impl<'a> TypeLocator<'a> { TypeLocator { base: BaseLocator::new(dex), parsed: false, - type_maps: HashMap::new(), + str_offsets: Vec::new(), + type_idxs: Vec::new(), } } @@ -30,24 +35,51 @@ impl<'a> TypeLocator<'a> { _debug_log("TypeLocator", stage, t_start, count, self.base.debug); } + fn types_of(&self, str_idx: u32) -> &[u32] { + let index = str_idx as usize; + match (self.str_offsets.get(index), self.str_offsets.get(index + 1)) { + (Some(&start), Some(&end)) => &self.type_idxs[start as usize..end as usize], + _ => &[], + } + } + pub fn _build_map(&mut self) { if self.parsed { return; } let t_start = if self.base.debug { Some(Instant::now()) } else { None }; let (type_ids_off, type_ids_size) = self.base.dex.header.types; - let mut type_ids_off = type_ids_off as usize; - let type_maps = &mut self.type_maps; + let (_, string_ids_size) = self.base.dex.header.strings; let buf = self.base.buf; + let mut type_ids_off = type_ids_off as usize; + + self.str_offsets = vec![0u32; string_ids_size as usize + 1]; + self.type_idxs = Vec::new(); + for _ in 0..type_ids_size { + let str_idx = _STRUCT_I::unpack_from(buf, type_ids_off); + type_ids_off += 4; + if let Some(slot) = self.str_offsets.get_mut(str_idx as usize + 1) { + *slot += 1; + } + } + for index in 1..self.str_offsets.len() { + self.str_offsets[index] += self.str_offsets[index - 1]; + } + self.type_idxs = vec![0u32; self.str_offsets.last().copied().unwrap_or(0) as usize]; + + let mut cursor = self.str_offsets.clone(); + let mut type_ids_off = self.base.dex.header.types.0 as usize; for type_idx in 0..type_ids_size { let str_idx = _STRUCT_I::unpack_from(buf, type_ids_off); type_ids_off += 4; - type_maps.entry(str_idx).or_insert_with(HashSet::new).insert(type_idx); + if let Some(slot) = cursor.get_mut(str_idx as usize) { + self.type_idxs[*slot as usize] = type_idx; + *slot += 1; + } } self.parsed = true; - let count = type_maps.len(); - self._debug_log("build_map", t_start, Some(count)); + self._debug_log("build_map", t_start, Some(type_ids_size as usize)); } pub fn locate(&self, type_str: &str, str_locator: &StringLocator) -> HashSet { @@ -59,13 +91,10 @@ impl<'a> TypeLocator<'a> { let str_idxs = str_locator.locate(type_str); let mut ret = HashSet::new(); - let type_maps = &self.type_maps; for str_idx in str_idxs { - if let Some(types) = type_maps.get(&str_idx) { - ret.extend(types); - } + ret.extend(self.types_of(str_idx)); } self._debug_log("locate", t_start, Some(ret.len())); ret } -} \ No newline at end of file +} diff --git a/src/asc_core/findrefs/scan/code_item_scan.rs b/src/asc_core/findrefs/scan/code_item_scan.rs index 83ba4f9..3b1a07d 100644 --- a/src/asc_core/findrefs/scan/code_item_scan.rs +++ b/src/asc_core/findrefs/scan/code_item_scan.rs @@ -123,6 +123,22 @@ fn _mask_table(values: &[u8]) -> [u8; 256] { table } +/// 65536-bit occupancy bitmap over index values. +fn _idx_bits(idxs: &HashSet) -> Vec { + let mut bits = vec![0u64; (u16::MAX as usize + 1) / 64]; + for &idx in idxs { + if idx <= u16::MAX as u32 { + bits[(idx as usize) >> 6] |= 1u64 << (idx & 63); + } + } + bits +} + +#[inline(always)] +fn _idx_bit(bits: &[u64], idx: usize) -> bool { + bits[idx >> 6] & (1u64 << (idx & 63)) != 0 +} + fn _idx_mask(idxs: &HashSet) -> Vec { let mut table = vec![0u8; u16::MAX as usize + 1]; for &idx in idxs { @@ -218,16 +234,36 @@ impl<'a> CodeItemScanner<'a> { } } _ => { - // >3 opcode values: mask-table loop, still use hi_mask pre-filter for speed + // >3 opcode values: mask-table loop, still use hi_mask pre-filter for speed. + // DEX instructions are 16-bit aligned and every instruction length is a + // multiple of 2 bytes, so opcode bytes only ever appear at even offsets. + // Odd-offset hits can never survive `InsnLocator::locate` (its verifier + // walks instruction boundaries from an even method start), so scanning + // half the positions produces byte-identical results for half the work. let op_mask = _mask_table(opcode_values); - for start in 0..raw_len.saturating_sub(3) { + if idx_width == 2 { + // Test the referenced index first: the query index set is a + // handful of values out of 65536, so the bitmap rejects + // nearly every position and the opcode table is touched + // only on real candidates. Same predicate, same output. + let idx_bits = _idx_bits(idxs); + let mut start = 0usize; + while start + 4 <= raw_len { + let idx = u16::from_le_bytes([raw[start + 2], raw[start + 3]]) as usize; + if _idx_bit(&idx_bits, idx) && op_mask[raw[start] as usize] != 0 { + matched_offset.push((start + off_start) as u32); + if mark { + mark_idx.push(idx as u32); + } + } + start += 2; + } + return; + } + for start in (0..raw_len.saturating_sub(3)).step_by(2) { if op_mask[raw[start] as usize] == 0 { continue; } if start + 2 + idx_width > raw_len { continue; } - let idx = if idx_width == 2 { - u16::from_le_bytes([raw[start + 2], raw[start + 3]]) as u32 - } else { - u32::from_le_bytes([raw[start + 2], raw[start + 3], raw[start + 4], raw[start + 5]]) - }; + let idx = u32::from_le_bytes([raw[start + 2], raw[start + 3], raw[start + 4], raw[start + 5]]); if idx_mask[idx as usize] == 0 { continue; } matched_offset.push((start + off_start) as u32); if mark { mark_idx.push(idx); } @@ -283,8 +319,21 @@ impl<'a> CodeItemScanner<'a> { .flatten() .copied() .collect(); + let profile = std::env::var_os("RASC_PROFILE").is_some(); + let ts0 = profile.then(std::time::Instant::now); let insn_offs = self._scan_code_item(&type_key, &idxs_set, mark); + let ts1 = profile.then(std::time::Instant::now); let mids = insn_locator.locate(&insn_offs); + let ts2 = profile.then(std::time::Instant::now); + if profile { + eprintln!( + "[PROF] {} bytescan={:.3}ms insn_locate={:.3}ms hits={}", + type_key, + (ts1.unwrap() - ts0.unwrap()).as_secs_f64() * 1e3, + (ts2.unwrap() - ts1.unwrap()).as_secs_f64() * 1e3, + insn_offs.len() + ); + } scan.insert(type_key, mids); } } diff --git a/src/asc_core/findrefs/scan/method_id_scan.rs b/src/asc_core/findrefs/scan/method_id_scan.rs index c1ab853..fac4acb 100644 --- a/src/asc_core/findrefs/scan/method_id_scan.rs +++ b/src/asc_core/findrefs/scan/method_id_scan.rs @@ -268,6 +268,7 @@ pub(crate) fn render_dex_methods( }; let class_bytes = class.map(cesu8::to_java_cesu8); let mut text = String::new(); + let mut descriptor = String::new(); let mut rendered_count = 0; for class_index in 0..reader.classes.count { let entry = reader.classes.entry(class_index)?; @@ -277,7 +278,10 @@ pub(crate) fn render_dex_methods( { continue; } - let descriptor = reader.type_name(type_index)?.to_owned(); + // Reuse one buffer: this runs once per class (hundreds of thousands on + // a large DEX) and a fresh `String` per iteration is pure allocator churn. + descriptor.clear(); + descriptor.push_str(reader.type_name(type_index)?); let mut cursor = u32_at(buf, entry + 24)? as usize; if cursor == 0 { continue; diff --git a/src/asc_core/utils/tinydex.rs b/src/asc_core/utils/tinydex.rs index 8cc11fa..82cc772 100644 --- a/src/asc_core/utils/tinydex.rs +++ b/src/asc_core/utils/tinydex.rs @@ -696,6 +696,19 @@ impl DEX { Self::from_vec(buf.to_vec(), name) } + /// Recover the backing buffer so a worker can reuse its pages for the next + /// DEX instead of allocating fresh multi-MB memory per entry. + pub fn into_vec(mut self) -> Vec { + let buf = std::mem::replace(&mut self.buf, std::rc::Rc::new(Vec::new())); + // Dropping the DEX releases the context's second reference, leaving + // `buf` uniquely owned in the common case (no copy). + drop(self); + match std::rc::Rc::try_unwrap(buf) { + Ok(vec) => vec, + Err(shared) => shared.as_ref().clone(), + } + } + pub fn from_vec(buf: Vec, name: &str) -> Self { let buf = Rc::new(buf); let header = DEXHeader::new(&buf); diff --git a/src/cli.rs b/src/cli.rs index 5d42285..c4f3e3d 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -155,17 +155,21 @@ fn _handle_listclass(args: &ArgMatches) -> Result<(), String> { let debug = args.get_flag("debug"); let thread = _get_worker_count(args)?; let prefix = args.get_one::("prefix").map(String::as_str); - let names = ApkHandler::new(apk_path, debug, thread).list_classes(prefix)?; - let text = if names.is_empty() { - String::new() - } else { - names.join("\n") + "\n" - }; - + let handler = ApkHandler::new(apk_path, debug, thread); if let Some(output) = args.get_one::("output") { - fs::write(output, text).map_err(|error| format!("cannot write {output}: {error}"))?; + let mut file = File::create(output) + .map_err(|error| format!("cannot write {output}: {error}"))?; + handler.for_each_class_chunk(prefix, |chunk| { + file.write_all(chunk) + .map_err(|error| format!("cannot write {output}: {error}")) + })?; } else { - print!("{text}"); + let stdout = io::stdout(); + let mut out = stdout.lock(); + handler.for_each_class_chunk(prefix, |chunk| { + out.write_all(chunk) + .map_err(|error| format!("cannot write stdout: {error}")) + })?; } if debug { @@ -184,19 +188,21 @@ fn _handle_listmethod(args: &ArgMatches) -> Result<(), String> { let thread = _get_worker_count(args)?; let class = args.get_one::("class_name").map(String::as_str); let filter = args.get_one::("filter").map(String::as_str); - let chunks = ApkHandler::new(apk_path, debug, thread).list_method_chunks(class, filter)?; - let write_chunks = |writer: &mut dyn Write| { - chunks - .iter() - .try_for_each(|chunk| writer.write_all(chunk.as_bytes())) - }; + let handler = ApkHandler::new(apk_path, debug, thread); if let Some(output) = args.get_one::("output") { - File::create(output) - .and_then(|mut file| write_chunks(&mut file)) + let mut file = File::create(output) .map_err(|error| format!("cannot write {output}: {error}"))?; + handler.for_each_method_chunk(class, filter, |chunk| { + file.write_all(chunk) + .map_err(|error| format!("cannot write {output}: {error}")) + })?; } else { - write_chunks(&mut io::stdout().lock()) - .map_err(|error| format!("cannot write stdout: {error}"))?; + let stdout = io::stdout(); + let mut out = stdout.lock(); + handler.for_each_method_chunk(class, filter, |chunk| { + out.write_all(chunk) + .map_err(|error| format!("cannot write stdout: {error}")) + })?; } if debug { println!( diff --git a/tests/field_query.rs b/tests/field_query.rs new file mode 100644 index 0000000..02157f6 --- /dev/null +++ b/tests/field_query.rs @@ -0,0 +1,132 @@ +use rust_ASC::asc_core::findrefs::locator::field_locator::{FieldFind, FieldLocator}; +use rust_ASC::asc_core::findrefs::locator::string_locator::StringLocator; +use rust_ASC::asc_core::findrefs::locator::type_locator::TypeLocator; +use rust_ASC::asc_core::utils::tinydex::DEX; +use std::collections::HashSet; + +fn put_u32(buf: &mut [u8], offset: usize, value: u32) { + buf[offset..offset + 4].copy_from_slice(&value.to_le_bytes()); +} + +// Only the string/type/field tables are needed for locator queries. Descriptors +// are sorted for exact-class binary search; member names include prefix +// collisions, repeated names across owners, and non-ASCII strings. +fn fixture(reverse_fields: bool, reverse_string_data: bool) -> Vec { + let strings = [ + "FIRST", "I", "Lalpha/A;", "Lalpha/B;", "Lbeta/C;", "SECOND", "SECONDARY", + "Ωname", "名字", + ]; + let string_off = 0x70; + let type_off = string_off + strings.len() * 4; + let fields_off = type_off + 4 * 4; + let mut rows = vec![ + (1u16, 0u16, 0u32), + (1, 0, 5), + (2, 0, 5), + (2, 0, 7), + (3, 0, 6), + (3, 0, 8), + ]; + if reverse_fields { + rows.reverse(); + } + let mut buf = vec![0u8; fields_off + rows.len() * 8]; + buf[..8].copy_from_slice(b"dex\n035\0"); + for (i, (count, offset)) in [ + (strings.len(), string_off), (4, type_off), (0, 0), (rows.len(), fields_off), + (0, 0), (0, 0), + ].into_iter().enumerate() { + put_u32(&mut buf, 0x38 + i * 8, count as u32); + put_u32(&mut buf, 0x3c + i * 8, offset as u32); + } + for (i, string_idx) in [1u32, 2, 3, 4].into_iter().enumerate() { + put_u32(&mut buf, type_off + i * 4, string_idx); + } + for (i, (owner, ty, name)) in rows.into_iter().enumerate() { + let off = fields_off + i * 8; + buf[off..off + 2].copy_from_slice(&owner.to_le_bytes()); + buf[off + 2..off + 4].copy_from_slice(&ty.to_le_bytes()); + put_u32(&mut buf, off + 4, name); + } + let mut order: Vec<_> = (0..strings.len()).collect(); + if reverse_string_data { + order.reverse(); + } + for i in order { + let offset = buf.len(); + put_u32(&mut buf, string_off + i * 4, offset as u32); + let value = strings[i]; + buf.push(value.encode_utf16().count() as u8); + buf.extend_from_slice(value.as_bytes()); + buf.push(0); + } + let length = buf.len() as u32; + put_u32(&mut buf, 0x20, length); + put_u32(&mut buf, 0x24, 0x70); + buf +} + +#[test] +fn direct_field_queries_match_indexed_queries_for_all_match_modes() { + for reverse_fields in [false, true] { + for reverse_strings in [false, true] { + let data = fixture(reverse_fields, reverse_strings); + let dex = DEX::parse(&data, "fields.dex"); + let mut strings = StringLocator::new(&dex); + strings._build_map(); + let mut types = TypeLocator::new(&dex); + types._build_map(); + let direct = FieldLocator::new(&dex); + let mut indexed = FieldLocator::new(&dex); + indexed._build_map(); + + for class in [None, Some(""), Some("Lalpha/A;"), Some("Lalpha/B;"), + Some("Lalpha/"), Some("Lbeta/C;"), Some("Lmissing/Nope;")] + { + for class_precise in [false, true] { + for name in [None, Some(""), Some("FIRST"), Some("SEC"), Some("SECOND"), + Some("SECONDARY"), Some("Ωname"), Some("名字"), Some("absent")] + { + for field_precise in [false, true] { + let query = FieldFind { + class: class.map(|name| (name.to_owned(), class_precise)), + field: name.map(str::to_owned), + field_precise, + }; + assert_eq!( + direct.locate_unindexed(&query, &strings, &types), + indexed.locate(&query, &strings, &types), + "class={class:?}, class_precise={class_precise}, name={name:?}, \ + field_precise={field_precise}, reverse_fields={reverse_fields}, \ + reverse_strings={reverse_strings}", + ); + } + } + } + } + assert!(!direct.parsed, "one-shot queries must not build full maps"); + assert!(direct.clz_maps.is_empty()); + assert!(direct.field_maps.is_empty()); + } + } +} + +#[test] +fn direct_field_queries_preserve_exact_names_and_repeated_owners() { + let data = fixture(false, false); + let dex = DEX::parse(&data, "fields.dex"); + let mut strings = StringLocator::new(&dex); + strings._build_map(); + let mut types = TypeLocator::new(&dex); + types._build_map(); + let locator = FieldLocator::new(&dex); + for (query, expected) in [ + (FieldFind::new_exact(None, Some("SECOND".to_owned())), HashSet::from([1, 2])), + (FieldFind::new_exact(None, Some("SEC".to_owned())), HashSet::new()), + (FieldFind::new_exact(None, Some("名字".to_owned())), HashSet::from([5])), + (FieldFind::new(Some(("Lalpha/".to_owned(), false)), None), HashSet::from([0, 1, 2, 3])), + (FieldFind::new(Some(("Lbeta/C;".to_owned(), true)), Some("SEC".to_owned())), HashSet::from([4])), + ] { + assert_eq!(locator.locate_unindexed(&query, &strings, &types), expected); + } +} diff --git a/tests/listclass_api.rs b/tests/listclass_api.rs new file mode 100644 index 0000000..78e1281 --- /dev/null +++ b/tests/listclass_api.rs @@ -0,0 +1,108 @@ +mod dex_fixture; +mod zip_fixture; + +use dex_fixture::{make_class_only_dex, make_dex041_container}; +use rust_ASC::asc_client::apk_handler::ApkHandler; +use tempfile::tempdir; +use zip_fixture::{CompressionMethod, write_apk}; + +#[test] +fn row_api_returns_strings_in_dex_order_and_preserves_duplicates() { + let dir = tempdir().unwrap(); + let apk = dir.path().join("classes.apk"); + let first = make_class_only_dex(&["Lalpha/First;", "Lbeta/Other;"], 0, b"035"); + let second = make_class_only_dex(&["Lalpha/Second;", "Lalpha/First;"], 0, b"035"); + for compression in [CompressionMethod::Stored, CompressionMethod::Deflated] { + write_apk( + &apk, + &[("classes.dex", &first), ("classes2.dex", &second)], + compression, + ); + for threads in [1, 2] { + let handler = ApkHandler::new(apk.to_str().unwrap(), false, threads); + // The explicit type catches public API changes at compile time. + let rows: Vec = handler.list_classes(None).unwrap(); + assert_eq!( + rows, + [ + "Lalpha/First;", + "Lbeta/Other;", + "Lalpha/Second;", + "Lalpha/First;", + ] + ); + for prefix in ["alpha.", "alpha/", "Lalpha/"] { + let rows: Vec = handler.list_classes(Some(prefix)).unwrap(); + assert_eq!(rows, ["Lalpha/First;", "Lalpha/Second;", "Lalpha/First;"]); + } + assert!(handler.list_classes(Some("missing")).unwrap().is_empty()); + } + } +} + +#[test] +fn row_api_does_not_reparse_names_as_newline_delimited_output() { + // The original row API does not validate descriptor syntax. Preserve each + // string as a record, even if it cannot be safely framed as a CLI line. + let descriptors = [ + "Lalpha/With\nNewline;", + "Lalpha/With\rReturn;", + "Lalpha/名字;", + ]; + let dex = make_class_only_dex(&descriptors, 0, b"035"); + let dir = tempdir().unwrap(); + let apk = dir.path().join("records.apk"); + for compression in [CompressionMethod::Stored, CompressionMethod::Deflated] { + write_apk(&apk, &[("classes.dex", &dex)], compression); + let handler = ApkHandler::new(apk.to_str().unwrap(), false, 1); + let rows: Vec = handler.list_classes(None).unwrap(); + assert_eq!(rows, descriptors); + } +} + +#[test] +fn row_api_lists_every_dex041_member() { + let dex = make_dex041_container(); + let dir = tempdir().unwrap(); + let apk = dir.path().join("container.apk"); + for compression in [CompressionMethod::Stored, CompressionMethod::Deflated] { + write_apk(&apk, &[("classes.dex", &dex)], compression); + let handler = ApkHandler::new(apk.to_str().unwrap(), false, 2); + let rows: Vec = handler.list_classes(None).unwrap(); + assert_eq!( + rows, + [ + "Lalpha/First;", + "Lshared/Duplicate;", + "Lbeta/Second;", + "Lshared/Duplicate;", + ] + ); + assert_eq!( + handler.list_classes(Some("beta")).unwrap(), + ["Lbeta/Second;"] + ); + } +} + +#[test] +fn row_api_preserves_validation_errors() { + let dir = tempdir().unwrap(); + let apk = dir.path().join("classes.apk"); + let dex = make_class_only_dex(&["Lalpha/First;"], 0, b"035"); + write_apk(&apk, &[("classes.dex", &dex)], CompressionMethod::Stored); + assert_eq!( + ApkHandler::new(apk.to_str().unwrap(), false, 0).list_classes(Some(" ")), + Err("Worker count must be greater than zero".to_owned()), + ); + assert_eq!( + ApkHandler::new(apk.to_str().unwrap(), false, 1).list_classes(Some(" ")), + Err("Class prefix cannot be empty".to_owned()), + ); + let missing = dir.path().join("missing.apk"); + assert!( + ApkHandler::new(missing.to_str().unwrap(), false, 1) + .list_classes(None) + .is_err() + ); +} diff --git a/tests/method_query.rs b/tests/method_query.rs new file mode 100644 index 0000000..ec728f2 --- /dev/null +++ b/tests/method_query.rs @@ -0,0 +1,130 @@ +use rust_ASC::asc_core::findrefs::locator::method_locator::{MethodFind, MethodLocator}; +use rust_ASC::asc_core::findrefs::locator::string_locator::StringLocator; +use rust_ASC::asc_core::findrefs::locator::type_locator::TypeLocator; +use rust_ASC::asc_core::utils::tinydex::DEX; +use std::collections::HashSet; + +fn put_u32(buf: &mut [u8], offset: usize, value: u32) { + buf[offset..offset + 4].copy_from_slice(&value.to_le_bytes()); +} + +// Only the string/type/method tables are needed for locator queries. Descriptors +// are sorted for exact-class binary search; member names include prefix +// collisions, repeated names across owners, and non-ASCII strings. +fn fixture(reverse_methods: bool, reverse_string_data: bool) -> Vec { + let strings = [ + "FIRST", "I", "Lalpha/A;", "Lalpha/B;", "Lbeta/C;", "SECOND", "SECONDARY", + "Ωname", "名字", + ]; + let string_off = 0x70; + let type_off = string_off + strings.len() * 4; + let methods_off = type_off + 4 * 4; + let mut rows = vec![ + (1u16, 0u16, 0u32), + (1, 0, 5), + (2, 0, 5), + (2, 0, 7), + (3, 0, 6), + (3, 0, 8), + ]; + if reverse_methods { + rows.reverse(); + } + let mut buf = vec![0u8; methods_off + rows.len() * 8]; + buf[..8].copy_from_slice(b"dex\n035\0"); + for (i, (count, offset)) in [ + (strings.len(), string_off), (4, type_off), (0, 0), (0, 0), + (rows.len(), methods_off), (0, 0), + ].into_iter().enumerate() { + put_u32(&mut buf, 0x38 + i * 8, count as u32); + put_u32(&mut buf, 0x3c + i * 8, offset as u32); + } + for (i, string_idx) in [1u32, 2, 3, 4].into_iter().enumerate() { + put_u32(&mut buf, type_off + i * 4, string_idx); + } + for (i, (owner, ty, name)) in rows.into_iter().enumerate() { + let off = methods_off + i * 8; + buf[off..off + 2].copy_from_slice(&owner.to_le_bytes()); + buf[off + 2..off + 4].copy_from_slice(&ty.to_le_bytes()); + put_u32(&mut buf, off + 4, name); + } + let mut order: Vec<_> = (0..strings.len()).collect(); + if reverse_string_data { + order.reverse(); + } + for i in order { + let offset = buf.len(); + put_u32(&mut buf, string_off + i * 4, offset as u32); + let value = strings[i]; + buf.push(value.encode_utf16().count() as u8); + buf.extend_from_slice(value.as_bytes()); + buf.push(0); + } + let length = buf.len() as u32; + put_u32(&mut buf, 0x20, length); + put_u32(&mut buf, 0x24, 0x70); + buf +} + +#[test] +fn direct_method_queries_match_indexed_queries_for_all_match_modes() { + for reverse_methods in [false, true] { + for reverse_strings in [false, true] { + let data = fixture(reverse_methods, reverse_strings); + let dex = DEX::parse(&data, "methods.dex"); + let mut strings = StringLocator::new(&dex); + strings._build_map(); + let mut types = TypeLocator::new(&dex); + types._build_map(); + let direct = MethodLocator::new(&dex); + let mut indexed = MethodLocator::new(&dex); + indexed._build_map(); + + for class in [None, Some(""), Some("Lalpha/A;"), Some("Lalpha/B;"), + Some("Lalpha/"), Some("Lbeta/C;"), Some("Lmissing/Nope;")] + { + for class_precise in [false, true] { + for name in [None, Some(""), Some("FIRST"), Some("SEC"), Some("SECOND"), + Some("SECONDARY"), Some("Ωname"), Some("名字"), Some("absent")] + { + for method_precise in [false, true] { + let query = MethodFind { + class: class.map(|name| (name.to_owned(), class_precise)), + method: name.map(str::to_owned), + method_precise, + }; + assert_eq!( + direct.locate_unindexed(&query, &strings, &types), + indexed.locate(&query, &strings, &types), + "class={class:?}, class_precise={class_precise}, name={name:?}, \ + method_precise={method_precise}, reverse_methods={reverse_methods}, \ + reverse_strings={reverse_strings}", + ); + } + } + } + } + assert!(!direct.parsed, "one-shot queries must not build full maps"); + } + } +} + +#[test] +fn direct_method_queries_preserve_exact_names_and_repeated_owners() { + let data = fixture(false, false); + let dex = DEX::parse(&data, "methods.dex"); + let mut strings = StringLocator::new(&dex); + strings._build_map(); + let mut types = TypeLocator::new(&dex); + types._build_map(); + let locator = MethodLocator::new(&dex); + for (query, expected) in [ + (MethodFind::new_exact(None, Some("SECOND".to_owned())), HashSet::from([1, 2])), + (MethodFind::new_exact(None, Some("SEC".to_owned())), HashSet::new()), + (MethodFind::new_exact(None, Some("名字".to_owned())), HashSet::from([5])), + (MethodFind::new(Some(("Lalpha/".to_owned(), false)), None), HashSet::from([0, 1, 2, 3])), + (MethodFind::new(Some(("Lbeta/C;".to_owned(), true)), Some("SEC".to_owned())), HashSet::from([4])), + ] { + assert_eq!(locator.locate_unindexed(&query, &strings, &types), expected); + } +} diff --git a/tests/prefix_layout.rs b/tests/prefix_layout.rs new file mode 100644 index 0000000..2a7157d --- /dev/null +++ b/tests/prefix_layout.rs @@ -0,0 +1,80 @@ +mod dex_fixture; +mod zip_fixture; + +use dex_fixture::{get_u32, make_class_only_dex, put_u32}; +use std::process::Command; +use tempfile::tempdir; +use zip_fixture::{write_apk, CompressionMethod}; + +// Relocate class descriptor data without changing its logical content. Large +// incompressible padding exercises both the streaming and full-decode choices. +fn dex_with_strings_at(percent: usize) -> Vec { + let mut dex = make_class_only_dex(&["Lalpha/First;", "Lbeta/Second;"], 0, b"035"); + let strings = get_u32(&dex, 0x3c) as usize; + let data_start = get_u32(&dex, strings) as usize; + let payload = dex[data_start..].to_vec(); + let offset = 4 * 1024 * 1024 * percent / 100; + let old_len = dex.len(); + dex.resize(4 * 1024 * 1024, 0); + let mut state = 0x1234_5678_u32; + for byte in &mut dex[old_len..] { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + *byte = state as u8; + } + dex[offset..offset + payload.len()].copy_from_slice(&payload); + for index in 0..2 { + let old_offset = get_u32(&dex, strings + index * 4) as usize; + put_u32(&mut dex, strings + index * 4, (offset + old_offset - data_start) as u32); + } + let len = dex.len() as u32; + put_u32(&mut dex, 0x20, len); + dex +} + +#[test] +fn descriptor_layout_does_not_change_listing_results() { + let dir = tempdir().unwrap(); + let apk = dir.path().join("layout.apk"); + for percent in [10, 55, 82, 98] { + let dex = dex_with_strings_at(percent); + for compression in [CompressionMethod::Stored, CompressionMethod::Deflated] { + write_apk(&apk, &[("classes.dex", &dex)], compression); + let output = Command::new(env!("CARGO_BIN_EXE_rust_ASC")) + .args(["listclass", "--thread", "1"]) + .arg(&apk) + .output() + .unwrap(); + assert!(output.status.success(), "{}", String::from_utf8_lossy(&output.stderr)); + assert_eq!(output.stdout, b"Lalpha/First;\nLbeta/Second;\n", "layout {percent}%"); + } + } +} + +#[test] +fn a_late_unused_string_does_not_change_class_selection() { + let mut dex = dex_with_strings_at(10); + let strings = get_u32(&dex, 0x3c) as usize; + let old_offset = get_u32(&dex, strings) as usize; + let end = old_offset + 1 + "Lalpha/First;".len() + 1; + let payload = dex[old_offset..end].to_vec(); + let late_offset = dex.len() * 9 / 10; + dex[late_offset..late_offset + payload.len()].copy_from_slice(&payload); + put_u32(&mut dex, strings, late_offset as u32); + // The first string is no longer used by a defined class. It can only be a + // decoder-selection hint, not a required-prefix/correctness assumption. + put_u32(&mut dex, 0x60, 1); + let classes = get_u32(&dex, 0x64) as usize; + put_u32(&mut dex, classes, 1); + let dir = tempdir().unwrap(); + let apk = dir.path().join("unordered.apk"); + write_apk(&apk, &[("classes.dex", &dex)], CompressionMethod::Deflated); + let output = Command::new(env!("CARGO_BIN_EXE_rust_ASC")) + .arg("listclass") + .arg(&apk) + .output() + .unwrap(); + assert!(output.status.success(), "{}", String::from_utf8_lossy(&output.stderr)); + assert_eq!(output.stdout, b"Lbeta/Second;\n"); +}