From cd91dbed9c5b7da8da866516ffa38c8b07cb790d Mon Sep 17 00:00:00 2001 From: Tsingshui Date: Mon, 19 Oct 2026 14:41:16 +0800 Subject: [PATCH] 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() + ); +}