From cd91dbed9c5b7da8da866516ffa38c8b07cb790d Mon Sep 17 00:00:00 2001 From: Tsingshui Date: Mon, 19 Oct 2026 14:41:16 +0800 Subject: [PATCH 1/4] feat(mcp): add default-enabled stdio query server Expose getclass, listclass, getmanifest and findrefs as typed, read-only MCP tools using direct Rust queries. Support stateless requests, bounded concurrency and cooperative cancellation without changing CLI result allocation paths. Share query input normalization between CLI and MCP. Add protocol and output regression tests, feature gating and concise Pi CodeMode documentation. --- Cargo.lock | 492 +++++++++++++++++++++++++++ Cargo.toml | 10 + README.md | 45 +++ src/asc_client/apk/loader.rs | 69 +++- src/asc_client/asc_handler.rs | 109 ++++-- src/asc_client/manifest_handler.rs | 7 +- src/asc_client/mod.rs | 2 + src/asc_client/query.rs | 354 ++++++++++++++++++++ src/asc_client/query_input.rs | 66 ++++ src/cli.rs | 76 +---- src/lib.rs | 2 + src/mcp/executor.rs | 175 ++++++++++ src/mcp/mod.rs | 66 ++++ src/mcp/schema.rs | 199 +++++++++++ src/mcp/server.rs | 79 +++++ tests/cli_thread.rs | 16 + tests/dex_fixture/mod.rs | 20 ++ tests/mcp_stdio.rs | 519 +++++++++++++++++++++++++++++ tests/reference_output.rs | 151 +++++++++ 19 files changed, 2358 insertions(+), 99 deletions(-) create mode 100644 src/asc_client/query.rs create mode 100644 src/asc_client/query_input.rs create mode 100644 src/mcp/executor.rs create mode 100644 src/mcp/mod.rs create mode 100644 src/mcp/schema.rs create mode 100644 src/mcp/server.rs create mode 100644 tests/mcp_stdio.rs create mode 100644 tests/reference_output.rs 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() + ); +} From 01dc748f39ad0e3ae20cbe0407e77fe3adce20f6 Mon Sep 17 00:00:00 2001 From: MG1937 <2586364982@qq.com> Date: Mon, 19 Oct 2026 20:46:09 +0800 Subject: [PATCH 2/4] support mcp feature --- Cargo.lock | 2 - Cargo.toml | 5 +- README.md | 96 +++-- src/asc_client/apk/loader.rs | 69 +--- src/asc_client/asc_handler.rs | 109 ++---- src/asc_client/manifest_handler.rs | 7 +- src/asc_client/mod.rs | 2 - src/asc_client/query.rs | 354 ----------------- src/asc_client/query_input.rs | 66 ---- src/cli.rs | 69 +++- src/mcp/executor.rs | 175 --------- src/mcp/mod.rs | 62 +-- src/mcp/schema.rs | 303 ++++++++------- src/mcp/server.rs | 42 +-- tests/cli_thread.rs | 7 +- tests/dex_fixture/mod.rs | 20 - tests/mcp_stdio.rs | 587 ++++++++++++++--------------- tests/reference_output.rs | 151 -------- 18 files changed, 660 insertions(+), 1466 deletions(-) delete mode 100644 src/asc_client/query.rs delete mode 100644 src/asc_client/query_input.rs delete mode 100644 src/mcp/executor.rs delete mode 100644 tests/reference_output.rs diff --git a/Cargo.lock b/Cargo.lock index af379aa..05bf4b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -746,7 +746,6 @@ dependencies = [ "sha1", "tempfile", "tokio", - "tokio-util", ] [[package]] @@ -969,7 +968,6 @@ dependencies = [ "bytes", "futures-core", "futures-sink", - "futures-util", "libc", "pin-project-lite", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 48c8a3d..e8706dd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,7 +6,7 @@ license = "Apache-2.0" [features] default = ["mcp"] -mcp = ["dep:rmcp", "dep:tokio", "dep:tokio-util", "dep:serde", "dep:serde_json", "dep:schemars"] +mcp = ["dep:rmcp", "dep:tokio", "dep:serde", "dep:serde_json", "dep:schemars"] [profile.release] opt-level = 3 # max speed @@ -21,8 +21,7 @@ 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"] } +tokio = { version = "1", optional = true, features = ["rt", "io-std", "time"] } serde = { version = "1", optional = true, features = ["derive"] } serde_json = { version = "1", optional = true } schemars = { version = "1", optional = true } diff --git a/README.md b/README.md index c12d0ea..fb50dd5 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,10 @@ # rust_ASC -`rust_ASC` is a one-to-one Rust port of [MG1937/ASC](https://github.com/MG1937/ASC), the Python implementation of ASC. The Rust implementation preserves the corresponding source layout, data structures, control flow, function roles, and runtime behavior while targeting substantially lower command-line latency and memory overhead. +**rust_ASC is a stateless Android APK and DEX decompiler with millisecond cross-reference search, written in Rust.** It queries the compiled artifact directly instead of building a global index, so there is no preprocessing step, no persistent cache, and no multi-gigabyte memory footprint. -On the reported CLI workloads below, `rust_ASC` is approximately **10× faster** than the Python ASC implementation. It remains stateless: no preprocessing is required before a query, and it writes no persistent cache to disk. +It is the Rust implementation of [MG1937/ASC](https://github.com/MG1937/ASC), the Python reference project, and preserves the same source layout, data structures, control flow and runtime behavior. On the CLI workloads below it is approximately **10× faster** than the Python version. + +Presented at Black Hat Europe Arsenal: [Droid ASC: R8 Compiler Optimization as a DeCompiler Primitive](https://blackhat.com/europe/arsenal/schedule/index.html#droid-asc-r8-compiler-optimization-as-a-decompiler-primitive-54834). Repository links: @@ -20,21 +22,39 @@ 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. +- Expose CLI queries as MCP tools over stdio. + +## Why it is fast + +- Stateless by design: no preprocessing, no global index build, no disk cache. +- The APK is memory-mapped, and only the DEX entries a query needs are inflated. +- Cross-DEX scans run in parallel, so a search over dozens of DEX files stays in the millisecond range. +- R8 output is already highly structured. Searches jump to candidate offsets directly instead of rebuilding a bloated cross-reference database. +- On a hit, rust_ASC rebuilds a minimal self-consistent DEX in memory and decompiles it immediately. ## Build +`radguard` is a Git submodule. Clone it with the main repository: + +```bash +git clone --recurse-submodules https://github.com/MG1937/rust_ASC.git +cd rust_ASC +cargo build --release +``` + +For an existing checkout: + ```bash git submodule update --init --recursive cargo build --release +``` -# Build without MCP and its dedicated dependencies. +MCP is enabled by default. Build only the original CLI with: + +```bash 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 @@ -79,29 +99,9 @@ rust_ASC export --thread 10 --debug app.apk output ## 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. +`rust_ASC mcp` runs a stdio MCP server that exposes the CLI commands as tools: `getclass`, `listclass`, `getmanifest`, and `findrefs`. Each call spawns a fresh CLI process, so behavior always matches the command line. -
-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: +Configure it in your MCP client: ```json { @@ -114,7 +114,43 @@ and follow your client's MCP configuration instructions: } ``` -## CLI Benchmark +## Benchmark + +### rust_ASC vs JADX + +Same APKs, same workloads. JADX figures are from the ASC benchmark suite. + +**Global string cross-reference search** + +| Tool | TelegramX (59 MB) | WhatsApp (130 MB) | Grab (228 MB) | +| --- | ---: | ---: | ---: | +| rust_ASC | 55 ms | 67 ms | 118 ms | +| JADX | 20 s | 32 s | 2m14s | + +**Single-class decompile** + +| Tool | TelegramX (59 MB) | WhatsApp (130 MB) | Grab (228 MB) | +| --- | ---: | ---: | ---: | +| rust_ASC | 79 ms | 70 ms | 111 ms | +| JADX | 6 s | 15 s | 37 s | + +**CLI peak memory** + +| Tool | TelegramX (59 MB) | WhatsApp (130 MB) | Grab (228 MB) | +| --- | ---: | ---: | ---: | +| rust_ASC | 32 MB | 116 MB | 358 MB | +| JADX | 1,016 MB | 2.1 GB | 7.1 GB | + +**Disk cache** + +| Tool | TelegramX | WhatsApp | Grab | +| --- | ---: | ---: | ---: | +| rust_ASC | 0 MB | 0 MB | 0 MB | +| JADX | 119 MB | 30 MB | 127 MB | + +JADX spends seconds to minutes per operation and gigabytes of memory. rust_ASC answers in tens to low hundreds of milliseconds and keeps a small resident set. Because every query is cheap and stateless, a user or an AI agent can move across dozens or hundreds of APKs on a weak machine, where JADX would need several GB just to preload one. + +### rust_ASC CLI details 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 99d622d..e68d459 100644 --- a/src/asc_client/apk/loader.rs +++ b/src/asc_client/apk/loader.rs @@ -139,63 +139,6 @@ 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); @@ -288,14 +231,6 @@ 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; @@ -327,7 +262,7 @@ pub(crate) fn try_extract_zip_entry( local_header_off: _STRUCT_I::unpack_from(buf, ptr + 42), comp_method: _STRUCT_H::unpack_from(buf, ptr + 10), }; - return _inflate_dex(buf, &entry, None); + return _inflate_dex(buf, &entry, None)?.ok_or_else(|| "inflate stopped".to_owned()); } ptr = name_end @@ -336,5 +271,5 @@ pub(crate) fn try_extract_zip_entry( .ok_or_else(|| "bad central directory entry range".to_owned())?; } - Ok(None) + Err(format!("{target_name} not found in APK")) } diff --git a/src/asc_client/asc_handler.rs b/src/asc_client/asc_handler.rs index c8d0820..0e4eb48 100644 --- a/src/asc_client/asc_handler.rs +++ b/src/asc_client/asc_handler.rs @@ -6,57 +6,14 @@ use std::vec; use crate::asc_core::core::dex::dex_manager::DexManager; use crate::asc_core::findrefs::findrefs_manager::{FindRefInput, FindRefManager}; +use crate::asc_core::findrefs::locator::field_locator::FieldFind; +use crate::asc_core::findrefs::locator::method_locator::MethodFind; use crate::asc_core::utils::decompiler::{ - MemberReference, decompile_dex_bytes, decompile_dex_bytes_with_metadata, + decompile_dex_bytes, decompile_dex_bytes_with_metadata, MemberReference, }; use crate::asc_core::utils::tinydex::DEX; -/// One caller and its matched targets. Built from DEX records, never CLI text. -#[derive(Debug, Clone, PartialEq, Eq)] -#[cfg_attr(feature = "mcp", derive(serde::Serialize, schemars::JsonSchema))] -pub struct Reference { - pub dex_name: String, - pub class_name: String, - pub method_name: String, - pub matched: Vec, -} - -// 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 { @@ -85,13 +42,9 @@ impl AscHandler { decompile_dex_bytes_with_metadata(&new_dex_bytes, dalvik_class) } - /// 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_method(&self, dex: &DEX, midx: usize) -> String { + let method = dex.get_method(midx as u32); + format!("{}->{}", method.cls().fullname, method.name()) } fn _format_matched_name(&self, dex: &DEX, find_type: &str, idx: usize) -> String { @@ -120,7 +73,7 @@ impl AscHandler { aggregate: bool, ) -> Vec { let dex = DEX::parse(dex_buf, dex_name); - self.findrefs_output::(dex_name, &dex, find_type, find, aggregate) + self._findrefs(dex_name, &dex, find_type, find, aggregate) } pub fn findrefs_vec( @@ -132,28 +85,17 @@ impl AscHandler { aggregate: bool, ) -> Vec { let dex = DEX::from_vec(dex_buf, dex_name); - 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) + self._findrefs(dex_name, &dex, find_type, find, aggregate) } - fn findrefs_output( + fn _findrefs( &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); @@ -179,12 +121,15 @@ 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(R::single(dex_name, dex, mid, &matched)); + ret.push(format!( + "{} | {} | matched=({})", + dex_name, + self._format_method(&dex, mid as usize), + matched + )); } } return ret; @@ -194,9 +139,7 @@ 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); } @@ -211,9 +154,15 @@ impl AscHandler { sorted_idxs.sort(); let matched = sorted_idxs .iter() - .map(|&&idx| self._format_matched_name(dex, find_type, idx)) - .collect::>(); - ret.push(R::aggregate(dex_name, dex, *mid as u32, matched)); + .map(|&&idx| self._format_matched_name(&dex, find_type, idx)) + .collect::>() + .join("; "); + ret.push(format!( + "{} | {} | matched=({})", + dex_name, + self._format_method(&dex, *mid), + matched + )); } ret } diff --git a/src/asc_client/manifest_handler.rs b/src/asc_client/manifest_handler.rs index 04aab5c..833deab 100644 --- a/src/asc_client/manifest_handler.rs +++ b/src/asc_client/manifest_handler.rs @@ -46,13 +46,8 @@ 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; + let mut input = manifest_data.as_slice(); 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 cdf138f..1ff390c 100644 --- a/src/asc_client/mod.rs +++ b/src/asc_client/mod.rs @@ -3,5 +3,3 @@ 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 deleted file mode 100644 index 2711a2b..0000000 --- a/src/asc_client/query.rs +++ /dev/null @@ -1,354 +0,0 @@ -//! 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 deleted file mode 100644 index 8a724de..0000000 --- a/src/asc_client/query_input.rs +++ /dev/null @@ -1,66 +0,0 @@ -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 13a0f8f..31581f4 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -9,8 +9,72 @@ 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 @@ -404,8 +468,7 @@ pub fn _build_main_parser() -> Command { ) .subcommand(findrefs); #[cfg(feature = "mcp")] - let command = command - .subcommand(Command::new("mcp").about("Serve read-only MCP tools over stdio.")); + let command = command.subcommand(Command::new("mcp").about("Serve MCP tools over stdio.")); command } diff --git a/src/mcp/executor.rs b/src/mcp/executor.rs deleted file mode 100644 index 1367a57..0000000 --- a/src/mcp/executor.rs +++ /dev/null @@ -1,175 +0,0 @@ -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 index 35f35f7..a2a9c7d 100644 --- a/src/mcp/mod.rs +++ b/src/mcp/mod.rs @@ -1,66 +1,20 @@ -//! Optional stdio-only MCP adapter. No listener or CLI subprocess is created. -mod executor; +//! Stdio MCP adapter for the existing command-line interface. 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() + .enable_time() .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 + let service = server::Server::new() + .serve((tokio::io::stdin(), tokio::io::stdout())) + .await + .map_err(|error| error.to_string())?; + service.waiting().await.map_err(|error| error.to_string())?; + Ok(()) }) } - -/// 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 index 6b793a9..d43e516 100644 --- a/src/mcp/schema.rs +++ b/src/mcp/schema.rs @@ -1,123 +1,101 @@ -use rmcp::model::{CallToolResult, Tool, ToolAnnotations}; -use schemars::JsonSchema; -use serde::{Deserialize, Serialize}; -use serde_json::{Value, json}; +use std::ffi::OsString; +use std::process::Command; -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; +use rmcp::model::{Tool, ToolAnnotations}; +use schemars::JsonSchema; +use serde::Deserialize; +use serde_json::Value; #[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, +struct GetClass { + /// Path to the APK to analyze. + apk_path: String, + /// Dotted class name or Dalvik descriptor. + class_name: String, + /// Positive worker count. Defaults to available CPU count. + thread: Option, + /// Optional source output file. Existing files are overwritten. Content is also returned. + output: Option, } #[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, +struct ListClass { + /// Path to the APK to analyze. + apk_path: String, + /// Optional case sensitive class or package prefix. + prefix: Option, + /// Positive worker count. Defaults to available CPU count. + thread: Option, + /// Optional class list output file. When set, stdout is empty. Existing files are overwritten. + output: 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, +struct GetManifest { + /// Path to the APK to analyze. + apk_path: String, + /// Optional XML output file. Existing files are overwritten. Content is also returned. + output: Option, } #[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, +struct FindRefs { + /// Path to the APK to analyze. + apk_path: String, + /// Reference target and match mode. + query: RefQuery, + /// Positive worker count. Defaults to available CPU count. + thread: Option, + /// Optional reference output file. Existing files are overwritten. Content is also returned. + output: Option, } #[derive(Deserialize, JsonSchema)] #[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] -pub(super) enum RefQuery { +enum RefQuery { + /// Find string constant references. String { + /// Case sensitive string constant substring. value: String, }, + /// Find type references. Type { + /// Case sensitive Dalvik type descriptor substring. value: String, }, + /// Find method call references. Method { - /// Fuzzy method name. Supply name, class_name, or both. + /// Case sensitive method name substring. Supply name, class_name, or both. name: Option, + /// Dotted name or Dalvik descriptor of the declaring class. class_name: Option, - /// If false, class_name is an exact descriptor/dotted class name. + /// Match class_name as a case sensitive descriptor substring. #[serde(default)] fuzzy_class: bool, }, + /// Find field access references. Field { - /// Fuzzy field name. Supply name, class_name, or both. + /// Case sensitive field name substring. Supply name, class_name, or both. name: Option, + /// Dotted name or Dalvik descriptor of the declaring class. class_name: Option, + /// Match class_name as a case sensitive descriptor substring. #[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 { +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) + .read_only(false) + .destructive(true) .idempotent(true) .open_world(false), ) @@ -125,75 +103,148 @@ fn tool( 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::( + tool::("getclass", "Decompile one APK class to Java like source."), + 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.", + "List class descriptors from every DEX in an APK.", ), - tool::( + 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.", + "Decode AndroidManifest.xml to formatted XML.", ), - tool::( + 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.", + "Find methods that reference strings, types, methods, or fields.", ), ] } -pub(super) enum Query { - Class(GetClass), - Classes(ListClass), - Manifest(GetManifest), - References(FindRefs), +pub(super) struct Query { + args: Vec, } 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")), + pub(super) fn parse(name: &str, args: Value) -> Result { + let args = match name { + "getclass" => getclass(serde_json::from_value(args).map_err(invalid_args)?), + "listclass" => listclass(serde_json::from_value(args).map_err(invalid_args)?), + "getmanifest" => getmanifest(serde_json::from_value(args).map_err(invalid_args)?), + "findrefs" => findrefs(serde_json::from_value(args).map_err(invalid_args)?), + _ => return Err("unknown tool".to_owned()), }; - result.map_err(|error| QueryError::invalid(error.to_string())) + Ok(Self { args }) + } +} + +fn invalid_args(error: serde_json::Error) -> String { + format!("invalid arguments: {error}") +} + +fn push_thread(args: &mut Vec, thread: Option) { + if let Some(thread) = thread { + args.push("--thread".into()); + args.push(thread.to_string().into()); + } +} + +fn getclass(input: GetClass) -> Vec { + let mut args = vec!["getclass".into()]; + push_thread(&mut args, input.thread); + args.push(input.apk_path.into()); + args.push(input.class_name.into()); + push_output(&mut args, input.output); + args +} + +fn listclass(input: ListClass) -> Vec { + let mut args = vec!["listclass".into()]; + push_thread(&mut args, input.thread); + if let Some(prefix) = input.prefix { + args.push("--prefix".into()); + args.push(prefix.into()); } + args.push(input.apk_path.into()); + push_output(&mut args, input.output); + args +} - 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 getmanifest(input: GetManifest) -> Vec { + let mut args = vec!["getmanifest".into(), input.apk_path.into()]; + push_output(&mut args, input.output); + args +} + +fn findrefs(input: FindRefs) -> Vec { + let mut args = vec!["findrefs".into()]; + push_thread(&mut args, input.thread); + args.push(input.apk_path.into()); + match input.query { + RefQuery::String { value } => { + args.push("string".into()); + args.push(value.into()); } + RefQuery::Type { value } => { + args.push("type".into()); + args.push(value.into()); + } + RefQuery::Method { + name, + class_name, + fuzzy_class, + } => push_member(&mut args, "method", name, class_name, fuzzy_class), + RefQuery::Field { + name, + class_name, + fuzzy_class, + } => push_member(&mut args, "field", name, class_name, fuzzy_class), } + push_output(&mut args, input.output); + args } -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(), - ) - }) +fn push_output(args: &mut Vec, output: Option) { + if let Some(output) = output { + args.push("-o".into()); + args.push(output.into()); + } } -pub(super) fn error_result(error: QueryError) -> CallToolResult { - CallToolResult::structured_error(json!({"code": error.code.as_str(), "message": error.message})) +fn push_member( + args: &mut Vec, + kind: &str, + name: Option, + class_name: Option, + fuzzy_class: bool, +) { + args.push(kind.into()); + if let Some(name) = name { + args.push(name.into()); + } + if let Some(class_name) = class_name { + args.push("--class".into()); + args.push(class_name.into()); + } + if fuzzy_class { + args.push("--fuzzy-class".into()); + } +} + +pub(super) fn execute(query: Query) -> Result { + let executable = std::env::current_exe() + .map_err(|error| format!("cannot resolve rust_ASC executable: {error}"))?; + let output = Command::new(executable) + .args(query.args) + .output() + .map_err(|error| format!("cannot start rust_ASC CLI: {error}"))?; + if output.status.success() { + return Ok(String::from_utf8_lossy(&output.stdout).into_owned()); + } + + let stderr = String::from_utf8_lossy(&output.stderr); + let message = stderr.trim(); + if message.is_empty() { + Err(format!("rust_ASC CLI exited with {}", output.status)) + } else { + Err(message.to_owned()) + } } diff --git a/src/mcp/server.rs b/src/mcp/server.rs index 9e454f7..ceb08ad 100644 --- a/src/mcp/server.rs +++ b/src/mcp/server.rs @@ -1,22 +1,20 @@ use rmcp::model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, Implementation, ListToolsResult, - PaginatedRequestParams, ServerCapabilities, ServerConfig, Tool, + CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, Implementation, + ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerConfig, Tool, }; use rmcp::service::RequestContext; use rmcp::{ErrorData, RoleServer, ServerHandler}; use serde_json::Value; -use super::{executor::Executor, schema}; +use super::schema; pub(super) struct Server { - pub executor: Executor, tools: Vec, } impl Server { - pub(super) fn new(executor: Executor) -> Self { + pub(super) fn new() -> Self { Self { - executor, tools: schema::tools(), } } @@ -26,7 +24,10 @@ 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.") + .with_instructions( + "Inspect local Android APK files. Calls are stateless. Paths use the server \ + filesystem. Output files may be overwritten. Failures return CLI stderr.", + ) } async fn list_tools( @@ -50,7 +51,7 @@ impl ServerHandler for Server { async fn call_tool( &self, request: CallToolRequestParams, - context: RequestContext, + _context: RequestContext, ) -> Result { if self.get_tool(&request.name).is_none() { return Err(ErrorData::invalid_params( @@ -61,19 +62,18 @@ impl ServerHandler for Server { 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()), + Err(error) => { + return Ok(CallToolResult::error(vec![ContentBlock::text(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()) + let result = tokio::task::spawn_blocking(move || schema::execute(query)).await; + let result = match result { + Ok(Ok(stdout)) => CallToolResult::success(vec![ContentBlock::text(stdout)]), + Ok(Err(error)) => CallToolResult::error(vec![ContentBlock::text(error)]), + Err(error) => CallToolResult::error(vec![ContentBlock::text(format!( + "CLI task failed: {error}" + ))]), + }; + Ok(result.into()) } } diff --git a/tests/cli_thread.rs b/tests/cli_thread.rs index a34d71a..7a23710 100644 --- a/tests/cli_thread.rs +++ b/tests/cli_thread.rs @@ -27,12 +27,7 @@ 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}"); + assert!(text.contains("Serve MCP tools over stdio."), "{text}"); } else { assert!(text.contains("unrecognized subcommand 'mcp'"), "{text}"); } diff --git a/tests/dex_fixture/mod.rs b/tests/dex_fixture/mod.rs index 1a18fab..e06013b 100644 --- a/tests/dex_fixture/mod.rs +++ b/tests/dex_fixture/mod.rs @@ -203,26 +203,6 @@ 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 index 9ecdc3e..9fc48e7 100644 --- a/tests/mcp_stdio.rs +++ b/tests/mcp_stdio.rs @@ -3,16 +3,14 @@ mod dex_fixture; mod zip_fixture; -use std::collections::BTreeMap; +use std::ffi::OsString; use std::io::{BufRead, BufReader, Read, Write}; -use std::process::{Child, ChildStdin, Command, Stdio}; +use std::process::{Child, ChildStdin, Command, Output, 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 dex_fixture::{make_dex, make_instance_field_dex, make_invoke_dex}; use serde_json::{Value, json}; use tempfile::tempdir; use zip_fixture::{CompressionMethod, write_apk}; @@ -23,14 +21,11 @@ struct Client { output: Receiver>, stdout: Option>, stderr: Option>, + next_id: u64, } 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()) @@ -49,9 +44,9 @@ impl Client { 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| { + let result = line.map_err(|error| error.to_string()).and_then(|line| { serde_json::from_str(&line) - .map_err(|e| format!("non-protocol stdout: {line}: {e}")) + .map_err(|error| format!("non-protocol stdout: {line}: {error}")) }); if send.send(result).is_err() { break; @@ -64,17 +59,17 @@ impl Client { output, stdout: Some(stdout), stderr: Some(stderr), + next_id: 1, }; - client.send(json!({"jsonrpc":"2.0", "id":0, "method":"initialize", "params":{ - "protocolVersion":version, "capabilities":{}, "clientInfo":{"name":"test", "version":"1"} - }})); + client.send( + json!({"jsonrpc":"2.0", "id":0, "method":"initialize", "params":{ + "protocolVersion":"2025-11-25", + "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 @@ -88,37 +83,23 @@ impl Client { fn recv(&self) -> Value { self.output - .recv_timeout(Duration::from_secs(20)) + .recv_timeout(Duration::from_secs(30)) .expect("MCP reply timeout/disconnect") .expect("stdout must contain only JSON-RPC") } fn call(&mut self, name: &str, args: Value) -> Value { + let id = self.next_id; + self.next_id += 1; self.send( - json!({"jsonrpc":"2.0", "id":1, "method":"tools/call", "params":{ + json!({"jsonrpc":"2.0", "id":id, "method":"tools/call", "params":{ "name":name, "arguments":args }}), ); let response = self.recv(); - assert_eq!(response["id"], 1, "{response}"); + assert_eq!(response["id"], id, "{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() - ); + response["result"].clone() } fn finish(mut self) { @@ -138,9 +119,6 @@ impl Client { 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:?}"); - } } } @@ -152,38 +130,127 @@ impl Drop for Client { } } -fn success(result: &Value) -> &Value { - assert_eq!(result["isError"], false, "{result}"); - &result["structuredContent"] +fn text(result: &Value) -> &str { + result["content"][0]["text"].as_str().unwrap() +} + +fn run_cli(args: impl IntoIterator) -> Output { + Command::new(env!("CARGO_BIN_EXE_rust_ASC")) + .args(args) + .output() + .unwrap() +} + +fn cli_stdout(args: impl IntoIterator) -> String { + let output = run_cli(args); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + String::from_utf8(output.stdout).unwrap() } #[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); +fn exposes_four_documented_cli_wrappers() { + let mut client = Client::new(); + client.send(json!({"jsonrpc":"2.0", "id":99, "method":"tools/list"})); + let response = client.recv(); + let tools = response["result"]["tools"].as_array().unwrap(); + let mut names: Vec<_> = tools + .iter() + .map(|tool| tool["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["annotations"]["readOnlyHint"], false); + assert_eq!(tool["annotations"]["destructiveHint"], true); + assert!(tool.get("outputSchema").is_none()); + let description = tool["description"].as_str().unwrap(); + assert!( + (20..=80).contains(&description.len()), + "tool description is not concise: {description}" + ); + for forbidden in ["-", "—", "{", "apk_path", "class_name", "fuzzy_class"] { + assert!( + !description.contains(forbidden), + "tool description contains {forbidden}: {description}" + ); + } + assert_described_properties(&tool["inputSchema"]); + assert_concise_descriptions(&tool["inputSchema"]); + let output = tool["inputSchema"]["properties"]["output"]["description"] + .as_str() + .unwrap(); + if tool["name"] == "listclass" { + assert!(output.contains("Returned text is empty.")); + } else { + assert!(output.contains("Content is also returned.")); + } + } + client.send( + json!({"jsonrpc":"2.0", "id":100, "method":"tools/call", "params":{ + "name":"export", "arguments":{} + }}), + ); + assert_eq!(client.recv()["error"]["code"], -32602); + client.finish(); +} + +fn assert_described_properties(schema: &Value) { + if let Some(properties) = schema.get("properties").and_then(Value::as_object) { + for (name, property) in properties { + if name != "kind" { + assert!( + property + .get("description") + .and_then(Value::as_str) + .is_some_and(|description| !description.is_empty()), + "parameter {name} has no description: {property}" + ); + } + } + } + if let Some(object) = schema.as_object() { + for value in object.values() { + assert_described_properties(value); + } + } else if let Some(array) = schema.as_array() { + for value in array { + assert_described_properties(value); + } + } +} + +fn assert_concise_descriptions(schema: &Value) { + if let Some(object) = schema.as_object() { + if let Some(description) = object.get("description").and_then(Value::as_str) { + assert!( + description.len() <= 120, + "description is too long: {description}" + ); + assert!( + !description.contains('-') && !description.contains('—'), + "description contains a dash: {description}" + ); + } + for value in object.values() { + assert_concise_descriptions(value); + } + } else if let Some(array) = schema.as_array() { + for value in array { + assert_concise_descriptions(value); } - 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() { +fn tool_results_are_exact_cli_stdout() { let dir = tempdir().unwrap(); - let apk = dir.path().join("fixture.apk"); + let apk = dir.path().join("fixture with space.apk"); let dex = make_dex(); let manifest = minimal_manifest(); write_apk( @@ -191,221 +258,133 @@ fn four_queries_return_full_structured_data_and_do_not_write_files() { &[("classes.dex", &dex), ("AndroidManifest.xml", &manifest)], CompressionMethod::Deflated, ); + let path = apk.to_string_lossy().into_owned(); 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" + text(&result), + cli_stdout( + ["listclass", "--thread", "1", "--prefix", "example", &path,].map(OsString::from) + ) ); - 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( + let result = client.call( "getclass", - json!({"apk_path":apk,"class_name":"not.Found"}), - "NOT_FOUND", + json!({"apk_path":path, "class_name":"example.Test", "thread":1}), ); - client.error("getmanifest", json!({"apk_path":apk}), "NOT_FOUND"); - client.error( - "findrefs", - json!({"apk_path":apk,"query":{"kind":"method"}}), - "INVALID_ARGUMENT", + assert_eq!(result["isError"], false); + assert_eq!( + text(&result), + cli_stdout(["getclass", "--thread", "1", &path, "example.Test"].map(OsString::from)) ); - client.error( - "findrefs", - json!({"apk_path":apk,"query":{"kind":"invalid"}}), - "INVALID_ARGUMENT", + + let result = client.call("getmanifest", json!({"apk_path":path})); + assert_eq!(result["isError"], false); + assert_eq!( + text(&result), + cli_stdout(["getmanifest", &path].map(OsString::from)) ); - 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"}}), + json!({ + "apk_path":path, + "thread":1, + "query":{"kind":"string", "value":"token"} + }), + ); + assert_eq!(result["isError"], false); + assert_eq!( + text(&result), + cli_stdout(["findrefs", "--thread", "1", &path, "string", "token"].map(OsString::from)) ); - assert_eq!(success(&result)["references"], json!([])); + assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 1); client.finish(); } #[test] -fn malformed_inputs_are_not_reported_as_empty_or_partial_success() { +fn optional_output_writes_files_with_cli_semantics() { 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", - ); - } + let dex = make_dex(); + let manifest = minimal_manifest(); write_apk( &apk, - &[("classes.dex", &make_dex()), ("classes2.dex", b"not a dex")], + &[("classes.dex", &dex), ("AndroidManifest.xml", &manifest)], 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, + let path = apk.to_string_lossy().into_owned(); + let mut client = Client::new(); + + let class_output = dir.path().join("Test.java"); + let result = client.call( + "getclass", + json!({ + "apk_path":path, + "class_name":"example.Test", + "thread":1, + "output":class_output + }), ); - client.error( - "findrefs", - json!({"apk_path":apk,"query":{"kind":"string","value":"token"}}), - "INVALID_FORMAT", + assert_eq!(result["isError"], false, "{result}"); + assert_eq!( + std::fs::read_to_string(&class_output).unwrap().trim_end(), + text(&result).trim_end() ); - let mut broken_container = make_dex041_container(); - broken_container.truncate(broken_container.len() - 8); - write_apk( - &apk, - &[("classes.dex", &broken_container)], - CompressionMethod::Stored, + + let classes_output = dir.path().join("classes.txt"); + let result = client.call( + "listclass", + json!({"apk_path":path, "thread":1, "output":classes_output}), ); - client.error("listclass", json!({"apk_path":apk}), "INVALID_FORMAT"); - write_apk( - &apk, - &[("AndroidManifest.xml", b"bad xml")], - CompressionMethod::Stored, + assert_eq!(result["isError"], false, "{result}"); + assert_eq!( + std::fs::read_to_string(&classes_output).unwrap(), + "Lexample/Test;\n" ); - 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(); -} + assert_eq!(text(&result), ""); -#[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 manifest_output = dir.path().join("AndroidManifest.xml"); + let result = client.call( + "getmanifest", + json!({"apk_path":path, "output":manifest_output}), ); - let result = client.call("listclass", json!({"apk_path":apk})); + assert_eq!(result["isError"], false, "{result}"); assert_eq!( - success(&result)["classes"], - json!([ - "Lalpha/First;", - "Lshared/Duplicate;", - "Lbeta/Second;", - "Lshared/Duplicate;" - ]) + std::fs::read_to_string(&manifest_output) + .unwrap() + .trim_end(), + text(&result).trim_end() ); + + let refs_output = dir.path().join("references.txt"); let result = client.call( - "getclass", - json!({"apk_path":apk,"class_name":"beta.Second"}), + "findrefs", + json!({ + "apk_path":path, + "thread":1, + "query":{"kind":"string", "value":"token"}, + "output":refs_output + }), + ); + assert_eq!(result["isError"], false, "{result}"); + assert_eq!( + std::fs::read_to_string(&refs_output).unwrap(), + text(&result) ); - assert_eq!(success(&result)["dex_name"], "classes.dex!classes2.dex"); client.finish(); } #[test] -fn all_reference_kinds_have_structured_results() { +fn findrefs_query_variants_map_to_cli_arguments() { 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 apk = dir.path().join("references.apk"); 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") @@ -414,87 +393,95 @@ fn all_reference_kinds_have_structured_results() { write_apk( &apk, &[ - ("classes.dex", &invoke_dex), - ("classes2.dex", &make_instance_field_dex()), - ("classes3.dex", &type_dex), + ("classes.dex", &type_dex), + ("classes2.dex", &make_invoke_dex()), + ("classes3.dex", &make_instance_field_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") - ); - } + let path = apk.to_string_lossy().into_owned(); + let mut client = Client::new(); + let cases = [ + ( + json!({"kind":"type", "value":"example"}), + vec!["type", "example"], + ), + ( + json!({ + "kind":"method", + "name":"callToMethod", + "class_name":"example.Target" + }), + vec!["method", "callToMethod", "--class", "example.Target"], + ), + ( + json!({ + "kind":"field", + "name":"a", + "class_name":"example", + "fuzzy_class":true + }), + vec!["field", "a", "--class", "example", "--fuzzy-class"], + ), + ]; + for (query, query_args) in cases { + let result = client.call("findrefs", json!({"apk_path":path, "query":query})); + assert_eq!(result["isError"], false, "{result}"); + let mut cli_args = vec![OsString::from("findrefs"), OsString::from(&path)]; + cli_args.extend(query_args.into_iter().map(OsString::from)); + assert_eq!(text(&result), cli_stdout(cli_args)); } client.finish(); } #[test] -fn concurrent_requests_have_independent_ids_and_complete_untruncated_results() { +fn cli_failures_are_isolated_and_the_server_remains_usable() { 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); + write_apk( + &apk, + &[("classes.dex", &make_dex())], + CompressionMethod::Stored, + ); + let path = apk.to_string_lossy().into_owned(); 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()); + + let result = client.call("getclass", json!({"apk_path":path})); + assert_eq!(result["isError"], true); + assert!(text(&result).contains("invalid arguments"), "{result}"); + + let cli = run_cli( + ["getclass", &path, "example.Missing"] + .map(OsString::from) + .into_iter(), + ); + assert!(!cli.status.success()); + let result = client.call( + "getclass", + json!({"apk_path":path, "class_name":"example.Missing"}), + ); + assert_eq!(result["isError"], true); + assert_eq!(text(&result), String::from_utf8_lossy(&cli.stderr).trim()); + + let result = client.call("listclass", json!({"apk_path":path})); + assert_eq!(result["isError"], false); + assert_eq!(text(&result), "Lexample/Test;\n"); 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 u16_(out: &mut Vec, value: u16) { + out.extend(value.to_le_bytes()); } - fn u32_(out: &mut Vec, n: u32) { - out.extend(n.to_le_bytes()); + fn u32_(out: &mut Vec, value: u32) { + out.extend(value.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); diff --git a/tests/reference_output.rs b/tests/reference_output.rs deleted file mode 100644 index fb26221..0000000 --- a/tests/reference_output.rs +++ /dev/null @@ -1,151 +0,0 @@ -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() - ); -} From 32b6fbb61ea96dff3969ea9108fe3abac36fc05a Mon Sep 17 00:00:00 2001 From: MG1937 <2586364982@qq.com> Date: Mon, 19 Oct 2026 20:49:07 +0800 Subject: [PATCH 3/4] add test for dev --- .github/workflows/test.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index df4a4d5..1aa6463 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -4,6 +4,7 @@ on: push: branches: - main + - dev pull_request: workflow_dispatch: From 343c329caf4ea6f86bc0a8845728f7efe494a54b Mon Sep 17 00:00:00 2001 From: MG1937 <2586364982@qq.com> Date: Mon, 19 Oct 2026 21:07:31 +0800 Subject: [PATCH 4/4] add test for dev --- tests/mcp_stdio.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/mcp_stdio.rs b/tests/mcp_stdio.rs index 9fc48e7..8dc62eb 100644 --- a/tests/mcp_stdio.rs +++ b/tests/mcp_stdio.rs @@ -186,7 +186,7 @@ fn exposes_four_documented_cli_wrappers() { .as_str() .unwrap(); if tool["name"] == "listclass" { - assert!(output.contains("Returned text is empty.")); + assert!(output.contains("stdout is empty.")); } else { assert!(output.contains("Content is also returned.")); }