diff --git a/Cargo.lock b/Cargo.lock index 40e5350..af379aa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "android_system_properties" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" +dependencies = [ + "libc", +] + [[package]] name = "anstream" version = "1.0.0" @@ -90,6 +99,12 @@ version = "1.0.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9a868f137d0dbaade83da5c4700fd9aef2dfe439fac1763fdb925e5054452a65" +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + [[package]] name = "bitflags" version = "2.13.2" @@ -105,6 +120,18 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + [[package]] name = "cc" version = "1.5.1" @@ -127,6 +154,18 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "num-traits", + "serde", + "windows-link", +] + [[package]] name = "clap" version = "4.6.7" @@ -160,6 +199,12 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + [[package]] name = "cpufeatures" version = "0.2.17" @@ -189,6 +234,18 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + +[[package]] +name = "equivalent" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" + [[package]] name = "errno" version = "0.3.14" @@ -220,6 +277,94 @@ dependencies = [ "zlib-rs", ] +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.5", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -241,6 +386,48 @@ dependencies = [ "r-efi", ] +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown", + "serde", + "serde_core", +] + [[package]] name = "is_terminal_polyfill" version = "1.70.2" @@ -253,6 +440,17 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "js-sys" +version = "0.3.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7883d941dae510fb2d978fc3fe018c71c9e2892fd38854de3e8b92c2e5ad9cc5" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -328,6 +526,15 @@ dependencies = [ "libmimalloc-sys", ] +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -340,6 +547,12 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "pastey" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" + [[package]] name = "phf" version = "0.14.0" @@ -392,6 +605,12 @@ dependencies = [ "siphasher", ] +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + [[package]] name = "proc-macro2" version = "1.0.107" @@ -433,6 +652,26 @@ dependencies = [ "mimalloc", ] +[[package]] +name = "ref-cast" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e440fb4e4b4147295338efb76001ab9e4efc0e5839df2c47fc5ac2381d365c3" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ecd8964f8453721699a1ed72037b0db49ce2f5a5138486ee89bed6f67cdf3a" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.5", +] + [[package]] name = "regex" version = "1.13.1" @@ -462,6 +701,27 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "rmcp" +version = "3.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab5023804cbe0304cbe06197ac16e645a0a0d0ffa2953bf735c15c2fa0213350" +dependencies = [ + "chrono", + "futures", + "indexmap", + "pastey", + "pin-project-lite", + "schemars", + "serde", + "serde_json", + "thiserror", + "tokio", + "tokio-util", + "tracing", + "uuid", +] + [[package]] name = "rust_ASC" version = "0.1.0" @@ -479,8 +739,14 @@ dependencies = [ "quick-xml", "radguard", "regex", + "rmcp", + "schemars", + "serde", + "serde_json", "sha1", "tempfile", + "tokio", + "tokio-util", ] [[package]] @@ -496,6 +762,38 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "schemars" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab508826f74a77ca9d5aba6ff19b522583ee3eaf28a19384ff3d0e5835fadf6e" +dependencies = [ + "chrono", + "dyn-clone", + "ref-cast", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1c3a92094fa7d61aa124645844facb6b554dfc797136d0f5fd1f890e2bffc69" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn 3.0.5", +] + [[package]] name = "serde" version = "1.0.229" @@ -526,6 +824,17 @@ dependencies = [ "syn 3.0.5", ] +[[package]] +name = "serde_derive_internals" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f852137cce035d6a4df67ccce505ff6b3e9fd3a10e3e52b24dc71e650bb1a9bd" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.5", +] + [[package]] name = "serde_json" version = "1.0.151" @@ -562,6 +871,12 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + [[package]] name = "strsim" version = "0.11.1" @@ -623,6 +938,74 @@ dependencies = [ "syn 3.0.5", ] +[[package]] +name = "tokio" +version = "1.53.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce3335fa71841cda333a58d7615b03901380ecf09d59b3296d21f8bbac0dde4e" +dependencies = [ + "bytes", + "pin-project-lite", + "tokio-macros", +] + +[[package]] +name = "tokio-macros" +version = "2.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.5", +] + +[[package]] +name = "tokio-util" +version = "0.7.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e464cf451ba96ebfc6f9b6542f17ee8b8956e33f1e40d9690624e59d7a7f8a4b" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "futures-util", + "libc", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + [[package]] name = "typenum" version = "1.20.1" @@ -641,18 +1024,127 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "uuid" +version = "1.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cc1186384beb7dd8eedea376413fd654937285ea6c9cfbb928dc3043ea4b606" +dependencies = [ + "getrandom", + "js-sys", + "wasm-bindgen", +] + [[package]] name = "version_check" version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "wasm-bindgen" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bb54f33acc68fd454578d9820b0bde1a1a3d17aa17bb7b6595806d02886d409" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e29d0c35b16e224a7eeb5cd2d25e3e1968fbd65604117b44d3b789d00ee8535" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f501a8bc3719dba86ef8ae4728879c08001bea749eb1333ac5b91e040e2a6b7" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 3.0.5", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23f0c9c52aa7cd7d77769a4cfe2a9adb1b331f489a41d912ce14513d5ab995c6" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.61.2" diff --git a/Cargo.toml b/Cargo.toml index a843dfe..48c8a3d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,10 @@ version = "0.1.0" edition = "2024" license = "Apache-2.0" +[features] +default = ["mcp"] +mcp = ["dep:rmcp", "dep:tokio", "dep:tokio-util", "dep:serde", "dep:serde_json", "dep:schemars"] + [profile.release] opt-level = 3 # max speed lto = "fat" # full link-time optimization @@ -16,6 +20,12 @@ debug = true opt-level = 0 [dependencies] +rmcp = { version = "3.5.1", optional = true, default-features = false, features = ["server", "transport-io"] } +tokio = { version = "1", optional = true, features = ["rt", "sync", "macros", "time"] } +tokio-util = { version = "0.7", optional = true, features = ["rt"] } +serde = { version = "1", optional = true, features = ["derive"] } +serde_json = { version = "1", optional = true } +schemars = { version = "1", optional = true } # getclass keeps chunked flate2 decompression for cooperative early cancellation. flate2 = { version = "1", default-features = false, features = ["zlib-rs"] } libdeflater = { version = "1.26.1", features = ["use_rust_alloc"] } diff --git a/README.md b/README.md index aae624e..c12d0ea 100644 --- a/README.md +++ b/README.md @@ -20,13 +20,21 @@ https://github.com/user-attachments/assets/f1f71a75-1eca-4647-89ce-406ed2c63c6b - Export pseudocode for every class in an APK with the public API of [MG1937/radguard](https://github.com/MG1937/radguard). - Handle multi-DEX APKs and DEX 041 logical containers. - Use memory-mapped APK access and no on-disk cache. +- Expose four read-only, structured MCP tools over stdio, suitable for Pi CodeMode. ## Build ```bash +git submodule update --init --recursive cargo build --release + +# Build without MCP and its dedicated dependencies. +cargo build --release --no-default-features ``` +MCP is enabled by default, including CI/releases. `--no-default-features` +removes only MCP and its dedicated dependencies; existing CLI commands remain. + The executable is written to: ```text @@ -69,6 +77,43 @@ rust_ASC export --thread 10 --debug app.apk output `export` writes one `.java` pseudocode file per class under package-relative paths in the output directory. Its multi-threaded APK export pipeline is supplied by [MG1937/radguard](https://github.com/MG1937/radguard). With `--debug`, it reports the APK export lifecycle and worker activity. +## MCP + +Read-only MCP tools over stdio, designed for CodeMode: +`getclass`, `listclass`, `getmanifest`, and `findrefs`. + +Configure: + +Ask your agent to install and configure ASC, or follow the manual setup below. + +
+Agent setup instructions + +1. Confirm the target MCP client and user- or project-level configuration scope. +2. Locate an MCP-enabled `rust_ASC` executable. If needed, run + `git submodule update --init --recursive` followed by `cargo install --path . --locked` + from this repository. MCP is enabled by default; the executable is normally installed + to `~/.cargo/bin/rust_ASC`. +3. Register `asc` as a stdio server using the executable's absolute path and + `args: ["mcp"]`. Use the client's native configuration format and preserve existing settings. +4. Reload or restart the client and verify that the four tools listed above are available. + +
+ +For manual setup, replace `/absolute/path/to/rust_ASC` with your executable path +and follow your client's MCP configuration instructions: + +```json +{ + "mcpServers": { + "asc": { + "command": "/absolute/path/to/rust_ASC", + "args": ["mcp"] + } + } +} +``` + ## CLI Benchmark The following measurements use CLI commands with no cache, no preprocessing, and direct searches of every code reference to the string `https://`. diff --git a/src/asc_client/apk/loader.rs b/src/asc_client/apk/loader.rs index e68d459..99d622d 100644 --- a/src/asc_client/apk/loader.rs +++ b/src/asc_client/apk/loader.rs @@ -139,6 +139,63 @@ pub(crate) fn _parse_cd_dex_entries(buf: &[u8]) -> Vec { entries } +/// Strict validation for long-lived query callers. The CLI's fast scanner is +/// retained, but malformed archives must not become successful empty results. +pub(crate) fn checked_dex_entries(buf: &[u8]) -> Result, String> { + let eocd = _find_eocd(buf).ok_or("EOCD not found")?; + if _STRUCT_H::unpack_from(buf, eocd + 4) != 0 + || _STRUCT_H::unpack_from(buf, eocd + 6) != 0 + || _STRUCT_H::unpack_from(buf, eocd + 8) != _STRUCT_H::unpack_from(buf, eocd + 10) + { + return Err("multi-disk ZIP is unsupported".to_owned()); + } + let count = _STRUCT_H::unpack_from(buf, eocd + 10) as usize; + if count == u16::MAX as usize { + return Err("ZIP64 is unsupported".to_owned()); + } + let comment = _STRUCT_H::unpack_from(buf, eocd + 20) as usize; + if eocd + 22 + comment != buf.len() { + return Err("bad ZIP comment length".to_owned()); + } + let start = _STRUCT_I::unpack_from(buf, eocd + 16) as usize; + let end = start + .checked_add(_STRUCT_I::unpack_from(buf, eocd + 12) as usize) + .filter(|&end| end <= eocd) + .ok_or("bad central directory range")?; + let mut pos = start; + let mut entries = Vec::new(); + for _ in 0..count { + if pos.checked_add(46).is_none_or(|next| next > end) || &buf[pos..pos + 4] != _CD_SIG { + return Err("bad central directory entry".to_owned()); + } + let name_len = _STRUCT_H::unpack_from(buf, pos + 28) as usize; + let extra_len = _STRUCT_H::unpack_from(buf, pos + 30) as usize; + let comment_len = _STRUCT_H::unpack_from(buf, pos + 32) as usize; + let next = pos + .checked_add(46 + name_len + extra_len + comment_len) + .filter(|&next| next <= end) + .ok_or("bad central directory entry range")?; + if _STRUCT_H::unpack_from(buf, pos + 8) & 1 != 0 { + return Err("encrypted ZIP entries are unsupported".to_owned()); + } + let name = &buf[pos + 46..pos + 46 + name_len]; + if name.ends_with(_DEX_SUFFIX) && !name.contains(&b'/') { + entries.push(DexEntry { + name: String::from_utf8_lossy(name).into_owned(), + uncomp_size: _STRUCT_I::unpack_from(buf, pos + 24), + comp_size: _STRUCT_I::unpack_from(buf, pos + 20), + local_header_off: _STRUCT_I::unpack_from(buf, pos + 42), + comp_method: _STRUCT_H::unpack_from(buf, pos + 10), + }); + } + pos = next; + } + if pos != end { + return Err("central directory size/count mismatch".to_owned()); + } + Ok(entries) +} + fn _inflate_deflate_chunks(compressed: &[u8], stop_event: &AtomicBool) -> Option> { let mut decoder = DeflateDecoder::new(compressed); let mut out = Vec::with_capacity(compressed.len() * 3); @@ -231,6 +288,14 @@ pub(crate) fn _inflate_dex( } pub(crate) fn _extract_zip_entry(buf: &[u8], target_name: &str) -> Result, String> { + try_extract_zip_entry(buf, target_name)? + .ok_or_else(|| format!("{target_name} not found in APK")) +} + +pub(crate) fn try_extract_zip_entry( + buf: &[u8], + target_name: &str, +) -> Result>, String> { let eocd_off = _find_eocd(buf).ok_or_else(|| "EOCD not found".to_owned())?; let cd_size = _STRUCT_I::unpack_from(buf, eocd_off + 12) as usize; let cd_off = _STRUCT_I::unpack_from(buf, eocd_off + 16) as usize; @@ -262,7 +327,7 @@ pub(crate) fn _extract_zip_entry(buf: &[u8], target_name: &str) -> Result Result, +} + +// Static dispatch keeps structured allocations off the CLI path. +trait ReferenceOutput: Sized { + fn single(dex_name: &str, dex: &DEX, mid: u32, matched: &str) -> Self; + fn aggregate(dex_name: &str, dex: &DEX, mid: u32, matched: Vec) -> Self; +} + +impl ReferenceOutput for String { + fn single(dex_name: &str, dex: &DEX, mid: u32, matched: &str) -> Self { + let method = dex.get_method(mid); + let caller = format!("{}->{}", method.cls().fullname, method.name()); + format!("{} | {} | matched=({})", dex_name, caller, matched) + } + + fn aggregate(dex_name: &str, dex: &DEX, mid: u32, matched: Vec) -> Self { + let joined = matched.join("; "); + // Release target buffers before formatting, as in the legacy CLI path. + drop(matched); + Self::single(dex_name, dex, mid, &joined) + } +} + +impl ReferenceOutput for Reference { + fn single(dex_name: &str, dex: &DEX, mid: u32, matched: &str) -> Self { + Self::aggregate(dex_name, dex, mid, vec![matched.to_owned()]) + } + + fn aggregate(dex_name: &str, dex: &DEX, mid: u32, matched: Vec) -> Self { + let method = dex.get_method(mid); + Self { + dex_name: dex_name.to_owned(), + class_name: method.cls().fullname, + method_name: method.name().to_owned(), + matched, + } + } +} /// Core orchestrator integrating DexManager, FindRefManager and decompiler pub struct AscHandler { @@ -42,9 +85,13 @@ impl AscHandler { decompile_dex_bytes_with_metadata(&new_dex_bytes, dalvik_class) } - fn _format_method(&self, dex: &DEX, midx: usize) -> String { - let method = dex.get_method(midx as u32); - format!("{}->{}", method.cls().fullname, method.name()) + /// Unlike the legacy text helper, decompiler failures remain errors. + pub fn try_getclass(&self, dex_buf: &[u8], dalvik_class: &str) -> Result { + let mut manager = DexManager::from_bytes(dex_buf.to_vec(), self.debug); + let new_dex_bytes = manager.extract_and_rebuild(dalvik_class); + radguard::decompile_dex_bytes(&new_dex_bytes, dalvik_class) + .map(|output| output.source) + .map_err(|error| error.to_string()) } fn _format_matched_name(&self, dex: &DEX, find_type: &str, idx: usize) -> String { @@ -73,7 +120,7 @@ impl AscHandler { aggregate: bool, ) -> Vec { let dex = DEX::parse(dex_buf, dex_name); - self._findrefs(dex_name, &dex, find_type, find, aggregate) + self.findrefs_output::(dex_name, &dex, find_type, find, aggregate) } pub fn findrefs_vec( @@ -85,17 +132,28 @@ impl AscHandler { aggregate: bool, ) -> Vec { let dex = DEX::from_vec(dex_buf, dex_name); - self._findrefs(dex_name, &dex, find_type, find, aggregate) + self.findrefs_output::(dex_name, &dex, find_type, find, aggregate) + } + + pub fn findrefs_records( + &self, + dex_name: &str, + dex_buf: Vec, + find_type: &str, + find: &FindRefInput, + ) -> Vec { + let dex = DEX::from_vec(dex_buf, dex_name); + self.findrefs_output::(dex_name, &dex, find_type, find, true) } - fn _findrefs( + fn findrefs_output( &self, dex_name: &str, dex: &DEX, find_type: &str, find: &FindRefInput, aggregate: bool, - ) -> Vec { + ) -> Vec { let mut ref_manager = FindRefManager::new(dex, self.debug); let output = ref_manager.find_ref(find, true); @@ -121,15 +179,12 @@ impl AscHandler { let mut ret = Vec::new(); for (i, mid_entry) in mids.iter().enumerate() { let idx = matched_idxs.get(i).copied().unwrap_or(0); - let matched = self._format_matched_name(&dex, find_type, idx); - let Some(mid_list) = mid_entry else { continue; }; + let matched = self._format_matched_name(dex, find_type, idx); + let Some(mid_list) = mid_entry else { + continue; + }; for &mid in mid_list { - ret.push(format!( - "{} | {} | matched=({})", - dex_name, - self._format_method(&dex, mid as usize), - matched - )); + ret.push(R::single(dex_name, dex, mid, &matched)); } } return ret; @@ -139,7 +194,9 @@ impl AscHandler { let mut grouped: HashMap> = HashMap::new(); for (i, mid_entry) in mids.iter().enumerate() { let idx = matched_idxs.get(i).copied().unwrap_or(0); - let Some(mid_list) = mid_entry else { continue; }; + let Some(mid_list) = mid_entry else { + continue; + }; for &mid in mid_list { grouped.entry(mid as usize).or_default().insert(idx); } @@ -154,15 +211,9 @@ impl AscHandler { sorted_idxs.sort(); let matched = sorted_idxs .iter() - .map(|&&idx| self._format_matched_name(&dex, find_type, idx)) - .collect::>() - .join("; "); - ret.push(format!( - "{} | {} | matched=({})", - dex_name, - self._format_method(&dex, *mid), - matched - )); + .map(|&&idx| self._format_matched_name(dex, find_type, idx)) + .collect::>(); + ret.push(R::aggregate(dex_name, dex, *mid as u32, matched)); } ret } diff --git a/src/asc_client/manifest_handler.rs b/src/asc_client/manifest_handler.rs index 833deab..04aab5c 100644 --- a/src/asc_client/manifest_handler.rs +++ b/src/asc_client/manifest_handler.rs @@ -46,8 +46,13 @@ pub fn get_manifest_xml(apk_path: &str, pretty: bool) -> Result let apk = fs::read(apk_path).map_err(|e| format!("cannot open APK: {}", e))?; let manifest_data = _extract_zip_entry(&apk, "AndroidManifest.xml")?; + decode_manifest_xml(&manifest_data, pretty) +} + +/// Decode already-extracted data so query callers can retain typed file errors. +pub(crate) fn decode_manifest_xml(manifest_data: &[u8], pretty: bool) -> Result { _quiet_loguru(); - let mut input = manifest_data.as_slice(); + let mut input = manifest_data; let axml = AXML::new(&mut input, None) .map_err(|e| format!("AndroidManifest.xml is not a valid binary XML: {}", e))?; _format_xml(&axml.get_xml_string(), pretty) diff --git a/src/asc_client/mod.rs b/src/asc_client/mod.rs index 1ff390c..cdf138f 100644 --- a/src/asc_client/mod.rs +++ b/src/asc_client/mod.rs @@ -3,3 +3,5 @@ pub mod apk_handler; pub mod asc_handler; pub mod dex_container; pub mod manifest_handler; +pub mod query; +pub(crate) mod query_input; diff --git a/src/asc_client/query.rs b/src/asc_client/query.rs new file mode 100644 index 0000000..2711a2b --- /dev/null +++ b/src/asc_client/query.rs @@ -0,0 +1,354 @@ +//! Request-local, fallible APK queries. +use std::fmt; +use std::fs::File; +use std::io; +use std::path::Path; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +use cesu8::to_java_cesu8; +use memmap2::{Mmap, MmapOptions}; + +use super::apk::dex_reader::{_dex_defines_class, list_dex_classes}; +use super::apk::loader::{_inflate_dex, DexEntry, checked_dex_entries, try_extract_zip_entry}; +use super::asc_handler::{AscHandler, Reference}; +use super::dex_container::into_logical_dex_buffers; +use super::manifest_handler::decode_manifest_xml; +use super::query_input::_format_class_name; +use crate::asc_core::findrefs::findrefs_manager::FindRefInput; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum QueryErrorCode { + InvalidArgument, + FileNotFound, + AccessDenied, + IoError, + InvalidFormat, + NotFound, + QueryFailed, + Cancelled, +} + +impl QueryErrorCode { + pub fn as_str(self) -> &'static str { + match self { + Self::InvalidArgument => "INVALID_ARGUMENT", + Self::FileNotFound => "FILE_NOT_FOUND", + Self::AccessDenied => "ACCESS_DENIED", + Self::IoError => "IO_ERROR", + Self::InvalidFormat => "INVALID_FORMAT", + Self::NotFound => "NOT_FOUND", + Self::QueryFailed => "QUERY_FAILED", + Self::Cancelled => "CANCELLED", + } + } +} + +#[derive(Debug)] +pub struct QueryError { + pub code: QueryErrorCode, + pub message: String, +} + +impl QueryError { + pub fn new(code: QueryErrorCode, message: impl Into) -> Self { + Self { + code, + message: message.into(), + } + } + + pub fn invalid(message: impl Into) -> Self { + Self::new(QueryErrorCode::InvalidArgument, message) + } + + fn format(message: impl Into) -> Self { + Self::new(QueryErrorCode::InvalidFormat, message) + } + + fn io(error: io::Error) -> Self { + let code = match error.kind() { + io::ErrorKind::NotFound => QueryErrorCode::FileNotFound, + io::ErrorKind::PermissionDenied => QueryErrorCode::AccessDenied, + _ => QueryErrorCode::IoError, + }; + Self::new(code, error.to_string()) + } +} + +impl fmt::Display for QueryError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}: {}", self.code.as_str(), self.message) + } +} +impl std::error::Error for QueryError {} + +#[derive(Debug)] +#[cfg_attr(feature = "mcp", derive(serde::Serialize, schemars::JsonSchema))] +pub struct ClassSource { + pub class_name: String, + pub dex_name: String, + pub source: String, +} + +/// The mapped input must remain unchanged until the query completes. +pub struct ApkQuery { + data: Mmap, + entries: Vec, + workers: usize, +} + +impl ApkQuery { + pub fn new(path: &str, workers: usize) -> Result { + if !Path::new(path).is_absolute() { + return Err(QueryError::invalid("apk_path must be an absolute path")); + } + if workers == 0 { + return Err(QueryError::invalid("worker count must be positive")); + } + // Reject FIFOs before open, which could otherwise block. + let metadata = std::fs::metadata(path).map_err(QueryError::io)?; + if !metadata.is_file() { + return Err(QueryError::invalid("apk_path must identify a regular file")); + } + if metadata.len() == 0 { + return Err(QueryError::format("empty APK")); + } + let file = File::open(path).map_err(QueryError::io)?; + // SAFETY: callers must not modify the mapped file. + let data = unsafe { MmapOptions::new().map(&file) }.map_err(QueryError::io)?; + let entries = checked_dex_entries(&data).map_err(QueryError::format)?; + Ok(Self { + data, + entries, + workers, + }) + } + + // Preserve archive order; only single-class lookup may stop early. + fn scan( + &self, + first_hit: bool, + visit: impl Fn(&DexEntry) -> Result, QueryError> + Sync, + ) -> Result, QueryError> { + let next = AtomicUsize::new(0); + let stopped = AtomicBool::new(false); + std::thread::scope(|scope| { + let mut handles = Vec::new(); + for _ in 0..self.workers.min(self.entries.len()) { + let visit = &visit; + let next = &next; + let stopped = &stopped; + handles.push(scope.spawn(move || { + let mut results = Vec::new(); + while !stopped.load(Ordering::Relaxed) { + let index = next.fetch_add(1, Ordering::Relaxed); + let Some(entry) = self.entries.get(index) else { + break; + }; + let result = visit(entry); + if first_hit && matches!(result, Ok(Some(_))) { + stopped.store(true, Ordering::Relaxed); + } + results.push((index, result)); + } + results + })); + } + let mut results = Vec::new(); + let mut failed = false; + for handle in handles { + match handle.join() { + Ok(items) => results.extend(items), + Err(_) => failed = true, + } + } + if failed { + return Err(QueryError::new( + QueryErrorCode::QueryFailed, + "query worker panicked", + )); + } + results.sort_by_key(|(index, _)| *index); + let mut output = Vec::new(); + for (_, result) in results { + if let Some(value) = result? { + output.push(value); + } + } + Ok(output) + }) + } + + fn dex_buffers(&self, entry: &DexEntry) -> Result)>, QueryError> { + let data = _inflate_dex(&self.data, entry, None) + .map_err(QueryError::format)? + .ok_or_else(|| QueryError::format("unexpected interrupted inflate"))?; + validate_dex(&data)?; + // Reject broken containers before the legacy single-DEX fallback. + if data.starts_with(b"dex\n041\0") { + let mut offset = 0; + while offset < data.len() { + let member = &data[offset..]; + validate_dex_header(member)?; + if !member.starts_with(b"dex\n041\0") { + return Err(QueryError::format("invalid DEX 041 member")); + } + offset += u32_at(member, 0x20) as usize; + } + } + let buffers = into_logical_dex_buffers(&entry.name, data); + for (_, data) in &buffers { + validate_dex(data)?; + } + Ok(buffers) + } + + pub fn getclass(&self, class_name: &str) -> Result { + let class_name = _format_class_name(class_name).map_err(QueryError::invalid)?; + let target = to_java_cesu8(&class_name); + let hits = self.scan(true, |entry| { + for (dex_name, data) in self.dex_buffers(entry)? { + if _dex_defines_class(&data, &target).map_err(QueryError::format)? { + return Ok(Some((dex_name, data))); + } + } + Ok(None) + })?; + let (dex_name, data) = hits.into_iter().next().ok_or_else(|| { + QueryError::new( + QueryErrorCode::NotFound, + format!("class {class_name} not found"), + ) + })?; + let source = AscHandler::new(false) + .try_getclass(&data, &class_name) + .map_err(|message| QueryError::new(QueryErrorCode::QueryFailed, message))?; + Ok(ClassSource { + class_name, + dex_name, + source, + }) + } + + pub fn listclass(&self, prefix: Option<&str>) -> Result, QueryError> { + let prefix = prefix + .map(str::trim) + .map(|prefix| { + if prefix.is_empty() { + return Err(QueryError::invalid("class prefix cannot be empty")); + } + Ok(if prefix.starts_with('L') { + prefix.to_owned() + } else { + format!("L{}", prefix.replace('.', "/")) + }) + }) + .transpose()?; + let bytes = prefix + .as_deref() + .filter(|p| p.is_ascii()) + .map(str::as_bytes); + let groups = self.scan(false, |entry| { + let mut names = Vec::new(); + for (_, data) in self.dex_buffers(entry)? { + names.extend( + list_dex_classes(&data, prefix.as_deref(), bytes) + .map_err(QueryError::format)?, + ); + } + Ok(Some(names)) + })?; + Ok(groups.into_iter().flatten().collect()) + } + + pub fn findrefs(&self, kind: &str, find: &FindRefInput) -> Result, QueryError> { + if !matches!(kind, "string" | "type" | "method" | "field") { + return Err(QueryError::invalid("unknown reference kind")); + } + let groups = self.scan(false, |entry| { + let mut records = Vec::new(); + for (name, data) in self.dex_buffers(entry)? { + records.extend(AscHandler::new(false).findrefs_records(&name, data, kind, find)); + } + Ok(Some(records)) + })?; + let mut records: Vec<_> = groups.into_iter().flatten().collect(); + records.sort_by(|a, b| a.dex_name.cmp(&b.dex_name)); + Ok(records) + } + + pub fn getmanifest(&self) -> Result { + let data = try_extract_zip_entry(&self.data, "AndroidManifest.xml") + .map_err(QueryError::format)? + .ok_or_else(|| { + QueryError::new(QueryErrorCode::NotFound, "AndroidManifest.xml not found") + })?; + decode_manifest_xml(&data, true).map_err(QueryError::format) + } +} + +fn u32_at(data: &[u8], offset: usize) -> u32 { + u32::from_le_bytes(data[offset..offset + 4].try_into().unwrap()) +} + +fn validate_dex_header(data: &[u8]) -> Result<(), QueryError> { + if data.len() < 0x70 + || !data.starts_with(b"dex\n") + || data[7] != 0 + || !data[4..7].iter().all(u8::is_ascii_digit) + { + return Err(QueryError::format("invalid DEX header")); + } + let size = u32_at(data, 0x20) as usize; + let header = u32_at(data, 0x24) as usize; + let expected = if &data[4..7] == b"041" { 0x78 } else { 0x70 }; + if header != expected || size < header || size > data.len() || u32_at(data, 0x28) != 0x1234_5678 + { + return Err(QueryError::format("invalid DEX size/header/endianness")); + } + Ok(()) +} + +fn validate_dex(data: &[u8]) -> Result<(), QueryError> { + validate_dex_header(data)?; + for (index, width) in [4usize, 4, 12, 8, 8, 32].into_iter().enumerate() { + let count = u32_at(data, 0x38 + index * 8) as usize; + let start = u32_at(data, 0x3c + index * 8) as usize; + if count != 0 + && (start == 0 + || count + .checked_mul(width) + .and_then(|size| start.checked_add(size)) + .is_none_or(|end| end > data.len())) + { + return Err(QueryError::format("invalid DEX table range")); + } + } + // Reject malformed strings instead of the legacy locators' empty fallback. + let strings = u32_at(data, 0x38) as usize; + let string_start = u32_at(data, 0x3c) as usize; + for index in 0..strings { + let offset = u32_at(data, string_start + index * 4) as usize; + let tail = data + .get(offset..) + .filter(|_| offset != 0) + .ok_or_else(|| QueryError::format("invalid DEX string offset"))?; + let prefix = tail + .iter() + .take(5) + .position(|byte| byte & 0x80 == 0) + .ok_or_else(|| QueryError::format("invalid DEX string length"))? + + 1; + if memchr::memchr(0, &tail[prefix..]).is_none() { + return Err(QueryError::format("unterminated DEX string")); + } + } + let types = u32_at(data, 0x40) as usize; + let type_start = u32_at(data, 0x44) as usize; + for index in 0..types { + if u32_at(data, type_start + index * 4) as usize >= strings { + return Err(QueryError::format("invalid DEX type string index")); + } + } + Ok(()) +} diff --git a/src/asc_client/query_input.rs b/src/asc_client/query_input.rs new file mode 100644 index 0000000..8a724de --- /dev/null +++ b/src/asc_client/query_input.rs @@ -0,0 +1,66 @@ +use crate::asc_core::findrefs::findrefs_manager::FindRefInput; +use crate::asc_core::findrefs::locator::field_locator::FieldFind; +use crate::asc_core::findrefs::locator::method_locator::MethodFind; + +pub(crate) fn _format_class_name(name: &str) -> Result { + if name.is_empty() { + return Err("Class name cannot be empty".to_owned()); + } + if name.starts_with('L') && name.ends_with(';') && name.contains('/') { + return Ok(name.to_owned()); + } + let mut name = name.replace('.', "/"); + if !name.starts_with('L') { + name = format!("L{name}"); + } + if !name.ends_with(';') { + name.push(';'); + } + Ok(name) +} + +fn _normalize_class_query(name: Option<&str>, fuzzy: bool) -> Result, String> { + let Some(name) = name else { + return Ok(None); + }; + if fuzzy { + return Ok(Some(if name.contains('.') && !name.contains('/') { + name.replace('.', "/") + } else { + name.to_owned() + })); + } + _format_class_name(name).map(Some) +} + +pub(crate) fn _build_member_find( + key: &str, + clz: Option<&str>, + clz_fuzzy: bool, + name: Option<&str>, +) -> Result { + let clz = clz.filter(|value| !value.is_empty()); + let name = name.filter(|value| !value.is_empty()); + if clz.is_none() && name.is_none() { + return Err(format!( + "{key} query needs at least one of class or {key} name" + )); + } + + let class = _normalize_class_query(clz, clz_fuzzy)?.map(|value| (value, !clz_fuzzy)); + match key { + "method" => Ok(FindRefInput { + string: None, + type_: None, + method: Some(MethodFind::new(class, name.map(str::to_owned))), + field: None, + }), + "field" => Ok(FindRefInput { + string: None, + type_: None, + method: None, + field: Some(FieldFind::new(class, name.map(str::to_owned))), + }), + _ => Err(format!("Unknown member type: {key}")), + } +} diff --git a/src/cli.rs b/src/cli.rs index 42b7c9b..13a0f8f 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -9,72 +9,8 @@ use radguard::cli::{ExportOptions, ExportProgress, export_to_directory}; use crate::asc_client::apk_handler::ApkHandler; use crate::asc_client::asc_handler::AscHandler; use crate::asc_client::manifest_handler::get_manifest_xml; +use crate::asc_client::query_input::{_build_member_find, _format_class_name}; use crate::asc_core::findrefs::findrefs_manager::FindRefInput; -use crate::asc_core::findrefs::locator::field_locator::FieldFind; -use crate::asc_core::findrefs::locator::method_locator::MethodFind; - -fn _format_class_name(name: &str) -> Result { - if name.is_empty() { - return Err("Class name cannot be empty".to_owned()); - } - if name.starts_with('L') && name.ends_with(';') && name.contains('/') { - return Ok(name.to_owned()); - } - let mut name = name.replace('.', "/"); - if !name.starts_with('L') { - name = format!("L{name}"); - } - if !name.ends_with(';') { - name.push(';'); - } - Ok(name) -} - -fn _normalize_class_query(name: Option<&str>, fuzzy: bool) -> Result, String> { - let Some(name) = name else { - return Ok(None); - }; - if fuzzy { - return Ok(Some(if name.contains('.') && !name.contains('/') { - name.replace('.', "/") - } else { - name.to_owned() - })); - } - _format_class_name(name).map(Some) -} - -fn _build_member_find( - key: &str, - clz: Option<&str>, - clz_fuzzy: bool, - name: Option<&str>, -) -> Result { - let clz = clz.filter(|value| !value.is_empty()); - let name = name.filter(|value| !value.is_empty()); - if clz.is_none() && name.is_none() { - return Err(format!( - "{key} query needs at least one of class or {key} name" - )); - } - - let class = _normalize_class_query(clz, clz_fuzzy)?.map(|value| (value, !clz_fuzzy)); - match key { - "method" => Ok(FindRefInput { - string: None, - type_: None, - method: Some(MethodFind::new(class, name.map(str::to_owned))), - field: None, - }), - "field" => Ok(FindRefInput { - string: None, - type_: None, - method: None, - field: Some(FieldFind::new(class, name.map(str::to_owned))), - }), - _ => Err(format!("Unknown member type: {key}")), - } -} fn _get_find_query(args: &ArgMatches) -> Result<(String, FindRefInput), String> { let (find_type, find_args) = args @@ -391,7 +327,7 @@ pub fn _build_main_parser() -> Command { )), ); - Command::new("rust_ASC") + let command = Command::new("rust_ASC") .about("ASC tooling entry.") .after_help( "examples:\n rust_ASC getclass app.apk Lcom/poc/Main; -o Main.java\n \ @@ -466,7 +402,11 @@ pub fn _build_main_parser() -> Command { ) .arg(Arg::new("apk_path").required(true).help("Path to the input APK file.")), ) - .subcommand(findrefs) + .subcommand(findrefs); + #[cfg(feature = "mcp")] + let command = command + .subcommand(Command::new("mcp").about("Serve read-only MCP tools over stdio.")); + command } pub fn run_from(args: I) -> Result<(), String> @@ -478,6 +418,8 @@ where .try_get_matches_from(args) .map_err(|error| error.to_string())?; match matches.subcommand() { + #[cfg(feature = "mcp")] + Some(("mcp", _)) => crate::mcp::run(), Some(("getclass", args)) => _handle_getclass(args), Some(("listclass", args)) => _handle_listclass(args), Some(("export", args)) => _handle_export(args), diff --git a/src/lib.rs b/src/lib.rs index 127b77d..d1dbeda 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,3 +3,5 @@ pub mod asc_client; pub mod asc_core; pub mod cli; +#[cfg(feature = "mcp")] +pub mod mcp; diff --git a/src/mcp/executor.rs b/src/mcp/executor.rs new file mode 100644 index 0000000..1367a57 --- /dev/null +++ b/src/mcp/executor.rs @@ -0,0 +1,175 @@ +use std::sync::Arc; +use tokio::sync::Semaphore; +use tokio_util::sync::CancellationToken; + +use crate::asc_client::query::{QueryError, QueryErrorCode}; + +#[derive(Clone)] +pub(super) struct Executor { + permits: Arc, + threads: usize, + shutdown: CancellationToken, +} + +impl Executor { + pub(super) fn new(parallelism: usize) -> Self { + let parallelism = parallelism.max(1); + let concurrency = parallelism.min(2); + Self { + permits: Arc::new(Semaphore::new(concurrency)), + threads: (parallelism / concurrency).max(1), + shutdown: CancellationToken::new(), + } + } + + pub(super) fn shutdown_token(&self) -> CancellationToken { + self.shutdown.clone() + } + + pub(super) async fn run( + &self, + cancelled: CancellationToken, + work: impl FnOnce(usize) -> Result + Send + 'static, + ) -> Result { + let permit = tokio::select! { + biased; + _ = self.shutdown.cancelled() => return Err(cancelled_error()), + _ = cancelled.cancelled() => return Err(cancelled_error()), + permit = self.permits.clone().acquire_owned() => permit.map_err(|_| cancelled_error())?, + }; + let threads = self.threads; + let work_cancelled = cancelled.clone(); + let shutdown = self.shutdown.clone(); + let task = tokio::task::spawn_blocking(move || { + // Cancellation must not release capacity while CPU work continues. + let _permit = permit; + if work_cancelled.is_cancelled() || shutdown.is_cancelled() { + return Err(cancelled_error()); + } + work(threads) + }); + tokio::select! { + biased; + _ = self.shutdown.cancelled() => Err(cancelled_error()), + _ = cancelled.cancelled() => Err(cancelled_error()), + result = task => result.map_err(|error| + QueryError::new(QueryErrorCode::QueryFailed, error.to_string()))?, + } + } +} + +fn cancelled_error() -> QueryError { + QueryError::new(QueryErrorCode::Cancelled, "query cancelled") +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + use tokio::sync::oneshot; + + #[test] + fn automatic_budget() { + for (p, c, t) in [(0, 1, 1), (1, 1, 1), (2, 2, 1), (3, 2, 1), (16, 2, 8)] { + let executor = Executor::new(p); + assert_eq!(executor.permits.available_permits(), c); + assert_eq!(executor.threads, t); + } + } + + #[tokio::test] + async fn running_cancellation_keeps_its_permit_and_queued_cancellation_never_runs() { + let executor = Executor::new(1); + let first_cancel = CancellationToken::new(); + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let first = tokio::spawn({ + let executor = executor.clone(); + let token = first_cancel.clone(); + async move { + executor + .run(token, move |_| { + let _ = started_tx.send(()); + let _ = release_rx.recv(); + Ok(()) + }) + .await + } + }); + tokio::time::timeout(Duration::from_secs(5), started_rx) + .await + .unwrap() + .unwrap(); + first_cancel.cancel(); + assert_eq!( + first.await.unwrap().unwrap_err().code, + QueryErrorCode::Cancelled + ); + assert_eq!(executor.permits.available_permits(), 0); + + let queued_cancel = CancellationToken::new(); + let queued = executor.run::<()>(queued_cancel.clone(), |_| panic!("queued query ran")); + tokio::pin!(queued); + // Poll once to ensure this request really entered the semaphore queue. + assert!(futures_poll_pending(queued.as_mut()).await); + queued_cancel.cancel(); + assert_eq!(queued.await.unwrap_err().code, QueryErrorCode::Cancelled); + release_tx.send(()).unwrap(); + tokio::time::timeout( + Duration::from_secs(5), + executor.run(CancellationToken::new(), |_| Ok(())), + ) + .await + .unwrap() + .unwrap(); + } + + async fn futures_poll_pending( + mut future: std::pin::Pin<&mut F>, + ) -> bool { + std::future::poll_fn(|cx| std::task::Poll::Ready(future.as_mut().poll(cx).is_pending())) + .await + } + + #[tokio::test] + async fn two_queries_can_run_but_a_third_waits_and_shutdown_cancels_it() { + let executor = Executor::new(8); + let mut releases = Vec::new(); + let mut running = Vec::new(); + for _ in 0..2 { + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + releases.push(release_tx); + running.push(tokio::spawn({ + let executor = executor.clone(); + async move { + executor + .run(CancellationToken::new(), move |threads| { + assert_eq!(threads, 4); + let _ = started_tx.send(()); + let _ = release_rx.recv(); + Ok(()) + }) + .await + } + })); + tokio::time::timeout(Duration::from_secs(5), started_rx) + .await + .unwrap() + .unwrap(); + } + let third = executor.run(CancellationToken::new(), |_| panic!("queued query ran")); + tokio::pin!(third); + assert!(futures_poll_pending(third.as_mut()).await); + executor.shutdown_token().cancel(); + let result: Result<(), QueryError> = third.await; + assert_eq!(result.unwrap_err().code, QueryErrorCode::Cancelled); + assert_eq!(executor.permits.available_permits(), 0); + for release in releases { + let _ = release.send(()); + } + for task in running { + let _ = task.await.unwrap(); + } + } +} diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs new file mode 100644 index 0000000..35f35f7 --- /dev/null +++ b/src/mcp/mod.rs @@ -0,0 +1,66 @@ +//! Optional stdio-only MCP adapter. No listener or CLI subprocess is created. +mod executor; +mod schema; +mod server; + +use rmcp::ServiceExt; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::io::{AsyncRead, ReadBuf}; +use tokio_util::sync::CancellationToken; + +/// Only `mcp` starts this runtime; drop waits for running queries. +pub fn run() -> Result<(), String> { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|error| error.to_string())?; + runtime.block_on(async { + let parallelism = std::thread::available_parallelism() + .map(usize::from) + .unwrap_or(1); + let executor = executor::Executor::new(parallelism); + let shutdown = executor.shutdown_token(); + let input = StopOnEof { + reader: tokio::io::stdin(), + shutdown: shutdown.clone(), + }; + let result = async { + let service = server::Server::new(executor) + .serve((input, tokio::io::stdout())) + .await + .map_err(|error| error.to_string())?; + service.waiting().await.map_err(|error| error.to_string())?; + Ok(()) + } + .await; + shutdown.cancel(); + result + }) +} + +/// Stop queued work as soon as stdin closes, before the SDK drains replies. +struct StopOnEof { + reader: R, + shutdown: CancellationToken, +} + +impl AsyncRead for StopOnEof { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let before = buf.filled().len(); + let remaining = buf.remaining(); + let result = Pin::new(&mut self.reader).poll_read(cx, buf); + match &result { + Poll::Ready(Err(_)) => self.shutdown.cancel(), + Poll::Ready(Ok(())) if remaining > 0 && buf.filled().len() == before => { + self.shutdown.cancel() + } + _ => {} + } + result + } +} diff --git a/src/mcp/schema.rs b/src/mcp/schema.rs new file mode 100644 index 0000000..6b793a9 --- /dev/null +++ b/src/mcp/schema.rs @@ -0,0 +1,199 @@ +use rmcp::model::{CallToolResult, Tool, ToolAnnotations}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; + +use crate::asc_client::asc_handler::Reference; +use crate::asc_client::query::{ApkQuery, ClassSource, QueryError}; +use crate::asc_client::query_input::_build_member_find; +use crate::asc_core::findrefs::findrefs_manager::FindRefInput; + +#[derive(Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub(super) struct GetClass { + /// Absolute path to an APK on the MCP server's local filesystem. + pub apk_path: String, + /// Dalvik descriptor (Lcom/example/Main;) or dotted class name. + pub class_name: String, +} + +#[derive(Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub(super) struct ListClass { + /// Absolute path to an APK on the MCP server's local filesystem. + pub apk_path: String, + /// Optional class/package prefix, such as com.example or Lcom/example. + pub prefix: Option, +} + +#[derive(Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub(super) struct GetManifest { + /// Absolute path to an APK on the MCP server's local filesystem. + pub apk_path: String, +} + +#[derive(Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub(super) struct FindRefs { + /// Absolute path to an APK on the MCP server's local filesystem. + pub apk_path: String, + pub query: RefQuery, +} + +#[derive(Deserialize, JsonSchema)] +#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] +pub(super) enum RefQuery { + String { + value: String, + }, + Type { + value: String, + }, + Method { + /// Fuzzy method name. Supply name, class_name, or both. + name: Option, + class_name: Option, + /// If false, class_name is an exact descriptor/dotted class name. + #[serde(default)] + fuzzy_class: bool, + }, + Field { + /// Fuzzy field name. Supply name, class_name, or both. + name: Option, + class_name: Option, + #[serde(default)] + fuzzy_class: bool, + }, +} + +impl RefQuery { + fn into_find(self) -> Result<(&'static str, FindRefInput), QueryError> { + let kind = match &self { + Self::String { .. } => "string", + Self::Type { .. } => "type", + Self::Method { .. } => "method", + Self::Field { .. } => "field", + }; + let find = match self { + Self::String { value } => FindRefInput::new(Some(value), None, None, None), + Self::Type { value } => FindRefInput::new(None, Some(value), None, None), + Self::Method { + name, + class_name, + fuzzy_class, + } + | Self::Field { + name, + class_name, + fuzzy_class, + } => _build_member_find(kind, class_name.as_deref(), fuzzy_class, name.as_deref()) + .map_err(QueryError::invalid)?, + }; + Ok((kind, find)) + } +} + +#[derive(Serialize, JsonSchema)] +struct Classes { + classes: Vec, +} +#[derive(Serialize, JsonSchema)] +struct Manifest { + xml: String, +} +#[derive(Serialize, JsonSchema)] +struct References { + references: Vec, +} + +fn tool( + name: &'static str, + description: &'static str, +) -> Tool { + Tool::new(name, description, serde_json::Map::new()) + .with_input_schema::() + .with_output_schema::() + .with_annotations( + ToolAnnotations::new() + .read_only(true) + .destructive(false) + .idempotent(true) + .open_world(false), + ) +} + +pub(super) fn tools() -> Vec { + vec![ + tool::( + "getclass", + "Decompile one class from a local APK. Accepts an absolute apk_path and a Dalvik descriptor or dotted class_name. Returns source, class_name and dex_name; missing classes are errors. No files are written.", + ), + tool::( + "listclass", + "List all classes in a local APK, optionally filtered by class/package prefix. Returns the complete classes array (Dalvik descriptors); no matches is an empty array. Requires an absolute apk_path. No pagination or file output.", + ), + tool::( + "getmanifest", + "Decode AndroidManifest.xml from a local APK. Requires an absolute apk_path. Returns the complete XML in xml; a missing or invalid manifest is an error. No files are written.", + ), + tool::( + "findrefs", + "Find callers referencing a string, type, method or field across all DEX files in a local APK. query.kind selects the search; string/type use fuzzy value, method/field use fuzzy name and/or class_name with optional fuzzy_class. Returns complete structured references with dex_name, class_name, method_name and matched targets. Empty matches succeed; partial scan failures are errors. Requires an absolute apk_path.", + ), + ] +} + +pub(super) enum Query { + Class(GetClass), + Classes(ListClass), + Manifest(GetManifest), + References(FindRefs), +} + +impl Query { + pub(super) fn parse(name: &str, args: Value) -> Result { + let result = match name { + "getclass" => serde_json::from_value(args).map(Self::Class), + "listclass" => serde_json::from_value(args).map(Self::Classes), + "getmanifest" => serde_json::from_value(args).map(Self::Manifest), + "findrefs" => serde_json::from_value(args).map(Self::References), + _ => return Err(QueryError::invalid("unknown tool")), + }; + result.map_err(|error| QueryError::invalid(error.to_string())) + } + + pub(super) fn execute(self, threads: usize) -> Result { + match self { + Self::Class(args) => { + encode(ApkQuery::new(&args.apk_path, threads)?.getclass(&args.class_name)?) + } + Self::Classes(args) => encode(Classes { + classes: ApkQuery::new(&args.apk_path, threads)? + .listclass(args.prefix.as_deref())?, + }), + Self::Manifest(args) => encode(Manifest { + xml: ApkQuery::new(&args.apk_path, threads)?.getmanifest()?, + }), + Self::References(args) => { + let (kind, find) = args.query.into_find()?; + encode(References { + references: ApkQuery::new(&args.apk_path, threads)?.findrefs(kind, &find)?, + }) + } + } + } +} + +fn encode(value: impl Serialize) -> Result { + serde_json::to_value(value).map_err(|error| { + QueryError::new( + crate::asc_client::query::QueryErrorCode::QueryFailed, + error.to_string(), + ) + }) +} + +pub(super) fn error_result(error: QueryError) -> CallToolResult { + CallToolResult::structured_error(json!({"code": error.code.as_str(), "message": error.message})) +} diff --git a/src/mcp/server.rs b/src/mcp/server.rs new file mode 100644 index 0000000..9e454f7 --- /dev/null +++ b/src/mcp/server.rs @@ -0,0 +1,79 @@ +use rmcp::model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, Implementation, ListToolsResult, + PaginatedRequestParams, ServerCapabilities, ServerConfig, Tool, +}; +use rmcp::service::RequestContext; +use rmcp::{ErrorData, RoleServer, ServerHandler}; +use serde_json::Value; + +use super::{executor::Executor, schema}; + +pub(super) struct Server { + pub executor: Executor, + tools: Vec, +} + +impl Server { + pub(super) fn new(executor: Executor) -> Self { + Self { + executor, + tools: schema::tools(), + } + } +} + +impl ServerHandler for Server { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info(Implementation::new("rust_ASC", env!("CARGO_PKG_VERSION"))) + .with_instructions("Read-only, stateless local APK queries: getclass, listclass, getmanifest and findrefs. Pass an absolute apk_path on this machine. No open/close, cache, file writes, pagination or JavaScript execution. Keep input files unchanged during each query. Results are complete and can be large. CodeMode callers should check isError and consume structuredContent; tool errors resolve rather than throw. Concurrent calls share an automatic CPU budget. Cancellation does not interrupt already-running computation. Parser aborts can terminate this process.") + } + + async fn list_tools( + &self, + request: Option, + _context: RequestContext, + ) -> Result { + if request.is_some_and(|request| request.cursor.is_some()) { + return Err(ErrorData::invalid_params( + "tool list has no continuation cursor", + None, + )); + } + Ok(ListToolsResult::with_all_items(self.tools.clone())) + } + + fn get_tool(&self, name: &str) -> Option { + self.tools.iter().find(|tool| tool.name == name).cloned() + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + if self.get_tool(&request.name).is_none() { + return Err(ErrorData::invalid_params( + format!("unknown tool: {}", request.name), + None, + )); + } + let args = Value::Object(request.arguments.unwrap_or_default()); + let query = match schema::Query::parse(&request.name, args) { + Ok(query) => query, + Err(error) => return Ok(schema::error_result(error).into()), + }; + let result = self + .executor + .run(context.ct, move |threads| { + // Keep large-result serialization off the protocol loop. + query.execute(threads).map(CallToolResult::structured) + }) + .await; + Ok(match result { + Ok(result) => result, + Err(error) => schema::error_result(error), + } + .into()) + } +} diff --git a/tests/cli_thread.rs b/tests/cli_thread.rs index 5adedd7..a34d71a 100644 --- a/tests/cli_thread.rs +++ b/tests/cli_thread.rs @@ -22,6 +22,22 @@ fn thread_option_is_singular_and_belongs_to_each_command() { } } +#[test] +fn mcp_command_follows_the_compile_time_feature() { + let output = run_cli(&["mcp", "--help"]); + let text = stderr(&output); + if cfg!(feature = "mcp") { + assert!( + text.contains("Serve read-only MCP tools over stdio."), + "{text}" + ); + assert!(!text.contains("--thread"), "{text}"); + assert!(!text.contains("--port"), "{text}"); + } else { + assert!(text.contains("unrecognized subcommand 'mcp'"), "{text}"); + } +} + #[test] fn plural_threads_option_is_rejected() { let output = run_cli(&[ diff --git a/tests/dex_fixture/mod.rs b/tests/dex_fixture/mod.rs index e06013b..1a18fab 100644 --- a/tests/dex_fixture/mod.rs +++ b/tests/dex_fixture/mod.rs @@ -203,6 +203,26 @@ pub fn make_dex() -> Vec { ) } +/// Build one caller with repeated string references, including a CLI delimiter. +/// Target index order intentionally differs from lexical and instruction order. +pub fn make_multi_string_dex() -> Vec { + make_member_dex( + &[ + b"Lexample/Test;", + b"Ljava/lang/Object;", + b"V", + b"run", + b"token-z", + b"token-a; token-b", + ], + &[0, 1, 2], + None, + &[(0, 0, 3)], + code_item(1, 0, 0, 7, b"\x1a\0\x05\0\x1a\0\x04\0\x1a\0\x05\0\x0e\0"), + b"\0\0\x01\0\0\x09", + ) +} + /// Build a class whose run method invokes example.Target.callToMethod. pub fn make_invoke_dex() -> Vec { let strings: &[&[u8]] = &[ diff --git a/tests/mcp_stdio.rs b/tests/mcp_stdio.rs new file mode 100644 index 0000000..9ecdc3e --- /dev/null +++ b/tests/mcp_stdio.rs @@ -0,0 +1,519 @@ +#![cfg(feature = "mcp")] + +mod dex_fixture; +mod zip_fixture; + +use std::collections::BTreeMap; +use std::io::{BufRead, BufReader, Read, Write}; +use std::process::{Child, ChildStdin, Command, Stdio}; +use std::sync::mpsc::{self, Receiver}; +use std::thread::JoinHandle; +use std::time::{Duration, Instant}; + +use dex_fixture::{ + make_class_only_dex, make_dex, make_dex041_container, make_instance_field_dex, make_invoke_dex, +}; +use serde_json::{Value, json}; +use tempfile::tempdir; +use zip_fixture::{CompressionMethod, write_apk}; + +struct Client { + child: Child, + input: Option, + output: Receiver>, + stdout: Option>, + stderr: Option>, +} + +impl Client { + fn new() -> Self { + Self::version("2025-11-25") + } + + fn version(version: &str) -> Self { + let mut child = Command::new(env!("CARGO_BIN_EXE_rust_ASC")) + .arg("mcp") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + let input = child.stdin.take(); + let stdout = child.stdout.take().unwrap(); + let mut stderr = child.stderr.take().unwrap(); + let stderr = std::thread::spawn(move || { + let mut text = String::new(); + stderr.read_to_string(&mut text).unwrap(); + text + }); + let (send, output) = mpsc::channel(); + let stdout = std::thread::spawn(move || { + for line in BufReader::new(stdout).lines() { + let result = line.map_err(|e| e.to_string()).and_then(|line| { + serde_json::from_str(&line) + .map_err(|e| format!("non-protocol stdout: {line}: {e}")) + }); + if send.send(result).is_err() { + break; + } + } + }); + let mut client = Self { + child, + input, + output, + stdout: Some(stdout), + stderr: Some(stderr), + }; + client.send(json!({"jsonrpc":"2.0", "id":0, "method":"initialize", "params":{ + "protocolVersion":version, "capabilities":{}, "clientInfo":{"name":"test", "version":"1"} + }})); + let response = client.recv(); + assert_eq!(response["id"], 0); + assert_eq!(response["result"]["protocolVersion"], version, "{response}"); + assert!( + response["result"]["capabilities"]["tools"].is_object(), + "{response}" + ); + assert_eq!(response["result"]["serverInfo"]["name"], "rust_ASC"); + client.send(json!({"jsonrpc":"2.0", "method":"notifications/initialized"})); + client + } + + fn send(&mut self, value: Value) { + let input = self.input.as_mut().unwrap(); + writeln!(input, "{value}").unwrap(); + input.flush().unwrap(); + } + + fn recv(&self) -> Value { + self.output + .recv_timeout(Duration::from_secs(20)) + .expect("MCP reply timeout/disconnect") + .expect("stdout must contain only JSON-RPC") + } + + fn call(&mut self, name: &str, args: Value) -> Value { + self.send( + json!({"jsonrpc":"2.0", "id":1, "method":"tools/call", "params":{ + "name":name, "arguments":args + }}), + ); + let response = self.recv(); + assert_eq!(response["id"], 1, "{response}"); + assert!(response.get("error").is_none(), "{response}"); + let result = &response["result"]; + let text: Value = + serde_json::from_str(result["content"][0]["text"].as_str().unwrap()).unwrap(); + assert_eq!(text, result["structuredContent"]); + result.clone() + } + + fn error(&mut self, name: &str, args: Value, code: &str) { + let result = self.call(name, args); + assert_eq!(result["isError"], true, "{result}"); + assert_eq!(result["structuredContent"]["code"], code, "{result}"); + assert!( + !result["structuredContent"]["message"] + .as_str() + .unwrap() + .is_empty() + ); + } + + fn finish(mut self) { + self.input.take(); + let start = Instant::now(); + loop { + if let Some(status) = self.child.try_wait().unwrap() { + let stderr = self.stderr.take().unwrap().join().unwrap(); + assert!(status.success(), "MCP exit: {status}; {stderr}"); + assert!(stderr.is_empty(), "unexpected stderr: {stderr}"); + break; + } + assert!( + start.elapsed() < Duration::from_secs(10), + "MCP did not exit on EOF" + ); + std::thread::sleep(Duration::from_millis(10)); + } + self.stdout.take().unwrap().join().unwrap(); + if let Some(remaining) = self.output.try_iter().next() { + panic!("unexpected trailing protocol output: {remaining:?}"); + } + } +} + +impl Drop for Client { + fn drop(&mut self) { + self.input.take(); + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +fn success(result: &Value) -> &Value { + assert_eq!(result["isError"], false, "{result}"); + &result["structuredContent"] +} + +#[test] +fn discovers_exactly_four_read_only_typed_tools_and_supports_legacy_clients() { + for version in ["2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25"] { + let mut client = Client::version(version); + client.send(json!({"jsonrpc":"2.0", "id":1, "method":"tools/list"})); + let response = client.recv(); + let tools = response["result"]["tools"].as_array().unwrap(); + let mut names: Vec<_> = tools.iter().map(|t| t["name"].as_str().unwrap()).collect(); + names.sort(); + assert_eq!(names, ["findrefs", "getclass", "getmanifest", "listclass"]); + for tool in tools { + assert_eq!(tool["inputSchema"]["type"], "object"); + assert_eq!(tool["inputSchema"]["additionalProperties"], false); + assert_eq!(tool["outputSchema"]["type"], "object"); + assert_eq!(tool["annotations"]["readOnlyHint"], true); + assert_eq!(tool["annotations"]["openWorldHint"], false); + } + client.send(json!({"jsonrpc":"2.0", "id":2, "method":"tools/call", "params":{"name":"export","arguments":{}}})); + assert_eq!(client.recv()["error"]["code"], -32602); + client.finish(); + } +} + +#[test] +fn four_queries_return_full_structured_data_and_do_not_write_files() { + let dir = tempdir().unwrap(); + let apk = dir.path().join("fixture.apk"); + let dex = make_dex(); + let manifest = minimal_manifest(); + write_apk( + &apk, + &[("classes.dex", &dex), ("AndroidManifest.xml", &manifest)], + CompressionMethod::Deflated, + ); + let mut client = Client::new(); + let result = client.call("listclass", json!({"apk_path":apk})); + assert_eq!(success(&result)["classes"], json!(["Lexample/Test;"])); + let result = client.call( + "getclass", + json!({"apk_path":apk,"class_name":"example.Test"}), + ); + let data = success(&result); + assert_eq!(data["class_name"], "Lexample/Test;"); + assert_eq!(data["dex_name"], "classes.dex"); + assert!(data["source"].as_str().unwrap().contains("class Test")); + let result = client.call("getmanifest", json!({"apk_path":apk})); + assert!( + success(&result)["xml"] + .as_str() + .unwrap() + .contains("first | matched=(token)\nclasses.dex | Lexample/Test;->second | matched=(token)\n" + ); + assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 1); + client.finish(); +} + +#[test] +fn failures_are_tool_errors_but_no_matches_are_successful() { + let dir = tempdir().unwrap(); + let apk = dir.path().join("fixture.apk"); + write_apk( + &apk, + &[("classes.dex", &make_dex())], + CompressionMethod::Stored, + ); + let mut client = Client::new(); + client.error( + "listclass", + json!({"apk_path":"relative.apk"}), + "INVALID_ARGUMENT", + ); + client.error( + "listclass", + json!({"apk_path":dir.path()}), + "INVALID_ARGUMENT", + ); + client.error( + "listclass", + json!({"apk_path":dir.path().join("missing.apk")}), + "FILE_NOT_FOUND", + ); + client.error( + "listclass", + json!({"apk_path":apk,"prefix":" "}), + "INVALID_ARGUMENT", + ); + client.error( + "listclass", + json!({"apk_path":apk,"output":"not-allowed"}), + "INVALID_ARGUMENT", + ); + client.error("getclass", json!({"apk_path":apk}), "INVALID_ARGUMENT"); + client.error( + "getclass", + json!({"apk_path":apk,"class_name":"not.Found"}), + "NOT_FOUND", + ); + client.error("getmanifest", json!({"apk_path":apk}), "NOT_FOUND"); + client.error( + "findrefs", + json!({"apk_path":apk,"query":{"kind":"method"}}), + "INVALID_ARGUMENT", + ); + client.error( + "findrefs", + json!({"apk_path":apk,"query":{"kind":"invalid"}}), + "INVALID_ARGUMENT", + ); + let result = client.call("listclass", json!({"apk_path":apk,"prefix":"missing"})); + assert_eq!(success(&result)["classes"], json!([])); + let result = client.call( + "findrefs", + json!({"apk_path":apk,"query":{"kind":"string","value":"missing"}}), + ); + assert_eq!(success(&result)["references"], json!([])); + client.finish(); +} + +#[test] +fn malformed_inputs_are_not_reported_as_empty_or_partial_success() { + let dir = tempdir().unwrap(); + let apk = dir.path().join("fixture.apk"); + let mut client = Client::new(); + for bytes in [b"not a zip".as_slice(), b""] { + std::fs::write(&apk, bytes).unwrap(); + client.error("listclass", json!({"apk_path":apk}), "INVALID_FORMAT"); + client.error( + "getclass", + json!({"apk_path":apk,"class_name":"example.Test"}), + "INVALID_FORMAT", + ); + } + write_apk( + &apk, + &[("classes.dex", &make_dex()), ("classes2.dex", b"not a dex")], + CompressionMethod::Deflated, + ); + client.error( + "findrefs", + json!({"apk_path":apk,"query":{"kind":"string","value":"token"}}), + "INVALID_FORMAT", + ); + client.error("listclass", json!({"apk_path":apk}), "INVALID_FORMAT"); + let mut bad_strings = make_dex(); + let string_ids = dex_fixture::get_u32(&bad_strings, 0x3c) as usize; + dex_fixture::put_u32(&mut bad_strings, string_ids, u32::MAX); + write_apk( + &apk, + &[("classes.dex", &bad_strings)], + CompressionMethod::Stored, + ); + client.error( + "findrefs", + json!({"apk_path":apk,"query":{"kind":"string","value":"token"}}), + "INVALID_FORMAT", + ); + let mut broken_container = make_dex041_container(); + broken_container.truncate(broken_container.len() - 8); + write_apk( + &apk, + &[("classes.dex", &broken_container)], + CompressionMethod::Stored, + ); + client.error("listclass", json!({"apk_path":apk}), "INVALID_FORMAT"); + write_apk( + &apk, + &[("AndroidManifest.xml", b"bad xml")], + CompressionMethod::Stored, + ); + client.error("getmanifest", json!({"apk_path":apk}), "INVALID_FORMAT"); + // A well-formed archive without DEX files is not a malformed archive. + let result = client.call("listclass", json!({"apk_path":apk})); + assert_eq!(success(&result)["classes"], json!([])); + // Corrupt the central-directory byte count, keeping the EOCD present. + let mut data = std::fs::read(&apk).unwrap(); + let offset = data.len() - 22 + 12; + dex_fixture::put_u32(&mut data, offset, u32::MAX); + std::fs::write(&apk, data).unwrap(); + client.error("listclass", json!({"apk_path":apk}), "INVALID_FORMAT"); + client.finish(); +} + +#[test] +fn repeated_calls_observe_new_files_and_multidex041_is_preserved() { + let dir = tempdir().unwrap(); + let apk = dir.path().join("fixture.apk"); + let mut client = Client::new(); + write_apk( + &apk, + &[("classes.dex", &make_dex())], + CompressionMethod::Stored, + ); + let result = client.call("listclass", json!({"apk_path":apk})); + assert_eq!(success(&result)["classes"], json!(["Lexample/Test;"])); + write_apk( + &apk, + &[("classes.dex", &make_dex041_container())], + CompressionMethod::Deflated, + ); + let result = client.call("listclass", json!({"apk_path":apk})); + assert_eq!( + success(&result)["classes"], + json!([ + "Lalpha/First;", + "Lshared/Duplicate;", + "Lbeta/Second;", + "Lshared/Duplicate;" + ]) + ); + let result = client.call( + "getclass", + json!({"apk_path":apk,"class_name":"beta.Second"}), + ); + assert_eq!(success(&result)["dex_name"], "classes.dex!classes2.dex"); + client.finish(); +} + +#[test] +fn all_reference_kinds_have_structured_results() { + let dir = tempdir().unwrap(); + let apk = dir.path().join("fixture.apk"); + let mut client = Client::new(); + let mut invoke_dex = make_invoke_dex(); + // Make the call self-referential. The shared fixture's trailing Target + // type is unsorted; Test is searchable by the exact-class binary search. + let methods = dex_fixture::get_u32(&invoke_dex, 0x5c) as usize; + invoke_dex[methods + 8..methods + 10].copy_from_slice(&0u16.to_le_bytes()); + let mut type_dex = make_dex(); + // Replace const-string v0, token with const-class v0, Lexample/Test;. + let instruction = type_dex + .windows(6) + .position(|bytes| bytes == b"\x1a\0\x05\0\x0e\0") + .unwrap(); + type_dex[instruction..instruction + 4].copy_from_slice(b"\x1c\0\0\0"); + write_apk( + &apk, + &[ + ("classes.dex", &invoke_dex), + ("classes2.dex", &make_instance_field_dex()), + ("classes3.dex", &type_dex), + ], + CompressionMethod::Stored, + ); + for query in [ + json!({"kind":"type","value":"example"}), + json!({"kind":"method","name":"callToMethod","class_name":"example.Test"}), + json!({"kind":"method","name":"callToMethod","class_name":"Test","fuzzy_class":true}), + json!({"kind":"field","name":"a","class_name":"example.Test"}), + ] { + let result = client.call("findrefs", json!({"apk_path":apk,"query":query})); + let references = success(&result)["references"].as_array().unwrap(); + assert!(!references.is_empty(), "{query}: {result}"); + for reference in references { + assert!(!reference["matched"].as_array().unwrap().is_empty()); + assert!( + reference["dex_name"] + .as_str() + .unwrap() + .starts_with("classes") + ); + } + } + client.finish(); +} + +#[test] +fn concurrent_requests_have_independent_ids_and_complete_untruncated_results() { + let dir = tempdir().unwrap(); + let apk = dir.path().join("fixture.apk"); + let names: Vec<_> = (0..4000) + .map(|i| format!("Lexample/Class{i:04};")) + .collect(); + let descriptors: Vec<_> = names.iter().map(String::as_str).collect(); + let dex = make_class_only_dex(&descriptors, 0, b"035"); + write_apk(&apk, &[("classes.dex", &dex)], CompressionMethod::Deflated); + let mut client = Client::new(); + for id in 1..=8 { + client.send( + json!({"jsonrpc":"2.0","id":id,"method":"tools/call","params":{ + "name":"listclass","arguments":{"apk_path":apk} + }}), + ); + } + client.send(json!({"jsonrpc":"2.0","id":9,"method":"ping"})); + let mut results = BTreeMap::new(); + for _ in 1..=9 { + let response = client.recv(); + assert!( + results + .insert(response["id"].as_u64().unwrap(), response) + .is_none() + ); + } + for id in 1..=8 { + let result = &results[&id]["result"]; + assert_eq!(success(result)["classes"], json!(names)); + let text: Value = + serde_json::from_str(result["content"][0]["text"].as_str().unwrap()).unwrap(); + assert_eq!(text, result["structuredContent"]); + } + assert!(results[&9].get("error").is_none()); + client.finish(); +} + +// Minimal valid binary XML: one UTF-8 string, an empty resource map, and a +// manifest start/end pair. Kept generated and reviewable instead of a binary blob. +fn minimal_manifest() -> Vec { + fn u16_(out: &mut Vec, n: u16) { + out.extend(n.to_le_bytes()); + } + fn u32_(out: &mut Vec, n: u32) { + out.extend(n.to_le_bytes()); + } + fn header(out: &mut Vec, kind: u16, header_size: u16, size: u32) { + u16_(out, kind); + u16_(out, header_size); + u32_(out, size); + } + let mut out = Vec::new(); + header(&mut out, 3, 8, 120); + header(&mut out, 1, 28, 44); + for value in [1, 0, 0x100, 32, 0, 0] { + u32_(&mut out, value); + } + out.extend(b"\x08\x08manifest\0\0"); + header(&mut out, 0x180, 8, 8); + header(&mut out, 0x102, 16, 36); + for value in [1, u32::MAX, u32::MAX, 0] { + u32_(&mut out, value); + } + for value in [20, 20, 0, 0, 0, 0] { + u16_(&mut out, value); + } + header(&mut out, 0x103, 16, 24); + for value in [1, u32::MAX, u32::MAX, 0] { + u32_(&mut out, value); + } + assert_eq!(out.len(), 120); + out +} diff --git a/tests/reference_output.rs b/tests/reference_output.rs new file mode 100644 index 0000000..fb26221 --- /dev/null +++ b/tests/reference_output.rs @@ -0,0 +1,151 @@ +mod dex_fixture; + +use dex_fixture::{make_dex, make_instance_field_dex, make_invoke_dex, make_multi_string_dex}; +use rust_ASC::asc_client::asc_handler::{AscHandler, Reference}; +use rust_ASC::asc_core::findrefs::findrefs_manager::FindRefInput; +use rust_ASC::asc_core::findrefs::locator::field_locator::FieldFind; +use rust_ASC::asc_core::findrefs::locator::method_locator::MethodFind; + +fn text(data: &[u8], kind: &str, query: &FindRefInput, aggregate: bool) -> Vec { + let handler = AscHandler::new(false); + let mut borrowed = handler.findrefs("fixture.dex", data, kind, query, aggregate); + let mut owned = handler.findrefs_vec("fixture.dex", data.to_vec(), kind, query, aggregate); + if !aggregate { + borrowed.sort(); + owned.sort(); + } + assert_eq!(borrowed, owned); + borrowed +} + +#[test] +fn aggregated_text_and_records_preserve_target_order_and_delimiters() { + let handler = AscHandler::new(false); + let data = make_multi_string_dex(); + let query = FindRefInput::new(Some("token".into()), None, None, None); + let expected = ["fixture.dex | Lexample/Test;->run | matched=(token-z; token-a; token-b)"]; + assert_eq!(text(&data, "string", &query, true), expected); + assert_eq!( + handler.findrefs_records("fixture.dex", data, "string", &query), + vec![Reference { + dex_name: "fixture.dex".into(), + class_name: "Lexample/Test;".into(), + method_name: "run".into(), + matched: vec!["token-z".into(), "token-a; token-b".into()], + }] + ); +} + +#[test] +fn non_aggregated_text_preserves_repeated_occurrences() { + let query = FindRefInput::new(Some("token".into()), None, None, None); + let expected = [ + "fixture.dex | Lexample/Test;->run | matched=(token-a; token-b)", + "fixture.dex | Lexample/Test;->run | matched=(token-a; token-b)", + "fixture.dex | Lexample/Test;->run | matched=(token-z)", + ]; + assert_eq!( + text(&make_multi_string_dex(), "string", &query, false), + expected + ); +} + +#[test] +fn shared_code_keeps_both_callers_in_text_and_records() { + let handler = AscHandler::new(false); + let query = FindRefInput::new(Some("token".into()), None, None, None); + let expected = [ + "fixture.dex | Lexample/Test;->first | matched=(token)", + "fixture.dex | Lexample/Test;->second | matched=(token)", + ]; + for aggregate in [false, true] { + assert_eq!(text(&make_dex(), "string", &query, aggregate), expected); + } + let records = handler.findrefs_records("fixture.dex", make_dex(), "string", &query); + assert_eq!( + records + .iter() + .map(|r| r.method_name.as_str()) + .collect::>(), + ["first", "second"] + ); + assert!(records.iter().all(|r| r.matched == ["token"])); +} + +#[test] +fn type_method_and_field_targets_keep_exact_cli_and_structured_formats() { + let handler = AscHandler::new(false); + let mut type_dex = make_dex(); + let insn = type_dex + .windows(6) + .position(|b| b == b"\x1a\0\x05\0\x0e\0") + .unwrap(); + type_dex[insn..insn + 4].copy_from_slice(b"\x1c\0\0\0"); + let cases = [ + ( + type_dex, + "type", + FindRefInput::new(None, Some("Lexample/Test;".into()), None, None), + vec!["first", "second"], + "Lexample/Test;", + ), + ( + make_invoke_dex(), + "method", + FindRefInput::new( + None, + None, + Some(MethodFind::new(None, Some("callToMethod".into()))), + None, + ), + vec!["run"], + "Lexample/Target;->callToMethod", + ), + ( + make_instance_field_dex(), + "field", + FindRefInput::new( + None, + None, + None, + Some(FieldFind::new(None, Some("a".into()))), + ), + vec!["getA"], + "Lexample/Test;->a", + ), + ]; + for (data, kind, query, callers, target) in cases { + let expected: Vec<_> = callers + .iter() + .map(|caller| format!("fixture.dex | Lexample/Test;->{caller} | matched=({target})")) + .collect(); + for aggregate in [false, true] { + assert_eq!(text(&data, kind, &query, aggregate), expected, "{kind}"); + } + let records = handler.findrefs_records("fixture.dex", data, kind, &query); + let expected: Vec<_> = callers + .into_iter() + .map(|caller| Reference { + dex_name: "fixture.dex".into(), + class_name: "Lexample/Test;".into(), + method_name: caller.into(), + matched: vec![target.into()], + }) + .collect(); + assert_eq!(records, expected, "{kind}"); + } +} + +#[test] +fn absent_targets_return_empty_results_in_both_representations() { + let handler = AscHandler::new(false); + let query = FindRefInput::new(Some("missing-target".into()), None, None, None); + for aggregate in [false, true] { + assert!(text(&make_dex(), "string", &query, aggregate).is_empty()); + } + assert!( + handler + .findrefs_records("fixture.dex", make_dex(), "string", &query) + .is_empty() + ); +}