diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 00000000..9a43f475 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,15 @@ +# Memory guard for test harnesses. +# +# Doc tests link one binary per fenced example, and `[profile.release]` sets `lto = "fat"`, +# so every one of those links is a whole-crate LTO job costing gigabytes of RAM. libtest +# defaults its concurrency to one thread per core, so a bare `cargo test --release --doc` +# on a many-core machine starts dozens of them at once and takes the host out of memory +# before any test reports. +# +# Capping here rather than only in the Makefile means the guard holds for anyone who runs +# cargo directly. Cargo's `[env]` does not override a variable that is already set, so the +# unit-test run — which is cheap per test and wants every core — opts back up to full +# parallelism by exporting `RUST_TEST_THREADS` itself; see the `test` target in the Makefile. +# An explicit `-- --test-threads=N` also still wins, which is what `make test-doc` and CI use. +[env] +RUST_TEST_THREADS = "8" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 336a41d7..7e553a1a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -94,22 +94,26 @@ jobs: - name: Run tests (Linux - full) if: runner.os == 'Linux' - run: cargo test --workspace --release --lib --features z3 --verbose + run: cargo test --workspace --release --features z3 --lib --bins --tests --verbose - name: Run tests (macOS/Windows - skip expensive tests) if: runner.os != 'Linux' - run: cargo test --workspace --release --lib --features z3,skip-expensive-tests --verbose + run: cargo test --workspace --release --features z3,skip-expensive-tests --lib --bins --tests --verbose - name: Run doc tests (Linux only) if: runner.os == 'Linux' - run: cargo test --workspace --release --doc --features z3 + run: cargo test --workspace --release --doc --features z3 -- --test-threads=4 - name: Check documentation + env: + RUSTDOCFLAGS: "-Dwarnings" run: cargo doc -p dotscope --all-features --no-deps - # Ensures the crate compiles and tests pass without legacy-crypto feature (Linux only) + # Compiles and tests without the legacy-crypto feature (Linux only). Pinned to the declared + # MSRV so this job also guards `rust-version`; the workspace check covers the default-feature + # path, which `--no-default-features` alone would not. minimal-features: - name: Minimal Features + name: Minimal Features & MSRV runs-on: ubuntu-latest steps: @@ -117,18 +121,21 @@ jobs: uses: actions/checkout@v7 - name: Install Rust toolchain - uses: dtolnay/rust-toolchain@stable + uses: dtolnay/rust-toolchain@1.95.0 - name: Setup Rust cache uses: Swatinem/rust-cache@v2 with: key: minimal-features + - name: Check compilation (default features, whole workspace) + run: cargo check --workspace --all-targets + - name: Check compilation (no default features) run: cargo check -p dotscope --no-default-features - name: Run tests (no default features) - run: cargo test -p dotscope --no-default-features --release --verbose + run: cargo test -p dotscope --no-default-features --release --lib --bins --tests --verbose # Fuzzing on pushes to main or PRs targeting main fuzzing: @@ -154,19 +161,41 @@ jobs: with: key: fuzz-quick - - name: Build fuzz target - run: cargo +nightly fuzz build --release cilobject + - name: Build fuzz targets + run: cargo +nightly fuzz build --release working-directory: dotscope/fuzz + # Every target gets the committed crash corpus. The inputs are PE-shaped, so they are + # directly meaningful to `cilobject`/`assemblyview` and act as structured starting material + # for the blob and body targets rather than starting them from scratch. + - name: Seed corpora from committed regressions + run: | + for t in cilobject assemblyview signatures customattributes methodbody emulation; do + mkdir -p "dotscope/fuzz/corpus/$t" + cp dotscope/tests/samples/fuzz-regressions/* "dotscope/fuzz/corpus/$t/" || true + done + - name: Run quick fuzz test - run: timeout 60 cargo +nightly fuzz run cilobject --release -- -max_total_time=50 || true + run: | + for t in cilobject assemblyview signatures customattributes methodbody emulation; do + echo "::group::fuzz $t" + timeout 60 cargo +nightly fuzz run "$t" --release -- \ + -max_total_time=40 -timeout=25 -rss_limit_mb=2048 || true + echo "::endgroup::" + done working-directory: dotscope/fuzz - name: Check for crashes run: | - if [ -d "dotscope/fuzz/artifacts/cilobject" ] && [ "$(ls -A dotscope/fuzz/artifacts/cilobject)" ]; then - echo "Fuzzing found crashes!" - ls -la dotscope/fuzz/artifacts/cilobject/ + found=0 + for t in cilobject assemblyview signatures customattributes methodbody emulation; do + if [ -d "dotscope/fuzz/artifacts/$t" ] && [ "$(ls -A "dotscope/fuzz/artifacts/$t")" ]; then + echo "Fuzzing found crashes in $t:" + ls -la "dotscope/fuzz/artifacts/$t" + found=1 + fi + done + if [ "$found" = "1" ]; then exit 1 fi diff --git a/CHANGELOG.md b/CHANGELOG.md index 60e98fa4..6a356d6f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,227 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.9.0] - 2026-08-15 + +A security and correctness release. dotscope parses, emulates and rewrites +hostile input, and this release closes the gap between what the resource limits +claimed to enforce and what they actually did, along with a set of +miscompilations in the SSA back end and layout defects in the PE writer. + +### Security + +- **Resource limits are enforced before the work happens, not after.** The + managed-heap ceiling was checked once the object was already materialised; + unmanaged allocation (`localloc`, `AllocHGlobal`, `AllocCoTaskMem`, + `VirtualAlloc`) had no budget at all, and `max_unmanaged_bytes` and + `max_heap_objects` were declared but never read. Allocation + now runs through a reservation that must succeed first, in-place mutation is + accounted, and forks inherit the ceiling instead of escaping it. +- **Unbounded and quadratic work on attacker input.** Fixed in the inheritance + walker (a cyclic `extends` graph caused an uncatchable native stack overflow), + the x86 traversal (O(n²) to end of file), method-body decoding (disassembled + past the declared `code_size`), exception-handler association (O(H²·B) at load + time), DEFLATE/GZIP/LZMA expansion, and the signature parser (a blob could + build a ~61 000-deep type whose recursive drop overflowed the stack). +- **Argument validation across the BCL hooks.** Negative or oversized lengths + reaching `Marshal.Copy`, `Stream.SetLength`, `StringBuilder.set_Length`, + `String.PadLeft`/`PadRight`, the `BinaryReader` readers and the PBKDF2 + constructors reserved `usize::MAX`, ran multi-billion-iteration loops, or + drove a ~4.3-billion-round KDF. They now reject the value and raise the .NET + exception. +- **Emulator forks were not isolated.** "Isolated" forks shared one mutable + runtime state, AppDomain and synthetic-method map while running concurrently. + `Assembly.Load(byte[])` is now bounded by `max_loaded_assemblies` and + `max_loaded_assembly_bytes`, and runtime-loaded assemblies parse with minimal + validation rather than the full pipeline over hostile bytes. +- **Memory protection flags are enforced** on read and write, faulting through a + new catchable `AccessViolationException`, and region mappings are overlap-checked. +- `deny(unsafe_code)` is enabled. One `unsafe` block remains, for the writer's + output mapping, with a targeted allow and a SAFETY note. +- `SECURITY.md` now states the supported version, the real `EmulationLimits` + defaults and what is actually run. The previous text listed DoS protections as + "ToDo" and claimed Valgrind testing that does not exist. + +### Fixed + +- **Malformed table rows silently truncated a table.** The row iterators + reported a parse failure as end-of-iteration, and because the writer rebuilds + tables by iterating them, an unreadable row became *missing output* rather + than an error. Iterators now yield `Result`, `get` returns + `Result>`, and `MetadataTable::new` validates and truncates to the + declared extent. +- **`MethodPtr`, `EventPtr` and `PropertyPtr` tokens used the wrong table id**, + so any assembly carrying a `*Ptr` table lost its method-bearing types. +- **Three back-end miscompilations.** Full inlining placed the return-value copy + before the instruction defining it; switch and conditional-branch phi + trampolines fell through into the next edge's copies. Critical edges are now + split into real blocks by a dedicated out-of-SSA pass. +- **Handler SSA used a "last block wins" snapshot** of try-scope definitions + because the CIL CFG carried no exception edges. Real EH edges make handler + entries ordinary join points. +- **Linear-scan allocation computed live intervals with no liveness solve**, so + a value live across a back edge could have its slot clobbered. +- **Four exception-unwind defects**: the caller's `finally` ran against the + grandparent frame, queued `finally` blocks were never drained once a catch was + selected, a `leave` out of nested `finally`s spun on `endfinally`, and a filter + returning zero terminated emulation instead of resuming the handler search. +- **PE writer layout.** Heap offsets were computed twice from different inputs, + so offsets baked into tables and IL disagreed with where data was written; + heap index widths were inherited from the input and truncated above 0xFFFF; + section `SizeOfRawData` came from the virtual extent; and the input's + certificate directory offset was applied to the output, zeroing live `.text` + before the checksum was computed over the damage. `Output` now writes to a + temp file and renames. +- **Cleanup deleted live metadata**: TypeRef liveness ignored `ResolutionScope`, + and the opaque-field pass folded any static-to-instance load and deleted the + owning type with no immutability precondition. +- Byte-offset slicing of string literals panicked on multi-byte UTF-8; + `clippy::string_slice` is now denied, which surfaced ten genuine sites. +- The fuzz crash-corpus regression test passed on any checkout without the + corpus, and CI ran `cargo test --lib`, so the integration tests never executed + on Windows or macOS. Both are fixed, and the 72 crash artifacts are committed. +- **An array signature's rank was never bounded**, and it was the only ceiling on + the lower-bound count that follows it, so a declared rank of 0x400000 made that + check permissive rather than protective and the dimension list grew to the + declared count before any read could run out of input. This accounted for every + out-of-memory artifact found by fuzzing. +- **Type-name validation rejected legitimate compiler-generated names.** It + matched a hand-written list of prefixes, so `{GUID}` failed on untouched + input as well as on rewritten output; the closed angle bracket the C# compiler + guarantees is the real invariant. Validation failures also reported only how + many validators failed, discarding the messages saying why. +- **Cleanup deleted enclosing types whose nested types were still referenced**, + leaving a NestedClass row pointing at a TypeDef that no longer existed. + Reachability now walks the nesting relation to a fixed point. +- **Reachability used the SSA call graph alone**, so every method without SSA + looked unreachable and the live set was under-approximated. SSA edges are now + preferred where they exist and the static graph fills in where they do not. +- **Opaque static fields were only folded when every write came from a `.cctor`.** + Obfuscators route initialization through helpers, so those fields stayed opaque + and their predicates survived. A write site now counts when every caller of the + writing method is itself initialization-only; a method with no known caller is + not admitted. .NET Reactor string samples go from 223 decryption failures to + none. +- **Parameters removed with their method left dangling references behind them.** + `Constant`, `FieldMarshal` and `CustomAttribute` rows name a parameter through + a coded index and are dropped by asking whether their parent was deleted, but a + parameter discarded along with its method never entered that record — what had + been deleted was the method. The rows outlived the parameters they named and + the output failed raw validation with an out-of-range `Param` RID. Removed + parameters are now cascaded to all three tables. +- **.NET Reactor NecroBit recovered nothing from full-protection binaries.** + Every encrypted body was lost on both such samples — 0 of 59 and 0 of 562 — + while necrobit-only binaries were unaffected. The cause was not in the + decryption: a protection that resolves `VirtualProtect` through + `LoadLibrary`/`GetProcAddress` and calls it through a delegate never reached + the hook that implements it, so the pages holding the method bodies stayed + read-only and the write-back faulted on the first body. Both samples now + restore every stub and validate. +- **A native function resolved at runtime never reached its hook.** Hook matching + required a declared P/Invoke, so any function obtained through `GetProcAddress` + and invoked through `Marshal.GetDelegateForFunctionPointer` bypassed it — the + delegate path answered from a small table of hardcoded return values instead, + reporting success without performing the call's effect. Such calls now carry + their arguments and dispatch through the ordinary hook path. `LoadLibrary` + hands out a distinct handle per module so the resolved function can be matched + against the library it came from. +- **A refused write was retried as a fresh mapping.** `Marshal`'s write path + treated "mapped, but not writable" the same as "not mapped" and tried to + materialise a window at the enclosing 64KB boundary. For an address inside a + loaded image that is the image base, so the attempt collided with the image and + reported an overlap — turning a recoverable permission error into a fatal one + that named the wrong cause. The two cases are now distinguished. +- **A failed body-decryption transform caused cleanup to delete the code it + could not decrypt.** A technique fills its cleanup request during detection, + before it knows whether the transform those deletions depend on will run. When + a byte transform fails, the bodies it was meant to restore stay encrypted; + such a method contributes no call edges, so everything it references reads as + unreachable and the type-level sweep removes it. One .NET Reactor sample fell + from 1181 methods to 87. Techniques now report what they could not restore + (`Technique::unrecovered_methods`), cleanup protects those methods, withholds + the failed technique's own request, and skips unreferenced-type removal for the + run — the call graph cannot tell unreachable from undecrypted. The same sample + now keeps 946 methods and validates. +- **Unflattening could emit a function that failed SSA validation**, which + aborted deobfuscation for the whole assembly rather than the method. Rewiring a + dispatcher edge can skip a definition that a surviving block still reads; the + guards that prevent this have gaps, so the rebuilt form is now checked and a + method that cannot be rewired safely is left flattened. + +### Performance + +- `EmValue` drops from 200 to 104 bytes on x86-64 by boxing `CilFlavor::FnPtr`, + halving every value in the interpreter. Pinned by a static assertion. +- `Method`, `CilType` and `Param` no longer eagerly allocate 8–11 `Arc>` + each; `LazyList` defers to first use. +- Type resolution: `get_by_fullname` no longer falls back to a linear scan of + every registered type (reachable once per custom-attribute argument), and + `fullname()` returns `Arc` instead of allocating a fresh `String` at every + call site. +- Declaring-type lookups are indexed rather than brute-force scans over every + type and member — they sat on the emulator's hottest paths. +- Table loaders no longer take a shared `Mutex` once per row inside the rayon + loop; an inherent `try_for_each` had been shadowing rayon's in every loader. +- Handler SSA no longer rebuilds a version-stack snapshot per exception + successor per block, and unmanaged access is a `BTreeMap` lookup rather than a + linear region scan that allocated a `Vec` for a 1–8 byte read. + +### Changed + +- **BREAKING**: **CFF unflattening resolves dispatcher edges from SSA instead of + enumerating execution paths.** The old tracer walked the method from entry and + forked at every conditional, which is exponential in the number of branches and + re-explored the whole method once per dispatcher. The state reaching a + dispatcher is a phi whose operands are indexed by predecessor, so the value on + each edge can simply be read; recovering it is linear in the number of edges. + Encodings that derive each state from the previous one are resolved by a fixed + point over states — one iteration per original block, not per path. + + On one .NET Reactor sample the tree cost 108.8 million nodes and 62 seconds + across 40 dispatchers; the same work now takes 0.16. `reactor_full` drops from + 615 to 154 seconds, of which unflattening is 1.1. The tracer and the patch-plan + reconstruction are deleted, roughly 3 400 lines net. + + Edges are only rewired when the answer is provable: the case index is obtained + by evaluating the dispatcher's own switch operand rather than a reconstructed + transform, arithmetic folds at the operand's width because state encodings rely + on int32 wraparound, and a block is skippable only when everything it computes + feeds the state machine and nothing else. An edge that cannot be resolved keeps + routing through the dispatcher, so coverage degrades rather than correctness, + and blocks holding a call, a store or a string are never removed on the strength + of an analysis that is allowed to be incomplete. + + `unflatten` and `unflatten_with_dispatchers` no longer take a config or an + assembly, `CffReconstructionPass::new` takes only the context, and the patch + plan API is gone. `UnflattenConfig` and `UnflatteningThresholds` lose the knobs + that drove path enumeration; the ones that remain are now actually applied by + detection, which previously built a config and then ignored it. +- The workspace declares `rust-version = "1.95"`, and the minimal-features CI job + is pinned to it — and extended with a default-feature workspace check — so the + MSRV is verified rather than merely stated. +- **BREAKING**: `Error` is `#[non_exhaustive]` and derives `Clone`. The previous + hand-written `Clone` rewrote most variants into `Error::Other(String)`, + destroying the taxonomy for any caller that cloned. +- **BREAKING**: `MetadataTable::get` returns `Result>` and the table + iterators yield `Result`. +- **BREAKING**: `CilType::fullname()` returns `Arc`; `UserStrings::get` + returns an owned `U16String`; `derive_pbkdf2_key` returns `Result` and errors + on an unavailable algorithm instead of silently substituting SHA-256 for SHA-1. +- `CaptureConfig` gains `max_items` and `max_total_bytes` ceilings (10 000 and + 256 MB), and buffer capture is off under a default config, honouring the + documented "no capture by default" contract. `CaptureContext::new()` sets it + explicitly, so the "capture what is useful" constructor is unchanged. +- Stale rustdoc `# Errors` contracts across the crate referenced `Error` variants + that had been deleted. They are rewritten to what the code returns, and + `deny(rustdoc::broken_intra_doc_links)` plus `RUSTDOCFLAGS: -Dwarnings` in CI + keeps them accurate — `RUSTFLAGS` does not reach rustdoc, which is why the + existing `-Dwarnings` never caught them. +- Doc tests run under a concurrency cap: each fenced example is a whole-crate + fat-LTO link, and one per core exhausts memory on a many-core machine. +- Five new fuzz targets beside `cilobject`, covering the assembly view, the + signature and custom-attribute blob parsers, method-body decode and bounded + emulation. + ## [0.8.5] - 2026-08-09 ### Fixed diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 9de019b3..20939f25 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -151,7 +151,7 @@ cargo fmt --all -- --check cargo doc --all-features --no-deps # For extended fuzzing (do this locally, not in CI): -cd fuzz && cargo +nightly fuzz run cilobject -- -max_total_time=1800 # 30 minutes +cd dotscope/fuzz && cargo +nightly fuzz run cilobject -- -max_total_time=1800 # 30 minutes ``` ### 4. Submit Pull Request @@ -188,7 +188,7 @@ cd fuzz && cargo +nightly fuzz run cilobject -- -max_total_time=1800 # 30 minut 3. **Fuzzing**: Test with random inputs ```bash - cd fuzz + cd dotscope/fuzz cargo +nightly fuzz run cilobject ``` diff --git a/Cargo.lock b/Cargo.lock index af49f288..f5635c1c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1564,7 +1564,7 @@ dependencies = [ [[package]] name = "dotscope" -version = "0.8.5" +version = "0.9.0" dependencies = [ "aes", "analyssa", @@ -1610,7 +1610,7 @@ dependencies = [ [[package]] name = "dotscope-cli" -version = "0.8.5" +version = "0.9.0" dependencies = [ "anyhow", "clap", diff --git a/Cargo.toml b/Cargo.toml index 2f1edc01..eb465c08 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ resolver = "2" [workspace.package] edition = "2021" +rust-version = "1.95" license = "Apache-2.0" repository = "https://github.com/ATRAPSLLC/dotscope" homepage = "https://github.com/ATRAPSLLC/dotscope" diff --git a/Makefile b/Makefile index c75d3ec5..db115c82 100644 --- a/Makefile +++ b/Makefile @@ -1,13 +1,31 @@ # Makefile for dotscope development # Provides convenient commands for common development tasks -.PHONY: help build test clean fmt clippy doc bench fuzz install coverage audit +# Doc tests link one binary per fenced example, and `[profile.release]` uses +# `lto = "fat"`, so each link is a whole-crate LTO job costing gigabytes of RAM. The +# harness defaults to one thread per core, which on a many-core machine means dozens of +# those at once and an out-of-memory freeze. Cap the concurrency instead; override with +# `make test-doc DOC_TEST_THREADS=n`. +# +# `.cargo/config.toml` sets `RUST_TEST_THREADS = "8"` so that guard also holds for anyone +# who runs cargo directly instead of going through this file. This default matches it: +# a value above the config's guard would silently defeat it, since an explicit +# `-- --test-threads=N` always wins over cargo's `[env]`. +# +# Unit tests are cheap per test and want every core, so the `test` target exports its own +# `RUST_TEST_THREADS` to opt back out -- cargo's `[env]` deliberately does not override a +# variable that is already set. +DOC_TEST_THREADS ?= 8 +TEST_THREADS ?= $(shell nproc 2>/dev/null || echo 8) + +.PHONY: help build test test-doc clean fmt clippy doc bench fuzz install coverage audit # Default target help: @echo "Available targets:" @echo " build - Build the project" @echo " test - Run all tests" + @echo " test-doc - Run doc tests (memory-bounded; see DOC_TEST_THREADS)" @echo " clean - Clean build artifacts" @echo " fmt - Format code" @echo " clippy - Run clippy lints" @@ -28,8 +46,20 @@ build-release: cargo build --release --all-features # Run tests +# +# Release profile is mandatory, not a speed preference: the sample-driven integration tests +# run the full detection -> SSA -> pass-pipeline -> codegen path over real packed binaries, +# which takes hours unoptimised and tens of seconds optimised. +# `--lib --bins --tests` excludes doc tests, which have their own target below because they +# need a concurrency cap this run does not. test: - cargo test --workspace --all-features --verbose + RUST_TEST_THREADS=$(TEST_THREADS) cargo test --workspace --release --all-features --lib --bins --tests --verbose + +# Run doc tests +# +# Split out from `test` and concurrency-capped: see DOC_TEST_THREADS at the top of this file. +test-doc: + cargo test --workspace --release --all-features --doc -- --test-threads=$(DOC_TEST_THREADS) # Run tests with coverage test-coverage: @@ -67,8 +97,20 @@ bench: cargo bench --all-features # Run fuzzing +# Runs every target in turn, each seeded from the committed crash corpus. Override the target +# with `make fuzz FUZZ_TARGETS=signatures`, or the duration with `FUZZ_TIME=300`. +FUZZ_TARGETS ?= cilobject assemblyview signatures customattributes methodbody emulation +FUZZ_TIME ?= 60 +FUZZ_RSS_LIMIT_MB ?= 4096 + fuzz: - cd fuzz && cargo +nightly fuzz run cilobject -- -max_total_time=60 + @for t in $(FUZZ_TARGETS); do \ + echo "=== fuzzing $$t ==="; \ + mkdir -p dotscope/fuzz/corpus/$$t; \ + cp dotscope/tests/samples/fuzz-regressions/* dotscope/fuzz/corpus/$$t/ 2>/dev/null || true; \ + (cd dotscope/fuzz && cargo +nightly fuzz run $$t -- \ + -max_total_time=$(FUZZ_TIME) -rss_limit_mb=$(FUZZ_RSS_LIMIT_MB)) || exit 1; \ + done # Install development tools install: @@ -89,7 +131,7 @@ outdated: cargo outdated # Run all checks -check-all: fmt-check clippy test audit +check-all: fmt-check clippy test test-doc audit @echo "All checks passed!" # Prepare for release @@ -102,5 +144,5 @@ dev: fmt clippy test @echo "Development cycle completed" # CI simulation (run what CI runs) -ci: fmt-check clippy test doc +ci: fmt-check clippy test test-doc doc @echo "CI simulation completed" diff --git a/README.md b/README.md index 06c2eab8..3a8cc46c 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,7 @@ Add `dotscope` to your `Cargo.toml`: ```toml [dependencies] -dotscope = "0.8.4" +dotscope = "0.9" ``` ### Raw Access Example @@ -339,7 +339,7 @@ See the [examples README](examples/README.md) for a recommended learning path. Security is a top priority: - **Memory Safety**: Built on Rust's memory safety guarantees -- **Fuzzing**: Continuous fuzzing with cargo-fuzz +- **Fuzzing**: `cargo-fuzz` targets run on demand (`make fuzz`), with a committed crash corpus replayed by the test suite - **Input Validation**: Strict validation of all inputs - **Audit Trail**: Regular dependency auditing @@ -385,7 +385,7 @@ make coverage # Generate coverage report make fuzz # Extended fuzzing (manual) -cd fuzz && cargo +nightly fuzz run cilobject --release -- -max_total_time=1800 +cd dotscope/fuzz && cargo +nightly fuzz run cilobject --release -- -max_total_time=1800 # All quality checks make check-all diff --git a/SECURITY.md b/SECURITY.md index 6c9796ca..a1d94ed8 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -6,7 +6,11 @@ We actively support the following versions of dotscope with security updates: | Version | Supported | | ------- | ------------------ | -| 0.1.x | :white_check_mark: | +| 0.9.x | :white_check_mark: | +| < 0.9 | :x: | + +dotscope is pre-1.0 and ships breaking changes in minor releases. Only the +latest minor version receives security fixes; there are no backports. ## Reporting a Vulnerability @@ -60,29 +64,53 @@ dotscope parses potentially untrusted .NET assemblies. We take several precautio - **Memory Safety**: Built on Rust's memory safety guarantees - **Bounds Checking**: All array and buffer accesses are bounds-checked -- **Fuzzing**: Continuous fuzzing with cargo-fuzz to find parsing edge cases +- **Fuzzing**: `cargo-fuzz` targets covering the object, view, signature, + custom-attribute, method-body and emulation paths, run on demand via + `make fuzz`. Crash artifacts are committed and replayed by the test suite - **Input Validation**: Strict validation of metadata structures and bytecode +Note that `ValidationConfig::disabled()` and the `lenient` presets do not +disable bounds checking -- what they give up is *semantic* rejection, so +incoherent metadata is analysed as if it were coherent. + ### Denial of Service Protection -- **Resource Limits**: ToDo -- **Timeout Handling**: ToDo +Emulation runs under `EmulationLimits`, enforced during execution rather than +after the fact. Defaults: + +| Limit | Default | +| ----- | ------- | +| `max_instructions` | 10,000,000 | +| `max_call_depth` | 1,000 | +| `max_heap_objects` | 100,000 | +| `max_heap_bytes` | 256 MB | +| `max_unmanaged_bytes` | 64 MB | +| `max_loaded_assemblies` | 64 | +| `max_loaded_assembly_bytes` | 32 MB | +| `timeout_ms` | 60,000 | + +- **Timeout Handling**: the wall-clock budget is checked between instructions - **Malformed Input**: Graceful handling of corrupted or crafted files ### Known Security Considerations 1. **Memory-Mapped Files**: We use memory mapping for performance, which requires careful handling -2. **Unsafe Code**: Limited use of `unsafe` code with careful review and testing +2. **Unsafe Code**: the crate builds under `deny(unsafe_code)`. One block carries + a targeted allow, for the memory mapping of the writer's output file. The + primary load path maps input through the `cowfile` dependency 3. **Dependency Chain**: Regular auditing of dependencies for vulnerabilities ## Security Testing Our security testing includes: -- **Continuous Fuzzing**: Automated fuzzing with various input types -- **Static Analysis**: Clippy and other static analysis tools +- **Fuzzing**: six `cargo-fuzz` targets, run on demand and in CI, seeded from + the committed crash corpus +- **Regression Corpus**: every crash artifact found by fuzzing is committed and + replayed by the test suite +- **Static Analysis**: Clippy with `panic`, `unwrap`, `expect`, `indexing_slicing`, + `arithmetic_side_effects` and `string_slice` denied in the library crate - **Dependency Auditing**: Regular `cargo audit` runs -- **Memory Safety**: Valgrind testing for memory leaks and corruption ## Acknowledgments diff --git a/docs/research/netreactor/necrobit.md b/docs/research/netreactor/necrobit.md index 48b911b1..78a6e362 100644 --- a/docs/research/netreactor/necrobit.md +++ b/docs/research/netreactor/necrobit.md @@ -2,12 +2,19 @@ Analysis of .NET Reactor 7.5.0 NecroBit protection based on reverse engineering `reactor_necrobit.exe` (70,144 bytes, 258 methods) against `original.exe` (14,336 bytes, -35 methods) using dotscope disassembly. +35 methods) using dotscope disassembly. The named tokens and RVAs throughout are +from that sample; `reactor_full.exe` and `reactor_virtualization_full.exe` +provide the full-protection layout, whose runtime is split across more methods +and whose storage differs (see *Two Operating Modes*). NecroBit is .NET Reactor's most critical protection. It encrypts all method bodies, replacing them with stubs. Without reversing NecroBit first, all other deobfuscation stages operate on encrypted bytecode and produce no results. +Both storage variants are fully supported — see *Support Status* for the +per-sample recovery counts, and *What the Emulator Must Provide* for what +reversing them depends on. + ## File-Level Changes | Property | Original | Protected | Delta | @@ -242,24 +249,23 @@ Instead, it resolves and calls VirtualProtect dynamically: - Signature: `int32 Invoke(native int, int32, int32, int32&)` — matches `BOOL VirtualProtect(LPVOID, SIZE_T, DWORD, PDWORD)` -### Emulation Impact of VirtualProtect +### Why the VirtualProtect Return Value Matters -The VirtualProtect delegate invocation is **critical for CFF path selection**. -The init method checks the VirtualProtect return value to determine which -CFF switch cases to execute. If the delegate invocation fails (returns -Symbolic instead of `TRUE`/`1`), the CFF flow takes incorrect branches and -skips the Hashtable population code path entirely. +The delegate's return value drives **CFF path selection** inside the init +method, not just the success of the memory-protection call. The init checks it to +decide which switch cases to execute: -Observed behavior when VirtualProtect delegate fails: -- The init method takes the **direct-write mode**: writes individual method - table fields via `Marshal.WriteInt32` (872 calls to 368 addresses) -- Skips the **Hashtable mode**: never creates the Hashtable, never calls - `Hashtable.Add`, never stores complete method bodies -- The direct-write mode writes to CLR runtime addresses that don't exist - in the emulated address space, producing no usable data -- The emulation terminates with `EndOfStreamException` from `BinaryReader` - after processing all entries (this is the normal loop termination signal) +- **Returns `TRUE`**: the init proceeds through the Hashtable-population path, + building the per-method records and writing the decrypted bodies back. +- **Returns anything else** (including a symbolic value, when the call is not + modelled): the flow takes different branches, skips Hashtable population + entirely, and falls into a direct-write path whose `Marshal.WriteInt32` calls + target CLR runtime addresses that do not exist in an emulated process. +So an emulator that cannot resolve this call does not merely lose the protection +change — it steers the protection down a path that produces no usable data at +all. This is why the resolution chain has to be modelled end to end rather than +short-circuited with a plausible return value. ## Injected Type Inventory (33 new types) @@ -318,9 +324,11 @@ At runtime, the NecroBit protection works as follows: ## Deobfuscation Strategy -### Emulation-Based Approach (dotscope) — Verified Working +### Emulation-Based Approach (dotscope) -Our approach uses emulation to let the protection's own code do the decryption: +Rather than reimplementing the cipher, dotscope runs the protection's own +initialization under emulation and reads the result out of the emulated process. +Both storage variants are recovered this way — see *Support Status*. 1. **Detect** via structural patterns (stub methods, .cctor fan-in, trial check pattern, body patcher pattern) — no hardcoded names @@ -330,13 +338,17 @@ Our approach uses emulation to let the protection's own code do the decryption: - Injected .cctors bypassed to prevent re-entrancy - Full BCL hook coverage for Marshal, Process, Module, crypto, streams - Transparent pinned array support for managed/native shared backing -3. **Extract decrypted bodies** via two strategies (best result wins): - - **Heap extraction**: Find a byte array on the managed heap matching the - NecroBit data format (see Data Format below). Parses Variant A (with group - entries) or Variant B (inline complete bodies). - - **PE image extraction**: Read patched method bodies directly from the PE - image in the address space. In full-protection binaries, the init method - writes decrypted bodies to PE RVAs via `Marshal.WriteInt32`. +3. **Extract decrypted bodies** via one of two strategies, chosen by variant — + not "best result wins": the heap path is tried first and the image path is + the fallback when it declines. + - **Heap extraction** (variant A): Find a byte array on the managed heap + matching the NecroBit data format (see Data Format below). + `find_variant_a_blob` skips `group_count == 0`, because variant B's + authoritative bodies are in the address space rather than the heap. + - **PE image extraction** (variant B): Read patched method bodies directly + from the PE image in the address space. The init method writes decrypted + bodies to PE RVAs via `Marshal.WriteInt32`, which requires the emulated + `VirtualProtect` to have made those pages writable. 4. **Store restored bodies**: Use `CilAssembly::store_method_body()` to replace stub RVAs with real method body data. 5. **Regenerate PE**: Rebuild the assembly. @@ -391,22 +403,96 @@ The init method reads a **data format flag** from the decrypted data stream via `IntPtr` addresses and `Marshal.WriteInt32`. Extraction reads from the PE image in the address space. -Both modes work generically — no version-specific logic is needed. - -### Critical Emulation Fixes - -Two bugs in `IntPtr` BCL hooks caused all method body writes to target address 0: - -1. **`IntPtr.ToInt64`** did not handle `ManagedPointer` as `this` (from - `ldloca + call` on a value type local). Fix: dereference the managed pointer - to extract the NativeInt value. - -2. **`IntPtr..ctor`** did not store the value when `this` was `ObjectRef` (from - `newobj` allocating a heap object). Fix: store the value as a synthetic field - on the heap object. - -Without these fixes, all `Marshal.WriteInt32` calls wrote to address 0 instead -of the correct PE image RVAs, and no method bodies were recoverable. +### Support Status + +Both variants are fully recovered. + +| Sample | Variant | Stubs | Restored | +|---|---|---|---| +| `reactor_necrobit.exe` | A | 45 | 45 | +| `reactor_necrobit_strings_cff.exe` | A | 56 | 56 | +| `reactor_full.exe` | B | 59 | 59 | +| `reactor_virtualization_full.exe` | B | 562 | 562 | + +Asserted per variant by `restores_every_stub_in_a_necrobit_only_binary` and +`restores_every_stub_in_a_full_protection_binary` in `necrobit.rs`. Both count +restored bodies rather than checking that the transform returned `Ok` — a +transform that recovers *some* bodies reports success, so a count is the only +assertion that distinguishes full recovery from partial, or from none. + +### What the Emulator Must Provide + +NecroBit is decrypted by running the protection's own code, so recovery depends +on the emulator being faithful in a few specific places. Each of these is load +bearing: without it, extraction yields nothing rather than less. + +**A working `VirtualProtect`, reached through dynamic resolution.** +Variant B's bodies are written into `.text`, which the PE loader maps +`READ|EXECUTE` from the section characteristics. The write only succeeds because +the protection calls `VirtualProtect` first — and it never declares that +P/Invoke, resolving it through `LoadLibrary`/`GetProcAddress` and calling it +through a delegate (see *VirtualProtect Resolution Chain*). Hook matching keyed +on declared P/Invoke alone therefore never sees it. + +dotscope preserves identity across the whole chain: `LoadLibrary` returns a +distinct handle per module, `GetProcAddress` records the resolved function +against that module, `GetDelegateForFunctionPointer` turns the fake address into +a token, and the delegate invocation reconstructs a native call context from it. +The call then runs through the ordinary hook path with its real arguments, so the +`kernel32!VirtualProtect` hook installs a page protection override — the *effect*, +not merely a `TRUE` return value. Returning success without applying the +protection is indistinguishable from working right up until the first write. + +**Writes into a mapped-but-unwritable page must report that.** +A write refused on protection is not the same as a write to unmapped memory. +Treating them alike — materialising a fresh region at the enclosing 64KB +boundary — maps over the loaded image, because for an image address that +boundary is the image base. The resulting overlap error names the wrong cause and +is fatal. `AccessViolation` is propagated; only genuinely unmapped addresses are +auto-allocated. + +**`IntPtr` values must survive both construction forms.** The body patcher builds +its target addresses through `IntPtr`, from a value-type local (`ldloca + call`) +and from `newobj`. See *IntPtr Construction Forms* below. + +**Cleanup must not run ahead of decryption.** A method whose body is still an +encrypted stub contributes no call edges, so everything it references reads as +unreachable. If the transform fails, deleting on that basis removes the original +code — the reason the sample was kept. Techniques report what they could not +restore via `Technique::unrecovered_methods`; cleanup protects those methods, +withholds the failed technique's own request, and skips type-level +unreferenced-type removal, since the call graph cannot then distinguish +unreachable from undecrypted. + +### On the Variant B Heap Blob + +A blob shaped like the NecroBit data format is present on the managed heap during +a variant-B run, and `is_necrobit_data_array` accepts it. `find_variant_a_blob` +discards it (`group_count == 0`) and extraction reads the address space instead, +which is correct: variant B's authoritative bodies are the ones written into the +image. + +Note that the variant-B branch of that shape check is weak — `len >= 36`, a +MethodDef token byte at offset 0, and one method-body header byte at offset 32. +It gates nothing today because variant B never parses from the heap. Anything +that starts parsing heap blobs should validate by walking the record chain +instead: each RVA resolving to a known method, each `MethodBody::from` succeeding, +and the walk terminating at the buffer's end. + +### IntPtr Construction Forms + +The body patcher computes its write targets through `IntPtr`, reaching the BCL +hooks in two shapes that must both be handled: + +1. **`IntPtr.ToInt64` with a `ManagedPointer` as `this`**, from `ldloca + call` + on a value-type local. The pointer has to be dereferenced to reach the + `NativeInt` value. + +2. **`IntPtr..ctor` with an `ObjectRef` as `this`**, from `newobj` allocating a + heap object. The value is stored as a synthetic field on that object. + +Handling only one of them leaves every `Marshal.WriteInt32` writing to address 0 +rather than a PE image RVA, and no bodies are recoverable. ### Comparison with NRS Approach diff --git a/docs/research/netreactor/virtualization.md b/docs/research/netreactor/virtualization.md index df5c8a77..23fb0530 100644 --- a/docs/research/netreactor/virtualization.md +++ b/docs/research/netreactor/virtualization.md @@ -1,12 +1,17 @@ # Code Virtualization (VM Protection) -Analysis of .NET Reactor 7.5.0 code virtualization based on reverse engineering -`reactor_virtualization.exe` (125,440 bytes, 851 methods) against `original.exe` -(14,336 bytes, 35 methods) using dotscope disassembly. +Analysis of .NET Reactor 7.5.0 code virtualization, based on `reactor_virtualization.exe` +(125,440 bytes, 851 methods) and `reactor_virtualization_full.exe` (361,984 bytes, +1181 methods) against `original.exe` (14,336 bytes, 35 methods). + +Tokens, field names and IL offsets throughout are from `reactor_virtualization.exe` +unless stated otherwise. Names are the obfuscated ones as they appear in the sample; +they are per-build and carry no meaning across samples — treat them as coordinates +for re-verification, not as identifiers to match on. This is .NET Reactor's most sophisticated protection. It converts CIL method bodies -into custom bytecode interpreted by an embedded virtual machine. 8 methods were -virtualized; the remaining methods are untouched. +into a custom instruction stream interpreted by an embedded interpreter. 8 methods +were virtualized; the remaining methods are untouched. ## File-Level Changes @@ -24,18 +29,61 @@ virtualized; the remaining methods are untouched. ## Virtualized Method Stub Format -All 8 virtualized methods follow an identical stub pattern: +All 8 virtualized methods follow an identical stub pattern. `Calculator::Add`, +verbatim: + +``` +IL_0000: ldc.i4 2 +IL_0005: newarr System.Object +IL_000a: stloc.0 +IL_000b: ldloc.0 +IL_000c: ldc.i4 0 +IL_0011: ldarg 1 +IL_0015: box System.Int32 +IL_001a: stelem.ref +IL_001b: ldloc.0 +IL_001c: ldc.i4 1 +IL_0021: ldarg 2 +IL_0025: box System.Int32 +IL_002a: stelem.ref +IL_002b: ldc.i4 0 // method id +IL_0030: ldloc.0 // boxed argument array +IL_0031: ldarg.0 // 'this' (null for static) +IL_0032: call object[] BD6lOYUCm3(int32, object[], object) /* 0x060000A9 */ +IL_0037: stloc.1 +IL_0038: ldloc.1 +IL_0039: ldc.i4.0 +IL_003a: ldelem.ref +IL_003b: unbox.any System.Int32 +IL_0040: ret +``` + +Two properties of this stub matter more than anything else in the VM: + +- **The MethodDef signature is untouched.** `Add` is still + `instance int32 Add(int32, int32)`. Unlike KoiVM, .NET Reactor destroys no + type information at the method boundary. +- **The `box`/`unbox.any` pairs name the exact parameter and return types.** + Even if the signature were stripped, the stub itself spells out the boundary + types. + +Together these give a devirtualizer a fully typed entry and exit for every +virtualized method, for free. See [Type Recovery](#type-recovery-is-cheap-here). + +### VM Entry Point ``` -ldc.i4 // unique integer identifying the VM bytecode -newarr object[] // create array for boxed arguments -// ... pack each argument into object[] via box + stelem.ref ... -ldarg.0 // push 'this' (or null for static) -call BD6lOYUCm3 // VM entry point (token 0x060000A9) -// ... unbox return value from returned object[] ... -ret +.method assembly static object[] BD6lOYUCm3(int32, object[], object) /* 0x060000A9 */ + ldc.i4.0; stloc.0 + ldarg.0; ldarg.1; ldarg.2; ldloca.s 0 + call object[] UhSaWDOYZTgwwgfSO6::j4XlTwXGiJ(int32, object[], object, !!0&) + ret /* MethodSpec 0x2B000001 */ ``` +**Detection**: find `call 0x060000A9` preceded by `ldc.i4 `. Structurally: +a static method whose only act is to forward all arguments to a generic method, +called from many small methods that box their arguments into an `object[]`. + ### Method ID Mapping | ID | Original Method | Notes | @@ -49,11 +97,9 @@ ret | 6 | `SecretHolder::DecryptSecret` | Try/catch + crypto | | 7 | `SecretHolder::XorEncrypt` | Loop + char ops | -**Non-virtualized methods** (Subtract, Multiply, Divide, DemoLoop, etc.) retain +**Non-virtualized methods** (`Subtract`, `Multiply`, `Divide`, `DemoLoop`, …) retain their original IL completely unmodified. -**Detection**: Find `call 0x060000A9` (the VM entry point) preceded by `ldc.i4 `. - ## VM Architecture @@ -61,42 +107,147 @@ their original IL completely unmodified. ``` Virtualized stub - -> BD6lOYUCm3 (entry point) - -> MethodSpec 0x2B000001 (generic VM entry) - -> j4XlTwXGiJ (body loader: loads/decrypts bytecode, creates context) - -> lIKxeQZNPA (entry wrapper) - -> vjuxAdNiYK (execution loop) - -> XGtxjqudOH (opcode dispatcher: 176-case switch) + -> BD6lOYUCm3 0x060000A9 entry point + -> j4XlTwXGiJ MethodSpec 0x2B000001 + body loader: decrypt, decode, build context + -> lIKxeQZNPA 0x060002D7 entry wrapper ← the harvest seam + -> vjuxAdNiYK step/execute loop + -> XGtxjqudOH 0x060002DE dispatcher: 175-case switch +``` + +`j4XlTwXGiJ` calls `lIKxeQZNPA()` exactly once, at `IL_08db`. Everything before +that call is decode; everything after is execution. That single seam is what makes +static-free extraction practical — see +[Extraction Strategy](#extraction-strategy-harvest-the-decoded-program). + +### The Program Is an Object Graph, Not a Byte Stream + +This is the single most important structural fact about .NET Reactor's VM, and it +is what separates it from KoiVM, EazVM and VirtualGuard. + +After decode, a virtualized method is a +`List` — a list of **instruction objects**. The instruction +type (`0x0200002f`) has exactly two fields: + +| Field | Token | Type | Role | +|-------|-------|------|------| +| `qLOHqjHEu6` | `0x0400006E` | `guRiCRPRexpb1c3DuN0` (enum, `0x02000045`) | **opcode** | +| `x8xH1QMtT2` | `0x0400006F` | `object` | **operand** | + +The operand is a boxed CLR object, and by the time it is stored it has already +been *materialized*: the loader resolves metadata into live `System.Type`, +`MethodBase` and `MethodInfo` instances (122 `System.Type` and 20 `MethodBase` +references in `j4XlTwXGiJ` alone), decodes constants with the custom varint reader +`eKLl9Pieuj`, and stores `null` for operand-less opcodes. A tag byte read with +`ldelem.u1` selects the per-kind decode path. + +So there is no persistent "VM bytecode" to disassemble. There is an encrypted +resource, and there is a decoded object graph. Nothing in between survives. + +### Fetch–Dispatch + +`vjuxAdNiYK` fetches and dispatches: + +``` +IL_0080: ldfld MwSjwlZbVA /* 0x04000093 */ // VM context -> program holder +IL_0085: ldfld u18xPWpKQV /* 0x0400008B */ // -> List +IL_008a: ldfld QULjDgrIqA /* 0x0400009A */ // program counter +IL_0090: callvirt List::get_Item(int32) +IL_0095: stloc.0 // fetched instruction +IL_0096: ldloc.0 +IL_0098: ldfld x8xH1QMtT2 /* 0x0400006F */ // operand +IL_009d: stfld W6nj3gfLbK /* 0x0400009D */ // -> VM context operand slot +.try { +IL_00a4: call XGtxjqudOH(instruction) /* 0x060002DE */ +``` + +The operand is stashed on the VM context *before* dispatch; handlers read it from +there rather than from their argument. + +| Field | Token | Purpose | +|-------|-------|---------| +| `QULjDgrIqA` | `0x0400009A` | Program counter | +| `pRejhltIYn` | `0x0400009B` | Previous/saved PC | +| `W6nj3gfLbK` | `0x0400009D` | Current operand | +| `RKyjTXlhNj` | `0x0400009E` | Branch flag | +| `qTqjOcsEDb` | `0x0400009F` | Return flag | +| `i2Ij2Wrl3j` | `0x040000A0` | Halt flag | +| `qeGj0vFdv1` | `0x04000096` | Operand stack (type `C6jel6PFv1y17TI7U6B`) | + +The three flag fields are read-and-clear: the loop tests each, resets it, and +either returns or continues. `ItKxx2Gt1k(int32, int32)` performs the branch +bookkeeping when `RKyjTXlhNj` is set. Exception dispatch is handled around the +call site with `List` handler lists and explicit +`TargetInvocationException` unwrapping. + +### Handlers Are Regions, Not Methods + +`XGtxjqudOH` (`0x060002DE`, 3,908 lines of disassembly, 39 locals) switches on the +opcode enum read directly off the instruction object: + +``` +IL_0000: ldarg.1 +IL_0001: ldfld guRiCRPRexpb1c3DuN0 pmhFGRPKUN0cv5gbsq2::qLOHqjHEu6 /* 0x0400006E */ +IL_0006: stloc.0 +IL_0007: ldloc.0 +IL_0008: switch ( IL_145c, IL_198d, IL_20f6, … ) // 175 entries +``` + +**175 case entries, plus the default path.** Duplicate targets are common +(`IL_27f6` appears at indices 10, 14, 48, 49; `IL_145c` at 0, 44, 56 …), so the +~154 distinct targets from the earlier count are consistent. + +Every handler body is **inline in this one method**. There is no handler table, no +array of delegates, and no per-opcode method. A "handler" is a region of basic +blocks inside `XGtxjqudOH` dominated by one switch-case target. + +This defeats the usual detection heuristic — "a type containing many small methods +with similar signatures" — completely. .NET Reactor's VM has zero handler methods. + +### Semantics Live Behind Virtual Dispatch + +Case 5 (`IL_046d`), a comparison handler, in full: + +``` +IL_046d: ldarg.0; ldfld qeGj0vFdv1 // operand stack +IL_0473: callvirt jvj3LnPpvxU2erTMeM4 C6jel6PFv1y17TI7U6B::Ysk3wvblFN() /* 0x06000344 = pop */ +IL_0478: stloc.s 4 +IL_047a: ldarg.0; ldfld qeGj0vFdv1 +IL_0480: callvirt Ysk3wvblFN() /* pop */ +IL_0485: call wm4NTEPOv5wQhC99DbT dIB6JIPIiI7GlyxJGUd::On9h7vfNJU(jvj3LnPpvxU2erTMeM4) +IL_048a: ldloc.s 4 +IL_048c: callvirt bool wm4NTEPOv5wQhC99DbT::ojAOxCx0yc(jvj3LnPpvxU2erTMeM4) +IL_0491: brfalse IL_04a8 +IL_0496: ldarg.0; ldfld qeGj0vFdv1 +IL_049d: newobj KTvJuZPhZohptT4uVf6::.ctor(int32) // push 1 +IL_04a2: callvirt void C6jel6PFv1y17TI7U6B::hYG3VSV4XB(jvj3LnPpvxU2erTMeM4) /* 0x06000342 = push */ +IL_04a7: ret +IL_04a8: … newobj KTvJuZPhZohptT4uVf6::.ctor(int32) // push 0 +IL_04b4: callvirt hYG3VSV4XB(…) +IL_04b9: ret ``` -### Execution Loop: `vjuxAdNiYK` (308 lines, 12 locals) +Read structurally, this region says: *pop two values, call something, push 0 or 1*. +It does **not** say which comparison. The actual operation lives in +`ojAOxCx0yc`, resolved at runtime by the concrete type of the popped VM value — +one of the four parallel value types. 152 of the 983 calls in the dispatcher go +directly to those four types. -| Field | Purpose | -|-------|---------| -| `0x0400009A` | Program counter (instruction pointer) | -| `0x0400009B` | Previous/saved PC | -| `0x0400009E` | Branch flag | -| `0x0400009F` | Return flag (exits loop when set) | -| `0x040000A0` | Halt flag | -| `0x04000093 -> 0x0400008B` | Bytecode array (indexed by PC) | -| `0x0400006F` | Opcode ID field (on fetched instruction) | -| `0x04000096` | Virtual operand stack | +**Consequence:** SSA pattern matching on a handler region cannot classify .NET +Reactor's opcodes. Structure yields the *stack effect*; it does not yield the +*operation*. Any classifier that works here must be behavioural. -The loop: -1. Reads instruction from bytecode array at current PC -2. Reads opcode ID from the instruction's opcode field -3. Dispatches via `XGtxjqudOH` -4. Wraps execution in try/catch for VM-level exception handling +### Execution Is Reflection-Driven -### Opcode Dispatcher: `XGtxjqudOH` (4,781 lines, 39 locals) +The dispatcher performs CIL-level operations through reflection: +`MethodBase.Invoke` (2 sites), `FieldInfo.GetValue` (14), `FieldInfo.SetValue` +(4), and `Activator`/`ConstructorInfo.CreateInstance` (7). Locals of type +`ConstructorInfo`, `MethodInfo`, `MethodInfo[]` and `List` are +declared at the top of `XGtxjqudOH`. -- **176-case switch** statement on the opcode ID -- Maps to **~154 unique handler targets** (some opcodes share handlers) -- Stack operations via polymorphic calls on field `0x04000096`: - - **Pop**: `callvirt 0x06000344` (row 836) - - **Push**: `callvirt 0x06000342` (row 834) - - **Stack Count**: `callvirt 0x06000340` (row 832) - - **Stack Top**: `callvirt 0x06000341` (row 833) +For a devirtualizer this means a call operand is a *live `MethodInfo`*, not a +metadata token. Emitting `call`/`callvirt`/`newobj`/`ldfld` requires mapping the +reflection object back to its defining token. ## VM Type System @@ -118,10 +269,13 @@ The type classifier `HEWjGoIVXq` (245 lines) maps .NET types to VM type IDs: | 8 | UInt64 | 18 | Null | | 9 | Single (Float32) | | | +Every value on the virtual stack carries one of these tags at runtime. That is a +per-value type oracle a devirtualizer can read directly out of emulated state. + ### 4 Parallel VM Value Types Four large types implement the polymorphic VM value abstraction, each specialized -for different data widths: +for different data widths. `jvj3LnPpvxU2erTMeM4` is the common base. | Type | Methods | Unique Methods | Likely Specialization | |------|---------|----------------|----------------------| @@ -131,20 +285,23 @@ for different data widths: | `aqAvMsP2XIn3GUhf0si` | 88 | 0 | Minimal | Each shares ~70 identical method signatures (`JfS3vO6qin`, `Th73tWZsCb`, -`FHPTbr88Ae`, `Add`, `ToString`, arithmetic/comparison handlers). -Arithmetic is dispatched polymorphically through these types. +`FHPTbr88Ae`, `Add`, `ToString`, arithmetic/comparison handlers). Arithmetic and +comparison are dispatched polymorphically through these types — which is precisely +what makes the dispatcher's own IL semantics-free. ### Value Boxing -`jVsj6JAdhh` (855 lines, 19-case switch): Boxes .NET values into the appropriate -VM value type based on the type ID from the classifier. +`jVsj6JAdhh` (855 lines, 19-case switch): boxes .NET values into the appropriate +VM value type based on the type ID from the classifier. Called 3 times from the +loader, when marshalling the incoming `object[]` arguments onto the virtual stack. ## Bytecode Storage and Decryption ### Embedded Resource -Bytecode stored as an embedded assembly resource, loaded via: +The encrypted program is stored as an embedded assembly resource, loaded via: + ``` ldtoken TypeDef(row 18) // UhSaWDOYZTgwwgfSO6 call Type.GetTypeFromHandle() @@ -158,19 +315,22 @@ callvirt Assembly.GetManifestResourceStream(string) `oYiEKuNrl` (866 lines, 20 locals): - AES-based decryption with custom key derivation - Key scheduling: calls rows 57/58 sixteen times each (round key setup) -- Custom padding: `(448 - len*8) % 512` (Merkle-Damgard style) +- Custom padding: `(448 - len*8) % 512` (Merkle–Damgård style) - Key derivation in `tRelL85we1`: 32-byte key from seed, `ICryptoTransform` -### Method Body Loading: `j4XlTwXGiJ` (1,141 lines, 46 locals) +### Method Body Loading: `j4XlTwXGiJ` (933 lines of disassembly, 46 locals) + +1. Reads from a static cache array (field row 71) indexed by method ID +2. If not cached, reads from the decrypted resource stream (field row 74) +3. Uses a **custom compressed integer encoding** (`eKLl9Pieuj`, 16 call sites — + 6-bit base plus continuation bits; *not* .NET's standard compressed integer) +4. Reads parameter types, local variable types, instruction count, instruction data +5. Materializes each instruction as an object: opcode enum + resolved operand +6. **Token resolution** via `osalebjCgR`: masks with `0x0FFFFFFF`, indexes a + pre-built token array (field row 79) populated during initialization -1. Reads from static cache array (field row 71) indexed by method ID -2. If not cached, reads from decrypted resource stream (field row 74) -3. Uses **compressed integer encoding** (`eKLl9Pieuj` — 6-bit base + continuation bits, - custom format, not .NET's standard compressed integer) -4. Reads: parameter types, local variable types, instruction count, instruction data -5. Each instruction has an **opcode field** (0-175) and operand data -6. **Token resolution** via `osalebjCgR`: masks with `0x0FFFFFFF`, looks up in - pre-built token array (field row 79) +Step 1 means the program for method *id* is only decoded when *id* is first +invoked. Harvesting all eight requires driving all eight ids. ## Injected Type Inventory (79 new types) @@ -179,10 +339,10 @@ callvirt Assembly.GetManifestResourceStream(string) | Type | Namespace | Methods | Role | |------|-----------|---------|------| -| `UhSaWDOYZTgwwgfSO6` | `Dinih72WZsCb9wcqjy` | 23 | VM bootstrap, entry point, bytecode loader | -| `dIB6JIPIiI7GlyxJGUd` | — | 28 | VM execution engine, 176-case dispatcher | -| `oMu6jVbdhHEH79DDhU` | `AoIBWWlDJbaf7LijnA` | 49 | Bytecode decryptor, crypto | -| `jvj3LnPpvxU2erTMeM4` | — | 17 | VM type system, value boxing | +| `UhSaWDOYZTgwwgfSO6` | `Dinih72WZsCb9wcqjy` | 23 | VM bootstrap, entry point, body loader | +| `dIB6JIPIiI7GlyxJGUd` | — | 28 | VM context + 175-case dispatcher | +| `oMu6jVbdhHEH79DDhU` | `AoIBWWlDJbaf7LijnA` | 49 | Resource decryptor, crypto | +| `jvj3LnPpvxU2erTMeM4` | — | 17 | VM value base type, boxing | | `uy4ZXuP8hhYiKuNrl4` | `AsG4wKEPrjKTCY31dc` | 3 | Token resolver | | `gttro5yuWySr2hbdEM` | `UWxvxUSU2ZrCqT9K8B` | 2 | License date check | | `gEHfEJ9aJKgHNTQig9` | — | 2 | Crypto helper (AES/SHA) | @@ -190,12 +350,8 @@ callvirt Assembly.GetManifestResourceStream(string) ### VM Value Types (4 parallel instantiations) -| Type | Methods | -|------|---------| -| `KFLZn1PTyq5gG2AjYpx` | 104 | -| `KTvJuZPhZohptT4uVf6` | 93 | -| `ShYNHsP3JZ1gO8BRpVH` | 92 | -| `aqAvMsP2XIn3GUhf0si` | 88 | +`KFLZn1PTyq5gG2AjYpx` (104), `KTvJuZPhZohptT4uVf6` (93), +`ShYNHsP3JZ1gO8BRpVH` (92), `aqAvMsP2XIn3GUhf0si` (88). ### VM Stack/Reader Variants (5 types, 11 methods each) @@ -206,92 +362,99 @@ callvirt Assembly.GetManifestResourceStream(string) | Type | Methods | Role | |------|---------|------| -| `LbcypsPYVeerQlaEm3O` | 6 | Instruction key (Equals, GetHashCode) | -| `NWOcXrPi1WKXw4SWh1B`1` | 8 | Nullable wrapper | -| `C6jel6PFv1y17TI7U6B` | 6 | Opcode descriptor | +| `C6jel6PFv1y17TI7U6B` | 6 | **Operand stack** — `hYG3VSV4XB` push (`0x06000342`), `Ysk3wvblFN` pop (`0x06000344`) | +| `pmhFGRPKUN0cv5gbsq2` | — | Instruction: opcode enum + `object` operand | +| `guRiCRPRexpb1c3DuN0` | — | Opcode enum (`0x02000045`) | +| `LbcypsPYVeerQlaEm3O` | 6 | Instruction key (`Equals`, `GetHashCode`) | +| `NWOcXrPi1WKXw4SWh1B\`1` | 8 | Nullable wrapper | | `FwrX5yPtqhsabjCgRnP` | 3 | Static initializer | | `d8DE92F8305BE09E` | 11 | String interpolation handler | | 10+ small types | 1-2 each | Enums, structs, exception types | -## Key Findings for Devirtualization - -### 1. Uniform Stub Detection - -All virtualized methods call the same entry point (`BD6lOYUCm3`, token `0x060000A9`) -with a numeric method ID. Detection is trivial: find `call ` preceded by -`ldc.i4 `. - -### 2. Single Central Dispatcher - -One 176-case switch handles all VM opcodes. Each case is a distinct handler block -with predictable structure — amenable to semantic analysis. - -### 3. Stack-Based VM - -Classic operand stack model (push/pop through virtual methods). Operations are -polymorphic calls on the 4 VM value types. This mirrors CIL's own stack model, -which simplifies lifting back to CIL. - -### 4. Encrypted Bytecode +## What This Means for Devirtualization -Requires emulating AES decryption + key derivation to extract raw bytecode. -The emulation engine's AES BCL support can handle RijndaelManaged. +### Extraction Strategy: Harvest the Decoded Program -### 5. Runtime Token Resolution +The obvious approach — reimplement AES key derivation, the custom varint reader, +the resource format and the token array in Rust — is the wrong one. It is a +per-version reimplementation of machinery the sample already contains, and .NET +Reactor changes it between releases. -Token resolver masks with `0x0FFFFFFF` and indexes into a pre-built array. -The array is populated during initialization from the encrypted resource. +The alternative uses infrastructure dotscope already has, and mirrors what already +works for NecroBit variant A: **emulate the sample's own loader, then read the +decoded program out of the emulator heap.** -### 6. Custom Instruction Encoding +1. Hook `lIKxeQZNPA` (`0x060002D7`) to return immediately. Decode completes; + execution never starts. +2. Emulate `BD6lOYUCm3(id, args, null)` once per method id, with `args` populated + from the stub's `box` types. +3. Walk `process.address_space().managed_heap()` for `HeapObject::Object` entries + whose `type_token` is the instruction type, reading field `0x0400006E` + (opcode) and `0x0400006F` (operand). Order comes from the backing array of the + `List` reachable from the VM context. -Variable-length: 6-bit base + continuation bits (`eKLl9Pieuj`). Not .NET's -standard compressed integer format — needs custom parser. +Every primitive this needs already exists — `HeapObject::Object { type_token, +fields: HashMap }` gives field-token-keyed access, and +`find_variant_a_blob` in `necrobit.rs` is the same harvest applied to an array. -### 7. Complexity Assessment +The residual risk is emulator coverage of the crypto and reflection the loader +performs, not the design. That is measurable up front: emulate one method id and +see whether the list materializes. -176 opcodes, ~154 unique handlers, 4,781 lines of dispatcher IL, 39 locals. -This is a **full-featured VM**, not a simple wrapper. Devirtualization requires -the generic VM framework described in -[design/vm_devirtualization.md](../../design/vm_devirtualization.md). +### Classification Strategy: Probe the VM, Don't Read the Handlers +Since handler regions are semantics-free (see +[Semantics Live Behind Virtual Dispatch](#semantics-live-behind-virtual-dispatch)), +opcodes must be classified behaviourally. -## Devirtualization Strategy +The cheapest behavioural method does not require entering a handler region at all. +Synthesize a **one-instruction program** — a `List` containing a single +instruction object with opcode *k* and a chosen operand — install it on a VM +context with a known virtual stack, run one dispatch step, and diff the resulting +state: -Following the layered architecture from `vm_devirtualization.md`: +- stack depth delta and the type tags of pushed values → stack effect +- pushed value for known inputs → which arithmetic/comparison operation +- PC delta and branch-flag writes → control flow role +- reflection calls observed during the step → `call` / `newobj` / field access -### Layer 1: Detection & Extraction +This exercises the real prologue, the real value-type dispatch, and the real +operand plumbing, and it needs no new emulator capability — no mid-method entry, +no synthetic frames. It also generalizes past .NET Reactor: any VM whose loader +and dispatcher can be emulated can be probed this way, regardless of whether its +handlers are methods, regions, or delegates, and regardless of opcode +randomization. -1. **Detect VM stubs**: Find methods calling the VM entry point with `ldc.i4 ` -2. **Extract encrypted bytecode**: Locate the embedded resource via the `ldtoken` + - `GetManifestResourceStream` pattern -3. **Decrypt**: Emulate the decryption pipeline (`oYiEKuNrl` + `tRelL85we1`) +### Type Recovery Is Cheap Here -### Layer 2: Handler Analysis +Three independent sources of type information survive virtualization: -1. **Map the 176-case switch**: Each case is a handler block — extract the opcode-to-handler mapping -2. **Classify handlers semantically**: Use SSA pattern matching on each handler to - determine what CIL operation it implements (add, sub, load, store, branch, call, etc.) -3. **Handle the 4 value types**: Recognize the polymorphic dispatch through the - value type hierarchy +1. The MethodDef signature, untouched. +2. The stub's `box`/`unbox.any` operands, naming each boundary type exactly. +3. The runtime type tag (0–18) carried by every value on the virtual stack. -### Layer 3: SSA Lifting +Compare KoiVM, where all type information is destroyed and must be reconstructed +by data flow analysis. .NET Reactor's virtualization is significantly friendlier +in this one respect, and it matters: dotscope's `SsaFunctionBuilder::build_with` +rejects any used variable still typed `SsaType::Unknown`, so a lifter must type +every value it creates. -1. **Parse bytecode**: Implement the custom compressed integer encoding -2. **Resolve tokens**: Reconstruct the token array from the encrypted resource -3. **Build CFG**: Use branch/return flags to identify basic blocks -4. **Lift to SSA**: Map VM opcodes to `SsaFunction` operations, leveraging the - semantic classification from Layer 2 +### Ordering Constraint -### Layer 4: Integration +In `reactor_virtualization_full.exe`, NecroBit encrypts every method body — +including the VM loader and dispatcher. Devirtualization there is strictly +downstream of NecroBit body decryption: until those 562 stubs are recovered, the +VM is not merely unanalysable, it is not present in readable form. -1. **Run through compiler pipeline**: The lifted SSA gets all 21 optimization passes -2. **Emit CIL**: Replace the stub with the devirtualized method body -3. **Cleanup**: Remove VM infrastructure types (79 types, 816 methods) +### Complexity Assessment -### dotscope Infrastructure Leverage +175 switch cases, ~154 distinct handler regions, 3,908 lines of dispatcher IL, 39 +locals, semantics behind polymorphic dispatch on four value types, execution via +reflection. This is a full-featured VM, and it is a poor fit for the byte-stream, +handler-table model that KoiVM and EazVM share. -- **Emulation engine**: Decrypt the bytecode resource (AES BCL support) -- **SSA framework**: Direct-to-SSA lifting avoids intermediate IR -- **Compiler passes**: All 21 optimization passes apply to lifted code -- **Cleanup pipeline**: Removes the massive VM infrastructure +The generic framework needed to accommodate it is described in +[design/vm_devirtualization.md](../../design/vm_devirtualization.md); the +NET Reactor findings above are what motivate that document's extraction and +classification boundaries. diff --git a/dotscope-cli/Cargo.toml b/dotscope-cli/Cargo.toml index 018931cf..5fcdc805 100644 --- a/dotscope-cli/Cargo.toml +++ b/dotscope-cli/Cargo.toml @@ -1,7 +1,8 @@ [package] name = "dotscope-cli" -version = "0.8.5" +version = "0.9.0" edition.workspace = true +rust-version.workspace = true license.workspace = true repository.workspace = true homepage.workspace = true diff --git a/dotscope-cli/src/commands/attrs.rs b/dotscope-cli/src/commands/attrs.rs index 9ecfa63c..58e31e65 100644 --- a/dotscope-cli/src/commands/attrs.rs +++ b/dotscope-cli/src/commands/attrs.rs @@ -61,6 +61,9 @@ pub fn run(path: &Path, owner_filter: Option<&str>, opts: &GlobalOptions) -> any let mut entries: Vec = Vec::new(); for row in ca_table { + let Ok(row) = row else { + continue; + }; let (owner_kind, owner_name) = resolve_owner(&assembly, &row.parent, strings); let attr_type = resolve_constructor_type(&assembly, &row.constructor, strings); @@ -206,6 +209,9 @@ fn resolve_string_from_raw_table( "field" => { if let Some(table) = tables.table::() { for row in table { + let Ok(row) = row else { + continue; + }; if row.rid == ci.row { if let Ok(s) = strings.get(row.name as usize) { return s.to_string(); @@ -217,6 +223,9 @@ fn resolve_string_from_raw_table( "param" => { if let Some(table) = tables.table::() { for row in table { + let Ok(row) = row else { + continue; + }; if row.rid == ci.row { if let Ok(s) = strings.get(row.name as usize) { return s.to_string(); @@ -228,6 +237,9 @@ fn resolve_string_from_raw_table( "property" => { if let Some(table) = tables.table::() { for row in table { + let Ok(row) = row else { + continue; + }; if row.rid == ci.row { if let Ok(s) = strings.get(row.name as usize) { return s.to_string(); @@ -239,6 +251,9 @@ fn resolve_string_from_raw_table( "event" => { if let Some(table) = tables.table::() { for row in table { + let Ok(row) = row else { + continue; + }; if row.rid == ci.row { if let Ok(s) = strings.get(row.name as usize) { return s.to_string(); @@ -303,12 +318,18 @@ fn resolve_memberref_class_name( }; for row in mr_table { + let Ok(row) = row else { + continue; + }; if row.rid == ci.row { // The class field is a MemberRefParent coded index if row.class.tag == TableId::TypeRef { // Resolve the TypeRef if let Some(tr_table) = tables.table::() { for tr in tr_table { + let Ok(tr) = tr else { + continue; + }; if tr.rid == row.class.row { let ns = strings.get(tr.type_namespace as usize).unwrap_or("?"); let name = strings.get(tr.type_name as usize).unwrap_or("?"); diff --git a/dotscope-cli/src/commands/heaps.rs b/dotscope-cli/src/commands/heaps.rs index ee8f88c4..3381845b 100644 --- a/dotscope-cli/src/commands/heaps.rs +++ b/dotscope-cli/src/commands/heaps.rs @@ -186,7 +186,7 @@ fn lookup_userstring( .userstrings() .with_context(|| "assembly has no #US heap")?; us.get(off) - .map(widestring::U16Str::to_string_lossy) + .map(|s| s.to_string_lossy()) .with_context(|| format!("no user string at offset 0x{off:x}")) }) } diff --git a/dotscope-cli/src/commands/tables.rs b/dotscope-cli/src/commands/tables.rs index 0c5241ad..642149fd 100644 --- a/dotscope-cli/src/commands/tables.rs +++ b/dotscope-cli/src/commands/tables.rs @@ -220,6 +220,9 @@ macro_rules! table_formatter { .collect(); let mut rows = Vec::new(); for row in table { + let Ok(row) = row else { + continue; + }; let mut vals = vec![row.rid.to_string(), row.token.to_string()]; let extra: Vec = ($row_fn)(&row); vals.extend(extra); @@ -246,6 +249,9 @@ macro_rules! table_formatter_with { ["RID", "Token"].iter().chain($cols.iter()).map(ToString::to_string).collect(); let mut rows = Vec::new(); for row in table { + let Ok(row) = row else { + continue; + }; let mut vals = vec![row.rid.to_string(), row.token.to_string()]; let extra: Vec = ($row_fn)(&row, $($param),+); vals.extend(extra); @@ -285,6 +291,9 @@ fn format_module( .collect(); let mut rows = Vec::new(); for row in table { + let Ok(row) = row else { + continue; + }; rows.push(vec![ row.rid.to_string(), row.token.to_string(), @@ -605,6 +614,9 @@ fn format_enclog(tables: &TablesHeader<'_>) -> anyhow::Result .collect(); let mut rows = Vec::new(); for row in table { + let Ok(row) = row else { + continue; + }; rows.push(vec![ row.rid.to_string(), row.token.to_string(), @@ -630,6 +642,9 @@ fn format_encmap(tables: &TablesHeader<'_>) -> anyhow::Result .collect(); let mut rows = Vec::new(); for row in table { + let Ok(row) = row else { + continue; + }; rows.push(vec![ row.rid.to_string(), row.token.to_string(), @@ -669,6 +684,9 @@ fn format_assembly( .collect(); let mut rows = Vec::new(); for row in table { + let Ok(row) = row else { + continue; + }; rows.push(vec![ row.rid.to_string(), row.token.to_string(), @@ -736,6 +754,9 @@ fn format_assemblyref( .collect(); let mut rows = Vec::new(); for row in table { + let Ok(row) = row else { + continue; + }; rows.push(vec![ row.rid.to_string(), row.token.to_string(), @@ -886,6 +907,9 @@ fn format_methoddebuginformation( .collect(); let mut rows = Vec::new(); for row in table { + let Ok(row) = row else { + continue; + }; rows.push(vec![ row.rid.to_string(), row.token.to_string(), @@ -920,6 +944,9 @@ fn format_localscope(tables: &TablesHeader<'_>) -> anyhow::Result) -> anyhow::Result Result<()> { // Verify our new type was added if let Some(tables) = verify_view.tables() { if let Some(typedef_table) = tables.table::() { - let found_type = typedef_table.iter().any(|t| { + // A row that fails to parse is not a match; the iterator reports it as `Err`. + let found_type = typedef_table.iter().flatten().any(|t| { verify_view.strings().is_some_and(|s| { s.get(t.type_name as usize) .is_ok_and(|name| name == "DotScopeModifiedClass") diff --git a/dotscope/fuzz/Cargo.lock b/dotscope/fuzz/Cargo.lock index c334cd02..313876ea 100644 --- a/dotscope/fuzz/Cargo.lock +++ b/dotscope/fuzz/Cargo.lock @@ -27,11 +27,12 @@ checksum = "250f629c0161ad8107cf89319e990051fae62832fd343083bea452d93e2205fd" [[package]] name = "analyssa" -version = "0.2.0" +version = "0.5.0" dependencies = [ "boxcar", "dashmap", "log", + "num_enum", "rayon", "thiserror", ] @@ -158,9 +159,9 @@ checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" [[package]] name = "cowfile" -version = "0.2.2" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b0c9a99dda5d60063c8daece5d8c7dc1a7b1cd8a3695fb0f4be1df2193ed138" +checksum = "62c10f4a5e79d5e0a109ffeadb808578e24f82123292efd4c2412a21e9f9b4ff" dependencies = [ "memmap2", "thiserror", @@ -295,7 +296,7 @@ dependencies = [ [[package]] name = "dotscope" -version = "0.8.0" +version = "0.9.0" dependencies = [ "aes", "analyssa", @@ -324,6 +325,7 @@ dependencies = [ "rustc-hash", "sha1", "sha2", + "smallvec", "strum", "tempfile", "thiserror", @@ -335,6 +337,7 @@ dependencies = [ name = "dotscope-fuzz" version = "0.0.0" dependencies = [ + "boxcar", "dotscope", "libfuzzer-sys", ] @@ -490,12 +493,13 @@ checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" [[package]] name = "imbl" -version = "7.0.0" +version = "7.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e525189e5f603908d0c6e0d402cb5de9c4b2c8866151fabc4ebd771ed2630a2e" +checksum = "43ea8d4c37ee560727e824d62804183624d371e632019b0e9e3532bce64a33e5" dependencies = [ "archery", "bitmaps", + "equivalent", "imbl-sized-chunks", "rand_core", "rand_xoshiro", @@ -569,9 +573,9 @@ checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "libfuzzer-sys" -version = "0.4.12" +version = "0.4.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f12a681b7dd8ce12bff52488013ba614b869148d54dd79836ab85aafdd53f08d" +checksum = "a9fd2f41a1cba099f79a0b6b6c35656cf7c03351a7bae8ff0f28f25270f929d2" dependencies = [ "arbitrary", "cc", @@ -594,9 +598,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.31" +version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "113b30b4cd05f7c06868fdb2854f66a7b9fece9a48425351cd532e810d74024f" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" [[package]] name = "lzma-rs" @@ -626,9 +630,9 @@ checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" [[package]] name = "memmap2" -version = "0.9.10" +version = "0.9.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "714098028fe011992e1c3962653c96b2d578c4b4bce9036e15ff220319b1e0e3" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" dependencies = [ "libc", ] @@ -645,9 +649,9 @@ dependencies = [ [[package]] name = "num-bigint" -version = "0.4.6" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" dependencies = [ "num-integer", "num-traits", @@ -671,6 +675,28 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_enum" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26" +dependencies = [ + "num_enum_derive", + "rustversion", +] + +[[package]] +name = "num_enum_derive" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 2.0.115", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -698,7 +724,7 @@ dependencies = [ "proc-macro2", "proc-macro2-diagnostics", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -737,7 +763,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn", + "syn 2.0.115", +] + +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", ] [[package]] @@ -757,16 +792,16 @@ checksum = "af066a9c399a26e020ada66a034357a868728e72cd426f3adcd35f80d88d88c8" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", "version_check", "yansi", ] [[package]] name = "quick-xml" -version = "0.40.1" +version = "0.41.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2474bd2e5029e7ccb6abb2ba48cf2383a333851dedf495901544281590c7da7f" +checksum = "e660451e55124f798a69a5af3f49ccfbefbd41910eefd25caf2393e1f3473ec1" dependencies = [ "memchr", ] @@ -832,9 +867,9 @@ dependencies = [ [[package]] name = "rustc-hash" -version = "2.1.2" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" [[package]] name = "rustix" @@ -849,6 +884,12 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + [[package]] name = "safe_arch" version = "0.7.4" @@ -881,7 +922,7 @@ checksum = "22fc4f90c27b57691bbaf11d8ecc7cfbfe98a4da6dbe60226115d322aa80c06e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -916,7 +957,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -996,7 +1037,7 @@ dependencies = [ "heck 0.5.0", "proc-macro2", "quote", - "syn", + "syn 2.0.115", ] [[package]] @@ -1010,6 +1051,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "tempfile" version = "3.27.0" @@ -1025,22 +1077,52 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.18" +version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.18" +version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.3", +] + +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.13+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.3+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" +dependencies = [ + "winnow", ] [[package]] @@ -1156,6 +1238,15 @@ dependencies = [ "windows-link", ] +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + [[package]] name = "wit-bindgen" version = "0.51.0" @@ -1186,7 +1277,7 @@ dependencies = [ "heck 0.5.0", "indexmap", "prettyplease", - "syn", + "syn 2.0.115", "wasm-metadata", "wit-bindgen-core", "wit-component", @@ -1202,7 +1293,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn", + "syn 2.0.115", "wit-bindgen-core", "wit-bindgen-rust", ] @@ -1255,3 +1346,59 @@ name = "zmij" version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + +[[patch.unused]] +name = "analysir" +version = "0.1.0" + +[[patch.unused]] +name = "autoit-rs" +version = "0.1.1" + +[[patch.unused]] +name = "darwinscope" +version = "0.1.1" + +[[patch.unused]] +name = "gobin" +version = "0.4.1" + +[[patch.unused]] +name = "innospect" +version = "0.1.3" + +[[patch.unused]] +name = "mallabel" +version = "0.1.0" + +[[patch.unused]] +name = "nimrod" +version = "0.3.1" + +[[patch.unused]] +name = "nsis" +version = "0.3.1" + +[[patch.unused]] +name = "pascalscript" +version = "0.1.2" + +[[patch.unused]] +name = "securs-fleet" +version = "0.1.0" + +[[patch.unused]] +name = "securs-spec" +version = "0.1.0" + +[[patch.unused]] +name = "securs-wg" +version = "0.1.0" + +[[patch.unused]] +name = "undelphi" +version = "0.3.2" + +[[patch.unused]] +name = "visualbasic" +version = "0.3.1" diff --git a/dotscope/fuzz/Cargo.toml b/dotscope/fuzz/Cargo.toml index 3b515c22..76132dec 100644 --- a/dotscope/fuzz/Cargo.toml +++ b/dotscope/fuzz/Cargo.toml @@ -13,6 +13,8 @@ cargo-fuzz = true [dependencies] libfuzzer-sys = "0.4.13" +# The custom-attribute parser takes the constructor parameter list as `Arc>`. +boxcar = "0.2.14" [dependencies.dotscope] path = ".." @@ -23,3 +25,38 @@ path = "fuzz_targets/cilobject.rs" test = false doc = false bench = false + +[[bin]] +name = "assemblyview" +path = "fuzz_targets/assemblyview.rs" +test = false +doc = false +bench = false + +[[bin]] +name = "signatures" +path = "fuzz_targets/signatures.rs" +test = false +doc = false +bench = false + +[[bin]] +name = "customattributes" +path = "fuzz_targets/customattributes.rs" +test = false +doc = false +bench = false + +[[bin]] +name = "methodbody" +path = "fuzz_targets/methodbody.rs" +test = false +doc = false +bench = false + +[[bin]] +name = "emulation" +path = "fuzz_targets/emulation.rs" +test = false +doc = false +bench = false diff --git a/dotscope/fuzz/fuzz_targets/assemblyview.rs b/dotscope/fuzz/fuzz_targets/assemblyview.rs new file mode 100644 index 00000000..dbd8cbaf --- /dev/null +++ b/dotscope/fuzz/fuzz_targets/assemblyview.rs @@ -0,0 +1,18 @@ +//! Fuzzes the raw PE + metadata-stream view. +//! +//! `CilAssemblyView` is the layer beneath `CilObject`: it parses the PE headers, locates the +//! CLI directory and slices the metadata heaps, all before any table or signature parsing +//! happens. Fuzzing it directly reaches header and stream-slicing arithmetic that the +//! `cilobject` target only exercises when the outer parse gets far enough to call it. +//! +//! Malformed input is expected; an `Err` is a pass. The property is that the call returns +//! rather than panicking, aborting, or reading out of bounds. + +#![no_main] + +use dotscope::metadata::cilassemblyview::CilAssemblyView; +use libfuzzer_sys::fuzz_target; + +fuzz_target!(|data: &[u8]| { + let _ = CilAssemblyView::from_mem(data.to_vec()); +}); diff --git a/dotscope/fuzz/fuzz_targets/customattributes.rs b/dotscope/fuzz/fuzz_targets/customattributes.rs new file mode 100644 index 00000000..c551af69 --- /dev/null +++ b/dotscope/fuzz/fuzz_targets/customattributes.rs @@ -0,0 +1,22 @@ +//! Fuzzes the custom-attribute value-blob parser. +//! +//! CA blobs are parsed for every row of the CustomAttribute table during the default load, on +//! rayon workers, before any validation stage runs. The blob is fully attacker-controlled and +//! its grammar is self-describing — element-type tags select the parse, so a dozen bytes can +//! request an enormous amount of work or an enormous reservation. +//! +//! The constructor parameter list is empty on purpose. Parameters only supply expected types +//! for the fixed arguments; leaving them empty drives the *tag-driven* path, which is the one +//! that lacked the array bound. + +#![no_main] + +use std::sync::Arc; + +use dotscope::metadata::customattributes::parse_custom_attribute_data; +use libfuzzer_sys::fuzz_target; + +fuzz_target!(|data: &[u8]| { + let params = Arc::new(boxcar::Vec::new()); + let _ = parse_custom_attribute_data(data, ¶ms); +}); diff --git a/dotscope/fuzz/fuzz_targets/emulation.rs b/dotscope/fuzz/fuzz_targets/emulation.rs new file mode 100644 index 00000000..8ac79904 --- /dev/null +++ b/dotscope/fuzz/fuzz_targets/emulation.rs @@ -0,0 +1,78 @@ +//! Fuzzes the CIL interpreter under tight resource limits. +//! +//! This is the only target that executes attacker-controlled *instructions* rather than just +//! parsing attacker-controlled bytes, so it reaches the interpreter, the BCL hooks, the managed +//! heap and the exception-unwind machinery — the surfaces where unbounded allocation and +//! unwind-ordering defects live. +//! +//! # Why it is structured this way +//! +//! Random bytes almost never form a loadable assembly, so most iterations stop at +//! `CilObject::from_mem`. That is expected and still useful: the fuzzer's coverage feedback +//! drives it toward inputs that get further, and the committed corpus seeds it with real +//! PE-shaped material. The alternative — synthesising a valid PE wrapper around fuzzed IL — +//! would test a narrower, more artificial surface. +//! +//! # Limits +//! +//! Every budget is set far below the production default. A fuzz iteration must finish in +//! milliseconds, and an input that merely *runs for a long time* is not the bug class this +//! target hunts — unbounded work that ignores the budget entirely is. Keeping the budget tiny +//! makes a hang stand out instead of blending into normal execution time. + +#![no_main] + +use dotscope::{ + emulation::ProcessBuilder, + metadata::{cilobject::CilObject, token::Token}, +}; +use libfuzzer_sys::fuzz_target; + +/// Instruction budget per emulated method. Production defaults to ~10M. +const MAX_INSTRUCTIONS: u64 = 20_000; + +/// Call-depth budget. Deep enough to exercise unwinding across frames, shallow enough to stay +/// fast. +const MAX_CALL_DEPTH: usize = 32; + +/// Wall-clock budget per emulated method. +const TIMEOUT_MS: u64 = 500; + +/// How many methods to try per input, so one assembly cannot dominate the run. +const MAX_METHODS: usize = 8; + +fuzz_target!(|data: &[u8]| { + let Ok(assembly) = CilObject::from_mem(data.to_vec()) else { + return; + }; + + // Collect a few method tokens before building the process: `methods()` borrows the + // assembly, which the builder takes by value. + let tokens: Vec = assembly + .methods() + .iter() + .take(MAX_METHODS) + .map(|entry| *entry.key()) + .collect(); + + if tokens.is_empty() { + return; + } + + let Ok(process) = ProcessBuilder::new() + .assembly(assembly) + .with_max_instructions(MAX_INSTRUCTIONS) + .with_max_call_depth(MAX_CALL_DEPTH) + .with_timeout_ms(TIMEOUT_MS) + .build() + else { + return; + }; + + for token in tokens { + // Errors and limit-reached outcomes are both fine. The property under test is that the + // call returns at all — rather than panicking, aborting, exhausting host memory, or + // running past its budget. + let _ = process.execute_method(token, Vec::new()); + } +}); diff --git a/dotscope/fuzz/fuzz_targets/methodbody.rs b/dotscope/fuzz/fuzz_targets/methodbody.rs new file mode 100644 index 00000000..9210c42a --- /dev/null +++ b/dotscope/fuzz/fuzz_targets/methodbody.rs @@ -0,0 +1,33 @@ +//! Fuzzes method-body header parsing and CIL block decoding. +//! +//! Two stages, both attacker-controlled: +//! +//! 1. `MethodBody::from` parses the tiny/fat header, the `size_code` field and the exception +//! section chain. The fat header's size-in-dwords nibble decides where the header ends, so +//! a hostile value can make it overlap the code it precedes. +//! 2. `decode_blocks` walks the IL into basic blocks. `switch` reads an attacker-chosen case +//! count, and decoding must stay inside the method's declared `size_code` rather +//! than run to the end of the buffer. +//! +//! The decode runs on the raw input rather than on the body's code window on purpose: it is +//! the linear-disassembly entry point, and bounding it is the caller's job — which is the +//! contract this target covers. + +#![no_main] + +use dotscope::assembly::decode_blocks; +use libfuzzer_sys::fuzz_target; + +/// Keeps a single iteration cheap enough that the fuzzer explores rather than grinds. +const MAX_DECODE_BYTES: usize = 64 * 1024; + +fuzz_target!(|data: &[u8]| { + // Stage 1: header + exception section chain. + let _ = dotscope::metadata::method::MethodBody::from(data); + + // Stage 2: block decoding, bounded so one pathological input cannot dominate the run. + if !data.is_empty() { + let limit = data.len().min(MAX_DECODE_BYTES); + let _ = decode_blocks(data, 0, 0x2000, Some(limit)); + } +}); diff --git a/dotscope/fuzz/fuzz_targets/signatures.rs b/dotscope/fuzz/fuzz_targets/signatures.rs new file mode 100644 index 00000000..65dfe1b3 --- /dev/null +++ b/dotscope/fuzz/fuzz_targets/signatures.rs @@ -0,0 +1,29 @@ +//! Fuzzes every signature-blob parser against the same input. +//! +//! Signature blobs come straight from the `#Blob` heap and are the crate's most recursive +//! parsing surface — the one that produced a deeply-nested `TypeSignature` whose +//! *drop glue* overflowed the stack. Driving all six entry points from one input is +//! deliberate: they share `SignatureParser`, so a single crafted blob usually reaches several +//! of them and any one may be the arm that mishandles it. +//! +//! Note this also exercises teardown. The returned values are dropped at the end of each +//! iteration, which is where a recursive `Drop` would fail rather than in the parser. + +#![no_main] + +use dotscope::metadata::signatures::{ + parse_field_signature, parse_local_var_signature, parse_method_signature, + parse_method_spec_signature, parse_property_signature, parse_type_spec_signature, +}; +use libfuzzer_sys::fuzz_target; + +fuzz_target!(|data: &[u8]| { + // Each parser is independent; a failure in one says nothing about the others, so all six + // run on every input rather than short-circuiting. + let _ = parse_method_signature(data); + let _ = parse_field_signature(data); + let _ = parse_property_signature(data); + let _ = parse_local_var_signature(data); + let _ = parse_type_spec_signature(data); + let _ = parse_method_spec_signature(data); +}); diff --git a/dotscope/src/analysis/cfg/graph.rs b/dotscope/src/analysis/cfg/graph.rs index 61f5ea8e..d3e57791 100644 --- a/dotscope/src/analysis/cfg/graph.rs +++ b/dotscope/src/analysis/cfg/graph.rs @@ -130,6 +130,7 @@ impl ControlFlowGraph<'static> { )) })?; let successors = block.successors.clone(); + let exception_successors = block.exception_successors.clone(); let last_instruction = block.instructions.last(); // Determine edge kinds based on the terminating instruction @@ -153,6 +154,43 @@ impl ControlFlowGraph<'static> { graph.add_edge(*node_id, target_node, edge)?; } + + // Exception edges. Any instruction in a protected region can raise, so control + // can leave this block for its handler entry at any point; the edge is modelled + // at block granularity. Without these, handler blocks are unreachable islands — + // outside the dominator tree and invisible to every dataflow analysis. + for &handler_idx in &exception_successors { + if handler_idx >= block_count { + return Err(GraphError(format!( + "Block {} has exception successor index {} which exceeds block count {}", + node_id.index(), + handler_idx, + block_count + ))); + } + + // `add_edge` does not deduplicate and `predecessors` yields one entry per + // incoming edge, so a parallel edge would list this block twice among the + // handler's predecessors and desynchronise phi operand count from + // predecessor count. Skip a handler that is already a normal successor. + if successors.contains(&handler_idx) { + continue; + } + + let target_node = *node_ids + .get(handler_idx) + .ok_or_else(|| GraphError(format!("missing node id {handler_idx}")))?; + + // The caught class token is not reachable from here: `exception_successors` + // holds block indices, and `HandlerEntryInfo` carries the handler's *kind*, + // not its type, which lives in the method's exception table. `None` is the + // documented catch-all/finally encoding. + graph.add_edge( + *node_id, + target_node, + CfgEdge::exception_handler(handler_idx, None), + )?; + } } // Identify entry and exit blocks @@ -251,6 +289,26 @@ impl<'a> ControlFlowGraph<'a> { graph.add_edge(node_id, target_node, edge)?; } + + // See `from_basic_blocks` for why these edges exist, why an already-present + // normal successor is skipped, and why the exception type is `None`. + for &handler_idx in &block.exception_successors { + if handler_idx >= block_count { + return Err(GraphError(format!( + "Block {block_idx} has exception successor index {handler_idx} which exceeds block count {block_count}" + ))); + } + + if block.successors.contains(&handler_idx) { + continue; + } + + graph.add_edge( + node_id, + NodeId::new(handler_idx), + CfgEdge::exception_handler(handler_idx, None), + )?; + } } // Identify entry and exit blocks @@ -1126,4 +1184,94 @@ mod tests { assert_eq!(self_loop.latches.len(), 1); assert_eq!(self_loop.latches[0], NodeId::new(1)); // Self back edge (latch) } + + /// Blocks 0 and 1 form a protected region; block 2 is the handler entry. + /// + /// Layout: 0 -> 1 -> 3 (normal flow), with 0 and 1 both protected by handler 2. + fn eh_blocks() -> Vec { + let mut blocks = vec![ + make_block(0, vec![1], FlowType::UnconditionalBranch), + make_block(1, vec![3], FlowType::UnconditionalBranch), + make_block(2, vec![3], FlowType::UnconditionalBranch), + make_block(3, vec![], FlowType::Return), + ]; + for idx in 0..2 { + if let Some(block) = blocks.get_mut(idx) { + block.exception_successors.push(2); + } + } + blocks + } + + /// Without exception edges a handler is an unreachable island — outside + /// the dominator tree and invisible to every dataflow analysis. + #[test] + fn exception_edges_make_the_handler_reachable() { + let cfg = ControlFlowGraph::from_basic_blocks(eh_blocks()).unwrap(); + + // Every protected block reaches the handler. + let mut handler_preds: Vec = cfg + .predecessors(NodeId::new(2)) + .map(NodeId::index) + .collect(); + handler_preds.sort_unstable(); + assert_eq!(handler_preds, vec![0, 1]); + + // And the handler is a successor of each of them, tagged as an EH edge. + for protected in [0usize, 1] { + let succs: Vec = cfg + .successors(NodeId::new(protected)) + .map(NodeId::index) + .collect(); + assert!(succs.contains(&2), "block {protected} must reach handler 2"); + } + } + + /// `add_edge` does not deduplicate and `predecessors` yields one entry per incoming + /// edge, so a handler that is also a normal successor must not be wired twice. + #[test] + fn exception_edge_is_not_duplicated_when_already_a_normal_successor() { + let mut blocks = eh_blocks(); + // Make block 1 fall through to the handler as ordinary control flow too. + if let Some(block) = blocks.get_mut(1) { + block.successors = vec![2]; + } + + let cfg = ControlFlowGraph::from_basic_blocks(blocks).unwrap(); + + let preds: Vec = cfg + .predecessors(NodeId::new(2)) + .map(NodeId::index) + .collect(); + assert_eq!( + preds.iter().filter(|&&p| p == 1).count(), + 1, + "block 1 must appear exactly once among the handler's predecessors, got {preds:?}" + ); + } + + /// An out-of-range handler index is a malformed input, not a panic. + #[test] + fn exception_successor_out_of_range_is_rejected() { + let mut blocks = eh_blocks(); + if let Some(block) = blocks.get_mut(0) { + block.exception_successors = vec![99]; + } + + assert!(ControlFlowGraph::from_basic_blocks(blocks).is_err()); + } + + /// The borrowed constructor must wire the same edges as the owning one. + #[test] + fn from_blocks_ref_wires_exception_edges_too() { + let blocks = eh_blocks(); + let cfg = ControlFlowGraph::from_blocks_ref(&blocks).unwrap(); + + let mut handler_preds: Vec = cfg + .predecessors(NodeId::new(2)) + .map(NodeId::index) + .collect(); + handler_preds.sort_unstable(); + assert_eq!(handler_preds, vec![0, 1]); + } } diff --git a/dotscope/src/analysis/ssa/converter.rs b/dotscope/src/analysis/ssa/converter.rs index fc6c3afc..0179eaad 100644 --- a/dotscope/src/analysis/ssa/converter.rs +++ b/dotscope/src/analysis/ssa/converter.rs @@ -170,17 +170,6 @@ pub struct SsaConverter<'a, 'cfg> { /// variables (those not covered by `infer_origin`, e.g. add, box, call results). var_stack_positions: BTreeMap, - /// Version stack snapshots at try_start_block entries. - /// - /// Saved during the main rename pass so that handler blocks can inherit - /// the reaching definitions from their try scope. Without this, handler - /// blocks get renamed with empty version stacks and resolve locals to - /// their initial (null/0) values instead of the values stored by the - /// try setup code. - /// - /// Maps handler_start_block → snapshot of current_def for each group. - handler_scope_defs: BTreeMap>, - /// Type provider for assigning types during SSA construction. /// /// Variables are assigned correct types at creation time based on method @@ -892,7 +881,6 @@ impl<'a, 'cfg> SsaConverter<'a, 'cfg> { entry_stacks: BTreeMap::new(), indirect_stores: BTreeMap::new(), var_stack_positions: BTreeMap::new(), - handler_scope_defs: BTreeMap::new(), type_provider, }; @@ -1586,7 +1574,7 @@ impl<'a, 'cfg> SsaConverter<'a, 'cfg> { } let tables = assembly.tables()?; let table = tables.table::()?; - let raw = table.get(token.row())?; + let raw = table.get(token.row()).ok().flatten()?; let blob = assembly.blob()?; let owned = raw.to_owned(blob).ok()?; match &owned.parsed_signature { @@ -1822,6 +1810,22 @@ impl<'a, 'cfg> SsaConverter<'a, 'cfg> { continue; } + // Exception handler entries never inherit the protected region's evaluation + // stack. Per ECMA-335 the stack at handler entry holds only the exception + // object the runtime pushes (catch/filter) or nothing at all (finally/fault), + // so the try blocks' exit stacks are not values that merge here. Now that the + // CFG carries exception edges these blocks do have multiple predecessors, and + // without this guard every one of those exit-stack slots would grow a phi — + // including slot 0, which `rename_block_process` defines directly as the + // exception object. + if self + .cfg + .block(node_id) + .is_some_and(|block| block.handler_entry.is_some()) + { + continue; + } + // Get enhanced stacks from all predecessors let pred_stacks: Vec>> = predecessors .iter() @@ -1999,13 +2003,15 @@ impl<'a, 'cfg> SsaConverter<'a, 'cfg> { let mut rename_map = BTreeMap::new(); self.rename_block(self.cfg.entry().index(), dom_tree, &mut rename_map)?; - // Also rename exception handler blocks that aren't reachable via dominator tree. - // These blocks are only entered via exception flow, not normal control flow, - // so they won't be visited during the dominator tree traversal. + // Fallback for handler blocks the dominator traversal did not reach. // - // The dominator tree is built from the method's entry block, so blocks only - // reachable via exception handlers won't appear as children in the tree. - // We need to explicitly traverse all blocks reachable from handler entries. + // Now that the CFG carries exception edges, a handler covering a reachable try + // region is itself reachable from the method entry, so it is renamed by the main + // traversal above and this loop skips it. What remains are handlers with no + // incoming exception edge at all: `wire_exception_edges` drops a handler whose + // RVA does not resolve to a block start, and a handler over an unreachable try + // region has no reachable predecessor. Those blocks still need renaming, and they + // have no try scope to inherit from, so they simply start from the entry stacks. let entry_idx = self.cfg.entry().index(); // Collect handler entry blocks @@ -2037,24 +2043,8 @@ impl<'a, 'cfg> SsaConverter<'a, 'cfg> { } } - // Now process handler regions - blocks that aren't reachable from main entry. - // Before renaming each handler, push the try-scope version stacks so that - // handler blocks resolve locals to the correct reaching definitions. + // Now process handler regions that the main traversal did not reach. for handler_entry in handler_entries { - // Push try-scope version stacks for this handler entry - let scope_defs = self.handler_scope_defs.get(&handler_entry).cloned(); - let mut scope_pushed: BTreeMap = BTreeMap::new(); - if let Some(ref defs) = scope_defs { - for (&group, &var_id) in defs { - self.version_stacks - .entry(group) - .or_default() - .push((0, var_id)); - let entry = scope_pushed.entry(group).or_insert(0); - *entry = entry.saturating_add(1); - } - } - // BFS through all blocks reachable from this handler let mut worklist: Vec = vec![handler_entry]; while let Some(block_idx) = worklist.pop() { @@ -2073,15 +2063,6 @@ impl<'a, 'cfg> SsaConverter<'a, 'cfg> { } } } - - // Pop the try-scope definitions - for (group, count) in scope_pushed { - if let Some(stack) = self.version_stacks.get_mut(&group) { - for _ in 0..count { - stack.pop(); - } - } - } } // Resolve phi types from their operands @@ -2647,28 +2628,14 @@ impl<'a, 'cfg> SsaConverter<'a, 'cfg> { dom_tree: &DominatorTree, rename_map: &mut BTreeMap, ) -> Result> { - // Save version stack snapshot for exception handler blocks. - // Handler blocks are renamed separately (they're not in the dominator - // tree), so they need the version stacks from the try scope to correctly - // resolve locals. We save the snapshot for each handler entry that this - // block can reach via exception dispatch. + // Handler blocks need no seeding of their own. An exception can be raised at any + // instruction in a protected region, so a handler's reaching definition for a local + // is the *join* over every program point in that region — which a snapshot of one + // try-body block's version stacks cannot express. // - // Last block wins: each try-body block overwrites the previous snapshot. - // The dominator tree traversal processes deeper blocks last, so the - // deepest reaching definitions (e.g., from inside CFF case blocks that - // modify the state variable) are captured. This ensures handler blocks - // see the state variable from the CFF dispatcher scope (via its phi), - // not just the initial value from the try entry. - if let Some(cfg_block) = self.cfg.block(NodeId::new(block_idx)) { - for &handler_idx in &cfg_block.exception_successors { - let snapshot: BTreeMap = self - .version_stacks - .iter() - .filter_map(|(&group, stack)| stack.last().map(|(_, var_id)| (group, *var_id))) - .collect(); - self.handler_scope_defs.insert(handler_idx, snapshot); - } - } + // The CFG carries real exception edges, which makes handler entries ordinary join + // points: `place_phi_nodes` puts a phi there via the dominance frontier, and Step 3 + // below fills one operand per try-region predecessor. let mut pushed_counts: BTreeMap = BTreeMap::new(); @@ -2734,9 +2701,10 @@ impl<'a, 'cfg> SsaConverter<'a, 'cfg> { // ECMA-335 §I.12.4.2.5: When control transfers to a catch/filter handler, // the runtime pushes the exception object on the stack. This is an implicit // definition that doesn't flow from any predecessor — it's created by the - // runtime. Exception handler entry blocks typically have 0 predecessors in - // the CFG (exception edges aren't modeled as regular edges), so no phi is - // placed and predecessor-based resolution will fail. + // runtime. Handler entries do have predecessors now that the CFG carries + // exception edges, but they are the *try* blocks, whose exit stacks say nothing + // about the exception object; `place_stack_phi_nodes` deliberately places no + // stack phi here, so predecessor-based resolution must not be relied on. // // We create a new SSA variable with the correct exception type here, before // the normal entry stack resolution loop, so the exception object has proper diff --git a/dotscope/src/analysis/ssa/types.rs b/dotscope/src/analysis/ssa/types.rs index c57a2b70..5c6348cf 100644 --- a/dotscope/src/analysis/ssa/types.rs +++ b/dotscope/src/analysis/ssa/types.rs @@ -973,7 +973,7 @@ pub fn resolve_corelib_valuetype(assembly: &CilObject, fullname: &str) -> SsaTyp let registry = assembly.types(); for entry in registry.iter() { let ty = entry.value(); - if ty.token.table() == 0x01 && ty.fullname() == fullname { + if ty.token.table() == 0x01 && &*ty.fullname() == fullname { return SsaType::ValueType(TypeRef::new(ty.token)); } } @@ -1266,7 +1266,7 @@ impl<'a> TypeContext<'a> { let Some(table) = tables.table::() else { return SsaType::Unknown; }; - let Some(raw) = table.get(sig_token.row()) else { + let Some(raw) = table.get(sig_token.row()).ok().flatten() else { return SsaType::Unknown; }; let Some(blob) = self.assembly.blob() else { @@ -1382,7 +1382,10 @@ impl<'a> TypeContext<'a> { } let tables = self.assembly.tables()?; let table = tables.table::()?; - let raw = table.get(body.local_var_sig_token & 0x00FF_FFFF)?; + let raw = table + .get(body.local_var_sig_token & 0x00FF_FFFF) + .ok() + .flatten()?; let blob = self.assembly.blob()?; let sig_data = blob.get(raw.signature as usize).ok()?; let locals_sig = parse_local_var_signature(sig_data).ok()?; diff --git a/dotscope/src/analysis/x86/cfg.rs b/dotscope/src/analysis/x86/cfg.rs index df01db9b..90b1ee92 100644 --- a/dotscope/src/analysis/x86/cfg.rs +++ b/dotscope/src/analysis/x86/cfg.rs @@ -328,41 +328,23 @@ impl X86Function { /// /// A CFG is reducible if every back edge goes to a loop header that dominates /// the source of the back edge. + /// + /// # Implementation + /// + /// The depth-first search is iterative, with an explicit stack carrying an enter/exit + /// marker so `in_stack` is cleared when a node's successors are exhausted — the point a + /// recursive version does on return. + /// + /// Recursion is not an option here. Its depth equals the longest simple path, which for a + /// chain of basic blocks equals the block count, and the blocks come from decoding an + /// attacker-supplied native region. Exhausting the stack is a SIGSEGV, which neither + /// `deny(panic)` nor a caller's `catch_unwind` can do anything about. This function is + /// public API, so an embedder can reach it directly with such a region. pub fn is_reducible(&self) -> bool { - fn dfs_check( - node: NodeId, - func: &X86Function, - doms: &DominatorTree, - visited: &mut [bool], - in_stack: &mut [bool], - ) -> bool { - let idx = node.index(); - if let Some(slot) = visited.get_mut(idx) { - *slot = true; - } - if let Some(slot) = in_stack.get_mut(idx) { - *slot = true; - } - - for succ in func.graph.successors(node) { - let succ_idx = succ.index(); - if in_stack.get(succ_idx).copied().unwrap_or(false) { - // This is a back edge (n -> succ where succ is on the stack) - // For reducibility, succ must dominate node - if !doms.dominates(succ, node) { - return false; - } - } else if !visited.get(succ_idx).copied().unwrap_or(false) - && !dfs_check(succ, func, doms, visited, in_stack) - { - return false; - } - } - - if let Some(slot) = in_stack.get_mut(idx) { - *slot = false; - } - true + /// One step of the iterative walk: visit a node's successors, or leave it. + enum Step { + Enter(NodeId), + Exit(NodeId), } if self.block_count() == 0 { @@ -371,16 +353,49 @@ impl X86Function { let doms = self.dominators(); - // A CFG is reducible if for every back edge (n -> h), - // h dominates n. A back edge is one where h dominates n. - // By definition, this is always true for back edges, so we check - // for edges that form cycles but aren't proper back edges. - - // Use DFS to find back edges + // A CFG is reducible if for every back edge (n -> h), h dominates n. A back edge is + // an edge to a node currently on the DFS stack. let mut visited = vec![false; self.block_count()]; let mut in_stack = vec![false; self.block_count()]; + let mut work = vec![Step::Enter(self.entry)]; + + while let Some(step) = work.pop() { + match step { + Step::Exit(node) => { + if let Some(slot) = in_stack.get_mut(node.index()) { + *slot = false; + } + } + Step::Enter(node) => { + let idx = node.index(); + if visited.get(idx).copied().unwrap_or(false) { + continue; + } + if let Some(slot) = visited.get_mut(idx) { + *slot = true; + } + if let Some(slot) = in_stack.get_mut(idx) { + *slot = true; + } + + // Popped only once every successor pushed below has been processed. + work.push(Step::Exit(node)); + + for succ in self.graph.successors(node) { + let succ_idx = succ.index(); + if in_stack.get(succ_idx).copied().unwrap_or(false) { + if !doms.dominates(succ, node) { + return false; + } + } else if !visited.get(succ_idx).copied().unwrap_or(false) { + work.push(Step::Enter(succ)); + } + } + } + } + } - dfs_check(self.entry, self, doms, &mut visited, &mut in_stack) + true } } diff --git a/dotscope/src/analysis/x86/decoder.rs b/dotscope/src/analysis/x86/decoder.rs index d4aedd17..5a7bffd5 100644 --- a/dotscope/src/analysis/x86/decoder.rs +++ b/dotscope/src/analysis/x86/decoder.rs @@ -3,11 +3,22 @@ //! This module provides a thin wrapper around iced-x86 that converts its //! instruction representation to our simplified [`X86Instruction`] types. -use std::collections::VecDeque; +use std::collections::{BTreeMap, VecDeque}; -use iced_x86::{Decoder, DecoderOptions, Instruction, Mnemonic, OpKind, Register}; +use iced_x86::{Decoder, DecoderOptions, FlowControl, Instruction, Mnemonic, OpKind, Register}; use rustc_hash::FxHashSet; +/// Maximum instructions a single x86 decode may produce. +/// +/// The x86 subsystem is reached from mixed-mode assemblies, where the region handed to it is +/// derived from attacker-controlled headers and can extend to end-of-file. Both +/// the traversal and the SSA translation built on it scale super-linearly in instruction count, +/// so an unbounded decode is a DoS primitive rather than merely slow. +/// +/// A native stub in a real mixed-mode assembly is kilobytes of code — low tens of thousands of +/// instructions at the very outside. +pub const MAX_DECODED_INSTRUCTIONS: usize = 250_000; + use crate::{ analysis::x86::types::{ X86Condition, X86DecodedInstruction, X86EpilogueInfo, X86Instruction, X86Memory, @@ -43,6 +54,25 @@ pub fn x86_decode_all( bytes: &[u8], bitness: u32, base_address: u64, +) -> Result> { + x86_decode_all_limited(bytes, bitness, base_address, MAX_DECODED_INSTRUCTIONS) +} + +/// Linear-sweep decode with an explicit instruction budget. +/// +/// See [`MAX_DECODED_INSTRUCTIONS`]. Callers that hand over a region derived from file +/// headers — rather than a stub whose extent they have already established — should use this +/// and pass a bound appropriate to what they expect. +/// +/// # Errors +/// +/// Returns [`crate::Error::X86Error`] on empty input, invalid bitness, an invalid instruction, +/// or when `max_instructions` is exhausted. +pub fn x86_decode_all_limited( + bytes: &[u8], + bitness: u32, + base_address: u64, + max_instructions: usize, ) -> Result> { if bytes.is_empty() { return Err(Error::X86Error("Empty input".to_string())); @@ -70,6 +100,14 @@ pub fn x86_decode_all( ))); } + // A linear sweep runs to the end of whatever slice it was given, and that slice can be + // the entire remainder of the file. + if instructions.len() >= max_instructions { + return Err(Error::X86Error(format!( + "x86 linear sweep exceeded {max_instructions} instructions" + ))); + } + let converted = convert_instruction(&instr, base_address)?; let is_ret = matches!(converted, X86Instruction::Ret); instructions.push(X86DecodedInstruction { @@ -166,6 +204,30 @@ pub fn x86_decode_traversal( bitness: u32, base_address: u64, entry_offset: u64, +) -> Result { + x86_decode_traversal_limited( + bytes, + bitness, + base_address, + entry_offset, + MAX_DECODED_INSTRUCTIONS, + ) +} + +/// Recursive-traversal decode with an explicit instruction budget. +/// +/// See [`MAX_DECODED_INSTRUCTIONS`] for why a budget is mandatory rather than advisory. +/// +/// # Errors +/// +/// Returns [`crate::Error::X86Error`] on empty input, invalid bitness, address overflow, or +/// when `max_instructions` is exhausted. +pub fn x86_decode_traversal_limited( + bytes: &[u8], + bitness: u32, + base_address: u64, + entry_offset: u64, + max_instructions: usize, ) -> Result { if bytes.is_empty() { return Err(Error::X86Error("Empty input".to_string())); @@ -187,6 +249,8 @@ pub fn x86_decode_traversal( let mut visited: FxHashSet = FxHashSet::default(); // Decoded instructions by offset let mut instructions: Vec = Vec::new(); + // Byte spans already claimed, as `offset -> length`, for O(log n) overlap rejection. + let mut decoded_spans: BTreeMap = BTreeMap::new(); // Targets we couldn't resolve let mut unresolved_targets: Vec = Vec::new(); let mut has_indirect = false; @@ -226,21 +290,36 @@ pub fn x86_decode_traversal( let offset = addr.saturating_sub(base_address); let length = instr.len(); - // Check if we've already decoded an instruction that overlaps - // (this can happen with certain obfuscation tricks) - let overlaps = instructions.iter().any(|existing| { - let existing_start = existing.offset; - let existing_end = existing.offset.saturating_add(existing.length as u64); - let new_start = offset; - let new_end = offset.saturating_add(length as u64); - // Check for overlap - new_start < existing_end && new_end > existing_start - }); - - if overlaps { + // Reject overlapping decodes (overlapping instructions are a known obfuscation + // trick) using an interval map rather than a linear scan of everything decoded so + // far, which made traversal O(n²) in decoded instructions over a region that could + // run to end-of-file. + // + // Intervals do not overlap each other by construction, so only two candidates can + // intersect `[offset, new_end)`: the nearest interval starting at or before + // `offset`, and the nearest starting after it. + let new_end = offset.saturating_add(length as u64); + let overlaps_before = decoded_spans + .range(..=offset) + .next_back() + .is_some_and(|(start, len)| start.saturating_add(*len as u64) > offset); + let overlaps_after = decoded_spans + .range(offset..) + .next() + .is_some_and(|(start, _)| *start < new_end); + + if overlaps_before || overlaps_after { continue; } + if instructions.len() >= max_instructions { + return Err(Error::X86Error(format!( + "x86 traversal exceeded {max_instructions} instructions" + ))); + } + + decoded_spans.insert(offset, length); + // Convert the instruction let converted = match convert_instruction(&instr, base_address) { Ok(i) => i, @@ -293,13 +372,40 @@ pub fn x86_decode_traversal( } } X86Instruction::Unsupported { .. } => { - // Check if this might be an indirect jump/call - if instr.mnemonic() == Mnemonic::Jmp || instr.mnemonic() == Mnemonic::Call { + // Classify by control flow, not by mnemonic. + // + // An unsupported *control transfer* is a path terminator: an indirect + // `jmp` (a jump table, or a tail call) does not fall through, and neither + // does `ret`/`retf`/`iret`. Enqueuing `next_addr` for them walks straight + // off the end of the function into alignment padding and then into whatever + // follows — which is why traversal never stopped at a function boundary. + // + // `flow_control()` classifies every transfer class, where a mnemonic match + // only caught `Jmp`/`Call` and silently treated the rest as sequential. + let flow = instr.flow_control(); + let is_indirect_transfer = matches!( + flow, + FlowControl::IndirectBranch | FlowControl::IndirectCall + ); + + if is_indirect_transfer { has_indirect = true; unresolved_targets.push(addr); } - // Try to continue to next instruction anyway - if next_addr < code_end && visited.insert(next_addr) { + + let terminates_path = matches!( + flow, + FlowControl::Return + | FlowControl::IndirectBranch + | FlowControl::UnconditionalBranch + | FlowControl::Interrupt + | FlowControl::Exception + | FlowControl::XbeginXabortXend + ); + + // Calls and conditional branches still fall through to the return address + // or the not-taken edge. + if !terminates_path && next_addr < code_end && visited.insert(next_addr) { worklist.push_back(next_addr); } } diff --git a/dotscope/src/analysis/x86/ssa.rs b/dotscope/src/analysis/x86/ssa.rs index 22012605..8e4b9024 100644 --- a/dotscope/src/analysis/x86/ssa.rs +++ b/dotscope/src/analysis/x86/ssa.rs @@ -60,6 +60,13 @@ use crate::{ Error, Result, }; +/// Maximum basic blocks accepted for x86-to-SSA translation. +/// +/// Real native stubs in mixed-mode assemblies have block counts in the tens. This bounds the +/// per-block register state and the phi worklist against input whose block count is ultimately +/// attacker-controlled. +const MAX_TRANSLATABLE_BLOCKS: usize = 65_536; + /// Number of registers tracked (0-15 GPRs for x64, 16-21 segment registers). const MAX_REGISTERS: usize = 22; @@ -222,6 +229,18 @@ impl<'a> X86ToSsaTranslator<'a> { return Err(Error::X86Error("Empty function".to_string())); } + // Translation allocates per-block register state and runs a worklist per register, so + // cost grows with block count — and block count comes from decoding a region derived + // from attacker-controlled file headers. The decoder's instruction budget bounds this + // indirectly; this is the direct statement of the limit. + if self.func.block_count() > MAX_TRANSLATABLE_BLOCKS { + return Err(Error::X86Error(format!( + "x86 function has {} blocks, exceeding the {} block translation limit", + self.func.block_count(), + MAX_TRANSLATABLE_BLOCKS + ))); + } + // Step 1: Analyze which blocks define which registers self.analyze_definitions(); @@ -280,6 +299,43 @@ impl<'a> X86ToSsaTranslator<'a> { .ok_or_else(|| Error::SsaError("place_phi_nodes: block_exit_states is empty".into()))? .register_count(); + // Precompute the dominance frontier once (Cytron et al.), rather than rediscovering it + // by scanning every block for every register on every worklist pop. + // + // The previous shape was O(registers × blocks² × in-degree) with two `dominates` queries + // per candidate, over input whose block count is attacker-controlled — the x86 subsystem + // decodes regions derived from file headers. This is one pass over the join points. + // + // DF[runner] gains `b` for every runner on the idom chain from each predecessor of a + // join point `b` up to (but excluding) idom(b). + let mut frontier: Vec> = vec![FxHashSet::default(); block_count]; + for block_idx in 0..block_count { + let node = NodeId::new(block_idx); + let preds: Vec = self.func.predecessors(node).collect(); + if preds.len() < 2 { + continue; + } + + let idom_of_node = doms.immediate_dominator(node); + for pred in preds { + let mut runner = pred; + while Some(runner) != idom_of_node { + let Some(slot) = frontier.get_mut(runner.index()) else { + break; + }; + slot.insert(block_idx); + + let Some(next) = doms.immediate_dominator(runner) else { + break; + }; + if next == runner { + break; + } + runner = next; + } + } + } + // For each register that is defined somewhere for reg_idx in 0..register_count { // Clone the def_blocks to avoid borrow issues @@ -299,33 +355,24 @@ impl<'a> X86ToSsaTranslator<'a> { let mut phi_blocks: FxHashSet = FxHashSet::default(); while let Some(def_block) = worklist.pop() { - // For each block in the dominance frontier of def_block - for block_idx in 0..block_count { - // Check if block_idx is in the dominance frontier of def_block - // DF(X) = {Y : ∃ pred of Y s.t. X dominates pred but X doesn't strictly dominate Y} - let node = NodeId::new(block_idx); - let def_node = NodeId::new(def_block); - - let in_frontier = self.func.predecessors(node).any(|pred| { - doms.dominates(def_node, pred) - && (def_node == node || !doms.dominates(def_node, node)) - }); - - if in_frontier && !phi_blocks.contains(&block_idx) { - phi_blocks.insert(block_idx); - - // Create phi variable - let phi_var = self.create_variable( - VariableOrigin::Phi, - DefSite::phi(block_idx), - bitness, - ); - self.phi_placement.set(block_idx, reg_idx, phi_var); + let Some(df) = frontier.get(def_block) else { + continue; + }; - // Phi defines the register, so add to worklist - if !def_blocks.contains(&block_idx) { - worklist.push(block_idx); - } + for &block_idx in df { + if phi_blocks.contains(&block_idx) { + continue; + } + phi_blocks.insert(block_idx); + + // Create phi variable + let phi_var = + self.create_variable(VariableOrigin::Phi, DefSite::phi(block_idx), bitness); + self.phi_placement.set(block_idx, reg_idx, phi_var); + + // Phi defines the register, so add to worklist + if !def_blocks.contains(&block_idx) { + worklist.push(block_idx); } } } diff --git a/dotscope/src/assembly/decoder.rs b/dotscope/src/assembly/decoder.rs index 6b26d3e8..0e7ae911 100644 --- a/dotscope/src/assembly/decoder.rs +++ b/dotscope/src/assembly/decoder.rs @@ -52,7 +52,7 @@ //! - [`crate::metadata::method`] - Supports method-level disassembly and caching use std::{ - collections::{HashMap, HashSet}, + collections::{BTreeMap, HashMap, HashSet}, sync::Arc, }; @@ -95,6 +95,27 @@ struct Decoder<'a> { offset_start: usize, /// Starting relative virtual address rva_start: usize, + /// Index from each block's start RVA to its index in `blocks`. + /// + /// `find_block_containing_rva` is called once per branch target. A scan over every block + /// would make block construction quadratic in a method whose branches are dense, and block + /// count is bounded only by the size of the method body. This `BTreeMap` answers the + /// containment query in O(log B) via `range(..=rva).next_back()`. + /// + /// Indices stay valid because blocks are only ever *appended*: `split_block_at` pushes the + /// tail as a new block rather than inserting it, so no existing index shifts. + block_starts: BTreeMap, + /// Exclusive upper bound on offsets this decoder may read. + /// + /// The parser spans the whole file image so that offsets stay file-absolute — the + /// `VisitedMap` is shared across methods and keyed by those offsets, so a per-method + /// window would make different methods collide at the same relative offset. Decoding is + /// instead bounded here, at the end of *this* method's declared `size_code`. + /// + /// Without it every bound in this decoder was the file length, so a method whose last + /// in-range instruction is not a terminator kept decoding into the next method, into the + /// metadata streams and into resources. + max_offset: usize, } impl<'a> Decoder<'a> { @@ -119,7 +140,7 @@ impl<'a> Decoder<'a> { /// /// # Errors /// - /// Returns [`crate::Error::OutOfBounds`] if the offset exceeds the parser's data length. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if the offset exceeds the parser's data length. /// /// # Thread Safety /// @@ -130,11 +151,14 @@ impl<'a> Decoder<'a> { rva: usize, exceptions: Option<&'a [ExceptionHandler]>, visited: Arc, + max_offset: usize, ) -> Result { if offset > parser.len() { return Err(out_of_bounds_error!()); } + let max_offset = max_offset.min(parser.len()); + Ok(Decoder { blocks: Vec::new(), exceptions, @@ -143,6 +167,8 @@ impl<'a> Decoder<'a> { block_id: 0, offset_start: offset, rva_start: rva, + block_starts: BTreeMap::new(), + max_offset, }) } @@ -246,6 +272,7 @@ impl<'a> Decoder<'a> { let mut entry_points: HashSet = HashSet::new(); // Create the first block at method entry + self.block_starts.insert(self.rva_start as u64, 0); self.blocks .push(BasicBlock::new(0, self.rva_start as u64, self.offset_start)); entry_points.insert(self.rva_start as u64); @@ -283,7 +310,8 @@ impl<'a> Decoder<'a> { .offset_start .checked_add(entry_offset_u32 as usize) .ok_or(out_of_bounds_error!())?; - if entry_offset < self.parser.len() && !self.visited.get(entry_offset) { + if entry_offset < self.max_offset && !self.visited.get(entry_offset) { + self.block_starts.insert(entry_rva, self.blocks.len()); self.blocks .push(BasicBlock::new(self.blocks.len(), entry_rva, entry_offset)); entry_points.insert(entry_rva); @@ -350,7 +378,7 @@ impl<'a> Decoder<'a> { (block.offset, block.rva) }; - if block_offset > self.parser.len() { + if block_offset > self.max_offset { return Err(out_of_bounds_error!()); } @@ -364,7 +392,7 @@ impl<'a> Decoder<'a> { let mut current_rva = block_rva; loop { - if current_offset >= self.parser.len() { + if current_offset >= self.max_offset { break; } @@ -489,7 +517,7 @@ impl<'a> Decoder<'a> { let Some(offset) = self.offset_start.checked_add(relative_offset) else { return; }; - if offset >= self.parser.len() { + if offset >= self.max_offset { return; } @@ -497,6 +525,7 @@ impl<'a> Decoder<'a> { self.split_block_at(block_idx, split_instr_idx, rva, offset); } else { let new_block = BasicBlock::new(self.blocks.len(), rva, offset); + self.block_starts.insert(rva, self.blocks.len()); self.blocks.push(new_block); } @@ -521,20 +550,25 @@ impl<'a> Decoder<'a> { /// /// * `rva` - The relative virtual address to search for fn find_block_containing_rva(&self, rva: u64) -> Option<(usize, usize)> { - for (block_idx, block) in self.blocks.iter().enumerate() { - if block.rva == rva { - return None; - } + // Already a block boundary — nothing to split. + if self.block_starts.contains_key(&rva) { + return None; + } - let block_end_rva = block.rva.checked_add(block.size as u64)?; - if rva > block.rva && rva < block_end_rva { - for (instr_idx, instr) in block.instructions.iter().enumerate() { - if instr.rva == rva { - return Some((block_idx, instr_idx)); - } + // The only block that can contain `rva` in its interior is the one with the greatest + // start below it, since blocks do not overlap. + let (_, &block_idx) = self.block_starts.range(..rva).next_back()?; + let block = self.blocks.get(block_idx)?; + + let block_end_rva = block.rva.checked_add(block.size as u64)?; + if rva > block.rva && rva < block_end_rva { + for (instr_idx, instr) in block.instructions.iter().enumerate() { + if instr.rva == rva { + return Some((block_idx, instr_idx)); } } } + None } @@ -574,6 +608,7 @@ impl<'a> Decoder<'a> { } // Create new block with instructions from split point onwards + self.block_starts.insert(rva, self.blocks.len()); let mut new_block = BasicBlock::new(self.blocks.len(), rva, offset); let Some(orig) = self.blocks.get(block_idx) else { return; @@ -665,9 +700,15 @@ impl<'a> Decoder<'a> { .checked_add(u64::from(handler.try_length)) .ok_or_else(|| malformed_error!("try region end RVA overflow"))?; - // Mark blocks in the try region - for block in &mut self.blocks { - if block.rva >= try_start && block.rva < try_end { + // Mark blocks in the try region. + // + // `self.blocks` is sorted by RVA, so binary-search the covered range instead of + // scanning every block per handler: with H handlers whose try ranges span the + // method, the scan was O(H·B), and B is bounded only by the method body's size. + let lo = self.blocks.partition_point(|b| b.rva < try_start); + let hi = self.blocks.partition_point(|b| b.rva < try_end); + if let Some(covered) = self.blocks.get_mut(lo..hi) { + for block in covered { block.exceptions.push(handler_idx); } } @@ -745,15 +786,26 @@ impl<'a> Decoder<'a> { .checked_add(u64::from(handler.try_length)) .ok_or_else(|| malformed_error!("try region end RVA overflow"))?; - for block in &mut self.blocks { - if block.rva >= try_start && block.rva < try_end { - // Add exception successor if not already present - if !block.exception_successors.contains(&handler_block_idx) { - block.exception_successors.push(handler_block_idx); - } + // Same binary-searched range as `process_exception_handlers`. Duplicates are + // removed in one pass afterwards rather than by a linear `contains` per push, + // which made this O(B·H²) — each push scanned successors already accumulated. + let lo = self.blocks.partition_point(|b| b.rva < try_start); + let hi = self.blocks.partition_point(|b| b.rva < try_end); + if let Some(covered) = self.blocks.get_mut(lo..hi) { + for block in covered { + block.exception_successors.push(handler_block_idx); } } } + + // Deduplicate while preserving first-occurrence order, which encodes handler priority — + // sorting would reorder the handler search. + let mut seen: HashSet = HashSet::new(); + for block in &mut self.blocks { + seen.clear(); + block.exception_successors.retain(|idx| seen.insert(*idx)); + } + Ok(()) } @@ -994,6 +1046,16 @@ pub(crate) fn decode_method( return Ok(()); } + // Bound decoding to this method's declared code window. + // + // The parser deliberately still spans the whole file image: offsets must stay + // file-absolute because `shared_visited` is shared across every method and keyed by + // them. Slicing to a per-method window would restart every method's offsets at 0, so + // methods would collide in that map and all but the first would decode nothing. + let code_end = code_start + .checked_add(body.size_code) + .ok_or_else(|| malformed_error!("code_start + size_code overflow"))?; + let mut parser = Parser::new(file.data()); let rva_start = rva .checked_add(body.size_header) @@ -1004,6 +1066,7 @@ pub(crate) fn decode_method( rva_start, Some(&body.exception_handlers), shared_visited, + code_end, )?; decoder.decode_blocks()?; @@ -1092,7 +1155,10 @@ pub fn decode_blocks( let mut parser = Parser::new(effective_data); let visited = Arc::new(VisitedMap::new(effective_data.len())); - let mut decoder = Decoder::new(&mut parser, 0, rva, None, visited)?; + // This entry point already slices `data` to the requested window and builds a private + // `VisitedMap` over it, so the whole buffer is in bounds. + let max_offset = effective_data.len(); + let mut decoder = Decoder::new(&mut parser, 0, rva, None, visited, max_offset)?; decoder.decode_blocks()?; @@ -1296,10 +1362,24 @@ pub fn decode_instruction(parser: &mut Parser, rva: u64) -> Result OperandType::Float64 => Operand::Immediate(Immediate::Float64(parser.read_le::()?)), OperandType::Token => Operand::Token(Token::new(parser.read_le::()?)), OperandType::Switch => { - let case_count = parser.read_le::()?; + let case_count = parser.read_le::()? as usize; + + // Each case is a 4-byte target that must actually be present, so the remaining + // input is a hard upper bound on the count. Without this check the raw `u32` sizes + // the reservation directly: five bytes (`0x45` plus `0xFFFFFFFF`) request ~16 GiB, + // and a failed allocation *aborts* rather than unwinding, so it escapes the + // crate's no-panic policy entirely. + let available_targets = parser.len().saturating_sub(parser.pos()) / 4; + if case_count > available_targets { + return Err(malformed_error!( + "switch declares {} cases but only {} target slots remain in the method body", + case_count, + available_targets + )); + } - let mut targets = Vec::with_capacity(case_count as usize); - for _ in 0..case_count as usize { + let mut targets = Vec::with_capacity(case_count); + for _ in 0..case_count { // Switch offsets are SIGNED 32-bit integers (can be negative for backward jumps) targets.push(parser.read_le::()?); } diff --git a/dotscope/src/assembly/encoder.rs b/dotscope/src/assembly/encoder.rs index 1694fce3..a66858be 100644 --- a/dotscope/src/assembly/encoder.rs +++ b/dotscope/src/assembly/encoder.rs @@ -1084,7 +1084,7 @@ impl InstructionEncoder { /// Returns an error if: /// - [`crate::Error::UndefinedLabel`] - Any referenced labels are undefined /// - [`crate::Error::InvalidBranch`] - Branch offsets exceed the allowed range for their instruction type - /// - [`crate::Error::Malformed`] - Stack underflow occurred during encoding (negative stack depth) + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] - Stack underflow occurred during encoding (negative stack depth) /// /// # Examples /// diff --git a/dotscope/src/cilassembly/cleanup/analysis.rs b/dotscope/src/cilassembly/cleanup/analysis.rs index 3c567aca..0bfec1ec 100644 --- a/dotscope/src/cilassembly/cleanup/analysis.rs +++ b/dotscope/src/cilassembly/cleanup/analysis.rs @@ -253,6 +253,13 @@ pub fn find_unreferenced_types( if let Some(tables) = assembly.tables() { if let Some(attr_table) = tables.table::() { for row in attr_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if row.constructor.token.is_table(TableId::MethodDef) { if let Some(&ctor_type) = method_to_type.get(&row.constructor.token) { if candidates.remove(&ctor_type) { diff --git a/dotscope/src/cilassembly/cleanup/compaction.rs b/dotscope/src/cilassembly/cleanup/compaction.rs index 1a57ba2d..ff4d4c85 100644 --- a/dotscope/src/cilassembly/cleanup/compaction.rs +++ b/dotscope/src/cilassembly/cleanup/compaction.rs @@ -306,7 +306,7 @@ fn collect_referenced_heap_entries( } // Get the row and serialize to bytes - let Some(row) = table.get(rid) else { + let Some(row) = table.get(rid).ok().flatten() else { continue; }; @@ -371,7 +371,7 @@ fn collect_referenced_heap_entries( for rid in 1..=row_count { // Get the row and serialize to bytes - let Some(row) = table.get(rid) else { + let Some(row) = table.get(rid).ok().flatten() else { continue; }; diff --git a/dotscope/src/cilassembly/cleanup/executor.rs b/dotscope/src/cilassembly/cleanup/executor.rs index b537997e..c9c33e25 100644 --- a/dotscope/src/cilassembly/cleanup/executor.rs +++ b/dotscope/src/cilassembly/cleanup/executor.rs @@ -10,7 +10,7 @@ //! The executor ensures deletions are applied in the correct order to maintain //! referential integrity and avoid RID shifting issues. -use std::collections::{BTreeSet, HashSet}; +use std::collections::{BTreeSet, HashMap, HashSet}; use crate::{ cilassembly::{ @@ -28,8 +28,8 @@ use crate::{ }, metadata::{ tables::{ - CustomAttributeRaw, FieldRaw, InterfaceImplRaw, MethodDefRaw, MethodImplRaw, - MethodSemanticsRaw, MethodSpecRaw, TableId, TypeDefRaw, + skip_unreadable, CustomAttributeRaw, FieldRaw, InterfaceImplRaw, MethodDefRaw, + MethodImplRaw, MethodSemanticsRaw, MethodSpecRaw, NestedClassRaw, TableId, TypeDefRaw, }, token::Token, }, @@ -121,7 +121,7 @@ pub fn execute_cleanup( if let Some(tables) = view.tables() { if let Some(attr_table) = tables.table::() { for attr_token in request.attributes() { - if let Some(attr) = attr_table.get(attr_token.row()) { + if let Some(attr) = attr_table.get(attr_token.row())? { pre_refs.il_tokens.insert(attr.constructor.token); } } @@ -183,6 +183,15 @@ pub fn execute_cleanup( .map(|t| t.row()) .collect(); + // Keeping a nested type alive means keeping its enclosing chain alive with it: a nested + // type cannot exist without its enclosing type (ECMA-335 II.22.32), and the cascade below + // deletes any nested type whose enclosing type went away. Without this expansion the + // enclosing type is deleted here, the cascade takes the nested type with it, and the + // signature that named the nested type is left pointing at a row that no longer exists -- + // which the writer cannot repair, because `typedef_remap` deliberately declines to remap + // references to deleted types onto surviving rows. + let sig_referenced_typedefs = expand_with_enclosing_types(assembly, sig_referenced_typedefs); + // 2d: Remove types (in descending RID order) for type_token in request.types() { if request.is_protected(*type_token) { @@ -439,6 +448,45 @@ pub fn execute_cleanup( Ok(stats) } +/// Adds the transitive enclosing chain of every nested type in `rids`. +/// +/// Nesting is transitive, so the walk runs to a fixed point: a type nested two levels deep +/// keeps both of its ancestors alive, not just its immediate parent. +/// +/// # Arguments +/// +/// * `assembly` - The assembly to read the `NestedClass` table from +/// * `rids` - TypeDef RIDs that must survive +/// +/// # Returns +/// +/// `rids` plus every TypeDef RID that encloses one of them. +fn expand_with_enclosing_types(assembly: &CilAssembly, mut rids: HashSet) -> HashSet { + let view = assembly.view(); + let Some(tables) = view.tables() else { + return rids; + }; + let Some(nested_table) = tables.table::() else { + return rids; + }; + + let mut enclosing_of: HashMap = HashMap::new(); + for nested in nested_table.iter().filter_map(skip_unreadable) { + enclosing_of.insert(nested.nested_class, nested.enclosing_class); + } + + let mut worklist: Vec = rids.iter().copied().collect(); + while let Some(rid) = worklist.pop() { + if let Some(&enclosing) = enclosing_of.get(&rid) { + if rids.insert(enclosing) { + worklist.push(enclosing); + } + } + } + + rids +} + /// Expands type deletions to include all their members. /// /// For each type marked for deletion, collects all its methods and fields @@ -475,13 +523,13 @@ fn expand_type_members( for type_token in request.types() { let type_rid = type_token.row(); - let Some(typedef) = typedef_table.get(type_rid) else { + let Some(typedef) = typedef_table.get(type_rid).ok().flatten() else { continue; }; // Get method range for this type let method_range = list_range(type_rid, type_count, methoddef_count, |rid| { - typedef_table.get(rid).map(|t| t.method_list) + typedef_table.get(rid).ok().flatten().map(|t| t.method_list) }); // Override start with actual typedef's method_list for method_rid in typedef.method_list..method_range.end { @@ -490,7 +538,7 @@ fn expand_type_members( // Get field range for this type let field_range = list_range(type_rid, type_count, field_count, |rid| { - typedef_table.get(rid).map(|t| t.field_list) + typedef_table.get(rid).ok().flatten().map(|t| t.field_list) }); // Override start with actual typedef's field_list for field_rid in typedef.field_list..field_range.end { @@ -560,7 +608,7 @@ fn remove_empty_types( let mut empty = Vec::new(); for type_rid in 1..=type_count { - let Some(typedef) = typedef_table.get(type_rid) else { + let Some(typedef) = typedef_table.get(type_rid).ok().flatten() else { continue; }; @@ -591,7 +639,7 @@ fn remove_empty_types( // incorrect after deletions: a type whose methods were all deleted still // has a non-zero range, so we must check each row individually. let method_range = list_range(type_rid, type_count, methoddef_count, |rid| { - typedef_table.get(rid).map(|t| t.method_list) + typedef_table.get(rid).ok().flatten().map(|t| t.method_list) }); let live_method_count = (typedef.method_list..method_range.end) .filter(|&rid| !assembly.changes().is_row_deleted(TableId::MethodDef, rid)) @@ -599,7 +647,7 @@ fn remove_empty_types( // Calculate field count for this type — same logic. let field_range = list_range(type_rid, type_count, field_count, |rid| { - typedef_table.get(rid).map(|t| t.field_list) + typedef_table.get(rid).ok().flatten().map(|t| t.field_list) }); let live_field_count = (typedef.field_list..field_range.end) .filter(|&rid| !assembly.changes().is_row_deleted(TableId::Field, rid)) @@ -618,19 +666,22 @@ fn remove_empty_types( // Skip types that are base classes of other surviving types. // Abstract base classes may have no direct members but provide // type hierarchy structure that must be preserved. - let is_base_class = typedef_table.iter().any(|other| { - other.rid != type_rid - && !empty.contains(&other.rid) - && other.extends.tag == TableId::TypeDef - && other.extends.row == type_rid - }); + let is_base_class = typedef_table + .iter() + .filter_map(skip_unreadable) + .any(|other| { + other.rid != type_rid + && !empty.contains(&other.rid) + && other.extends.tag == TableId::TypeDef + && other.extends.row == type_rid + }); if is_base_class { continue; } // Skip types that appear in InterfaceImpl as the interface being implemented. if let Some(iface_impl) = tables.table::() { - let is_implemented = iface_impl.iter().any(|row| { + let is_implemented = iface_impl.iter().filter_map(skip_unreadable).any(|row| { row.interface.tag == TableId::TypeDef && row.interface.row == type_rid }); if is_implemented { @@ -692,6 +743,13 @@ fn collect_alive_method_tokens(assembly: &CilAssembly) -> HashSet { // the underlying MethodDef is alive. if let Some(methodspec_table) = tables.table::() { for row in methodspec_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::MethodSpec, row.rid) @@ -709,6 +767,13 @@ fn collect_alive_method_tokens(assembly: &CilAssembly) -> HashSet { // its constructor method is alive. if let Some(attr_table) = tables.table::() { for row in attr_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::CustomAttribute, row.rid) @@ -725,6 +790,13 @@ fn collect_alive_method_tokens(assembly: &CilAssembly) -> HashSet { // These are alive if the row itself is not deleted. if let Some(sem_table) = tables.table::() { for row in sem_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::MethodSemantics, row.rid) @@ -739,6 +811,13 @@ fn collect_alive_method_tokens(assembly: &CilAssembly) -> HashSet { // MethodImpl.method_body / method_declaration → explicit overrides. if let Some(impl_table) = tables.table::() { for row in impl_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::MethodImpl, row.rid) diff --git a/dotscope/src/cilassembly/cleanup/orphans.rs b/dotscope/src/cilassembly/cleanup/orphans.rs index f3b47c80..f43fb777 100644 --- a/dotscope/src/cilassembly/cleanup/orphans.rs +++ b/dotscope/src/cilassembly/cleanup/orphans.rs @@ -143,7 +143,14 @@ where }; (1..=table.row_count) - .filter_map(|rid| table.get(rid).filter(|r| is_orphan(r)).map(|_| rid)) + .filter_map(|rid| { + table + .get(rid) + .ok() + .flatten() + .filter(|r| is_orphan(r)) + .map(|_| rid) + }) .collect() }; @@ -162,20 +169,33 @@ where /// /// Parameters belong to methods via the `param_list` field in MethodDef. /// When a method is deleted, its parameters become orphaned. -pub(super) fn remove_orphan_params(assembly: &mut CilAssembly, ctx: &DeletionContext) -> usize { +/// +/// Returns the removed Param RIDs alongside the count. Rows in other tables +/// point at parameters through a coded index — `Constant` for a default value, +/// `FieldMarshal` for marshalling information, `CustomAttribute` for anything +/// applied to the parameter — and those rows are dropped by asking +/// [`DeletionContext`] whether their parent went away. A parameter removed here +/// never enters that context: what was deleted is the *method*, and the context +/// is immutable by the time this runs. Without the RIDs, every such dependent +/// outlives the parameter it names and leaves a reference to a row that no +/// longer exists. +pub(super) fn remove_orphan_params( + assembly: &mut CilAssembly, + ctx: &DeletionContext, +) -> (usize, HashSet) { // Collect param RIDs that belong to deleted methods (immutable borrow scope) let mut orphan_params: Vec = { let view = assembly.view(); let Some(tables) = view.tables() else { - return 0; + return (0, HashSet::new()); }; let Some(methoddef_table) = tables.table::() else { - return 0; + return (0, HashSet::new()); }; if tables.table::().is_none() { - return 0; + return (0, HashSet::new()); } let method_count = methoddef_table.row_count; @@ -190,7 +210,11 @@ pub(super) fn remove_orphan_params(assembly: &mut CilAssembly, ctx: &DeletionCon } let range = list_range(method_rid, method_count, param_count, |rid| { - methoddef_table.get(rid).map(|m| m.param_list) + methoddef_table + .get(rid) + .ok() + .flatten() + .map(|m| m.param_list) }); params.extend(range); @@ -203,13 +227,24 @@ pub(super) fn remove_orphan_params(assembly: &mut CilAssembly, ctx: &DeletionCon orphan_params.dedup(); let mut removed_count: usize = 0; + let mut removed_rids: HashSet = HashSet::new(); for rid in orphan_params { if try_remove(assembly, TableId::Param, rid) { removed_count = removed_count.saturating_add(1); + removed_rids.insert(rid); } } - removed_count + (removed_count, removed_rids) +} + +/// Whether `parent` names one of the parameters removed by +/// [`remove_orphan_params`]. +/// +/// The coded index carries the table in the token, so a RID collision with +/// another table cannot be mistaken for a match. +fn names_removed_param(parent: Token, removed_params: &HashSet) -> bool { + parent.table() == TableId::Param.token_type() && removed_params.contains(&parent.row()) } /// Removes orphaned CustomAttribute entries for deleted tokens. @@ -217,12 +252,21 @@ pub(super) fn remove_orphan_params(assembly: &mut CilAssembly, ctx: &DeletionCon /// Custom attributes reference their parent via the `parent` coded index, /// and their constructor via the `constructor` coded index. /// When either the parent or constructor is deleted, the attribute becomes orphaned. -pub(super) fn remove_orphan_attributes(assembly: &mut CilAssembly, ctx: &DeletionContext) -> usize { +pub(super) fn remove_orphan_attributes( + assembly: &mut CilAssembly, + ctx: &DeletionContext, + removed_params: &HashSet, +) -> usize { remove_orphan_entries::(assembly, |attr| { // Remove if parent is deleted if ctx.is_deleted(attr.parent.token) { return true; } + // A parameter dropped with its method is gone without being recorded as + // deleted, so it has to be matched separately. + if names_removed_param(attr.parent.token, removed_params) { + return true; + } // Also remove if constructor method is deleted (MethodDef or MemberRef) ctx.is_deleted(attr.constructor.token) }) @@ -278,6 +322,13 @@ pub(super) fn collect_orphaned_nested_types( let mut orphaned = Vec::new(); for nested in nested_table.iter() { + let nested = match nested { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let enclosing_token = Token::from_parts(TableId::TypeDef, nested.enclosing_class); let nested_token = Token::from_parts(TableId::TypeDef, nested.nested_class); if ctx.is_type_deleted(enclosing_token) && !ctx.is_type_deleted(nested_token) { @@ -328,18 +379,27 @@ pub(super) fn remove_orphan_methodsemantics( }) } -/// Removes orphaned Constant entries for deleted fields. -pub(super) fn remove_orphan_constant(assembly: &mut CilAssembly, ctx: &DeletionContext) -> usize { - remove_orphan_entries::(assembly, |constant| ctx.is_deleted(constant.parent.token)) +/// Removes orphaned Constant entries for deleted fields, properties or params. +pub(super) fn remove_orphan_constant( + assembly: &mut CilAssembly, + ctx: &DeletionContext, + removed_params: &HashSet, +) -> usize { + remove_orphan_entries::(assembly, |constant| { + ctx.is_deleted(constant.parent.token) + || names_removed_param(constant.parent.token, removed_params) + }) } -/// Removes orphaned FieldMarshal entries for deleted fields. +/// Removes orphaned FieldMarshal entries for deleted fields or params. pub(super) fn remove_orphan_fieldmarshal( assembly: &mut CilAssembly, ctx: &DeletionContext, + removed_params: &HashSet, ) -> usize { remove_orphan_entries::(assembly, |marshal| { ctx.is_deleted(marshal.parent.token) + || names_removed_param(marshal.parent.token, removed_params) }) } @@ -374,7 +434,7 @@ pub(super) fn remove_orphan_genericparam( (1..=table.row_count) .filter_map(|rid| { - table.get(rid).and_then(|param| { + table.get(rid).ok().flatten().and_then(|param| { if ctx.is_deleted(param.owner.token) { Some(rid) } else { @@ -478,7 +538,7 @@ pub(super) fn remove_orphan_events( let mut events = Vec::new(); for map_rid in 1..=map_count { - let Some(eventmap) = eventmap_table.get(map_rid) else { + let Some(eventmap) = eventmap_table.get(map_rid).ok().flatten() else { continue; }; let parent_token = Token::from_parts(TableId::TypeDef, eventmap.parent); @@ -487,7 +547,7 @@ pub(super) fn remove_orphan_events( } let range = list_range(map_rid, map_count, event_count, |rid| { - eventmap_table.get(rid).map(|m| m.event_list) + eventmap_table.get(rid).ok().flatten().map(|m| m.event_list) }); events.extend(range); @@ -539,7 +599,7 @@ pub(super) fn remove_orphan_properties( let mut props = Vec::new(); for map_rid in 1..=map_count { - let Some(propertymap) = propertymap_table.get(map_rid) else { + let Some(propertymap) = propertymap_table.get(map_rid).ok().flatten() else { continue; }; let parent_token = Token::from_parts(TableId::TypeDef, propertymap.parent); @@ -548,7 +608,11 @@ pub(super) fn remove_orphan_properties( } let range = list_range(map_rid, map_count, property_count, |rid| { - propertymap_table.get(rid).map(|m| m.property_list) + propertymap_table + .get(rid) + .ok() + .flatten() + .map(|m| m.property_list) }); props.extend(range); @@ -629,7 +693,7 @@ pub(super) fn remove_orphan_modulerefs( if assembly.changes().is_row_deleted(TableId::ImplMap, rid) { continue; } - if let Some(implmap) = implmap_table.get(rid) { + if let Some(implmap) = implmap_table.get(rid).ok().flatten() { alive.insert(implmap.import_scope); } } @@ -641,7 +705,7 @@ pub(super) fn remove_orphan_modulerefs( if assembly.changes().is_row_deleted(TableId::TypeRef, rid) { continue; } - if let Some(typeref) = typeref_table.get(rid) { + if let Some(typeref) = typeref_table.get(rid).ok().flatten() { if typeref.resolution_scope.tag == TableId::ModuleRef { alive.insert(typeref.resolution_scope.row); } @@ -686,7 +750,7 @@ pub(super) fn remove_orphan_assemblyrefs( if assembly.changes().is_row_deleted(TableId::TypeRef, rid) { continue; } - if let Some(typeref) = typeref_table.get(rid) { + if let Some(typeref) = typeref_table.get(rid).ok().flatten() { if typeref.resolution_scope.tag == TableId::AssemblyRef { alive.insert(typeref.resolution_scope.row); } @@ -723,7 +787,7 @@ pub(super) fn remove_orphan_exportedtypes(assembly: &mut CilAssembly) -> (usize, (1..=table.row_count) .filter_map(|rid| { - let entry = table.get(rid)?; + let entry = table.get(rid).ok().flatten()?; // Null implementation (row == 0) means type is defined in this module if entry.implementation.row == 0 { return None; @@ -772,7 +836,7 @@ pub(super) fn remove_orphan_manifestresources(assembly: &mut CilAssembly) -> (us (1..=table.row_count) .filter_map(|rid| { - let entry = table.get(rid)?; + let entry = table.get(rid).ok().flatten()?; // Null implementation (row == 0) = embedded resource, skip if entry.implementation.row == 0 { return None; @@ -827,7 +891,7 @@ pub(super) fn remove_orphan_files(assembly: &mut CilAssembly, candidates: &BTree { continue; } - if let Some(entry) = table.get(rid) { + if let Some(entry) = table.get(rid).ok().flatten() { if entry.implementation.token.is_table(TableId::File) { alive.insert(entry.implementation.token.row()); } @@ -844,7 +908,7 @@ pub(super) fn remove_orphan_files(assembly: &mut CilAssembly, candidates: &BTree { continue; } - if let Some(entry) = table.get(rid) { + if let Some(entry) = table.get(rid).ok().flatten() { if entry.implementation.token.is_table(TableId::File) { alive.insert(entry.implementation.token.row()); } @@ -882,7 +946,7 @@ pub(super) fn remove_type_dependents( ); stats.add( TableId::CustomAttribute, - remove_orphan_attributes(assembly, ctx), + remove_orphan_attributes(assembly, ctx, &HashSet::new()), ); stats.add( TableId::ClassLayout, @@ -929,12 +993,13 @@ pub(super) fn remove_parent_child_dependents( let mut stats = CleanupStats::new(); // 1. Params (depend on methods) - stats.add(TableId::Param, remove_orphan_params(assembly, ctx)); + let (params, removed_params) = remove_orphan_params(assembly, ctx); + stats.add(TableId::Param, params); - // 2. Custom attributes (can target anything) + // 2. Custom attributes (can target anything, parameters included) stats.add( TableId::CustomAttribute, - remove_orphan_attributes(assembly, ctx), + remove_orphan_attributes(assembly, ctx, &removed_params), ); // 3. Type-related tables @@ -1008,9 +1073,12 @@ pub(super) fn remove_parent_child_dependents( ); stats.add( TableId::FieldMarshal, - remove_orphan_fieldmarshal(assembly, ctx), + remove_orphan_fieldmarshal(assembly, ctx, &removed_params), + ); + stats.add( + TableId::Constant, + remove_orphan_constant(assembly, ctx, &removed_params), ); - stats.add(TableId::Constant, remove_orphan_constant(assembly, ctx)); // 6. Generic params (and cascade to constraints) let (genericparams, removed_gp_rids) = remove_orphan_genericparam(assembly, ctx); @@ -1053,6 +1121,13 @@ pub(super) fn cascade_reference_cleanup( if let Some(tables) = view.tables() { if let Some(memberref_table) = tables.table::() { for memberref in memberref_table { + let memberref = match memberref { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if !assembly .changes() .is_row_deleted(TableId::MemberRef, memberref.rid) @@ -1082,6 +1157,13 @@ pub(super) fn cascade_reference_cleanup( if let Some(tables) = view.tables() { if let Some(typespec_table) = tables.table::() { for typespec in typespec_table { + let typespec = match typespec { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if !assembly .changes() .is_row_deleted(TableId::TypeSpec, typespec.rid) @@ -1117,7 +1199,7 @@ pub(super) fn cascade_reference_cleanup( if let Some(tables) = view.tables() { if let Some(memberref_table) = tables.table::() { for &memberref_rid in &deleted_memberref_rids { - if let Some(memberref) = memberref_table.get(memberref_rid) { + if let Some(memberref) = memberref_table.get(memberref_rid).ok().flatten() { if memberref.class.token.is_table(TableId::TypeRef) { typeref_candidates.insert(memberref.class.token.row()); } @@ -1139,6 +1221,13 @@ pub(super) fn cascade_reference_cleanup( if let Some(tables) = view.tables() { if let Some(typeref_table) = tables.table::() { for typeref in typeref_table { + let typeref = match typeref { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if !assembly .changes() .is_row_deleted(TableId::TypeRef, typeref.rid) @@ -1166,7 +1255,7 @@ pub(super) fn cascade_reference_cleanup( // From cascade-deleted TypeRefs if let Some(typeref_table) = tables.table::() { for &typeref_rid in &deleted_typeref_rids { - if let Some(typeref) = typeref_table.get(typeref_rid) { + if let Some(typeref) = typeref_table.get(typeref_rid).ok().flatten() { match typeref.resolution_scope.tag { TableId::AssemblyRef => { assemblyref_candidates.insert(typeref.resolution_scope.row); @@ -1183,6 +1272,13 @@ pub(super) fn cascade_reference_cleanup( // From deleted ImplMaps (import_scope → ModuleRef) if let Some(implmap_table) = tables.table::() { for implmap in implmap_table { + let implmap = match implmap { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::ImplMap, implmap.rid) @@ -1200,6 +1296,13 @@ pub(super) fn cascade_reference_cleanup( if let Some(tables) = view.tables() { if let Some(table) = tables.table::() { for row in table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if !assembly .changes() .is_row_deleted(TableId::ModuleRef, row.rid) @@ -1210,6 +1313,13 @@ pub(super) fn cascade_reference_cleanup( } if let Some(table) = tables.table::() { for row in table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if !assembly .changes() .is_row_deleted(TableId::AssemblyRef, row.rid) @@ -1248,7 +1358,7 @@ pub(super) fn cascade_reference_cleanup( if let Some(tables) = view.tables() { if let Some(table) = tables.table::() { for &rid in &deleted_exportedtype_rids { - if let Some(entry) = table.get(rid) { + if let Some(entry) = table.get(rid).ok().flatten() { if entry.implementation.token.is_table(TableId::File) { file_candidates.insert(entry.implementation.token.row()); } @@ -1257,7 +1367,7 @@ pub(super) fn cascade_reference_cleanup( } if let Some(table) = tables.table::() { for &rid in &deleted_manifestresource_rids { - if let Some(entry) = table.get(rid) { + if let Some(entry) = table.get(rid).ok().flatten() { if entry.implementation.token.is_table(TableId::File) { file_candidates.insert(entry.implementation.token.row()); } diff --git a/dotscope/src/cilassembly/cleanup/references.rs b/dotscope/src/cilassembly/cleanup/references.rs index 89da2a78..a224f4d8 100644 --- a/dotscope/src/cilassembly/cleanup/references.rs +++ b/dotscope/src/cilassembly/cleanup/references.rs @@ -49,9 +49,9 @@ use crate::{ }, streams::Blob, tables::{ - CustomAttributeRaw, FieldRaw, GenericParamConstraintRaw, InterfaceImplRaw, - MemberRefRaw, MethodDefRaw, MethodSpecRaw, PropertyRaw, StandAloneSigRaw, - TableDataOwned, TableId, TypeDefRaw, TypeSpecRaw, + skip_unreadable, CustomAttributeRaw, FieldRaw, GenericParamConstraintRaw, + InterfaceImplRaw, MemberRefRaw, MethodDefRaw, MethodSpecRaw, PropertyRaw, + StandAloneSigRaw, TableDataOwned, TableId, TypeDefRaw, TypeRefRaw, TypeSpecRaw, }, token::Token, }, @@ -116,6 +116,13 @@ pub(super) fn collect_pre_deletion_references( // Scan method bodies and signatures of methods being deleted if let Some(methoddef_table) = tables.table::() { for methoddef in methoddef_table { + let methoddef = match methoddef { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let method_token = Token::from_parts(TableId::MethodDef, methoddef.rid); if !methods.contains(&method_token) { continue; @@ -147,6 +154,13 @@ pub(super) fn collect_pre_deletion_references( if let Some(field_table) = tables.table::() { if let Some(blob) = &blob_heap { for field in field_table { + let field = match field { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let field_token = Token::from_parts(TableId::Field, field.rid); if !fields.contains(&field_token) { continue; @@ -159,7 +173,7 @@ pub(super) fn collect_pre_deletion_references( // Scan extends clause of types being deleted if let Some(typedef_table) = tables.table::() { for type_token in types { - if let Some(typedef) = typedef_table.get(type_token.row()) { + if let Some(typedef) = typedef_table.get(type_token.row()).ok().flatten() { if typedef.extends.token.is_table(TableId::TypeRef) { typeref_rids.insert(typedef.extends.token.row()); } @@ -170,6 +184,13 @@ pub(super) fn collect_pre_deletion_references( // Scan CustomAttribute constructors whose parent is being deleted if let Some(attr_table) = tables.table::() { for attr in attr_table { + let attr = match attr { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let parent_token = attr.parent.token; let parent_deleted = types.contains(&parent_token) || methods.contains(&parent_token) @@ -381,6 +402,13 @@ pub(super) fn collect_typedefs_from_field_signatures(assembly: &CilAssembly) -> }; for field in field_table.iter() { + let field = match field { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly.changes().is_row_deleted(TableId::Field, field.rid) { continue; } @@ -415,6 +443,7 @@ pub(super) fn collect_referenced_standalonesig_rids(assembly: &CilAssembly) -> H methoddef_table .into_iter() + .filter_map(skip_unreadable) .filter(|m| !assembly.changes().is_row_deleted(TableId::MethodDef, m.rid)) .map(|m| get_effective_method_rva(assembly, m.rid, m.rva)) .filter(|&rva| rva != 0) @@ -457,6 +486,7 @@ pub(super) fn scan_method_body_tokens(assembly: &CilAssembly) -> HashSet methoddef_table .into_iter() + .filter_map(skip_unreadable) .filter(|m| !assembly.changes().is_row_deleted(TableId::MethodDef, m.rid)) .map(|m| get_effective_method_rva(assembly, m.rid, m.rva)) .filter(|&rva| rva != 0) @@ -487,7 +517,7 @@ pub(super) fn scan_method_body_tokens(assembly: &CilAssembly) -> HashSet if let (Some(tables), Some(blob_heap)) = (view.tables(), view.blobs()) { if let Some(sig_table) = tables.table::() { for &rid in &local_sig_rids { - if let Some(sig_row) = sig_table.get(rid) { + if let Some(sig_row) = sig_table.get(rid).ok().flatten() { if let Ok(blob_data) = blob_heap.get(sig_row.signature as usize) { collect_type_tokens_from_local_sig(blob_data, &mut referenced); } @@ -529,7 +559,7 @@ pub(super) fn collect_typerefs_from_deleted_memberref_sigs( }; for &rid in memberref_rids { - if let Some(memberref) = memberref_table.get(rid) { + if let Some(memberref) = memberref_table.get(rid).ok().flatten() { if !scan_method_signature_blob(blob_heap, memberref.signature, &mut result) { scan_field_signature_blob(blob_heap, memberref.signature, &mut result); } @@ -558,6 +588,13 @@ pub(super) fn scan_typeref_metadata_refs(assembly: &CilAssembly) -> HashSet // TypeDef.extends - base class references (skip deleted types) if let Some(typedef_table) = tables.table::() { for typedef in typedef_table { + let typedef = match typedef { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::TypeDef, typedef.rid) @@ -573,6 +610,13 @@ pub(super) fn scan_typeref_metadata_refs(assembly: &CilAssembly) -> HashSet // InterfaceImpl.interface (skip deleted rows) if let Some(interfaceimpl_table) = tables.table::() { for impl_ in interfaceimpl_table { + let impl_ = match impl_ { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::InterfaceImpl, impl_.rid) @@ -588,6 +632,13 @@ pub(super) fn scan_typeref_metadata_refs(assembly: &CilAssembly) -> HashSet // MemberRef.class - declaring type of member references (skip deleted rows) if let Some(memberref_table) = tables.table::() { for memberref in memberref_table { + let memberref = match memberref { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::MemberRef, memberref.rid) @@ -603,6 +654,13 @@ pub(super) fn scan_typeref_metadata_refs(assembly: &CilAssembly) -> HashSet // GenericParamConstraint - type constraints (skip deleted rows) if let Some(constraint_table) = tables.table::() { for constraint in constraint_table { + let constraint = match constraint { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::GenericParamConstraint, constraint.rid) @@ -619,6 +677,13 @@ pub(super) fn scan_typeref_metadata_refs(assembly: &CilAssembly) -> HashSet if let Some(attr_table) = tables.table::() { if let Some(memberref_table) = tables.table::() { for attr in attr_table { + let attr = match attr { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::CustomAttribute, attr.rid) @@ -633,7 +698,7 @@ pub(super) fn scan_typeref_metadata_refs(assembly: &CilAssembly) -> HashSet { continue; } - if let Some(memberref) = memberref_table.get(memberref_rid) { + if let Some(memberref) = memberref_table.get(memberref_rid).ok().flatten() { if memberref.class.token.is_table(TableId::TypeRef) { referenced_rids.insert(memberref.class.token.row()); } @@ -646,6 +711,62 @@ pub(super) fn scan_typeref_metadata_refs(assembly: &CilAssembly) -> HashSet referenced_rids } +/// The `nested TypeRef -> enclosing TypeRef` edges, as `(rid, enclosing_rid)` pairs. +/// +/// A nested type is encoded as a TypeRef whose `ResolutionScope` points at the *enclosing* +/// TypeRef, and nothing in the ordinary reference scans reads that column — so an enclosing +/// TypeRef referenced only from a surviving nested one looks orphaned, is deleted, and leaves +/// the survivor's scope dangling. +/// +/// Collected separately from [`close_over_typeref_nesting`] because the closure has to run +/// against the *complete* live set. Rows already marked deleted are excluded: a nested TypeRef +/// on its way out must not keep its parent alive. +fn typeref_nesting_edges(assembly: &CilAssembly) -> Vec<(u32, u32)> { + let view = assembly.view(); + let Some(tables) = view.tables() else { + return Vec::new(); + }; + let Some(typeref_table) = tables.table::() else { + return Vec::new(); + }; + + typeref_table + .into_iter() + .filter_map(skip_unreadable) + .filter(|type_ref| { + !assembly + .changes() + .is_row_deleted(TableId::TypeRef, type_ref.rid) + }) + .filter(|type_ref| type_ref.resolution_scope.token.is_table(TableId::TypeRef)) + .map(|type_ref| (type_ref.rid, type_ref.resolution_scope.token.row())) + .collect() +} + +/// Marks every enclosing TypeRef of an already-live nested TypeRef live, transitively. +/// +/// Must be called once the live set is **complete**. Running it earlier — while it could see +/// only the metadata-table references and not the IL body tokens, signature blobs, or +/// MemberRef-derived RIDs added afterwards — meant a nested TypeRef kept alive by an IL token +/// never marked its parent, which is the most common way a nested type stays live. +/// +/// Nesting is transitive (`A+B+C`), so this iterates to a fixed point: marking `B` live must +/// in turn mark `A` live. Bounded by the edge count, since each pass either marks at least one +/// new RID live or stops. +fn close_over_typeref_nesting(referenced_rids: &mut HashSet, nesting: &[(u32, u32)]) { + loop { + let mut changed = false; + for (rid, enclosing) in nesting { + if referenced_rids.contains(rid) && referenced_rids.insert(*enclosing) { + changed = true; + } + } + if !changed { + break; + } + } +} + /// Scans metadata tables to collect MemberRef RIDs that are referenced. /// /// Collects references from: @@ -662,6 +783,13 @@ pub(super) fn scan_memberref_metadata_refs(assembly: &CilAssembly) -> HashSet() { for attr in attr_table { + let attr = match attr { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::CustomAttribute, attr.rid) @@ -677,6 +805,13 @@ pub(super) fn scan_memberref_metadata_refs(assembly: &CilAssembly) -> HashSet() { for spec in methodspec_table { + let spec = match spec { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::MethodSpec, spec.rid) @@ -710,6 +845,13 @@ pub(super) fn scan_typespec_metadata_refs(assembly: &CilAssembly) -> HashSet() { for memberref in memberref_table { + let memberref = match memberref { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::MemberRef, memberref.rid) @@ -725,6 +867,13 @@ pub(super) fn scan_typespec_metadata_refs(assembly: &CilAssembly) -> HashSet() { for impl_ in interfaceimpl_table { + let impl_ = match impl_ { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::InterfaceImpl, impl_.rid) @@ -740,6 +889,13 @@ pub(super) fn scan_typespec_metadata_refs(assembly: &CilAssembly) -> HashSet() { for constraint in constraint_table { + let constraint = match constraint { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::GenericParamConstraint, constraint.rid) @@ -755,6 +911,13 @@ pub(super) fn scan_typespec_metadata_refs(assembly: &CilAssembly) -> HashSet() { for typedef in typedef_table { + let typedef = match typedef { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::TypeDef, typedef.rid) @@ -909,6 +1072,13 @@ pub(super) fn scan_signature_typeref_refs(assembly: &CilAssembly) -> HashSet() { for methoddef in methoddef_table { + let methoddef = match methoddef { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::MethodDef, methoddef.rid) @@ -922,6 +1092,13 @@ pub(super) fn scan_signature_typeref_refs(assembly: &CilAssembly) -> HashSet() { for field in field_table { + let field = match field { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly.changes().is_row_deleted(TableId::Field, field.rid) { continue; } @@ -932,6 +1109,13 @@ pub(super) fn scan_signature_typeref_refs(assembly: &CilAssembly) -> HashSet() { for memberref in memberref_table { + let memberref = match memberref { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::MemberRef, memberref.rid) @@ -951,6 +1135,13 @@ pub(super) fn scan_signature_typeref_refs(assembly: &CilAssembly) -> HashSet() { for sig in standalonesig_table { + let sig = match sig { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; // Only scan StandAloneSigs that are referenced by current method bodies if referenced_sigs.contains(&sig.rid) { scan_local_var_signature_blob(blob_heap, sig.signature, &mut referenced_rids); @@ -961,6 +1152,13 @@ pub(super) fn scan_signature_typeref_refs(assembly: &CilAssembly) -> HashSet() { for typespec in typespec_table { + let typespec = match typespec { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::TypeSpec, typespec.rid) @@ -974,6 +1172,13 @@ pub(super) fn scan_signature_typeref_refs(assembly: &CilAssembly) -> HashSet() { for property in property_table { + let property = match property { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::Property, property.rid) @@ -1139,6 +1344,13 @@ pub(super) fn remove_unreferenced_typerefs( if let Some(tables) = view.tables() { if let Some(memberref_table) = tables.table::() { for memberref in memberref_table { + let memberref = match memberref { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::MemberRef, memberref.rid) @@ -1155,6 +1367,14 @@ pub(super) fn remove_unreferenced_typerefs( } } + // Last, once every source above has contributed: a nested TypeRef that survives keeps its + // enclosing TypeRef alive. This has to see the *whole* live set — running it inside + // `scan_typeref_metadata_refs` meant it saw only that function's own five metadata sources, + // so a nested TypeRef kept alive by an IL body token, a signature blob, or a MemberRef + // never marked its parent, and the parent was deleted out from under it. + let nesting = typeref_nesting_edges(assembly); + close_over_typeref_nesting(&mut referenced_rids, &nesting); + remove_candidates_not_alive(assembly, TableId::TypeRef, candidates, &referenced_rids) } @@ -1200,6 +1420,13 @@ pub(super) fn remove_unreferenced_memberrefs( if let Some(tables) = view.tables() { if let Some(methodspec_table) = tables.table::() { for spec in methodspec_table { + let spec = match spec { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if assembly .changes() .is_row_deleted(TableId::MethodSpec, spec.rid) @@ -1254,6 +1481,7 @@ pub(super) fn remove_unreferenced_typespecs( #[cfg(test)] mod tests { + use super::*; use crate::{ cilassembly::cleanup::utils::PLACEHOLDER_RVA_THRESHOLD, metadata::{tables::TableId, token::Token}, @@ -1280,4 +1508,51 @@ mod tests { assert!(memberref_token.is_table(TableId::MemberRef)); assert!(typespec_token.is_table(TableId::TypeSpec)); } + + /// A nested TypeRef kept alive by *any* source must keep its enclosing TypeRef alive too. + /// + /// The closure used to run inside `scan_typeref_metadata_refs`, where it could see only + /// that function's own metadata sources — not the IL body tokens, signature blobs or + /// MemberRef-derived RIDs unioned in afterwards. Seeding here with a RID that only those + /// later sources would have supplied is what distinguishes the fixed version: under the + /// old ordering the enclosing RIDs stayed absent and the rows were deleted, leaving the + /// survivor's `ResolutionScope` pointing at whatever shifted into their place. + #[test] + fn nesting_closure_marks_enclosing_typerefs_of_a_live_nested_typeref() { + // 3 nested in 2, 2 nested in 1 — the transitive `A+B+C` shape. + let nesting = [(3u32, 2u32), (2u32, 1u32)]; + + // RID 3 is live only because an IL body token referenced it. + let mut live = HashSet::from([3u32]); + close_over_typeref_nesting(&mut live, &nesting); + + assert!(live.contains(&2), "the immediately enclosing TypeRef"); + assert!( + live.contains(&1), + "nesting is transitive, so the outermost TypeRef must be marked too" + ); + } + + /// The closure must not resurrect an enclosing TypeRef whose nested type is itself dead. + #[test] + fn nesting_closure_leaves_unreferenced_chains_alone() { + let nesting = [(3u32, 2u32), (2u32, 1u32)]; + + let mut live = HashSet::from([9u32]); + close_over_typeref_nesting(&mut live, &nesting); + + assert_eq!(live, HashSet::from([9u32])); + } + + /// A cycle in `ResolutionScope` must not spin: the fixed point stops once nothing new is + /// marked, so a self- or mutually-nested pair terminates rather than looping forever. + #[test] + fn nesting_closure_terminates_on_a_cycle() { + let nesting = [(1u32, 2u32), (2u32, 1u32)]; + + let mut live = HashSet::from([1u32]); + close_over_typeref_nesting(&mut live, &nesting); + + assert_eq!(live, HashSet::from([1u32, 2u32])); + } } diff --git a/dotscope/src/cilassembly/cleanup/utils.rs b/dotscope/src/cilassembly/cleanup/utils.rs index d8ad452d..faefa530 100644 --- a/dotscope/src/cilassembly/cleanup/utils.rs +++ b/dotscope/src/cilassembly/cleanup/utils.rs @@ -127,7 +127,7 @@ pub(crate) fn is_cctor_method(assembly: &CilAssembly, method_rid: u32) -> bool { let Some(method_table) = tables.table::() else { return false; }; - let Some(row) = method_table.get(method_rid) else { + let Some(row) = method_table.get(method_rid).ok().flatten() else { return false; }; let Some(strings) = view.strings() else { diff --git a/dotscope/src/cilassembly/mod.rs b/dotscope/src/cilassembly/mod.rs index c7088cda..2b26bf08 100644 --- a/dotscope/src/cilassembly/mod.rs +++ b/dotscope/src/cilassembly/mod.rs @@ -117,7 +117,7 @@ //! //! ```rust,no_run //! use dotscope::{CilAssemblyView, CilAssembly}; -//! use std::path::Path; +//! use std::{fmt, path::Path}; //! //! // Load and convert to mutable assembly //! let view = CilAssemblyView::from_path(Path::new("assembly.dll"))?; @@ -130,7 +130,7 @@ //! assembly.to_file(Path::new("modified.dll"))?; //! # Ok::<(), dotscope::Error>(()) //! ``` -use std::path::Path; +use std::{fmt, path::Path}; use crate::{ file::File, @@ -143,7 +143,9 @@ use crate::{ encode_property_signature, encode_typespec_signature, SignatureField, SignatureLocalVariables, SignatureMethod, SignatureProperty, SignatureTypeSpec, }, - tables::{AssemblyRefRaw, CodedIndex, CodedIndexType, TableDataOwned, TableId}, + tables::{ + skip_unreadable, AssemblyRefRaw, CodedIndex, CodedIndexType, TableDataOwned, TableId, + }, token::Token, }, CilObject, Error, Result, ValidationConfig, @@ -1699,13 +1701,12 @@ impl CilAssembly { self.view.tables()?.table::(), self.view.strings(), ) { - for (index, assemblyref) in assembly_ref_table.iter().enumerate() { + for assemblyref in assembly_ref_table.iter().filter_map(skip_unreadable) { if let Ok(assembly_name) = strings.get(assemblyref.name as usize) { if assembly_name == name { - // Convert 0-based index to 1-based RID return Some(CodedIndex::new( TableId::AssemblyRef, - u32::try_from(index.saturating_add(1)).unwrap_or(u32::MAX), + assemblyref.rid, CodedIndexType::Implementation, )); } @@ -1855,7 +1856,7 @@ impl From for CilAssembly { } } -impl std::fmt::Debug for CilAssembly { +impl fmt::Debug for CilAssembly { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("CilAssembly") .field("original_view", &"") diff --git a/dotscope/src/cilassembly/writer/context.rs b/dotscope/src/cilassembly/writer/context.rs index 159e0fa9..c3106854 100644 --- a/dotscope/src/cilassembly/writer/context.rs +++ b/dotscope/src/cilassembly/writer/context.rs @@ -67,8 +67,23 @@ pub struct SectionWriteInfo { pub data_offset: Option, /// RVA assigned to this section's data pub rva: Option, - /// Actual size of data written (virtual size) + /// The section's extent in address space — its `VirtualSize`, and what the next section's + /// RVA is placed after. + /// + /// For sections the writer regenerates (`.text`) this is the generated size. For sections + /// copied verbatim it is the input's `VirtualSize`, which may differ from the number of + /// bytes on disk in either direction. pub data_size: Option, + /// Number of bytes actually written to the file for this section. + /// + /// The source of truth for `SizeOfRawData`, which is otherwise derived from + /// [`data_size`](Self::data_size) — a different quantity. A section carrying uninitialised + /// data has `VirtualSize > SizeOfRawData`, so deriving one from the other emits a header + /// describing an extent that does not match the bytes present. + /// + /// Must not be conflated with `data_size`: that value also drives RVA placement, and + /// shrinking it to the written length overlaps the following section. + pub raw_size: Option, /// Whether this section should be removed (header zeroed, count decremented) pub removed: bool, } @@ -228,6 +243,14 @@ pub struct WriteContext<'a> { pub import_data_offset: Option, /// RVA of native import table. pub import_data_rva: Option, + + /// Byte sizes of the heaps as they will be emitted: `(#Strings, #GUID, #Blob)`. + /// + /// Recorded by the heap pre-pass so the tables stream can declare heap index widths that + /// match the heaps actually written. ECMA-335 II.24.2.6 selects a 4-byte index once a heap + /// exceeds 0xFFFF; inheriting the input's `HeapSizes` byte instead silently truncates every + /// offset past that point once the writer appends enough. + pub output_heap_sizes: Option<(usize, usize, usize)>, /// Size of native import table in bytes. pub import_data_size: Option, @@ -474,6 +497,7 @@ impl<'a> WriteContext<'a> { import_data_offset: None, import_data_rva: None, + output_heap_sizes: None, import_data_size: None, pending_imports: None, native_entry_rva: None, @@ -778,11 +802,19 @@ impl<'a> WriteContext<'a> { /// * `data_offset` - File offset where section data was written /// * `rva` - RVA assigned to the section /// * `data_size` - Size of data written - pub fn update_section(&mut self, index: usize, data_offset: u64, rva: u32, data_size: u32) { + pub fn update_section( + &mut self, + index: usize, + data_offset: u64, + rva: u32, + data_size: u32, + raw_size: u32, + ) { if let Some(section) = self.sections.get_mut(index) { section.data_offset = Some(data_offset); section.rva = Some(rva); section.data_size = Some(data_size); + section.raw_size = Some(raw_size); } } diff --git a/dotscope/src/cilassembly/writer/fields.rs b/dotscope/src/cilassembly/writer/fields.rs index cbab965e..5c25d498 100644 --- a/dotscope/src/cilassembly/writer/fields.rs +++ b/dotscope/src/cilassembly/writer/fields.rs @@ -92,6 +92,8 @@ fn calculate_field_size( })?; let field_row = field_table .iter() + .collect::>>()? + .into_iter() .find(|r| r.rid == field_index) .ok_or_else(|| { Error::ModificationInvalid(format!("Field {field_index} not found in Field table")) @@ -142,6 +144,7 @@ fn calculate_type_size( if is_typedef { if let Some(class_layout_table) = tables.table::() { for layout_row in class_layout_table { + let layout_row = layout_row?; if layout_row.parent == row { return Ok(layout_row.class_size as usize); } @@ -227,6 +230,7 @@ fn collect_original_fieldrva_data( let mut entries_to_process: Vec<(u32, u32, u32)> = Vec::new(); // (rva, rid, field_index) for row in fieldrva_table { + let row = row?; if deleted_rids.contains(&row.rid) || deleted_field_rids.contains(&row.field) || modified_rids.contains(&row.rid) diff --git a/dotscope/src/cilassembly/writer/fixups.rs b/dotscope/src/cilassembly/writer/fixups.rs index 8b3e4a78..6478aaf3 100644 --- a/dotscope/src/cilassembly/writer/fixups.rs +++ b/dotscope/src/cilassembly/writer/fixups.rs @@ -211,18 +211,20 @@ pub fn fixup_section_table(ctx: &mut WriteContext) -> Result<()> { continue; // Skip sections without data }; - // SizeOfRawData must be a multiple of FileAlignment per PE spec. - // This is required for all sections including the last one. - let file_size = u32::try_from(align_to( - u64::from(data_size), - u64::from(ctx.file_alignment), - )) - .map_err(|_| { - Error::LayoutFailed(format!( - "Section {} file size exceeds u32 range", - section.name - )) - })?; + // SizeOfRawData must be a multiple of FileAlignment per PE spec, and must describe the + // bytes actually written — not `data_size`, which is the section's *virtual* extent and + // also drives RVA placement. The two differ whenever a section carries uninitialised + // data (VirtualSize > SizeOfRawData) or trailing alignment padding (the reverse). + // + // Falls back to `data_size` only for sections that recorded no written length. + let raw_size = section.raw_size.unwrap_or(data_size); + let file_size = u32::try_from(align_to(u64::from(raw_size), u64::from(ctx.file_alignment))) + .map_err(|_| { + Error::LayoutFailed(format!( + "Section {} file size exceeds u32 range", + section.name + )) + })?; let offset_u32 = u32::try_from(data_offset).map_err(|_| { Error::LayoutFailed(format!("Section {} offset exceeds u32 range", section.name)) @@ -682,24 +684,18 @@ pub fn zero_stripped_data_regions(ctx: &mut WriteContext) -> Result<()> { // we don't copy those sections either - we only preserve .rsrc and .reloc. let _ = ctx.original_debug_dir; // Stored for reference but not used - // Certificate data handling: + // Certificate data is deliberately NOT scrubbed here. // - // Certificates use a FILE OFFSET (not RVA) in the data directory, and are - // typically appended after all sections. Since we truncate the output to - // `bytes_written`, certificate data that was beyond our content is naturally - // excluded. If somehow certificate data falls within our written bounds - // (unusual but possible), we zero it since the signature is invalid after - // any modification. - if let Some((cert_offset, cert_size)) = ctx.original_certificate_dir { - let cert_offset_u64 = u64::from(cert_offset); - let cert_end = cert_offset_u64 - .checked_add(u64::from(cert_size)) - .ok_or_else(|| Error::LayoutFailed("Certificate region offset overflow".to_string()))?; - if cert_end <= ctx.bytes_written { - let zeros = vec![0u8; cert_size as usize]; - ctx.write_at(cert_offset_u64, &zeros)?; - } - } + // `original_certificate_dir` is a **file offset into the input**, and this writes into the + // *output*, which has an entirely different layout — the offset that held a signature in the + // input names unrelated content here, typically live `.text` or metadata. Zeroing it + // corrupted the generated file rather than sanitising it, and it ran before + // `fixup_checksum`, so the checksum was computed over the damage and the result looked + // internally consistent. + // + // Nothing needs scrubbing: the output is rebuilt from scratch and never copies certificate + // bytes, and `write_optional_header` already zeroes the CertificateTable data-directory + // entry. This mirrors the reasoning already applied to the debug directory. Ok(()) } diff --git a/dotscope/src/cilassembly/writer/generator.rs b/dotscope/src/cilassembly/writer/generator.rs index ee491995..9012030f 100644 --- a/dotscope/src/cilassembly/writer/generator.rs +++ b/dotscope/src/cilassembly/writer/generator.rs @@ -73,8 +73,8 @@ use crate::{ streams::{Blob, Guid, StreamHeader, Strings, UserStrings}, tablefields::get_heap_fields, tables::{ - ManifestResourceRaw, MethodDefRaw, RowWritable, StandAloneSigRaw, TableDataOwned, - TableId, TableInfoRef, + skip_unreadable, ManifestResourceRaw, MethodDefRaw, RowWritable, StandAloneSigRaw, + TableDataOwned, TableId, TableInfoRef, }, token::Token, }, @@ -82,6 +82,13 @@ use crate::{ Error, Result, }; +/// Upper bound on how far the native-stub decoder will scan for a function end. +/// +/// The scan window is otherwise "this PE section", which for a small file can still be the +/// bulk of it. A native stub in a mixed-mode assembly is kilobytes of code; anything past this +/// is not a stub whose extent we failed to find, it is a region we should refuse to guess at. +const MAX_NATIVE_BODY_SCAN_BYTES: usize = 64 * 1024; + /// IAT (Import Address Table) size for .NET executables (8 bytes). const IAT_SIZE: u64 = 8; @@ -384,16 +391,19 @@ impl<'a> PeGenerator<'a> { } self.write_cor20_header(&mut ctx)?; - // Pre-compute heap offsets early - needed for method bodies that reference - // newly added userstrings (ldstr instructions) and other heap entries. - // This resolves ChangeRefs to their final heap offsets. - precompute_heap_offsets(self.assembly.view(), &mut ctx, changes)?; - // Resolve table ChangeRefs early - needed for method bodies that reference // newly added StandAloneSig entries (local variable signatures) Self::resolve_table_change_refs(changes); - // Build RID remapper early - needed to patch IL tokens when rows are deleted. + // Build RID remapper BEFORE the heap pre-pass. + // + // The pre-pass must see the same remap maps the real write pass will use, or the two + // disagree on every offset after the first remapped signature blob — and the tables are + // serialised from the pre-pass values. Building the remapper afterwards left the + // pre-pass with empty maps by construction. + // + // The remapper depends only on `changes` and the original table row counts, not on any + // heap offset, so it is safe to compute this early. // When TypeDef/MethodDef/etc rows are removed, subsequent rows shift down // and IL tokens must be updated accordingly. if let Some(tables) = self.assembly.view().tables() { @@ -433,6 +443,11 @@ impl<'a> PeGenerator<'a> { self.build_standalonesig_dedup(&mut ctx, changes); } + // Pre-compute heap offsets - needed for method bodies that reference newly added + // userstrings (ldstr instructions) and other heap entries. This resolves ChangeRefs to + // their final heap offsets, using the remap maps established above. + precompute_heap_offsets(self.assembly.view(), &mut ctx, changes)?; + // Write method bodies ctx.align_to_4(); ctx.method_bodies_offset = ctx.pos(); @@ -475,11 +490,14 @@ impl<'a> PeGenerator<'a> { .ok_or_else(|| Error::LayoutFailed(".text section size underflow".to_string()))?; let text_size_u32 = u32::try_from(ctx.text_section_size).unwrap_or(u32::MAX); if let Some(idx) = ctx.find_section_index(".text") { + // `.text` is regenerated, so its virtual extent and its written length are the + // same measured `pos()` delta. ctx.update_section( idx, ctx.text_section_offset, ctx.text_section_rva, text_size_u32, + text_size_u32, ); } @@ -794,6 +812,7 @@ impl<'a> PeGenerator<'a> { data_offset: None, rva: None, data_size: None, + raw_size: None, removed: false, }); } @@ -1070,6 +1089,7 @@ impl<'a> PeGenerator<'a> { let mut written_rvas: HashSet = HashSet::new(); for row in method_table { + let row = row?; if deleted_method_rids.contains(&row.rid) { continue; } @@ -1104,21 +1124,52 @@ impl<'a> PeGenerator<'a> { #[cfg(feature = "x86")] { let offset = file.rva_to_offset(original_rva as usize)?; - let available = file.data().len().saturating_sub(offset); + + // Clamp the scan window to the PE section containing the stub, + // and to an absolute ceiling. Handing the decoder "everything + // from here to EOF" let a single native method absorb the rest + // of the input file into the output as one body, and made the + // decode cost proportional to file size rather than to the stub. + let section_end = file + .sections() + .iter() + .find_map(|section| { + let start = section.pointer_to_raw_data as usize; + let end = + start.saturating_add(section.size_of_raw_data as usize); + (offset >= start && offset < end).then_some(end) + }) + .unwrap_or_else(|| file.data().len()) + .min(file.data().len()); + let available = section_end + .saturating_sub(offset) + .min(MAX_NATIVE_BODY_SCAN_BYTES); let scan_data = file.data_slice(offset, available)?; let body_size = x86_native_body_size(scan_data, file.pe().is_64bit); - if body_size > 0 { - let body_data = file.data_slice(offset, body_size)?; + if body_size == 0 || body_size >= available { + // Either no function end was found, or the decode ran to the + // clamp — in both cases the extent is unknown. Emitting the + // row anyway leaves its RVA unmapped, and the fallback + // resolution maps an unmapped RVA to *itself*, so the + // method would point at whatever now occupies that offset in + // the relaid-out file. + return Err(Error::LayoutFailed(format!( + "native method body at RVA 0x{original_rva:08X} has no \ + determinable extent (decoded {body_size} of {available} \ + bytes); refusing to emit a row with an unresolvable RVA" + ))); + } + + let body_data = file.data_slice(offset, body_size)?; - // Native code needs 4-byte alignment for consistency - ctx.align_to_4_with_padding()?; + // Native code needs 4-byte alignment for consistency + ctx.align_to_4_with_padding()?; - let new_rva = ctx.current_rva(); - ctx.method_body_rva_map.insert(original_rva, new_rva); - ctx.write(body_data)?; - } + let new_rva = ctx.current_rva(); + ctx.method_body_rva_map.insert(original_rva, new_rva); + ctx.write(body_data)?; } } else { // CIL method - parse, remap tokens, and rebuild @@ -1226,7 +1277,7 @@ impl<'a> PeGenerator<'a> { .tables() .and_then(|t| t.table::()) .is_some_and(|table| { - table.iter().any(|row| { + table.iter().filter_map(skip_unreadable).any(|row| { // Only embedded resources (implementation.row == 0) have data // in this section. row.implementation.row == 0 @@ -1239,6 +1290,7 @@ impl<'a> PeGenerator<'a> { if let Some(table) = view.tables().and_then(|t| t.table::()) { let mut new_offset = 0u32; for row in table { + let row = row?; // External resources don't have data in this section if row.implementation.row != 0 { continue; @@ -1529,11 +1581,21 @@ impl<'a> PeGenerator<'a> { // Create output table info with new row counts (same as in write_tables_stream) // This ensures we use the correct row sizes for reading the output - let output_table_info = Arc::new( - tables + // Row counts first, then heap index widths derived from the heaps actually written. + // Both must reflect the *output*: `with_modified_row_counts` copies the input's heap + // widths verbatim, so without the second step a writer that appends past 0xFFFF keeps + // emitting 2-byte indices and every offset beyond it is masked. + let output_table_info = { + let with_rows = tables .info - .with_modified_row_counts(new_row_counts.iter().map(|(k, v)| (*k, *v))), - ); + .with_modified_row_counts(new_row_counts.iter().map(|(k, v)| (*k, *v))); + Arc::new(match ctx.output_heap_sizes { + Some((strings, guids, blobs)) => { + with_rows.with_modified_heap_sizes(strings, guids, blobs) + } + None => with_rows, + }) + }; // Calculate where table data starts (after tables stream header) let header_size = (valid.count_ones() as usize) @@ -1579,7 +1641,7 @@ impl<'a> PeGenerator<'a> { table_id, &output_table_info, changes, - ); + )?; // Second: apply old→new heap offset remapping (from deduplication) // This applies to ALL rows, including updated ones. ChangeRef placeholders @@ -1624,12 +1686,17 @@ impl<'a> PeGenerator<'a> { /// * `table_id` - The table type (determines which fields are heap references) /// * `table_info` - Table size information for field offset calculation /// * `changes` - The assembly changes for ChangeRef lookup + /// + /// # Errors + /// + /// Returns [`crate::Error::LayoutFailed`] if a resolved heap offset does not fit the index + /// width the emitted `HeapSizes` byte declares. fn patch_row_change_ref_placeholders( row_buffer: &mut [u8], table_id: TableId, table_info: &TableInfoRef, changes: &AssemblyChanges, - ) { + ) -> Result<()> { // Get heap field positions using the centralized schema let heap_fields = get_heap_fields(table_id, table_info); @@ -1670,17 +1737,30 @@ impl<'a> PeGenerator<'a> { if field.size == 4 { field_slice_mut.copy_from_slice(&resolved.to_le_bytes()); } else { - // Truncate to u16 - this is safe because heap offsets in small - // metadata files fit in u16. Overflow would indicate a corrupted state. - #[allow(clippy::cast_possible_truncation)] - let small_value = - u16::try_from(resolved).unwrap_or((resolved & 0xFFFF) as u16); + // A 2-byte heap index cannot represent an offset past 0xFFFF. + // + // Masking to 16 bits would not produce a broken file — it would + // produce a *valid-looking* one whose names, signatures and + // attribute blobs point at arbitrary earlier heap positions, + // deterministically modulo 0x10000, and an analyst consumes those + // attributions as fact. Reaching this means the emitted HeapSizes + // byte disagrees with the heap that was actually written, which is + // a layout bug and must not be silently encoded into the output. + let small_value = u16::try_from(resolved).map_err(|_| { + Error::LayoutFailed(format!( + "heap offset 0x{resolved:X} exceeds the 2-byte index width \ + declared by HeapSizes; the output heap grew past 0xFFFF but \ + the index width was not promoted" + )) + })?; field_slice_mut.copy_from_slice(&small_value.to_le_bytes()); } } } } } + + Ok(()) } /// Estimates the size of the metadata root header. @@ -1920,11 +2000,21 @@ impl<'a> PeGenerator<'a> { // Create output table info with new row counts // This recalculates coded index sizes based on new row counts - let output_table_info = Arc::new( - tables + // Row counts first, then heap index widths derived from the heaps actually written. + // Both must reflect the *output*: `with_modified_row_counts` copies the input's heap + // widths verbatim, so without the second step a writer that appends past 0xFFFF keeps + // emitting 2-byte indices and every offset beyond it is masked. + let output_table_info = { + let with_rows = tables .info - .with_modified_row_counts(new_row_counts.iter().map(|(k, v)| (*k, *v))), - ); + .with_modified_row_counts(new_row_counts.iter().map(|(k, v)| (*k, *v))); + Arc::new(match ctx.output_heap_sizes { + Some((strings, guids, blobs)) => { + with_rows.with_modified_heap_sizes(strings, guids, blobs) + } + None => with_rows, + }) + }; // Write tables stream header using OUTPUT table info let mut header_buffer = Vec::new(); @@ -2192,7 +2282,7 @@ impl<'a> PeGenerator<'a> { rid, output_table_info, )?; - } else if let Some(mut row) = table.get(rid) { + } else if let Some(mut row) = table.get(rid)? { // Use parsed original row if needs_remapping { row.remap_references(remapper); @@ -2362,6 +2452,13 @@ impl<'a> PeGenerator<'a> { let mut dup_to_canonical: HashMap = HashMap::new(); for sig in sig_table { + let sig = match sig { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if deleted_rids.contains(&sig.rid) { continue; } @@ -2382,6 +2479,13 @@ impl<'a> PeGenerator<'a> { let mut rid_to_output: HashMap = HashMap::new(); for sig in sig_table { + let sig = match sig { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if deleted_rids.contains(&sig.rid) || ctx.standalonesig_skip.contains(&sig.rid) { continue; } @@ -2416,6 +2520,13 @@ impl<'a> PeGenerator<'a> { // Phase 3: Build token remapping for original rows for sig in sig_table { + let sig = match sig { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if deleted_rids.contains(&sig.rid) { continue; } @@ -2829,9 +2940,10 @@ impl<'a> PeGenerator<'a> { // Write resource section with relocation let data_offset = ctx.pos(); let data_size = self.write_rsrc_data(ctx, original_section, section_rva)?; + let raw_size = Self::bytes_written_since(ctx, data_offset)?; if data_size > 0 { - ctx.update_section(section_idx, data_offset, section_rva, data_size); + ctx.update_section(section_idx, data_offset, section_rva, data_size, raw_size); current_end_rva = u64::from(section_rva) .checked_add(u64::from(data_size)) .ok_or_else(|| { @@ -2848,8 +2960,10 @@ impl<'a> PeGenerator<'a> { original_text_end, )?; + let raw_size = Self::bytes_written_since(ctx, data_offset)?; + if let Some(data_size) = result { - ctx.update_section(section_idx, data_offset, section_rva, data_size); + ctx.update_section(section_idx, data_offset, section_rva, data_size, raw_size); current_end_rva = u64::from(section_rva) .checked_add(u64::from(data_size)) .ok_or_else(|| { @@ -2862,10 +2976,16 @@ impl<'a> PeGenerator<'a> { } else { // Copy other sections as-is let data_offset = ctx.pos(); + // `data_size` is the section's virtual extent; `raw_size` is what actually + // reached the file. For a copied section these differ whenever the input's + // `VirtualSize` and `SizeOfRawData` differ — uninitialised data in one + // direction, file-alignment padding in the other — and `SizeOfRawData` must + // describe the bytes, not the extent. let data_size = self.write_generic_section(ctx, original_section)?; + let raw_size = Self::bytes_written_since(ctx, data_offset)?; if data_size > 0 { - ctx.update_section(section_idx, data_offset, section_rva, data_size); + ctx.update_section(section_idx, data_offset, section_rva, data_size, raw_size); current_end_rva = u64::from(section_rva) .checked_add(u64::from(data_size)) .ok_or_else(|| { @@ -2979,6 +3099,23 @@ impl<'a> PeGenerator<'a> { /// Writes a generic section by copying data as-is. /// /// Returns the size of data written, or 0 if no data. + /// Returns how many bytes have been written since `start`. + /// + /// Measured from the output position rather than trusted from a return value, so + /// `SizeOfRawData` cannot drift from what the section writer actually emitted. + /// + /// # Errors + /// + /// Returns [`crate::Error::LayoutFailed`] if the position moved backwards or the delta + /// exceeds `u32`. + fn bytes_written_since(ctx: &WriteContext, start: u64) -> Result { + let written = ctx.pos().checked_sub(start).ok_or_else(|| { + Error::LayoutFailed("Section write position moved backwards".to_string()) + })?; + u32::try_from(written) + .map_err(|_| Error::LayoutFailed("Section written size exceeds u32 range".to_string())) + } + fn write_generic_section(&self, ctx: &mut WriteContext, section: &SectionTable) -> Result { let view = self.assembly.view(); let file = view.file(); diff --git a/dotscope/src/cilassembly/writer/heaps/mod.rs b/dotscope/src/cilassembly/writer/heaps/mod.rs index 23699287..d8bf3c25 100644 --- a/dotscope/src/cilassembly/writer/heaps/mod.rs +++ b/dotscope/src/cilassembly/writer/heaps/mod.rs @@ -191,12 +191,42 @@ pub fn precompute_heap_offsets( let guid_data = view.guids().map_or(empty, crate::Guid::data); let us_data = view.userstrings().map_or(empty, crate::UserStrings::data); - // Pre-compute offsets for each heap (this resolves ChangeRefs) - let strings_result = compute_strings_heap_offsets(strings_data, &changes.string_heap_changes)?; - let blob_result = compute_blob_heap_offsets(blob_data, &changes.blob_heap_changes)?; + // Pre-compute offsets for each heap (this resolves ChangeRefs). + // + // These must be given exactly the inputs the real write pass will use. Running the pre-pass + // with empty remap sets makes the two passes disagree: RID remapping rewrites TypeDefOrRef + // tokens as compressed uints, so a RID crossing the 0x80/0x4000 compression boundary changes + // a blob's encoded length — and two distinct blobs can collapse into one, changing the dedup + // outcome. Either shifts `pos` for every later entry, so appended blobs land at different + // offsets in the two passes. The `#Strings` heap diverges the same way, because a non-empty + // `referenced_offsets` makes `emit_orphaned_substrings` allocate extra entries. + // + // The tables are serialised from the pre-pass values, so any disagreement is baked into the + // emitted file — silently, since generation still returns `Ok`. + let strings_result = compute_strings_heap_offsets( + strings_data, + &changes.string_heap_changes, + &changes.referenced_string_offsets, + )?; + let blob_result = compute_blob_heap_offsets( + blob_data, + &changes.blob_heap_changes, + &ctx.typedef_rid_remap, + &ctx.typeref_rid_remap, + &ctx.typespec_rid_remap, + )?; let guid_result = compute_guid_heap_offsets(guid_data, &changes.guid_heap_changes)?; let us_result = compute_userstring_heap_offsets(us_data, &changes.userstring_heap_changes)?; + // Record the emitted heap sizes so the tables stream can pick matching index widths. + // `usize::try_from` cannot fail for a heap that fits in memory; saturating keeps the + // failure direction safe (a larger value only widens the index). + ctx.output_heap_sizes = Some(( + usize::try_from(strings_result.bytes_written).unwrap_or(usize::MAX), + usize::try_from(guid_result.bytes_written).unwrap_or(usize::MAX), + usize::try_from(blob_result.bytes_written).unwrap_or(usize::MAX), + )); + // Store the remapping for later patching of existing table rows ctx.heap_remapping.strings = strings_result.remapping; ctx.heap_remapping.blobs = blob_result.remapping; diff --git a/dotscope/src/cilassembly/writer/heaps/streaming.rs b/dotscope/src/cilassembly/writer/heaps/streaming.rs index 9b40ba4a..0896a8d0 100644 --- a/dotscope/src/cilassembly/writer/heaps/streaming.rs +++ b/dotscope/src/cilassembly/writer/heaps/streaming.rs @@ -128,9 +128,9 @@ pub fn stream_strings_heap( pub fn compute_strings_heap_offsets( source_data: &[u8], changes: &HeapChanges, + referenced_offsets: &HashSet, ) -> Result { - let empty = HashSet::new(); - process_strings_heap(None, 0, source_data, changes, &empty) + process_strings_heap(None, 0, source_data, changes, referenced_offsets) } /// Emits original substrings as standalone entries when their parent string was @@ -493,9 +493,19 @@ pub fn stream_blob_heap( pub fn compute_blob_heap_offsets( source_data: &[u8], changes: &HeapChanges>, + typedef_remap: &HashMap, + typeref_remap: &HashMap, + typespec_remap: &HashMap, ) -> Result { - let empty = HashMap::new(); - process_blob_heap(None, 0, source_data, changes, &empty, &empty, &empty) + process_blob_heap( + None, + 0, + source_data, + changes, + typedef_remap, + typeref_remap, + typespec_remap, + ) } /// Unified blob heap processor. diff --git a/dotscope/src/cilassembly/writer/output.rs b/dotscope/src/cilassembly/writer/output.rs index a48cd369..7328486e 100644 --- a/dotscope/src/cilassembly/writer/output.rs +++ b/dotscope/src/cilassembly/writer/output.rs @@ -66,11 +66,16 @@ use crate::{utils::write_compressed_uint, Error, Result}; /// and in-memory operations without unnecessary copies. enum OutputBacking { /// File-backed memory mapping for efficient large file I/O. + /// + /// Writes go to a temporary file alongside the target; [`Output::finalize`] renames it into + /// place. The destination is not touched until generation has fully succeeded. File { - /// The memory mapping of the file + /// The memory mapping of the temporary file mmap: MmapMut, - /// The target file path + /// Where the finished file will be renamed to target_path: PathBuf, + /// The temporary file being written, adjacent to `target_path` + temp_path: PathBuf, }, /// In-memory vector for zero-copy memory output. Memory { @@ -128,9 +133,10 @@ pub struct Output { impl Output { /// Creates a new file-backed memory-mapped output. /// - /// This creates a file directly at the target path and maps it into memory - /// for efficient writing operations. If finalization fails or the output - /// is dropped without being finalized, the file will be automatically cleaned up. + /// Writes go to a temporary file alongside `target_path`, which [`finalize`](Self::finalize) + /// renames into place. The destination is therefore never modified until generation has + /// fully succeeded; if finalization fails or the output is dropped without being finalized, + /// only the temporary file is removed and any existing artefact at `target_path` survives. /// /// # Arguments /// @@ -150,20 +156,48 @@ impl Output { pub fn create>(target_path: P, size: u64) -> Result { let target_path = target_path.as_ref().to_path_buf(); - // Create the file directly at the target location - let file = std::fs::OpenOptions::new() - .read(true) - .write(true) - .create(true) - .truncate(true) - .open(&target_path) - .map_err(|e| Error::MmapFailed(format!("Failed to create target file: {e}")))?; + // Write to a temporary file beside the target, renamed into place by `finalize`. + // + // Truncating the destination up front and deleting it on failure — the previous + // behaviour — destroys the existing artefact the moment generation starts, and leaves + // nothing behind if any later step fails. When the caller writes back over the file the + // assembly was loaded from, that is the input as well as the output. + // + // Same directory as the target so the final step is a rename within one filesystem, + // which is atomic; a temp file in the system temp dir could land on another mount and + // degrade to a copy. + let parent = target_path.parent().unwrap_or_else(|| Path::new(".")); + std::fs::create_dir_all(parent) + .map_err(|e| Error::MmapFailed(format!("Failed to create output directory: {e}")))?; + + let temp = tempfile::Builder::new() + .prefix(".dotscope-") + .suffix(".tmp") + .tempfile_in(parent) + .map_err(|e| Error::MmapFailed(format!("Failed to create temporary file: {e}")))?; + + // Keep the path and hand back the `File`; the guard is dropped so it does not delete + // the file behind us. Cleanup on failure is `Output::drop`'s job. + let (file, temp_path) = temp + .keep() + .map_err(|e| Error::MmapFailed(format!("Failed to persist temporary file: {e}")))?; // Set the file size file.set_len(size) .map_err(|e| Error::MmapFailed(format!("Failed to set file size: {e}")))?; - // Create memory mapping + // Create memory mapping. + // + // SAFETY: `map_mut` is unsafe because the mapping aliases the file's contents, and Rust + // cannot rule out another process mutating or truncating the file underneath us — doing + // so would let safe code observe a torn value or take SIGBUS. The file here is one this + // function just created via `OpenOptions::create` at `target_path` and sized with + // `set_len` above; it is not handed to anything else before the mapping is established, + // and the `Mmap` is owned by the returned `Output`, so its lifetime cannot outlive the + // `File` it borrows. Concurrent external modification of a caller-chosen output path is + // outside the threat model this crate defends against — an attacker who can write that + // path can simply replace the artefact. + #[allow(unsafe_code)] let mmap = unsafe { MmapOptions::new() .map_mut(&file) @@ -171,7 +205,11 @@ impl Output { }; Ok(Self { - backing: OutputBacking::File { mmap, target_path }, + backing: OutputBacking::File { + mmap, + target_path, + temp_path, + }, finalized: false, }) } @@ -763,21 +801,25 @@ impl Output { ); match backing { - OutputBacking::File { mmap, target_path } => { + OutputBacking::File { + mmap, + target_path, + temp_path, + } => { // Flush memory mapping mmap.flush().map_err(|e| { Error::FinalizationFailed(format!("Failed to flush memory mapping: {e}")) })?; + // Release the mapping before touching the file through the filesystem: on + // Windows a mapped file cannot be renamed. + drop(mmap); + // Truncate if requested if let Some(size) = actual_size { - // Drop the mmap to release file handle - drop(mmap); - - // Truncate the file let file = std::fs::OpenOptions::new() .write(true) - .open(&target_path) + .open(&temp_path) .map_err(|e| { Error::FinalizationFailed(format!( "Failed to reopen file for truncation: {e}" @@ -791,6 +833,17 @@ impl Output { })?; } + // Rename into place. This is the only point at which the destination changes, + // so a failure anywhere earlier leaves the previous artefact untouched. + std::fs::rename(&temp_path, &target_path).map_err(|e| { + // Leave the temp file for inspection only if removing it also fails. + let _ = std::fs::remove_file(&temp_path); + Error::FinalizationFailed(format!( + "Failed to move generated file into place at {}: {e}", + target_path.display() + )) + })?; + // Mark as finalized self.finalized = true; Ok(()) @@ -896,9 +949,11 @@ impl Drop for Output { // First try to flush any pending writes let _ = self.flush(); - // For file-backed outputs, delete the incomplete file - if let OutputBacking::File { target_path, .. } = &self.backing { - let _ = std::fs::remove_file(target_path); + // For file-backed outputs, delete the incomplete *temporary* file. The target is + // deliberately left alone: it was never written to, so an abandoned generation + // leaves whatever was already there intact. + if let OutputBacking::File { temp_path, .. } = &self.backing { + let _ = std::fs::remove_file(temp_path); } // For in-memory outputs, the Vec will be dropped automatically } diff --git a/dotscope/src/compiler/codegen/coalescing.rs b/dotscope/src/compiler/codegen/coalescing.rs index 40472bf1..fa747d35 100644 --- a/dotscope/src/compiler/codegen/coalescing.rs +++ b/dotscope/src/compiler/codegen/coalescing.rs @@ -28,7 +28,7 @@ use std::{ cmp::Reverse, - collections::{BTreeMap, BinaryHeap}, + collections::{BTreeMap, BTreeSet, BinaryHeap}, }; use analyssa::BitSet; @@ -463,19 +463,47 @@ impl LocalCoalescer { /// Computes live intervals for all variables in the SSA. /// - /// A live interval is the range [start, end) where a variable is live. - /// We use instruction indices within a linearized view of the CFG. + /// A live interval is the range [start, end) where a variable is live, in + /// instruction indices over a linearized view of the CFG. + /// + /// # Why this walks the CFG instead of scanning textually or running a dataflow solve + /// + /// A purely *textual* scan — `[first mention, last use + 1)` over blocks in `block_id` + /// order — under-approximates liveness across back edges. A value defined before a loop + /// and used inside it has its interval truncated at its last textual use, so its slot + /// returns to the free pool and can be handed to a variable defined later in the same + /// loop body; on the second iteration the earlier use reads a clobbered slot. Linear + /// scan only runs above [`LINEAR_SCAN_THRESHOLD`] variables, so that failure is + /// size-gated and invisible to small unit tests. + /// + /// The obvious fix — reuse the `LiveVariables` solve that [`Self::build_graph_coloring`] + /// runs — is the wrong tool *here*. That framework computes the live set of every + /// variable at every block simultaneously, storing a live-in and a live-out `BitSet` of + /// `var_capacity` bits per block: `2·B·V` bits. Graph colouring can afford that only + /// because it is gated to at most [`LINEAR_SCAN_THRESHOLD`] variables. Linear scan + /// exists precisely for the large, attacker-sized methods where `B·V` explodes, and the + /// CIL path has no method-size cap. + /// + /// Nothing here needs the full live sets. A [`LiveInterval`] is a single contiguous + /// range, so all that is required per variable is the earliest and latest point it is + /// live. SSA gives exactly one definition per variable, and it dominates every use, so + /// each variable's live range is found by the classic *up-and-mark* walk: start at each + /// use, walk up predecessors, stop at the defining block — no second definition can be + /// encountered on the way. Cost is `O(B)` memory for one reusable generation-stamp + /// array, and time proportional to the total live range, which is the size of the + /// answer rather than the size of a `B × V` matrix. Most SSA variables are block-local, + /// so their walk terminates immediately. fn compute_live_intervals(ssa: &SsaFunction) -> BTreeMap { + let block_count = ssa.block_count(); let mut intervals: BTreeMap = BTreeMap::new(); - // Phase 1: Build a map of block → end instruction index. - // PHI operands are semantically used at the END of the predecessor - // block (not at the PHI's block), so we need to know each block's - // end position to correctly extend operand intervals. - let mut block_end_idx: Vec = Vec::with_capacity(ssa.block_count()); + // Phase 1: number instructions, recording each block's [start, end) span. + let mut block_start_idx: Vec = Vec::with_capacity(block_count); + let mut block_end_idx: Vec = Vec::with_capacity(block_count); { let mut idx = 0usize; - for block_id in 0..ssa.block_count() { + for block_id in 0..block_count { + block_start_idx.push(idx); if let Some(block) = ssa.block(block_id) { idx = idx.saturating_add(block.instructions().len()); } @@ -483,61 +511,158 @@ impl LocalCoalescer { } } - // Phase 2: Assign instruction indices by walking blocks in order - let mut instr_idx = 0usize; + // Phase 2: locate every definition. This must complete before uses are processed, + // because a use can precede its definition in block order (that is what a back edge + // means) and the walk below keys its stop condition on the defining block. + let mut def_block: BTreeMap = BTreeMap::new(); - for block_id in 0..ssa.block_count() { + for block_id in 0..block_count { let Some(block) = ssa.block(block_id) else { continue; }; + let Some(&entry_pos) = block_start_idx.get(block_id) else { + continue; + }; - // Phi nodes define at block entry + // Phi nodes define at block entry. for phi in block.phi_nodes() { let def = phi.result(); - intervals + def_block.insert(def, block_id); + let interval = intervals .entry(def) - .or_insert_with(|| LiveInterval::new(instr_idx)) - .extend_start(instr_idx); - - // PHI operands are used at the END of their predecessor block. - // A variable v in `v<-B_pred` must be live from its definition - // through the end of B_pred. If B_pred comes AFTER the PHI's - // block in the linearized layout, extending only to the PHI - // position would leave a gap where the local slot can be reused. + .or_insert_with(|| LiveInterval::new(entry_pos)); + interval.extend_start(entry_pos); + interval.extend_end(entry_pos.saturating_add(1)); + } + + let mut pos = entry_pos; + for instr in block.instructions() { + if let Some(def) = instr.def() { + def_block.insert(def, block_id); + let interval = intervals + .entry(def) + .or_insert_with(|| LiveInterval::new(pos)); + interval.extend_start(pos); + // Cover the definition point itself. `extend_start` is a min and + // `extend_end` a max, so a variable used at 5 and defined at 10 would + // otherwise hold [5, 6) — an interval that expires before the value it + // describes even exists, freeing the slot while the def is still ahead. + interval.extend_end(pos.saturating_add(1)); + } + pos = pos.saturating_add(1); + } + } + + // Phase 3: record uses, and seed the backward walk with the blocks each variable is + // live *out* of. + let cfg = SsaCfg::from_ssa(ssa); + let mut live_out_seeds: BTreeMap> = BTreeMap::new(); + + for block_id in 0..block_count { + let Some(block) = ssa.block(block_id) else { + continue; + }; + let Some(&entry_pos) = block_start_idx.get(block_id) else { + continue; + }; + + // A phi operand is used at the END of its predecessor block, not here, so the + // value must stay live through that predecessor even when it is laid out after + // the phi's own block. + for phi in block.phi_nodes() { for operand in phi.operands() { let pred = operand.predecessor(); - let pred_end = block_end_idx - .get(pred) - .copied() - .unwrap_or_else(|| instr_idx.saturating_add(1)); - // Use the later of: PHI position or predecessor end - let use_point = pred_end.max(instr_idx.saturating_add(1)); + let Some(&pred_end) = block_end_idx.get(pred) else { + continue; + }; + let value = operand.value(); intervals - .entry(operand.value()) - .or_insert_with(|| LiveInterval::new(instr_idx)) - .extend_end(use_point); + .entry(value) + .or_insert_with(|| LiveInterval::new(pred_end)) + .extend_end(pred_end); + live_out_seeds.entry(value).or_default().push(pred); } } - // Process instructions + // Seeding is per (variable, block), not per use. The walk's generation stamps + // would dedupe the extra work anyway, but the seed vector itself would not: + // a variable used u times in a block with p predecessors would push u·p entries + // instead of p, and both factors are attacker-controlled. + let mut seeded_here: BTreeSet = BTreeSet::new(); + + let mut pos = entry_pos; for instr in block.instructions() { - // Uses extend the interval for &use_var in &instr.uses() { - intervals + let interval = intervals .entry(use_var) - .or_insert_with(|| LiveInterval::new(instr_idx)) - .extend_end(instr_idx.saturating_add(1)); + .or_insert_with(|| LiveInterval::new(pos)); + interval.extend_end(pos.saturating_add(1)); + + // A use in the defining block needs no walk — the value is live only + // from the definition to this point. Otherwise the value is live-in + // here, hence live-out of every predecessor. + if def_block.get(&use_var) != Some(&block_id) { + interval.extend_start(entry_pos); + if seeded_here.insert(use_var) { + live_out_seeds + .entry(use_var) + .or_default() + .extend(cfg.block_predecessors(block_id).iter().copied()); + } + } } + pos = pos.saturating_add(1); + } + } - // Definitions start the interval - if let Some(def) = instr.def() { - intervals - .entry(def) - .or_insert_with(|| LiveInterval::new(instr_idx)) - .extend_start(instr_idx); + // Phase 4: up-and-mark. For each variable, walk up from every block it is live out + // of, covering whole blocks, and stop at its definition. + // + // `seen` holds a generation stamp per block rather than a boolean, so resetting + // between variables is a single counter bump instead of an O(B) clear. The counter + // is `u64`, which cannot wrap within one function's variable count. + let mut seen: Vec = vec![0; block_count]; + let mut generation: u64 = 0; + let mut worklist: Vec = Vec::new(); + + for (var, seeds) in &live_out_seeds { + generation = generation.saturating_add(1); + let stop = def_block.get(var).copied(); + + worklist.clear(); + worklist.extend(seeds.iter().copied()); + + while let Some(block_id) = worklist.pop() { + let Some(stamp) = seen.get_mut(block_id) else { + continue; + }; + if *stamp == generation { + continue; + } + *stamp = generation; + + if let (Some(&block_start), Some(&block_end), Some(interval)) = ( + block_start_idx.get(block_id), + block_end_idx.get(block_id), + intervals.get_mut(var), + ) { + interval.extend_end(block_end.max(block_start.saturating_add(1))); + // In the defining block the value is live only from its definition + // onward, so the block's start must not pull the interval back. + if stop != Some(block_id) { + interval.extend_start(block_start); + } } - instr_idx = instr_idx.saturating_add(1); + // The definition dominates every use, so nothing above it is live. A + // variable with no recorded definition (an argument or entry live-in) has + // no stop block and correctly walks back to the entry, which has no + // predecessors. + if stop == Some(block_id) { + continue; + } + + worklist.extend(cfg.block_predecessors(block_id).iter().copied()); } } @@ -1056,6 +1181,71 @@ mod tests { assert_eq!(graph.degree(vars[4]), 0); } + /// A value defined before a loop and used inside it must stay live for the whole + /// loop, not just up to its last *textual* use. + /// + /// Under a textual interval scan `outer`'s interval ends at its single use + /// in the loop body, so the slot returns to the free pool and can be handed to `inner`, + /// which is defined later in that same body — and on the second iteration the read of + /// `outer` returns `inner`'s value. + #[test] + fn value_live_across_a_back_edge_keeps_its_slot() { + // block 0: outer = 7; jump 1 + // block 1: header, branch to body (2) or exit (3) + // block 2: body, reads `outer`, then defines `inner`; jumps back to 1 + // block 3: ret + let mut outer_id: Option = None; + let mut inner_id: Option = None; + + let ssa = SsaFunctionBuilder::new(1, 0) + .build_with(|f| { + let cond = f.arg(0, SsaType::Bool); + f.block(0, |b| { + outer_id = Some(b.const_i32(7)); + b.jump(1); + }); + f.block(1, |b| { + b.branch(cond, 2, 3); + }); + let outer = outer_id.expect("block 0 is built first"); + f.block(2, |b| { + // `outer` is read here — its last textual mention. + let _used = b.add(outer, outer); + // ...and `inner` is defined after it, in the same block. + let inner = b.const_i32(9); + let _also = b.add(inner, inner); + inner_id = Some(inner); + b.jump(1); + }); + f.block(3, |b| { + b.ret(); + }); + }) + .unwrap(); + + let outer = outer_id.expect("outer was defined"); + let inner = inner_id.expect("inner was defined"); + + let intervals = LocalCoalescer::compute_live_intervals(&ssa); + + let outer_interval = intervals.get(&outer).expect("outer must have an interval"); + let inner_interval = intervals.get(&inner).expect("inner must have an interval"); + + // The two must overlap, which is what stops the allocator reusing one slot for + // both. Overlap is the property that matters; exact indices are not contractual. + assert!( + outer_interval.start < inner_interval.end + && inner_interval.start < outer_interval.end, + "outer {outer_interval:?} and inner {inner_interval:?} must overlap across the back edge" + ); + + // And `outer` must survive past the point where `inner` is defined. + assert!( + outer_interval.end > inner_interval.start, + "outer {outer_interval:?} expired before inner {inner_interval:?} was defined" + ); + } + #[test] fn test_type_compatibility_same_class() { // Same types are compatible diff --git a/dotscope/src/compiler/codegen/mod.rs b/dotscope/src/compiler/codegen/mod.rs index 75db38f1..11d3896b 100644 --- a/dotscope/src/compiler/codegen/mod.rs +++ b/dotscope/src/compiler/codegen/mod.rs @@ -68,8 +68,8 @@ use crate::{ TypeSignature, }, tables::{ - ClassLayoutRaw, CodedIndex, CodedIndexType, FieldRaw, FieldRvaRaw, MemberRefRaw, - NestedClassRaw, TableDataOwned, TableId, TypeDefRaw, TypeRefRaw, + skip_unreadable, ClassLayoutRaw, CodedIndex, CodedIndexType, FieldRaw, FieldRvaRaw, + MemberRefRaw, NestedClassRaw, TableDataOwned, TableId, TypeDefRaw, TypeRefRaw, }, token::Token, }, @@ -133,6 +133,32 @@ impl TempPool { } } +/// One entry in the emission order. +/// +/// Out-of-SSA needs somewhere to put a phi copy for the edge `from → to`. When `from` has a +/// single successor the copy can simply precede its branch, but a multi-way terminator +/// (`Branch`, `BranchCmp`, `Switch`) has no such place: the copy belongs *on the edge*, and +/// each edge needs its own. +/// +/// Edge blocks are therefore ordinary entries in the emission order rather than labelled runs +/// spliced into branch emission. That makes "does code sit between this block and its +/// successor" a structural property the layout answers — whatever lands between two blocks is +/// something the layout placed — rather than a condition each branch emitter has to maintain +/// by predicting what will be emitted after it. Getting that prediction wrong is the class of +/// unsound fall-through elision this pass removes. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum LayoutEntry { + /// A real SSA block, identified by its block id. + Block(usize), + /// A synthetic block carrying the phi copies for one CFG edge. + SplitEdge { + /// Block the edge leaves. + from: usize, + /// Block the edge enters. + to: usize, + }, +} + /// SSA to CIL code generator. /// /// Converts SSA form back to executable CIL bytecode with optimizations. @@ -143,6 +169,12 @@ pub struct SsaCodeGenerator { next_local: u16, /// Map from block index to label name block_labels: BTreeMap, + /// Label for each edge that was given its own block, keyed by `(from, to)`. + /// + /// Present only for edges leaving a multi-way terminator that carry phi copies. A + /// branch whose target appears here jumps to the edge block instead of the target + /// block; the edge block does the copies and then continues to the real target. + split_edge_labels: BTreeMap<(usize, usize), String>, /// Variables that are currently on the stack (for optimization) stack_vars: Vec, /// Cache of interned decrypted strings (string content -> heap index) @@ -262,6 +294,7 @@ impl SsaCodeGenerator { var_storage: BTreeMap::new(), next_local: 0, block_labels: BTreeMap::new(), + split_edge_labels: BTreeMap::new(), stack_vars: Vec::new(), interned_strings: HashMap::new(), deferred_constants: BTreeMap::new(), @@ -657,6 +690,7 @@ impl SsaCodeGenerator { self.var_storage.clear(); self.next_local = 0; self.block_labels.clear(); + self.split_edge_labels.clear(); self.stack_vars.clear(); self.interned_strings.clear(); self.deferred_constants.clear(); @@ -787,15 +821,39 @@ impl SsaCodeGenerator { // Compute optimal block layout to minimize unnecessary branches. // This reorders blocks so that fall-through paths don't need explicit jumps. let block_ids = Self::compute_block_layout(ssa, &blocks_to_include); - let blocks_to_generate: Vec<_> = block_ids.iter().filter_map(|&id| ssa.block(id)).collect(); - for (idx, block) in blocks_to_generate.iter().enumerate() { - // Record block start offset for exception handler remapping - let pos_before = encoder.current_position(); - self.block_offsets.insert(block.id(), pos_before); + // Give every phi-carrying edge out of a multi-way terminator its own block, and + // register a label for it, so the copies are ordinary layout entries rather than + // runs spliced into branch emission. + let layout = self.expand_layout_with_edge_blocks(ssa, &block_ids)?; + + for idx in 0..layout.len() { + let Some(&entry) = layout.get(idx) else { + continue; + }; + + // Fall-through is available only to whatever the *next entry* is. Naming the + // entry rather than a block id is what makes the elision safe by construction: + // an edge block between two blocks is itself an entry, so it suppresses the + // elision on its own instead of each branch helper having to predict that code + // will be emitted after it. + let next_entry = layout.get(idx.saturating_add(1)).copied(); - let next_block_idx = block_ids.get(idx.saturating_add(1)).copied(); - self.generate_block(&mut encoder, ssa, block, block.id(), next_block_idx)?; + match entry { + LayoutEntry::Block(block_id) => { + let Some(block) = ssa.block(block_id) else { + continue; + }; + // Record block start offset for exception handler remapping + let pos_before = encoder.current_position(); + self.block_offsets.insert(block_id, pos_before); + + self.generate_block(&mut encoder, ssa, block, block_id, next_entry)?; + } + LayoutEntry::SplitEdge { from, to } => { + self.generate_edge_block(&mut encoder, ssa, from, to, next_entry)?; + } + } } // Phase 4: Finalize and resolve labels @@ -1008,6 +1066,13 @@ impl SsaCodeGenerator { let typeref_table = tables.table::()?; for row in typeref_table.iter() { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let Ok(name) = strings.get(row.type_name as usize) else { continue; }; @@ -1030,7 +1095,7 @@ impl SsaCodeGenerator { let tables = assembly.view().tables()?; let strings = assembly.view().strings()?; let member_refs = tables.table::()?; - for row in member_refs { + for row in member_refs.iter().filter_map(skip_unreadable) { if let Ok(name) = strings.get(row.name as usize) { if name == "InitializeArray" { return Some(row.token); @@ -1159,6 +1224,13 @@ impl SsaCodeGenerator { let typeref_table = tables.table::()?; for row in typeref_table.iter() { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let Ok(name) = strings.get(row.type_name as usize) else { continue; }; @@ -1524,6 +1596,89 @@ impl SsaCodeGenerator { layout } + /// Expands a block layout into the emission order, giving every phi-carrying edge out of + /// a multi-way terminator its own [`LayoutEntry::SplitEdge`] block. + /// + /// # Which edges get a block + /// + /// Only edges leaving `Branch`, `BranchCmp` and `Switch`. A block with a single + /// successor has a place for its copies already — immediately before its branch, where + /// `generate_branch_op` still emits them — and an edge out of a single-successor block + /// is never critical, so it cannot need one. + /// + /// # Where they go + /// + /// Directly after the block the edge leaves, preferred-successor edge first. That keeps + /// each edge block inside whatever EH region its source block is in (a protected region + /// can only be left by `leave`, so a two-way branch never crosses out of one), and lets + /// the source block fall through into the edge it most likely takes. + /// + /// The edges are registered in `split_edge_labels` as they are placed, which is what + /// makes [`Self::edge_label`] and [`Self::edge_destination`] agree with the layout: a + /// branch names the edge block exactly when the layout created one. + /// + /// # Errors + /// + /// Propagates the storage lookup in [`Self::successor_has_phi_from`]. + fn expand_layout_with_edge_blocks( + &mut self, + ssa: &SsaFunction, + block_ids: &[usize], + ) -> Result> { + let mut layout: Vec = Vec::with_capacity(block_ids.len()); + + for &block_id in block_ids { + layout.push(LayoutEntry::Block(block_id)); + + let Some(block) = ssa.block(block_id) else { + continue; + }; + + // Preferred successor first, matching `compute_block_layout::preferred_successor`, + // so the edge block that can be reached by fall-through is the adjacent one. + let edge_targets: Vec = match block.terminator_op() { + Some( + SsaOp::Branch { + true_target, + false_target, + .. + } + | SsaOp::BranchCmp { + true_target, + false_target, + .. + }, + ) => vec![*false_target, *true_target], + Some(SsaOp::Switch { + targets, default, .. + }) => { + let mut all = Vec::with_capacity(targets.len().saturating_add(1)); + all.push(*default); + all.extend(targets.iter().copied()); + all + } + _ => continue, + }; + + for to in edge_targets { + // A switch may name the same target more than once, and a two-way branch + // may name the same block on both edges. One edge, one block. + if self.split_edge_labels.contains_key(&(block_id, to)) { + continue; + } + if !self.successor_has_phi_from(ssa, block_id, to)? { + continue; + } + + self.split_edge_labels + .insert((block_id, to), format!("edge_{block_id}_{to}")); + layout.push(LayoutEntry::SplitEdge { from: block_id, to }); + } + } + + Ok(layout) + } + /// Allocates storage for all SSA variables using graph coloring. /// /// Arguments get mapped to their original argument slots (ldarg/starg). @@ -2561,7 +2716,7 @@ impl SsaCodeGenerator { ssa: &SsaFunction, block: &SsaBlock, block_idx: usize, - next_block_idx: Option, + next_entry: Option, ) -> Result<()> { // Exception handler entry blocks (catch, filter, finally, fault) are only // entered via CLR exception dispatch — never via fallthrough from the previous @@ -2750,7 +2905,7 @@ impl SsaCodeGenerator { def_map: &def_map, current_block_idx: block_idx, }; - self.generate_ops_iterative(encoder, &ctx, &roots, next_block_idx)?; + self.generate_ops_iterative(encoder, &ctx, &roots, next_entry)?; // Check if the block ends with a non-terminator (falls through to next block) // If so, spill remaining stack values and emit phi stores for the fallthrough @@ -2759,7 +2914,11 @@ impl SsaCodeGenerator { // Spill any remaining stack values before fallthrough self.spill_stack(encoder, ssa)?; - if let Some(next_idx) = next_block_idx { + // A block with no terminator has exactly one successor, so its edge is never + // critical and never gets a block of its own — which is also why the next entry + // is always a `Block` here: edge blocks are only ever placed directly after a + // multi-way terminator. + if let Some(LayoutEntry::Block(next_idx)) = next_entry { self.emit_phi_stores_for_successor(encoder, ssa, block_idx, next_idx)?; } } @@ -2782,7 +2941,7 @@ impl SsaCodeGenerator { encoder: &mut InstructionEncoder, ctx: &BlockCodegenContext<'_>, roots: &[usize], - next_block_idx: Option, + next_entry: Option, ) -> Result<()> { // Use global use counts to determine if variables need storage. // Variables used multiple times (across ANY blocks) must be stored to locals @@ -3002,7 +3161,7 @@ impl SsaCodeGenerator { ctx.ssa, ctx.current_block_idx, cur_op, - next_block_idx, + next_entry, )?; generated.insert(idx); @@ -3089,7 +3248,7 @@ impl SsaCodeGenerator { ctx.ssa, ctx.current_block_idx, root_op, - next_block_idx, + next_entry, )?; generated.insert(root_idx); } @@ -3357,7 +3516,7 @@ impl SsaCodeGenerator { ssa: &SsaFunction, current_block_idx: usize, op: &SsaOp, - next_block_idx: Option, + next_entry: Option, ) -> Result<()> { // Clear last_load tracking - after any operation, the stack state changes // and the "last loaded" value may no longer be on top of the stack. @@ -3458,7 +3617,7 @@ impl SsaCodeGenerator { | SsaOp::EndFilter { .. } | SsaOp::BranchCmp { .. } | SsaOp::BranchFlags { .. } => { - self.generate_branch_op(encoder, ssa, current_block_idx, op, next_block_idx)?; + self.generate_branch_op(encoder, ssa, current_block_idx, op, next_entry)?; } // Simple standalone operations @@ -4043,7 +4202,7 @@ impl SsaCodeGenerator { ssa: &SsaFunction, current_block_idx: usize, op: &SsaOp, - next_block_idx: Option, + next_entry: Option, ) -> Result<()> { match op { SsaOp::Return { .. } => { @@ -4062,15 +4221,12 @@ impl SsaCodeGenerator { // Spill remaining stack values before control flow transfer self.spill_stack(encoder, ssa)?; - // Emit phi stores for the target block before jumping + // A jump has one successor, so its copies have a place already: right here, + // before the branch. No edge block is created for it. self.emit_phi_stores_for_successor(encoder, ssa, current_block_idx, *target)?; - if Some(*target) != next_block_idx { - let label = self - .block_labels - .get(target) - .cloned() - .unwrap_or_else(|| format!("block_{target}")); + if next_entry != Some(LayoutEntry::Block(*target)) { + let label = self.block_label(*target); // emit_branch validates stack depth at target encoder.emit_branch("br", &label)?; } @@ -4085,14 +4241,12 @@ impl SsaCodeGenerator { // The condition is already on the stack from operand loading. self.spill_stack(encoder, ssa)?; - // Handle phi stores for both targets using intermediate blocks - self.emit_branch_with_phi_stores( + self.emit_conditional_branch( encoder, - ssa, current_block_idx, *true_target, *false_target, - next_block_idx, + next_entry, )?; } @@ -4102,15 +4256,7 @@ impl SsaCodeGenerator { // Spill any remaining stack values except the switch value. self.spill_stack(encoder, ssa)?; - // Handle phi stores for switch targets using intermediate blocks - self.emit_switch_with_phi_stores( - encoder, - ssa, - current_block_idx, - targets, - *default, - next_block_idx, - )?; + self.emit_switch_branch(encoder, current_block_idx, targets, *default, next_entry)?; } SsaOp::Throw { .. } => { @@ -4131,14 +4277,10 @@ impl SsaCodeGenerator { // Spill any remaining stack values before leaving protected region. self.spill_stack(encoder, ssa)?; - // Emit phi stores for the target block before leaving + // A leave has one successor, so its copies go here, before the transfer. self.emit_phi_stores_for_successor(encoder, ssa, current_block_idx, *target)?; - let label = self - .block_labels - .get(target) - .cloned() - .unwrap_or_else(|| format!("block_{target}")); + let label = self.block_label(*target); encoder.emit_branch("leave", &label)?; } @@ -4166,16 +4308,15 @@ impl SsaCodeGenerator { // Spill any remaining stack values except the comparison operands. self.spill_stack(encoder, ssa)?; - // Emit comparison branch with phi stores - operands are already on the stack - self.emit_branch_cmp_with_phi_stores( + // Operands are already on the stack. + self.emit_comparison_branch( encoder, - ssa, current_block_idx, *cmp, *unsigned, *true_target, *false_target, - next_block_idx, + next_entry, )?; } @@ -4777,317 +4918,213 @@ impl SsaCodeGenerator { Ok(false) } - /// Emits a conditional branch (brtrue/brfalse) with proper phi store handling. + /// Label of a real block. + fn block_label(&self, block_id: usize) -> String { + self.block_labels + .get(&block_id) + .cloned() + .unwrap_or_else(|| format!("block_{block_id}")) + } + + /// The layout entry the edge `from → to` actually lands on. /// - /// When branch targets have phi nodes, we emit intermediate blocks that: - /// 1. Execute the phi stores for that specific edge - /// 2. Jump to the actual target + /// This is the whole of the fall-through question. A branch may elide its `br` exactly + /// when the entry it transfers to is the next one in the emission order — and because + /// edge blocks are entries, "is there code in between" is answered by looking at the + /// layout rather than by each emitter re-deriving what the others will emit after it. + fn edge_destination(&self, from: usize, to: usize) -> LayoutEntry { + if self.split_edge_labels.contains_key(&(from, to)) { + LayoutEntry::SplitEdge { from, to } + } else { + LayoutEntry::Block(to) + } + } + + /// Label a branch along the edge `from → to` must name: the edge's own block when it has + /// one, otherwise the target block directly. + fn edge_label(&self, from: usize, to: usize) -> String { + self.split_edge_labels + .get(&(from, to)) + .cloned() + .unwrap_or_else(|| self.block_label(to)) + } + + /// Generates the block holding the phi copies for one edge. /// - /// This handles the "critical edge splitting" problem where different predecessors - /// need to provide different values to phi nodes. - fn emit_branch_with_phi_stores( + /// The body is just the parallel copy for that edge followed by a transfer to the real + /// target, which is elided when the target is the next entry. + /// + /// # Errors + /// + /// Returns an error if the edge has no registered label — which would mean the layout + /// and [`Self::edge_destination`] disagree — or if emission fails. + fn generate_edge_block( &mut self, encoder: &mut InstructionEncoder, ssa: &SsaFunction, + from: usize, + to: usize, + next_entry: Option, + ) -> Result<()> { + let label = self + .split_edge_labels + .get(&(from, to)) + .cloned() + .ok_or_else(|| { + Error::CodegenFailed(format!("Missing split-edge label for edge {from} -> {to}")) + })?; + + encoder.define_label(&label)?; + + // An edge block is reached by a branch or by falling out of its source block, and in + // both cases the source has already spilled: the evaluation stack is empty here, so + // nothing the source left in `stack_vars` is still on it. + self.stack_vars.clear(); + + self.emit_phi_stores_for_successor(encoder, ssa, from, to)?; + + if next_entry != Some(LayoutEntry::Block(to)) { + let target_label = self.block_label(to); + encoder.emit_branch("br", &target_label)?; + } + + Ok(()) + } + + /// Emits a conditional branch (`brtrue`/`brfalse`). + /// + /// Phi copies are not this function's concern: any edge that carries them was given its + /// own block by [`Self::expand_layout_with_edge_blocks`], so all that is emitted here is + /// the branch, naming the edge block where there is one. + fn emit_conditional_branch( + &mut self, + encoder: &mut InstructionEncoder, current_block_idx: usize, true_target: usize, false_target: usize, - next_block_idx: Option, + next_entry: Option, ) -> Result<()> { - // Same-target: condition is irrelevant, discard it and emit unconditional branch + // Same-target: the condition cannot change where control goes, so discard it. if true_target == false_target { encoder.emit_instruction("pop", None)?; - if self.successor_has_phi_from(ssa, current_block_idx, true_target)? { - self.emit_phi_stores_for_successor(encoder, ssa, current_block_idx, true_target)?; - } - if Some(true_target) != next_block_idx { - let label = self - .block_labels - .get(&true_target) - .cloned() - .unwrap_or_else(|| format!("block_{true_target}")); + if next_entry != Some(self.edge_destination(current_block_idx, true_target)) { + let label = self.edge_label(current_block_idx, true_target); encoder.emit_branch("br", &label)?; } return Ok(()); } - let true_has_phi = self.successor_has_phi_from(ssa, current_block_idx, true_target)?; - let false_has_phi = self.successor_has_phi_from(ssa, current_block_idx, false_target)?; - - let true_label = self - .block_labels - .get(&true_target) - .cloned() - .unwrap_or_else(|| format!("block_{true_target}")); - let false_label = self - .block_labels - .get(&false_target) - .cloned() - .unwrap_or_else(|| format!("block_{false_target}")); - - // Simple case: no phi stores needed for either target - if !true_has_phi && !false_has_phi { - if Some(false_target) == next_block_idx { - encoder.emit_branch("brtrue", &true_label)?; - } else if Some(true_target) == next_block_idx { - encoder.emit_branch("brfalse", &false_label)?; - } else { - encoder.emit_branch("brtrue", &true_label)?; - encoder.emit_branch("br", &false_label)?; - } - return Ok(()); - } - - // Generate unique intermediate label for false path (true path is inline) - let phi_false_label = format!("phi_false_{current_block_idx}_{false_target}"); - - // Emit the conditional branch - // Strategy: brfalse to false handling, then handle true path - if true_has_phi { - encoder.emit_branch( - "brfalse", - if false_has_phi { - &phi_false_label - } else { - &false_label - }, - )?; + let true_label = self.edge_label(current_block_idx, true_target); + let false_label = self.edge_label(current_block_idx, false_target); + let true_dest = self.edge_destination(current_block_idx, true_target); + let false_dest = self.edge_destination(current_block_idx, false_target); - // True path: emit phi stores then jump - self.emit_phi_stores_for_successor(encoder, ssa, current_block_idx, true_target)?; - if Some(true_target) != next_block_idx { - encoder.emit_branch("br", &true_label)?; - } + if next_entry == Some(false_dest) { + encoder.emit_branch("brtrue", &true_label)?; + } else if next_entry == Some(true_dest) { + // `brfalse` is the exact complement of `brtrue`: both test one value against + // zero/null, so there is no third outcome to mishandle. That is what makes the + // inversion safe here and *not* safe for the comparison branches below. + encoder.emit_branch("brfalse", &false_label)?; } else { - // No phi stores for true - just branch directly encoder.emit_branch("brtrue", &true_label)?; - } - - // False path handling - if false_has_phi { - // Define the intermediate label for false path - if true_has_phi { - encoder.define_label(&phi_false_label)?; - } - // Emit phi stores for false target - self.emit_phi_stores_for_successor(encoder, ssa, current_block_idx, false_target)?; - if Some(false_target) != next_block_idx { - encoder.emit_branch("br", &false_label)?; - } - } else if !true_has_phi { - // Neither has phi, but we need to handle fallthrough - if Some(false_target) != next_block_idx { - encoder.emit_branch("br", &false_label)?; - } + encoder.emit_branch("br", &false_label)?; } Ok(()) } - /// Emits a comparison branch (beq, blt, etc.) with proper phi store handling. - #[allow(clippy::too_many_arguments)] // Branch emission requires comparison type, targets, and context - fn emit_branch_cmp_with_phi_stores( + /// Emits a comparison branch (`beq`, `blt`, …). + /// + /// As with [`Self::emit_conditional_branch`], phi copies live in their own blocks and are + /// not emitted here. + // The comparison kind and its signedness select the mnemonic, and both edges are needed to + // decide which one can fall through; neither pair compresses into a type that would carry + // its own meaning. + #[allow(clippy::too_many_arguments)] + fn emit_comparison_branch( &mut self, encoder: &mut InstructionEncoder, - ssa: &SsaFunction, current_block_idx: usize, cmp: CmpKind, unsigned: bool, true_target: usize, false_target: usize, - next_block_idx: Option, + next_entry: Option, ) -> Result<()> { - // Same-target: comparison is irrelevant, discard both operands + // Same-target: the comparison cannot change where control goes, so discard both + // operands rather than evaluating it. if true_target == false_target { encoder.emit_instruction("pop", None)?; encoder.emit_instruction("pop", None)?; - if self.successor_has_phi_from(ssa, current_block_idx, true_target)? { - self.emit_phi_stores_for_successor(encoder, ssa, current_block_idx, true_target)?; - } - if Some(true_target) != next_block_idx { - let label = self - .block_labels - .get(&true_target) - .cloned() - .unwrap_or_else(|| format!("block_{true_target}")); + if next_entry != Some(self.edge_destination(current_block_idx, true_target)) { + let label = self.edge_label(current_block_idx, true_target); encoder.emit_branch("br", &label)?; } return Ok(()); } - let true_has_phi = self.successor_has_phi_from(ssa, current_block_idx, true_target)?; - let false_has_phi = self.successor_has_phi_from(ssa, current_block_idx, false_target)?; - - let true_label = self - .block_labels - .get(&true_target) - .cloned() - .unwrap_or_else(|| format!("block_{true_target}")); - let false_label = self - .block_labels - .get(&false_target) - .cloned() - .unwrap_or_else(|| format!("block_{false_target}")); - - // Get the comparison mnemonic and its inverse - let (mnemonic, inverse_mnemonic) = match (cmp, unsigned) { - (CmpKind::Eq, _) => ("beq", "bne.un"), - (CmpKind::Ne, _) => ("bne.un", "beq"), - (CmpKind::Lt, false) => ("blt", "bge"), - (CmpKind::Lt, true) => ("blt.un", "bge.un"), - (CmpKind::Le, false) => ("ble", "bgt"), - (CmpKind::Le, true) => ("ble.un", "bgt.un"), - (CmpKind::Gt, false) => ("bgt", "ble"), - (CmpKind::Gt, true) => ("bgt.un", "ble.un"), - (CmpKind::Ge, false) => ("bge", "blt"), - (CmpKind::Ge, true) => ("bge.un", "blt.un"), + let true_label = self.edge_label(current_block_idx, true_target); + let false_label = self.edge_label(current_block_idx, false_target); + let false_dest = self.edge_destination(current_block_idx, false_target); + + // Only the comparison as written is needed; there is no inverse-mnemonic table + // because the branch below never inverts. See the note there for why ordered + // complements are wrong for floats. + let mnemonic = match (cmp, unsigned) { + (CmpKind::Eq, _) => "beq", + (CmpKind::Ne, _) => "bne.un", + (CmpKind::Lt, false) => "blt", + (CmpKind::Lt, true) => "blt.un", + (CmpKind::Le, false) => "ble", + (CmpKind::Le, true) => "ble.un", + (CmpKind::Gt, false) => "bgt", + (CmpKind::Gt, true) => "bgt.un", + (CmpKind::Ge, false) => "bge", + (CmpKind::Ge, true) => "bge.un", }; - // Simple case: no phi stores needed for either target - if !true_has_phi && !false_has_phi { - encoder.emit_branch(mnemonic, &true_label)?; - if Some(false_target) != next_block_idx { - encoder.emit_branch("br", &false_label)?; - } - return Ok(()); - } - - // Generate unique intermediate labels - let phi_false_label = format!("phi_false_{current_block_idx}_{false_target}"); - - if true_has_phi { - // Use inverse condition to branch to false handling, then handle true inline - encoder.emit_branch( - inverse_mnemonic, - if false_has_phi { - &phi_false_label - } else { - &false_label - }, - )?; - - // True path: emit phi stores then jump - self.emit_phi_stores_for_successor(encoder, ssa, current_block_idx, true_target)?; - if Some(true_target) != next_block_idx { - encoder.emit_branch("br", &true_label)?; - } - } else { - // No phi stores for true - just branch directly - encoder.emit_branch(mnemonic, &true_label)?; - } - - // False path handling - if false_has_phi { - if true_has_phi { - encoder.define_label(&phi_false_label)?; - } - self.emit_phi_stores_for_successor(encoder, ssa, current_block_idx, false_target)?; - if Some(false_target) != next_block_idx { - encoder.emit_branch("br", &false_label)?; - } - } else if !true_has_phi && Some(false_target) != next_block_idx { + // Always branch on the comparison as written, never on its complement. + // + // A complement table built from *ordered* opposites (`blt`/`bge`, `ble`/`bgt`, …) is + // wrong for floats: with a NaN operand every ordered comparison is false, so both + // `blt` and `bge` are false and the inverted branch sends the NaN case down the true + // edge — precisely the edge the comparison rejected. The complement can only be + // stated correctly once the operand type is known, since `.un` means *unordered* for + // floats but *unsigned* for integers. Emitting the original mnemonic sidesteps the + // question entirely: no inversion, no type dependence. + encoder.emit_branch(mnemonic, &true_label)?; + if next_entry != Some(false_dest) { encoder.emit_branch("br", &false_label)?; } Ok(()) } - /// Emits a switch instruction with proper phi store handling. + /// Emits a switch instruction and the transfer along its default edge. /// - /// For each switch target that has phi nodes from the current block, - /// we create an intermediate block that executes the phi stores and - /// then jumps to the actual target. - fn emit_switch_with_phi_stores( + /// Each case edge names its own edge block where it has one, which keeps the switch a + /// plain jump table: nothing is woven between the instruction and the default branch. + fn emit_switch_branch( &mut self, encoder: &mut InstructionEncoder, - ssa: &SsaFunction, current_block_idx: usize, targets: &[usize], default: usize, - next_block_idx: Option, + next_entry: Option, ) -> Result<()> { - // Determine which targets need phi stores - let mut needs_intermediate: Vec = Vec::with_capacity(targets.len()); - for &target in targets { - needs_intermediate.push(self.successor_has_phi_from(ssa, current_block_idx, target)?); - } - let default_needs_intermediate = - self.successor_has_phi_from(ssa, current_block_idx, default)?; - - // Build the label list for the switch instruction - // Use intermediate labels for targets that need phi stores - let mut switch_labels: Vec = Vec::with_capacity(targets.len()); - for (i, &target) in targets.iter().enumerate() { - if needs_intermediate.get(i).copied().unwrap_or(false) { - switch_labels.push(format!("phi_switch_{current_block_idx}_{i}")); - } else { - switch_labels.push( - self.block_labels - .get(&target) - .cloned() - .unwrap_or_else(|| format!("block_{target}")), - ); - } - } + let switch_labels: Vec = targets + .iter() + .map(|&target| self.edge_label(current_block_idx, target)) + .collect(); - // Emit the switch instruction let label_refs: Vec<&str> = switch_labels.iter().map(String::as_str).collect(); encoder.emit_switch(&label_refs)?; - // Emit jump to default (or intermediate for default) - let default_label = self - .block_labels - .get(&default) - .cloned() - .unwrap_or_else(|| format!("block_{default}")); - - if default_needs_intermediate { - let default_intermediate = format!("phi_switch_{current_block_idx}_default"); - if Some(default) != next_block_idx { - encoder.emit_branch("br", &default_intermediate)?; - } - - // Emit intermediate blocks for targets that need phi stores - for (i, &target) in targets.iter().enumerate() { - if needs_intermediate.get(i).copied().unwrap_or(false) { - let intermediate_label = format!("phi_switch_{current_block_idx}_{i}"); - encoder.define_label(&intermediate_label)?; - self.emit_phi_stores_for_successor(encoder, ssa, current_block_idx, target)?; - let target_label = self - .block_labels - .get(&target) - .cloned() - .unwrap_or_else(|| format!("block_{target}")); - encoder.emit_branch("br", &target_label)?; - } - } - - // Emit intermediate block for default - encoder.define_label(&default_intermediate)?; - self.emit_phi_stores_for_successor(encoder, ssa, current_block_idx, default)?; - if Some(default) != next_block_idx { - encoder.emit_branch("br", &default_label)?; - } - } else { - // Default doesn't need phi stores - if Some(default) != next_block_idx { - encoder.emit_branch("br", &default_label)?; - } - - // Emit intermediate blocks for targets that need phi stores - for (i, &target) in targets.iter().enumerate() { - if needs_intermediate.get(i).copied().unwrap_or(false) { - let intermediate_label = format!("phi_switch_{current_block_idx}_{i}"); - encoder.define_label(&intermediate_label)?; - self.emit_phi_stores_for_successor(encoder, ssa, current_block_idx, target)?; - let target_label = self - .block_labels - .get(&target) - .cloned() - .unwrap_or_else(|| format!("block_{target}")); - encoder.emit_branch("br", &target_label)?; - } - } + // Falling out of a switch is the default edge. + if next_entry != Some(self.edge_destination(current_block_idx, default)) { + let default_label = self.edge_label(current_block_idx, default); + encoder.emit_branch("br", &default_label)?; } Ok(()) diff --git a/dotscope/src/compiler/codegen/tests.rs b/dotscope/src/compiler/codegen/tests.rs index b7853bf5..76306d73 100644 --- a/dotscope/src/compiler/codegen/tests.rs +++ b/dotscope/src/compiler/codegen/tests.rs @@ -2132,6 +2132,7 @@ fn test_fibonacci_deobfuscation_preserves_semantics() { let mut found = false; for row in method_table.iter() { + let row = row.expect("row parses"); let name = strings.get(row.name as usize).unwrap_or(""); if name != "Fibonacci" { continue; @@ -2252,6 +2253,7 @@ fn test_fibonacci_pass_combinations() { let strings = output.strings().expect("strings"); for row in method_table.iter() { + let row = row.expect("row parses"); let method_name = strings.get(row.name as usize).unwrap_or(""); if method_name != "Fibonacci" { continue; diff --git a/dotscope/src/compiler/context.rs b/dotscope/src/compiler/context.rs index 7f3d38d7..64fe4f7a 100644 --- a/dotscope/src/compiler/context.rs +++ b/dotscope/src/compiler/context.rs @@ -510,7 +510,10 @@ impl CompilerContext { /// transformations that may have removed or added calls since the /// initial static call graph was built. /// - /// Used by both dead method elimination and cleanup request building. + /// Note this graph only has entries for methods that *have* an SSA function. A method whose + /// body never converted contributes no edges at all, which reads as "calls nothing" rather + /// than "unknown". Reachability consumers must use [`Self::build_effective_call_graph`]; + /// this one is for consumers that specifically want the post-transformation view. #[must_use] pub fn build_ssa_call_graph(&self) -> BTreeMap> { let mut call_graph = BTreeMap::new(); @@ -538,4 +541,27 @@ impl CompilerContext { } call_graph } + + /// Builds a call graph safe to drive reachability and deletion from. + /// + /// [`Self::build_ssa_call_graph`] is the post-transformation view, and it is the right + /// input wherever inlining and devirtualization must be reflected. It is the wrong input + /// for deciding what to *delete*: it only has entries for methods that converted to SSA, + /// so a method whose body never converted -- an encrypted or native body, or one that + /// failed to decode -- contributes no edges, and everything it references reads as + /// unreachable. Deleting on that basis strips types the surviving body still names, which + /// produces a structurally invalid assembly rather than a smaller one. + /// + /// So SSA edges win where they exist, and the static call graph fills in for every method + /// that has none. This mirrors what `CtxWorld::callees` does for dead-method elimination. + #[must_use] + pub fn build_effective_call_graph(&self) -> BTreeMap> { + let mut call_graph = self.build_ssa_call_graph(); + for node in self.call_graph.nodes() { + call_graph + .entry(node.token) + .or_insert_with(|| self.call_graph.callees(node.token).into_iter().collect()); + } + call_graph + } } diff --git a/dotscope/src/compiler/passes/constants/mod.rs b/dotscope/src/compiler/passes/constants/mod.rs index ded14f1e..398c4e90 100644 --- a/dotscope/src/compiler/passes/constants/mod.rs +++ b/dotscope/src/compiler/passes/constants/mod.rs @@ -72,6 +72,86 @@ fn is_method_on_type(assembly: &CilObject, token: Token, type_name: &str) -> boo } } +/// An integer constant reduced to its CIL stack width and raw bit pattern. +/// +/// Folding an operation whose meaning depends on operand width and signedness — `rem`/`rem.un` +/// and `div`/`div.un` — cannot go through a single `as_i64` accessor. `as_i64` sign-extends, +/// so reinterpreting its result as unsigned turns `rem.un` of `I32(-1)` into +/// `0xFFFF_FFFF_FFFF_FFFF % r` rather than the `0xFFFF_FFFF % r` CIL specifies. `as_u64` is +/// no better here: it refuses negative values outright, so it cannot express the +/// reinterpretation at all. +/// +/// Per ECMA-335 III.1.1 the stack holds `int32` and `int64`; anything narrower is widened to +/// `int32` when loaded. This mirrors that: sub-word constants become 32-bit, and the bit +/// pattern is kept unsigned so each operation can choose how to read it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum IntOperand { + /// A 32-bit stack value, as raw bits. + W32(u32), + /// A 64-bit stack value, as raw bits. + W64(u64), +} + +impl IntOperand { + /// Classifies an integer constant, or `None` for non-integer constants. + /// + /// `NativeInt`/`NativeUInt` are deliberately excluded rather than assumed 64-bit: their + /// width is the target's pointer size, so folding them here would bake a host assumption + /// into the output. + fn from_const(value: &ConstValue) -> Option { + #[allow(clippy::cast_sign_loss)] + Some(match value { + ConstValue::I8(v) => Self::W32(i32::from(*v) as u32), + ConstValue::I16(v) => Self::W32(i32::from(*v) as u32), + ConstValue::I32(v) => Self::W32(*v as u32), + ConstValue::U8(v) => Self::W32(u32::from(*v)), + ConstValue::U16(v) => Self::W32(u32::from(*v)), + ConstValue::U32(v) => Self::W32(*v), + ConstValue::I64(v) => Self::W64(*v as u64), + ConstValue::U64(v) => Self::W64(*v), + ConstValue::True => Self::W32(1), + ConstValue::False => Self::W32(0), + _ => return None, + }) + } +} + +/// Folds `rem`/`rem.un` over two constant operands, or `None` when it must not be folded. +/// +/// Returns `None` — leaving the instruction in place — when: +/// +/// - the divisor is zero, since the operation raises `DivideByZeroException` at runtime; +/// - the division overflows (`i32::MIN % -1`), which raises `OverflowException`; +/// - the operands have different widths, which is unverifiable IL. Inventing a promotion +/// would mean choosing sign- or zero-extension on the crate's behalf and folding an input +/// the runtime would reject. +/// +/// The result keeps the operand width, so a 64-bit remainder stays [`ConstValue::I64`]. +/// Narrowing it to `I32` — as a single `ConstValue::I32(result as i32)` does — is worse than a +/// wrong number: codegen then emits `ldc.i4` where the stack type requires `ldc.i8`, which is +/// a verification failure rather than a miscomputation. +fn fold_rem(left: IntOperand, right: IntOperand, unsigned: bool) -> Option { + #[allow(clippy::cast_possible_wrap)] + match (left, right) { + (IntOperand::W32(l), IntOperand::W32(r)) => { + if unsigned { + Some(ConstValue::I32(l.checked_rem(r)? as i32)) + } else { + Some(ConstValue::I32((l as i32).checked_rem(r as i32)?)) + } + } + (IntOperand::W64(l), IntOperand::W64(r)) => { + if unsigned { + Some(ConstValue::I64(l.checked_rem(r)? as i64)) + } else { + Some(ConstValue::I64((l as i64).checked_rem(r as i64)?)) + } + } + // Mixed widths: see above. + _ => None, + } +} + /// Result of checking an algebraic identity. /// /// Either the operation simplifies to a constant value (absorbing elements) @@ -805,13 +885,30 @@ impl ConstantPropagationPass { } /// Checks if an overflow-checked operation can be folded. - #[allow(clippy::cast_sign_loss, clippy::cast_possible_wrap)] // Intentional bit reinterpretation for overflow checking + /// + /// `add.ovf`/`sub.ovf`/`mul.ovf` throw `OverflowException` when the result does not fit + /// the operand width, so folding one is only sound when the operation provably does *not* + /// overflow. Getting that test wrong turns a throwing path into a wrong constant, which is + /// worse than not folding at all. + /// + /// The width test must therefore happen at the *operand's* own width. + /// [`ConstValue::add_checked`] and its siblings do exactly that — `checked_*` per variant, + /// `None` on overflow, signed or unsigned per the instruction's flag — so the whole + /// decision is delegated rather than reimplemented here. The plain `add`/`sub`/`mul` are + /// deliberately *not* used: they wrap silently, so they cannot serve as their own overflow + /// oracle. Nor is `as_i64`, which sign-extends to 64 bits and so reports no overflow for + /// any pair of 32-bit operands — `i32::MAX + 1` folded to `i32::MIN` under it. fn check_overflow_op( op: &SsaOp, constants: &BTreeMap, ptr_size: PointerSize, ) -> Option<(SsaVarId, ConstValue)> { - match op { + // No `x * 0 == 0` short-circuit for `MulOvf`: a multiply by zero cannot overflow, so + // `mul_checked` already folds it. The short-circuit that used to sit here returned the + // *zero operand's own* variant, so `I32(0) * I64(5)` yielded `I32(0)` and + // `I64(0) * I32(5)` yielded `I64(0)` — the result width depended on which side + // happened to be zero. + let (dest, folded) = match op { SsaOp::AddOvf { dest, left, @@ -819,22 +916,8 @@ impl ConstantPropagationPass { unsigned, .. } => { - let l = constants.get(left)?; - let r = constants.get(right)?; - let (lv, rv) = (l.as_i64()?, r.as_i64()?); - - if *unsigned { - let (result, overflow) = (lv as u64).overflowing_add(rv as u64); - if !overflow { - return Some((*dest, ConstValue::I64(result as i64))); - } - } else { - let (_, overflow) = lv.overflowing_add(rv); - if !overflow { - return Some((*dest, l.add(r, ptr_size)?)); - } - } - None + let (l, r) = (constants.get(left)?, constants.get(right)?); + (dest, l.add_checked(r, *unsigned, ptr_size)) } SsaOp::SubOvf { @@ -844,22 +927,8 @@ impl ConstantPropagationPass { unsigned, .. } => { - let l = constants.get(left)?; - let r = constants.get(right)?; - let (lv, rv) = (l.as_i64()?, r.as_i64()?); - - if *unsigned { - let (result, overflow) = (lv as u64).overflowing_sub(rv as u64); - if !overflow { - return Some((*dest, ConstValue::I64(result as i64))); - } - } else { - let (_, overflow) = lv.overflowing_sub(rv); - if !overflow { - return Some((*dest, l.sub(r, ptr_size)?)); - } - } - None + let (l, r) = (constants.get(left)?, constants.get(right)?); + (dest, l.sub_checked(r, *unsigned, ptr_size)) } SsaOp::MulOvf { @@ -869,35 +938,14 @@ impl ConstantPropagationPass { unsigned, .. } => { - let l = constants.get(left)?; - let r = constants.get(right)?; - - // Special case: x * 0 = 0, even with overflow check - if l.is_zero() { - return Some((*dest, l.clone())); - } - if r.is_zero() { - return Some((*dest, r.clone())); - } - - let (lv, rv) = (l.as_i64()?, r.as_i64()?); - - if *unsigned { - let (result, overflow) = (lv as u64).overflowing_mul(rv as u64); - if !overflow { - return Some((*dest, ConstValue::I64(result as i64))); - } - } else { - let (_, overflow) = lv.overflowing_mul(rv); - if !overflow { - return Some((*dest, l.mul(r, ptr_size)?)); - } - } - None + let (l, r) = (constants.get(left)?, constants.get(right)?); + (dest, l.mul_checked(r, *unsigned, ptr_size)) } - _ => None, - } + _ => return None, + }; + + folded.map(|value| (*dest, value)) } /// Folds calls to pure methods with all-constant arguments. @@ -1015,25 +1063,17 @@ impl ConstantPropagationPass { Some(SsaOp::Const { value, .. }) => Some(value), _ => None, }) - .and_then(ConstValue::as_i64); + .and_then(IntOperand::from_const); let rval = constants .get(right) .or_else(|| match ssa.get_definition(*right) { Some(SsaOp::Const { value, .. }) => Some(value), _ => None, }) - .and_then(ConstValue::as_i64); + .and_then(IntOperand::from_const); if let (Some(l), Some(r)) = (lval, rval) { - if r != 0 { - #[allow(clippy::cast_sign_loss)] - let result = if *unsigned { - (l as u64).checked_rem(r as u64).unwrap_or(0) as i64 - } else { - l.checked_rem(r).unwrap_or(0) - }; - #[allow(clippy::cast_possible_truncation)] - let value = ConstValue::I32(result as i32); + if let Some(value) = fold_rem(l, r, *unsigned) { new_constants.push((block_idx, instr_idx, *dest, value)); } } diff --git a/dotscope/src/compiler/passes/constants/tests.rs b/dotscope/src/compiler/passes/constants/tests.rs index be4171ca..f02cf3e8 100644 --- a/dotscope/src/compiler/passes/constants/tests.rs +++ b/dotscope/src/compiler/passes/constants/tests.rs @@ -6,7 +6,7 @@ use crate::{ SsaVarId, }, compiler::{ - passes::constants::{AlgebraicResult, ConstantPropagationPass}, + passes::constants::{fold_rem, AlgebraicResult, ConstantPropagationPass, IntOperand}, CompilerContext, EventLog, SsaPass, }, metadata::{token::Token, typesystem::PointerSize}, @@ -1286,3 +1286,118 @@ fn test_fold_string_operations_with_decrypted_concat() { "Unresolvable method should not fold" ); } + +/// `rem.un` reads its operands as unsigned at their own width. +/// +/// Reaching them through a sign-extending accessor turns `I32(-1)` into +/// `0xFFFF_FFFF_FFFF_FFFF`, so `-1 rem.un 7` folds to `1` instead of the `3` that +/// `0xFFFF_FFFF % 7` gives. That is a wrong constant on the default deobfuscation path — +/// modulus-with-a-key is the canonical decryptor shape. +#[test] +fn rem_un_treats_i32_operands_as_unsigned_32_bit() { + let l = IntOperand::from_const(&ConstValue::I32(-1)).unwrap(); + let r = IntOperand::from_const(&ConstValue::I32(7)).unwrap(); + + assert_eq!(fold_rem(l, r, true), Some(ConstValue::I32(3))); + // Signed: Rust and CIL agree that the remainder takes the dividend's sign. + assert_eq!(fold_rem(l, r, false), Some(ConstValue::I32(-1))); +} + +/// The folded constant keeps the operands' width. +/// +/// Narrowing a 64-bit remainder to `I32` is worse than a wrong value: codegen then emits +/// `ldc.i4` where the stack type requires `ldc.i8`, which fails verification. +#[test] +fn rem_keeps_the_operand_width() { + let wide = IntOperand::from_const(&ConstValue::I64(0x1_0000_0000 + 5)).unwrap(); + let divisor = IntOperand::from_const(&ConstValue::I64(0x1_0000_0000)).unwrap(); + + assert_eq!(fold_rem(wide, divisor, false), Some(ConstValue::I64(5))); +} + +/// Cases the runtime turns into an exception must not be folded away. +#[test] +fn rem_refuses_to_fold_faulting_and_unverifiable_cases() { + let one = IntOperand::from_const(&ConstValue::I32(1)).unwrap(); + let zero = IntOperand::from_const(&ConstValue::I32(0)).unwrap(); + let min = IntOperand::from_const(&ConstValue::I32(i32::MIN)).unwrap(); + let neg_one = IntOperand::from_const(&ConstValue::I32(-1)).unwrap(); + let wide = IntOperand::from_const(&ConstValue::I64(1)).unwrap(); + + // DivideByZeroException + assert_eq!(fold_rem(one, zero, false), None); + assert_eq!(fold_rem(one, zero, true), None); + // OverflowException + assert_eq!(fold_rem(min, neg_one, false), None); + // Mixed widths are unverifiable IL; folding would invent a promotion. + assert_eq!(fold_rem(one, wide, false), None); +} + +/// An overflow-checked op must test overflow at the operands' width, signed or unsigned. +/// +/// Two 32-bit values can never overflow the 64-bit range, so reading them through a +/// sign-extending accessor reports success for a case that must raise `OverflowException`, +/// and yields an `I64` where the stack holds an `int32`. The unsigned half of this was fixed +/// first; the signed half kept the `as_i64` shape and is the regression this pins. +#[test] +fn overflow_checks_use_the_operand_width() { + let ptr = PointerSize::Bit64; + let one: ConstValue = ConstValue::I32(1); + let two: ConstValue = ConstValue::I32(2); + let wide_one: ConstValue = ConstValue::I64(1); + + // Unsigned: 0xFFFF_FFFF + 1 overflows int32, so this must not fold. + let all_ones: ConstValue = ConstValue::I32(-1); + assert_eq!(all_ones.add_checked(&one, true, ptr), None); + + // Signed: i32::MAX + 1 overflows int32. Under `as_i64` both operands widened to 64 bits, + // no overflow was reported, and the wrapping `add` folded this to `i32::MIN` — turning a + // throwing path into a wrong constant. + let max: ConstValue = ConstValue::I32(i32::MAX); + let min: ConstValue = ConstValue::I32(i32::MIN); + assert_eq!(max.add_checked(&one, false, ptr), None); + assert_eq!(min.sub_checked(&one, false, ptr), None); + assert_eq!(max.mul_checked(&two, false, ptr), None); + + // In range, and the result keeps the 32-bit width. + assert_eq!( + two.add_checked(&one, true, ptr), + Some(ConstValue::I32(3)), + "an in-range unsigned add must fold at int32 width" + ); + assert_eq!( + two.add_checked(&one, false, ptr), + Some(ConstValue::I32(3)), + "an in-range signed add must fold at int32 width" + ); + + // The same values at 64-bit width do fit. + let wide: ConstValue = ConstValue::I64(0xFFFF_FFFF); + assert_eq!( + wide.add_checked(&wide_one, true, ptr), + Some(ConstValue::I64(0x1_0000_0000)) + ); +} + +/// A multiply by zero folds through the ordinary overflow check, at the operands' width. +/// +/// The `x * 0 == 0` short-circuit that used to precede the check returned the *zero operand's +/// own* variant, so the result width depended on which side happened to be zero. +#[test] +fn mul_ovf_by_zero_folds_at_the_operand_width() { + let ptr = PointerSize::Bit64; + + let hundred: ConstValue = ConstValue::I32(100); + let zero: ConstValue = ConstValue::I32(0); + assert_eq!( + hundred.mul_checked(&zero, false, ptr), + Some(ConstValue::I32(0)) + ); + + let wide_zero: ConstValue = ConstValue::I64(0); + let wide_hundred: ConstValue = ConstValue::I64(100); + assert_eq!( + wide_zero.mul_checked(&wide_hundred, false, ptr), + Some(ConstValue::I64(0)) + ); +} diff --git a/dotscope/src/compiler/passes/inlining.rs b/dotscope/src/compiler/passes/inlining.rs index cde080c4..f08090fd 100644 --- a/dotscope/src/compiler/passes/inlining.rs +++ b/dotscope/src/compiler/passes/inlining.rs @@ -565,7 +565,11 @@ impl<'a> InliningContext<'a> { return false; } - // Insert remaining inlined ops + // Insert remaining inlined ops. + // + // Capture the count before `into_iter` consumes the vector: the return-value copy below + // must land *after* every spliced op, and it needs this length to know where that is. + let spliced_count = inlined_ops.len().saturating_sub(1); let instructions = block.instructions_mut(); let base = call_instr_idx.saturating_add(1); for (i, op) in inlined_ops.into_iter().skip(1).enumerate() { @@ -580,7 +584,15 @@ impl<'a> InliningContext<'a> { let Some(block) = self.caller_ssa.block_mut(call_block_idx) else { return false; }; - let insert_pos = call_instr_idx.saturating_add(1); + // *After* the spliced body, not at its start. Inserting at + // `call_instr_idx + 1` puts the copy between the callee's first op and the + // rest, so for any callee with two or more non-`Return` ops the copy reads + // `remapped_ret` before the op that defines it — a use-before-def that hands + // `dest_var` whatever the unwritten storage holds. Nothing downstream re-sorts + // block instructions before codegen. + let insert_pos = call_instr_idx + .saturating_add(1) + .saturating_add(spliced_count); block.instructions_mut().insert( insert_pos, SsaInstruction::synthetic(SsaOp::Copy { diff --git a/dotscope/src/deobfuscation/cleanup.rs b/dotscope/src/deobfuscation/cleanup.rs index ead26594..75f12332 100644 --- a/dotscope/src/deobfuscation/cleanup.rs +++ b/dotscope/src/deobfuscation/cleanup.rs @@ -42,11 +42,15 @@ use crate::{ }, compiler::{EventKind, ProxyDevirtualizationPass}, deobfuscation::{ - context::AnalysisContext, engine::DeobfuscationEngine, renamer, techniques::Detections, + context::AnalysisContext, + engine::DeobfuscationEngine, + renamer, + techniques::{Detections, Technique, TechniqueCapability}, }, metadata::{ tables::{ - AssemblyRaw, FieldRaw, ModuleRaw, NestedClassRaw, TableDataOwned, TableId, TypeDefRaw, + skip_unreadable, AssemblyRaw, FieldRaw, ModuleRaw, NestedClassRaw, TableDataOwned, + TableId, TypeDefRaw, }, token::Token, typesystem::wellknown, @@ -80,7 +84,6 @@ fn decryptors_still_called( } still_called } - /// Builds a complete cleanup request from detection results and analysis state. /// /// This consolidates all cleanup sources into a single request: @@ -110,7 +113,58 @@ pub(crate) fn build_cleanup_request( ) -> CleanupRequest { // Start with technique-merged cleanup let registry = engine.technique_registry(); - let mut request = detections.merged_cleanup(); + + // Byte transforms that were detected but did not succeed. + // + // A byte transform rewrites the assembly before SSA is built — decrypting + // method bodies, unpacking resources. When one fails, the assembly still + // holds the protected form, but every technique's cleanup request was + // already filled during detection on the assumption it would not. Merging + // those requests deletes the infrastructure the protected code still needs. + // + // Failure is recorded by omission: the pipeline calls `mark_transformed` + // only on success and logs the error, so "detected, declares a byte + // transform, not transformed" is the signal. Techniques without the + // capability are never marked and must not be caught by this. + let failed: Vec<&dyn Technique> = registry + .sorted_techniques(detections) + .into_iter() + .filter(|tech| { + tech.capabilities() + .contains(&TechniqueCapability::ByteTransform) + && detections.is_detected(tech.id()) + && !detections.is_transformed(tech.id()) + }) + .collect(); + let failed_byte_transforms: BTreeSet<&str> = failed.iter().map(|tech| tech.id()).collect(); + + let mut request = detections.merged_cleanup_excluding(&failed_byte_transforms); + + for tech in &failed { + log::warn!( + "Withholding cleanup for {}: it declares a byte transform that did not succeed, \ + so the metadata it would delete is still in use", + tech.id() + ); + + // Whatever the technique was meant to restore is still in its protected + // form. Such a method carries no calls, so nothing it references looks + // reachable — see the suppression of type-level deletion below — and it + // is itself real code that must survive for later analysis. + let Some(detection) = detections.get(tech.id()) else { + continue; + }; + let unrecovered = tech.unrecovered_methods(detection); + if !unrecovered.is_empty() { + log::warn!( + "Protecting {} method(s) {} could not restore", + unrecovered.len(), + tech.id() + ); + request.protect_tokens(unrecovered); + } + } + let still_called = decryptors_still_called(ctx, ssa_call_graph); for tech in registry.sorted_techniques(detections) { if !detections.is_detected(tech.id()) { @@ -179,9 +233,26 @@ pub(crate) fn build_cleanup_request( // This uses cluster analysis: candidate types that only reference each // other (and already-deleted entities) are detected as isolated // infrastructure and removed. - let unreferenced_types = find_unreferenced_types(assembly, ssa_call_graph, &request); - for type_token in unreferenced_types { - request.add_type(type_token); + // + // Skipped entirely when a byte transform failed. The analysis decides a type + // is unreferenced by asking the call graph who names it, and a method whose + // body is still encrypted contributes no edges at all — neither the SSA + // graph, which never converted it, nor the static graph, which was built + // from the same stub. Almost every non-public type then looks unrooted and + // the assembly is gutted rather than cleaned. Protection cannot save it + // either: this analysis does not consult the protected set, and deleting a + // type takes its members with it. + if failed_byte_transforms.is_empty() { + let unreferenced_types = find_unreferenced_types(assembly, ssa_call_graph, &request); + for type_token in unreferenced_types { + request.add_type(type_token); + } + } else { + log::warn!( + "Skipping unreferenced-type removal: {} byte transform(s) failed, so the call \ + graph cannot distinguish unreachable from undecrypted", + failed_byte_transforms.len() + ); } // Add dead methods from analysis (aggressive mode only) @@ -269,6 +340,17 @@ pub fn execute_cleanup( if types_count > 0 || methods_count > 0 || fields_count > 0 { log::info!("Cleanup: {types_count} types, {methods_count} methods, {fields_count} fields"); + // Counts alone cannot answer "why is this still here" or "what took that away", + // which is the question any cleanup investigation starts from. + if log::log_enabled!(log::Level::Debug) { + let fmt = |it: &mut dyn Iterator| { + it.map(|t| format!("0x{:08x}", t.value())) + .collect::>() + .join(", ") + }; + log::debug!("Cleanup types: {}", fmt(&mut request.types())); + log::debug!("Cleanup methods: {}", fmt(&mut request.methods())); + } } for section_name in request.excluded_sections() { @@ -481,7 +563,7 @@ fn repair_invalid_module_guids(cil_assembly: &mut CilAssembly, ctx: &AnalysisCon let Some(module_table) = tables.table::() else { return; }; - let Some(row) = module_table.get(1) else { + let Some(row) = module_table.get(1).ok().flatten() else { return; }; @@ -550,7 +632,7 @@ fn needs_metadata_repair(assembly: &CilObject) -> bool { if t.row_count > 1 { return true; } - if let Some(row) = t.get(1) { + if let Some(row) = t.get(1).ok().flatten() { let guid_count: u32 = assembly.guids().map_or(0, |g| (g.data().len() / 16) as u32); if (row.encid != 0 && row.encid > guid_count) || (row.encbaseid != 0 && row.encbaseid > guid_count) @@ -596,6 +678,7 @@ fn repair_duplicate_typedef_rows( .table::() .map(|rows| { rows.into_iter() + .filter_map(skip_unreadable) .map(|r| (r.nested_class, r.enclosing_class, r.rid)) .collect() }) @@ -609,7 +692,7 @@ fn repair_duplicate_typedef_rows( // Collect TypeDef rows to check for duplicates and emptiness let typedef_rows: Vec = tables .table::() - .map(|t| t.into_iter().collect()) + .map(|t| t.into_iter().filter_map(skip_unreadable).collect()) .unwrap_or_default(); let method_count = cil_assembly.original_table_row_count(TableId::MethodDef); @@ -713,57 +796,71 @@ fn repair_duplicate_typedef_rows( /// This corrects any invalid access to `Private` (1), which is the most /// restrictive valid option and unlikely to break anything. fn repair_global_field_visibility(cil_assembly: &mut CilAssembly, ctx: &AnalysisContext) { - let Some(tables) = cil_assembly.view().tables() else { - return; - }; - let Some(typedefs) = tables.table::() else { - return; - }; + // Collected in a first pass so the read borrow of the view is released before the writes + // below take `cil_assembly` mutably. + let pending: Vec<(u32, FieldRaw)> = { + let Some(tables) = cil_assembly.view().tables() else { + return; + }; + let Some(typedefs) = tables.table::() else { + return; + }; - // Find type (always RID 1) and its field range - let typedef_rows: Vec = typedefs.into_iter().collect(); - let Some(module_type) = typedef_rows.first() else { - return; - }; - let field_start = module_type.field_list; - let field_end = if let Some(next) = typedef_rows.get(1) { - next.field_list - } else { - // If there's only one type, the field end is the total field count + 1 - cil_assembly - .original_table_row_count(TableId::Field) - .saturating_add(1) - }; + // Find type (always RID 1) and its field range. + // + // Fetched by RID rather than by position in a collected `Vec`: `skip_unreadable` + // drops unparseable rows, so in a damaged assembly — exactly what this function + // exists to repair — position stops tracking RID and `` would be read from + // whatever type happened to survive first. `MetadataTable::get` derives each row's + // offset from its RID, so one unreadable row cannot displace any other. + let Some(module_type) = typedefs.get(1).ok().flatten() else { + return; + }; + let field_start = module_type.field_list; + let field_end = match typedefs.get(2) { + Ok(Some(next)) => next.field_list, + // Only one type, or TypeDef RID 2 is unreadable: fall back to the whole table. + _ => cil_assembly + .original_table_row_count(TableId::Field) + .saturating_add(1), + }; - if field_start >= field_end { - return; - } + if field_start >= field_end { + return; + } - // Check each field for invalid access - let Some(fields_table) = tables.table::() else { - return; + let Some(fields_table) = tables.table::() else { + return; + }; + + (field_start..field_end) + .filter_map(|rid| { + // Same reason as above, and it matters more here: the repaired row is written + // back under this `rid`, so reading the wrong row would copy one field's flags + // onto a different field. + let field = fields_table.get(rid).ok().flatten()?; + let access = field.flags & 0x0007; // FieldAccessMask per ECMA-335 §II.23.1.5 + + // Valid for global fields: CompilerControlled(0), Private(1), Public(6). + if matches!(access, 0x0000 | 0x0001 | 0x0006) { + return None; + } + + let mut fixed = field; + fixed.flags = (fixed.flags & !0x0007) | 0x0001; // Set to Private + Some((rid, fixed)) + }) + .collect() }; - let fields: Vec = fields_table.into_iter().collect(); - let mut repaired: usize = 0; - for rid in field_start..field_end { - let idx = rid.saturating_sub(1) as usize; - let Some(field) = fields.get(idx) else { - break; - }; - let access = field.flags & 0x0007; // FieldAccessMask per ECMA-335 §II.23.1.5 - - // Valid access for global fields: CompilerControlled(0), Private(1), or Public(6) - if !matches!(access, 0x0000 | 0x0001 | 0x0006) { - let mut fixed = field.clone(); - fixed.flags = (fixed.flags & !0x0007) | 0x0001; // Set to Private - if let Err(e) = - cil_assembly.table_row_update(TableId::Field, rid, TableDataOwned::Field(fixed)) - { - log::warn!("Failed to repair field {rid} visibility: {e}"); - } else { - repaired = repaired.saturating_add(1); - } + let mut repaired: usize = 0; + for (rid, fixed) in pending { + if let Err(e) = + cil_assembly.table_row_update(TableId::Field, rid, TableDataOwned::Field(fixed)) + { + log::warn!("Failed to repair field {rid} visibility: {e}"); + } else { + repaired = repaired.saturating_add(1); } } diff --git a/dotscope/src/deobfuscation/config.rs b/dotscope/src/deobfuscation/config.rs index e8cdbd29..8c9fa01c 100644 --- a/dotscope/src/deobfuscation/config.rs +++ b/dotscope/src/deobfuscation/config.rs @@ -5,7 +5,10 @@ use std::time::Duration; -use crate::{deobfuscation::SmartRenameConfig, emulation::TracingConfig}; +use crate::{ + deobfuscation::SmartRenameConfig, + emulation::{EmulationLimits, TracingConfig}, +}; /// Fixpoint iteration limits. #[derive(Debug, Clone)] @@ -123,6 +126,46 @@ pub struct EmulationConfig { pub warmup_timeout: Duration, /// Number of retry passes for warmup methods with dependency chains. pub warmup_retry_passes: usize, + /// Maximum emulations a single deobfuscation run may perform, or 0 for unlimited. + /// + /// [`max_instructions`](Self::max_instructions) and [`timeout`](Self::timeout) bound one + /// emulation. Neither bounds how many a run performs: the decryption pass emulates once + /// per distinct call site with constant arguments, and its value cache is keyed on + /// `(decryptor, args)`, so a loop calling a decryptor with a varying key misses the cache + /// every iteration. This is the ceiling on the run as a whole. + pub max_emulations: usize, +} + +impl EmulationConfig { + /// Limits for one per-method emulation. + /// + /// A template process is warmed up once under [`warmup_timeout`](Self::warmup_timeout), + /// which is deliberately generous because warmup executes module and type initialisers. + /// Every per-method execution is forked from that template, and is a different job with + /// its own budget: [`timeout`](Self::timeout). + /// + /// Forking shares the emulator's configuration `Arc` by default, so without applying + /// these explicitly a per-method execution silently inherits the warmup budget — which is + /// what left `timeout` written and never read, and let a single method run under a wall + /// clock up to 120x larger than documented. + /// + /// Only the two per-method fields are overridden; everything else is inherited from + /// `base`, which must be the limits of the template being forked + /// ([`EmulationProcess::limits`]). Defaulting the remainder instead would silently + /// *raise* the call-depth ceiling and *lower* the heap ceiling relative to the template + /// the fork inherits its warmed heap from — 1000 and 256 MiB against the 100 and 512 MiB + /// the template pool builds with — so a sample whose warmup allocates past 256 MiB would + /// fail every per-method decryption on a limit it was never intended to be held to. + /// + /// [`EmulationProcess::limits`]: crate::emulation::EmulationProcess::limits + #[must_use] + pub fn execution_limits(&self, base: &EmulationLimits) -> EmulationLimits { + EmulationLimits { + max_instructions: self.max_instructions, + timeout_ms: u64::try_from(self.timeout.as_millis()).unwrap_or(u64::MAX), + ..base.clone() + } + } } impl Default for EmulationConfig { @@ -133,6 +176,7 @@ impl Default for EmulationConfig { tracing: None, warmup_timeout: Duration::from_secs(60), warmup_retry_passes: 5, + max_emulations: 100_000, } } } @@ -164,14 +208,6 @@ impl Default for DecryptorHeuristics { /// CFF unflattening thresholds. #[derive(Debug, Clone)] pub struct UnflatteningThresholds { - /// Minimum switch cases to consider as potential flattening dispatcher. - pub min_switch_cases: usize, - /// Maximum states to enumerate per case when solving dispatcher. - pub max_states_per_case: usize, - /// Maximum iterations when tracing execution through flattened CFG. - pub max_trace_iterations: usize, - /// Threshold for large constant detection in state encoding. - pub large_constant_threshold: i64, /// Maximum BFS depth for back-edge transitive reachability check. pub max_backedge_depth: usize, /// Confidence scoring weights for CFF dispatcher detection. @@ -181,10 +217,6 @@ pub struct UnflatteningThresholds { impl Default for UnflatteningThresholds { fn default() -> Self { Self { - min_switch_cases: 4, - max_states_per_case: 15, - max_trace_iterations: 500, - large_constant_threshold: 100_000, max_backedge_depth: 10, confidence_weights: DetectionWeights::default(), } @@ -452,6 +484,7 @@ impl EngineConfig { timeout: Duration::from_millis(500), warmup_timeout: Duration::from_secs(30), warmup_retry_passes: 2, + max_emulations: 10_000, ..Default::default() }, resolution_strategies: vec![ResolutionStrategy::Static, ResolutionStrategy::Pattern], @@ -491,6 +524,7 @@ impl EngineConfig { timeout: Duration::from_secs(30), warmup_timeout: Duration::from_secs(120), warmup_retry_passes: 8, + max_emulations: 1_000_000, ..Default::default() }, unflattening: UnflatteningThresholds { @@ -865,4 +899,79 @@ mod tests { assert!(ResolutionStrategy::Emulation.requires_emulation()); assert!(!ResolutionStrategy::Static.requires_emulation()); } + /// A per-method execution must run under `timeout`, not the template's `warmup_timeout`. + /// + /// Forking shares the emulator's configuration `Arc`, so without applying these limits + /// explicitly the documented per-method budget is written and never read and the method + /// runs under the warmup budget instead. + #[test] + fn execution_limits_use_the_per_method_timeout() { + let config = EmulationConfig::default(); + let limits = config.execution_limits(&EmulationLimits::default()); + + assert_eq!( + limits.timeout_ms, + u64::try_from(config.timeout.as_millis()).unwrap() + ); + assert_ne!( + limits.timeout_ms, + u64::try_from(config.warmup_timeout.as_millis()).unwrap(), + "a per-method execution must not inherit the warmup budget" + ); + assert_eq!(limits.max_instructions, config.max_instructions); + } + + /// The `fast()` preset tightens the per-method budget; the derivation must follow the + /// configuration rather than any default. + #[test] + fn execution_limits_follow_the_preset() { + let fast = EngineConfig::fast(); + let base = EmulationLimits::default(); + + assert_eq!( + fast.emulation.execution_limits(&base).timeout_ms, + u64::try_from(fast.emulation.timeout.as_millis()).unwrap() + ); + assert_eq!( + fast.emulation.execution_limits(&base).max_instructions, + fast.emulation.max_instructions + ); + } + + /// Everything the per-method budget does not name must come from the template being + /// forked, not from [`EmulationLimits::default`]. + /// + /// The template pool builds with a call depth of 100 and a 512 MiB heap, and a fork + /// inherits the warmed heap by copy-on-write. Defaulting the unnamed fields instead + /// raised the depth to 1000 and *lowered* the heap to 256 MiB, so a sample whose warmup + /// allocated past 256 MiB failed every per-method decryption on a ceiling the template + /// had never been held to. + #[test] + fn execution_limits_inherit_every_unnamed_field_from_the_template() { + let config = EmulationConfig::default(); + let template = EmulationLimits { + max_call_depth: 100, + max_heap_bytes: 512 * 1024 * 1024, + max_unmanaged_bytes: 7 * 1024 * 1024, + ..EmulationLimits::default() + }; + + let limits = config.execution_limits(&template); + + assert_eq!(limits.max_call_depth, template.max_call_depth); + assert_eq!(limits.max_heap_bytes, template.max_heap_bytes); + assert_eq!(limits.max_unmanaged_bytes, template.max_unmanaged_bytes); + assert_ne!( + limits.max_heap_bytes, + EmulationLimits::default().max_heap_bytes, + "the template's heap ceiling must survive the derivation" + ); + + // Only the two per-method fields are overridden. + assert_eq!(limits.max_instructions, config.max_instructions); + assert_eq!( + limits.timeout_ms, + u64::try_from(config.timeout.as_millis()).unwrap() + ); + } } diff --git a/dotscope/src/deobfuscation/decryptors.rs b/dotscope/src/deobfuscation/decryptors.rs index 80676d43..1eee7de1 100644 --- a/dotscope/src/deobfuscation/decryptors.rs +++ b/dotscope/src/deobfuscation/decryptors.rs @@ -52,7 +52,10 @@ //! } //! ``` -use std::collections::HashSet; +use std::{ + collections::HashSet, + sync::atomic::{AtomicUsize, Ordering}, +}; use dashmap::{DashMap, DashSet}; @@ -99,6 +102,16 @@ pub struct DecryptorContext { /// Secondary index for O(1) successful decryption lookups, keyed by (caller, location). decrypted_index: DashSet<(Token, usize)>, + + /// Emulations performed so far in this run. + /// + /// The per-emulation instruction and wall-clock budgets bound one call; nothing bounds + /// how many calls a run makes. The decryption pass emulates once per distinct call site + /// with constant arguments, and [`Self::cache`] is keyed on `(decryptor, args)`, so + /// distinct arguments always miss — a method that calls a decryptor in a loop with a + /// varying key produces one emulation per iteration. This counter is what gives the run + /// as a whole a ceiling. + emulations_performed: AtomicUsize, } /// Record of a successfully decrypted call. @@ -130,6 +143,8 @@ pub enum FailureReason { NonConstantArgs, /// Emulation failed or timed out. EmulationFailed(String), + /// The run's total emulation budget was exhausted before this site was reached. + EmulationBudgetExhausted, /// Couldn't resolve the method target. UnresolvedTarget, /// Return value couldn't be converted to a constant. @@ -151,7 +166,9 @@ impl FailureReason { /// - `NonConstantArgs` - the arguments might become constant in later passes #[must_use] pub fn is_permanent(&self) -> bool { - !matches!(self, Self::NonConstantArgs) + // Budget exhaustion is a property of the run, not of the call site: with a larger + // budget the same site could succeed, so it must not be cached as permanent. + !matches!(self, Self::NonConstantArgs | Self::EmulationBudgetExhausted) } } @@ -163,6 +180,7 @@ impl std::fmt::Display for FailureReason { Self::UnresolvedTarget => write!(f, "unresolved call target"), Self::InvalidReturnValue => write!(f, "invalid return value"), Self::MethodNotFound => write!(f, "method not found"), + Self::EmulationBudgetExhausted => write!(f, "run emulation budget exhausted"), } } } @@ -212,6 +230,35 @@ impl DecryptorContext { Self::default() } + /// Reserves one emulation against the run budget. + /// + /// Returns `false` once `max` emulations have been performed, at which point callers + /// should stop emulating and report the remaining sites as failures rather than + /// continuing. A `max` of 0 means unlimited. + /// + /// The reservation is taken before the emulation runs, so a concurrent pair of callers + /// cannot both pass the check at the limit. + pub fn try_reserve_emulation(&self, max: usize) -> bool { + if max == 0 { + return true; + } + self.emulations_performed + .try_update(Ordering::Relaxed, Ordering::Relaxed, |n| { + if n >= max { + None + } else { + Some(n.saturating_add(1)) + } + }) + .is_ok() + } + + /// Returns how many emulations this run has performed. + #[must_use] + pub fn emulations_performed(&self) -> usize { + self.emulations_performed.load(Ordering::Relaxed) + } + /// Registers a method as a known decryptor. /// /// # Arguments @@ -938,4 +985,64 @@ mod tests { assert_eq!(ctx.decryptor_count(), 200); assert_eq!(ctx.total_decrypted(), 200); } + + /// The per-emulation instruction and wall-clock budgets bound one call; nothing bounded + /// how many calls a run makes, and the value cache is keyed on `(decryptor, args)` so a + /// varying key misses it every time. This is that ceiling. + #[test] + fn emulation_budget_stops_the_run_at_the_limit() { + let ctx = DecryptorContext::new(); + + assert!(ctx.try_reserve_emulation(2)); + assert!(ctx.try_reserve_emulation(2)); + assert!( + !ctx.try_reserve_emulation(2), + "third reservation must be refused" + ); + assert_eq!(ctx.emulations_performed(), 2); + } + + /// A limit of zero means unlimited, matching the config's documented encoding. + #[test] + fn emulation_budget_of_zero_is_unlimited() { + let ctx = DecryptorContext::new(); + + for _ in 0..1000 { + assert!(ctx.try_reserve_emulation(0)); + } + } + + /// Reservations must not race: with N threads competing for M slots, exactly M succeed. + #[test] + fn emulation_budget_is_not_oversubscribed_under_contention() { + let ctx = Arc::new(DecryptorContext::new()); + let granted = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + + let handles: Vec<_> = (0..8) + .map(|_| { + let ctx = Arc::clone(&ctx); + let granted = Arc::clone(&granted); + thread::spawn(move || { + for _ in 0..50 { + if ctx.try_reserve_emulation(100) { + granted.fetch_add(1, Ordering::Relaxed); + } + } + }) + }) + .collect(); + for h in handles { + h.join().unwrap(); + } + + assert_eq!(granted.load(Ordering::Relaxed), 100); + } + + /// Budget exhaustion is a property of the run, not the call site, so it must not be + /// cached as a permanent failure — a larger budget would let the same site succeed. + #[test] + fn budget_exhaustion_is_not_a_permanent_failure() { + assert!(!FailureReason::EmulationBudgetExhausted.is_permanent()); + assert!(FailureReason::MethodNotFound.is_permanent()); + } } diff --git a/dotscope/src/deobfuscation/engine/codegen.rs b/dotscope/src/deobfuscation/engine/codegen.rs index f3497d91..ad939d50 100644 --- a/dotscope/src/deobfuscation/engine/codegen.rs +++ b/dotscope/src/deobfuscation/engine/codegen.rs @@ -98,7 +98,7 @@ impl DeobfuscationEngine { .view() .tables() .and_then(|t| t.table::()) - .and_then(|table| table.get(rid)) + .and_then(|table| table.get(rid).ok().flatten()) .ok_or_else(|| { Error::ModificationInvalid(format!("MethodDef row {rid} not found")) })?; diff --git a/dotscope/src/deobfuscation/engine/pipeline.rs b/dotscope/src/deobfuscation/engine/pipeline.rs index eaaa77bc..56e394d6 100644 --- a/dotscope/src/deobfuscation/engine/pipeline.rs +++ b/dotscope/src/deobfuscation/engine/pipeline.rs @@ -571,7 +571,9 @@ impl<'a> PipelineRun<'a> { .map_err(|e| Error::SsaError(e.0))?; } - let ssa_call_graph = ctx.build_ssa_call_graph(); + // Effective, not SSA-only: this graph decides what gets deleted, and a method whose + // body never converted to SSA must read as "unknown callees", not "no callees". + let ssa_call_graph = ctx.build_effective_call_graph(); let mut merged_cleanup = build_cleanup_request( self.engine, &ctx, diff --git a/dotscope/src/deobfuscation/passes/bitmono/strings.rs b/dotscope/src/deobfuscation/passes/bitmono/strings.rs index e65f79ca..a8306f9e 100644 --- a/dotscope/src/deobfuscation/passes/bitmono/strings.rs +++ b/dotscope/src/deobfuscation/passes/bitmono/strings.rs @@ -135,7 +135,7 @@ impl StringDecryptionPass { let salt_bytes = get_field_rva_data(assembly, salt_token.row(), field_rva_map)?; let key_bytes = get_field_rva_data(assembly, key_token.row(), field_rva_map)?; - let derived = derive_key_iv(&key_bytes, &salt_bytes, &self.crypto_params); + let derived = derive_key_iv(&key_bytes, &salt_bytes, &self.crypto_params)?; self.key_cache .lock() @@ -420,6 +420,13 @@ fn build_field_rva_map(assembly: &CilObject) -> HashMap { // Build direct FieldRVA map (backing_field_rid -> (rva, size)) for row in fieldrva_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if row.rva == 0 { continue; } @@ -504,13 +511,13 @@ mod tests { let key = [0u8; 8]; let params = CryptoParameters::default(); - let (aes_key, aes_iv) = derive_key_iv(&key, &salt, ¶ms); + let (aes_key, aes_iv) = derive_key_iv(&key, &salt, ¶ms).unwrap(); assert_eq!(aes_key.len(), 32, "AES-256 key should be 32 bytes"); assert_eq!(aes_iv.len(), 16, "AES IV should be 16 bytes"); // Verify deterministic output with all-zero inputs - let (aes_key2, aes_iv2) = derive_key_iv(&key, &salt, ¶ms); + let (aes_key2, aes_iv2) = derive_key_iv(&key, &salt, ¶ms).unwrap(); assert_eq!(aes_key, aes_key2, "Key derivation should be deterministic"); assert_eq!(aes_iv, aes_iv2, "IV derivation should be deterministic"); } @@ -520,7 +527,7 @@ mod tests { let salt = [0u8; 8]; let key = [0u8; 8]; let params = CryptoParameters::default(); - let (aes_key, aes_iv) = derive_key_iv(&key, &salt, ¶ms); + let (aes_key, aes_iv) = derive_key_iv(&key, &salt, ¶ms).unwrap(); // Encrypt a test string using the shared crypto utility let original = "Hello, BitMono!"; @@ -538,7 +545,7 @@ mod tests { let salt = [0u8; 8]; let key = [0u8; 8]; let params = CryptoParameters::default(); - let (aes_key, aes_iv) = derive_key_iv(&key, &salt, ¶ms); + let (aes_key, aes_iv) = derive_key_iv(&key, &salt, ¶ms).unwrap(); let result = decrypt_string(&[], &aes_key, &aes_iv).unwrap(); assert_eq!(result, ""); @@ -549,7 +556,7 @@ mod tests { let salt = [0u8; 8]; let key = [0u8; 8]; let params = CryptoParameters::default(); - let (aes_key, aes_iv) = derive_key_iv(&key, &salt, ¶ms); + let (aes_key, aes_iv) = derive_key_iv(&key, &salt, ¶ms).unwrap(); // Non-multiple-of-16 data should fail let result = decrypt_string(&[1, 2, 3], &aes_key, &aes_iv); diff --git a/dotscope/src/deobfuscation/passes/decryption.rs b/dotscope/src/deobfuscation/passes/decryption.rs index c4072084..d9c0e658 100644 --- a/dotscope/src/deobfuscation/passes/decryption.rs +++ b/dotscope/src/deobfuscation/passes/decryption.rs @@ -88,7 +88,7 @@ use crate::{ EmValue, EmulationError, EmulationOutcome, EmulationProcess, EmulationThread, StepResult, }, metadata::{ - tables::TypeRefRaw, + tables::{skip_unreadable, TypeRefRaw}, token::Token, typesystem::{CilFlavor, CilPrimitive, CilPrimitiveKind, PointerSize}, }, @@ -125,8 +125,11 @@ pub struct DecryptionPass { decryptors: Arc, /// State machine providers for order-dependent decryption. statemachine_providers: Arc>>, - /// Maximum instructions per emulation call. - emulation_max_instructions: u64, + /// Ceiling on how many emulations the whole run may perform, or 0 for unlimited. + /// + /// The per-emulation instruction and wall-clock budgets are applied by the template pool + /// when it forks; this is the run-level budget they do not provide. + max_emulations: usize, } /// Owned CFG analysis info, storing dominator tree and predecessors. @@ -173,7 +176,7 @@ impl DecryptionPass { template_pool: ctx.template_pool.get().cloned(), decryptors: Arc::clone(&ctx.decryptors), statemachine_providers: Arc::clone(&ctx.statemachine_providers), - emulation_max_instructions: ctx.config.emulation.max_instructions, + max_emulations: ctx.config.emulation.max_emulations, } } @@ -261,6 +264,12 @@ impl DecryptionPass { method: Token, args: &[ConstValue], ) -> (Option, Option) { + // Reserve against the run budget before doing any work. Taking the reservation up + // front means two concurrent callers cannot both pass the check on the last slot. + if !self.decryptors.try_reserve_emulation(self.max_emulations) { + return (None, Some(FailureReason::EmulationBudgetExhausted)); + } + // Fork from the template process - O(1) due to CoW semantics. // The template is created lazily on first call with expensive PE loading, // hooks, etc. Subsequent calls share memory via structural sharing. @@ -397,7 +406,7 @@ impl DecryptionPass { let strings = asm.strings()?; let type_refs = tables.table::()?; - for row in type_refs { + for row in type_refs.iter().filter_map(skip_unreadable) { let row_name = strings.get(row.type_name as usize).ok()?; let row_ns = strings.get(row.type_namespace as usize).ok()?; if row_name == name && row_ns == namespace { diff --git a/dotscope/src/deobfuscation/passes/delegates.rs b/dotscope/src/deobfuscation/passes/delegates.rs index b3f360fe..8e676cdf 100644 --- a/dotscope/src/deobfuscation/passes/delegates.rs +++ b/dotscope/src/deobfuscation/passes/delegates.rs @@ -24,8 +24,13 @@ //! //! # Safety //! -//! If warmup fails or a delegate cannot be resolved, those call sites are -//! silently skipped — no false positives are possible. +//! If warmup fails or a delegate cannot be resolved, those call sites are skipped. +//! +//! A call site is also skipped unless the delegate is **open** (no bound receiver) and has +//! **exactly one** invocation-list entry. Neither restriction is incidental: rewriting a +//! multicast delegate to a single call deletes its other targets from the output, and +//! inlining a closed delegate drops the receiver the callee's signature requires. Both +//! skips are logged at debug level so the analyst can see the proxy was left intact. use std::{ collections::{HashMap, HashSet}, @@ -40,7 +45,7 @@ use crate::{ assembly::{FlowType, Instruction, Operand}, compiler::{CompilerContext, EventKind, ModificationScope, SsaPass}, deobfuscation::{utils::build_def_map, EmulationTemplatePool, ProcessCell}, - emulation::{tokens, EmValue, EmulationProcess, HeapObject}, + emulation::{tokens, DelegateEntry, EmValue, EmulationProcess, HeapObject}, metadata::token::Token, CilObject, Result, }; @@ -90,6 +95,53 @@ pub struct DelegateProxyResolutionPass { processed_methods: DashSet, } +/// Why a delegate could not be inlined into a direct call. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum NotInlinable { + /// The invocation list is empty, so there is no target to call. + NoTarget, + /// More than one entry: invoking the delegate invokes all of them. + Multicast(usize), + /// The entry carries a bound receiver that a direct call would have to pass. + ClosedInstance, +} + +impl std::fmt::Display for NotInlinable { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::NoTarget => write!(f, "delegate has no invocation targets"), + Self::Multicast(n) => write!(f, "multicast with {n} entries"), + Self::ClosedInstance => write!(f, "closed delegate with a bound receiver"), + } + } +} + +/// Returns the single entry a delegate can be inlined to, or why it cannot be. +/// +/// Only an **open, single-target** delegate qualifies, and neither restriction is incidental: +/// +/// - A multicast delegate invokes *every* entry. Rewriting the call site to one of them +/// deletes the others from the output — precisely the primitive an obfuscator wants, since +/// the surviving entry can be the benign one. (The emulator's own dispatcher starts at +/// `first()`, so picking `last()` did not even agree with the entry that actually runs +/// first.) +/// - A closed delegate carries its receiver in `entry.target`. The rewrite drops the delegate +/// operand and emits the remaining arguments, so inlining one yields a call with one fewer +/// operand than the callee's signature requires — invalid IL, not merely wrong output. +/// +/// Supporting the closed case later means materialising the receiver and prepending it to the +/// argument list, not dropping it. +pub(crate) fn inlinable_entry( + invocation_list: &[DelegateEntry], +) -> std::result::Result<&DelegateEntry, NotInlinable> { + match invocation_list { + [] => Err(NotInlinable::NoTarget), + [entry] if entry.target.is_none() => Ok(entry), + [_] => Err(NotInlinable::ClosedInstance), + many => Err(NotInlinable::Multicast(many.len())), + } +} + impl DelegateProxyResolutionPass { /// Creates a new delegate proxy resolution pass with pre-computed findings. /// @@ -215,7 +267,17 @@ impl DelegateProxyResolutionPass { invocation_list, .. } = obj { - if let Some(entry) = invocation_list.last() { + let entry = match inlinable_entry(&invocation_list) { + Ok(entry) => entry, + Err(reason) => { + log::debug!( + "Leaving delegate proxy for field {field_token} intact: {reason}" + ); + continue; + } + }; + + { let method_token = entry.method_token; // Resolve synthetic DynamicMethod tokens to real metadata tokens. @@ -608,3 +670,65 @@ impl SsaPass for DelegateProxyResolutionPass { self.lazy_process.clear().map_err(Into::into) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::emulation::HeapRef; + + fn open_entry(token: u32) -> DelegateEntry { + DelegateEntry { + target: None, + method_token: Token::new(token), + } + } + + fn closed_entry(token: u32) -> DelegateEntry { + DelegateEntry { + target: Some(HeapRef::new(1)), + method_token: Token::new(token), + } + } + + /// The only shape that can be inlined into a direct call. + #[test] + fn a_single_open_target_is_inlinable() { + let list = vec![open_entry(0x0600_0001)]; + + assert_eq!( + inlinable_entry(&list).map(|e| e.method_token).ok(), + Some(Token::new(0x0600_0001)) + ); + } + + /// Invoking a multicast delegate invokes every entry, so rewriting the site to one of + /// them deletes the others from the output — an anti-analysis primitive, since the + /// surviving entry can be the benign one. + #[test] + fn a_multicast_delegate_is_left_intact() { + let list = vec![open_entry(0x0600_0001), open_entry(0x0600_0002)]; + + assert_eq!( + inlinable_entry(&list).err(), + Some(NotInlinable::Multicast(2)) + ); + } + + /// A closed delegate carries its receiver in `target`. The rewrite drops the delegate + /// operand, so inlining one emits a call with one fewer operand than the callee's + /// signature requires — invalid IL. + #[test] + fn a_closed_delegate_is_left_intact() { + let list = vec![closed_entry(0x0600_0001)]; + + assert_eq!( + inlinable_entry(&list).err(), + Some(NotInlinable::ClosedInstance) + ); + } + + #[test] + fn an_empty_invocation_list_has_nothing_to_inline() { + assert_eq!(inlinable_entry(&[]).err(), Some(NotInlinable::NoTarget)); + } +} diff --git a/dotscope/src/deobfuscation/passes/native.rs b/dotscope/src/deobfuscation/passes/native.rs index bd9643e4..68a374fc 100644 --- a/dotscope/src/deobfuscation/passes/native.rs +++ b/dotscope/src/deobfuscation/passes/native.rs @@ -231,7 +231,7 @@ impl NativeMethodConversionPass { .view() .tables() .and_then(|t| t.table::()) - .and_then(|table| table.get(rid)) + .and_then(|table| table.get(rid).ok().flatten()) .ok_or_else(|| Error::X86Error(format!("MethodDef row {rid} not found for token")))?; // Verify this is actually a native method diff --git a/dotscope/src/deobfuscation/passes/unflattening/mod.rs b/dotscope/src/deobfuscation/passes/unflattening/mod.rs index 216f35c0..a57dc649 100644 --- a/dotscope/src/deobfuscation/passes/unflattening/mod.rs +++ b/dotscope/src/deobfuscation/passes/unflattening/mod.rs @@ -20,27 +20,26 @@ //! instruction backwards through the SSA to find the PHI node carrying state //! 3. **Dispatcher Classification** ([`dispatcher`]): Determine type (switch with //! optional XOR/modulo transform, if-else chain, computed jump) -//! 4. **Tracing** ([`tracer`]): Evaluate the method from entry, following state -//! transitions through the dispatcher while forking at user branches to build -//! a tree of all execution paths -//! 5. **Reconstruction** ([`reconstruction`]): Extract a patch plan from the trace -//! tree (redirects, block clones, state instruction removal) and apply it to -//! the SSA, eliminating the dispatcher +//! 4. **Resolution** ([`resolve`]): Read the state travelling along each edge into +//! the dispatcher out of the SSA phi graph, and rewire that edge straight to +//! the block the dispatcher would have selected //! //! # Design Principles //! //! - **Structure-based detection**: Uses graph properties, not opcode patterns -//! - **Concrete evaluation**: SSA evaluator resolves state transitions at trace -//! time — no solver needed for standard arithmetic encodings -//! - **Graceful degradation**: Works partially even when full recovery isn't possible +//! - **Edges, not paths**: The state on an edge is a property of the edge, which +//! SSA already records per predecessor. Recovering it costs one pass over the +//! dispatcher's edges rather than an enumeration of execution paths, whose +//! number grows exponentially with the method's conditionals +//! - **Graceful degradation**: An edge whose state cannot be determined keeps +//! routing through the dispatcher, so partial recovery is always safe //! - **Clean separation**: Each phase is isolated for testability mod detection; mod dispatcher; -mod reconstruction; +mod resolve; mod spill; mod statevar; -mod tracer; use std::{collections::HashMap, sync::Arc}; @@ -48,44 +47,30 @@ use dashmap::DashSet; pub use detection::CffDetector; pub use dispatcher::Dispatcher; use rayon::prelude::*; -pub use reconstruction::{apply_patch_plan, extract_patch_plan, merge_patch_plans}; use crate::{ analysis::{CilTarget, MethodRef, SsaFunction}, compiler::{CompilerContext, PassCapability, SsaPass}, - deobfuscation::{ - config::DetectionWeights, - context::AnalysisContext, - passes::unflattening::tracer::{trace_for_dispatcher, TracedDispatcher}, - }, + deobfuscation::{config::DetectionWeights, context::AnalysisContext}, metadata::{token::Token, typesystem::PointerSize}, - CilObject, }; -/// High-level API: Unflatten a method using tree-based tracing and patching. +/// High-level API: unflatten a method by rewiring its dispatcher edges. /// -/// This is the main entry point for the trace-based unflattening approach. -/// It detects ALL CFF dispatchers in the method (there may be independent -/// dispatchers in different exception handler regions), traces each one -/// independently, merges the patch plans, and applies them all at once -/// before a single `rebuild_ssa()`. This avoids the block renumbering -/// corruption that occurs when dispatchers are processed iteratively -/// with intermediate SSA rebuilds. +/// Detects every CFF dispatcher in the method — there may be independent ones +/// per exception handler region — resolves all of their edges against the +/// unmodified function, and applies the whole set at once so a single +/// `rebuild_ssa()` sees the final graph. /// /// # Arguments /// /// * `ssa` - The SSA function to unflatten. -/// * `config` - Configuration controlling tracing limits and behavior. -/// * `assembly` - Optional assembly for predicate method resolution during tracing. +/// * `config` - Detection thresholds and target pointer size. /// /// # Returns -/// - `Some(ssa)` if unflattening succeeded (patched clone) -/// - `None` if method doesn't appear to be CFF-protected -pub fn unflatten( - ssa: &SsaFunction, - config: &UnflattenConfig, - assembly: Option<&CilObject>, -) -> Option { +/// - `Some(ssa)` if at least one edge was rewired (patched clone) +/// - `None` if the method doesn't appear to be CFF-protected +pub fn unflatten(ssa: &SsaFunction, config: &UnflattenConfig) -> Option { let mut detector = CffDetector::with_config(ssa, config); let dispatchers: Vec<_> = detector .detect_all_dispatchers() @@ -93,7 +78,7 @@ pub fn unflatten( .filter(|d| d.confidence >= config.min_confidence) .collect(); - unflatten_with_dispatchers(ssa, config, assembly, dispatchers) + unflatten_with_dispatchers(ssa, dispatchers) } /// Unflatten a method using pre-detected dispatchers. @@ -109,114 +94,87 @@ pub fn unflatten( /// # Arguments /// /// * `ssa` - The SSA function to unflatten. -/// * `config` - Configuration controlling tracing limits and behavior. -/// * `assembly` - Optional assembly for predicate method resolution during tracing. /// * `dispatchers` - Pre-detected dispatchers (already confidence-filtered). /// /// # Returns -/// - `Some(ssa)` if unflattening succeeded (patched clone) -/// - `None` if no dispatchers provided or tracing produced no changes +/// - `Some(ssa)` if at least one edge was rewired (patched clone) +/// - `None` if no dispatchers were given or none of their edges resolved pub fn unflatten_with_dispatchers( ssa: &SsaFunction, - config: &UnflattenConfig, - assembly: Option<&CilObject>, dispatchers: Vec, ) -> Option { if dispatchers.is_empty() { return None; } - // Trace dispatchers in parallel and extract patch plans. - // Each dispatcher trace is independent (shared &SsaFunction, own evaluator). - // Pass other dispatcher block indices so forks at foreign dispatchers - // don't consume tree depth budget (Problem A from §15.5). - let all_dispatcher_blocks: Vec = dispatchers.iter().map(|d| d.block).collect(); - - let plans: Vec<_> = dispatchers + // Resolve each dispatcher's edges against the unmodified function. This is + // a pure read, so the dispatchers are independent and can run in parallel. + let per_dispatcher: Vec> = dispatchers .par_iter() - .filter_map(|d| { - let traced = TracedDispatcher { - block: d.block, - switch_var: d.switch_var, - targets: d.cases.clone(), - default: d.default, - state_var: d.state_phi, - initial_state: d.initial_state, - }; - - let others: Vec = all_dispatcher_blocks - .iter() - .copied() - .filter(|&b| b != d.block) - .collect(); - let tree = trace_for_dispatcher(ssa, config, assembly, traced, &others); - - tree.dispatcher.as_ref()?; - - extract_patch_plan(&tree, ssa).filter(|plan| plan.state_transitions_removed > 0) + .map(|d| { + let (rewires, stats) = resolve::resolve_dispatch_edges(ssa, d); + log::debug!( + "CFF resolve b{}: {} edge(s) resolved, {} unresolved, {} conflicting \ + (table {} incl. {} overflow)", + d.block, + stats.resolved, + stats.unresolved, + stats.conflicts, + stats.table_size, + stats.overflow_entries + ); + if stats.unresolved > 0 { + log::debug!( + "CFF reasons b{}: no_target={} not_constant={} impure_chain={}", + d.block, + stats.reasons.no_target, + stats.reasons.not_constant, + stats.reasons.impure_chain + ); + } + rewires }) .collect(); - if plans.is_empty() { - return None; + let (rewires, conflicts) = resolve::merge_rewires(per_dispatcher); + if conflicts > 0 { + log::debug!("CFF resolve: {conflicts} edge(s) dropped across dispatchers"); } - - // Step 3: Merge all patch plans into a single combined plan - let merged = merge_patch_plans(plans); - - if merged.state_transitions_removed == 0 { + if rewires.is_empty() { return None; } - // Step 4: Clone the SSA, give every cross-block value a storage location, - // then apply the combined patches once. + // Give every cross-block value a storage location before the CFG changes. // - // The promotion must happen before any terminator is rewired. Patching - // invalidates the phi nodes that record which SSA names denote the same - // value across a merge, and for edges the patch creates no phi ever recorded - // anything at all — so a value carried between blocks in a stack temporary - // becomes unreconstructible the moment the CFG changes. Promoting those - // values to local slots first gives `rebuild_ssa` a storage location to - // resolve them against, which is what makes the rebuild well-defined for an - // arbitrarily rewired CFG. + // Rewiring replaces `case -> dispatcher -> next` with `case -> next`, so a + // value that used to be merged by a phi at the dispatcher now arrives along + // an edge no phi ever described. `rebuild_ssa` can reconstruct such a value + // only if it has an identity independent of control flow — a local slot. + // Promoting first is what makes the rebuild well-defined for the rewired + // graph. let mut patched = ssa.clone(); let spilled = spill::promote_cross_block_values(&mut patched); if spilled > 0 { log::debug!("Promoted {spilled} cross-block value(s) to locals before unflattening"); } - let _result = apply_patch_plan(&mut patched, &merged); - // Note: we do NOT reject based on dispatcher_still_needed. With multiple - // dispatchers, some may be fully resolved while others are only partial - // (e.g., handler CFF that depends on the outer CFF state). apply_patch_plan - // already handles this correctly: it only clears dispatchers that are fully - // resolved, leaving partial ones intact for later passes or as harmless - // residual. Rejecting the entire result would prevent the fully-resolved - // main body CFF from being applied. + let applied = resolve::apply_rewires(&mut patched, &rewires); + if applied == 0 { + return None; + } + let cleared = resolve::clear_unreachable(&mut patched); + log::debug!("CFF: rewired {applied} edge(s) past the dispatcher, {cleared} block(s) now dead"); + // Dispatchers whose edges did not all resolve keep their remaining + // predecessors and stay in the graph — partial recovery, not corruption. + // The state constants and the dispatcher itself become unreachable once + // every edge is rewired, and the ordinary dead-code passes remove them. Some(patched) } /// Configuration for the CFF reconstruction pass. #[derive(Debug, Clone)] pub struct UnflattenConfig { - /// Maximum states to explore before giving up. - /// - /// CFF typically has dozens to hundreds of states. If we exceed this - /// limit, the method likely has unusual structure or isn't CFF. - pub max_states: usize, - - /// Reserved: enable solver for complex state encodings. - /// - /// Placeholder for future solver integration. Currently unused — state - /// resolution relies entirely on concrete evaluation via the SSA evaluator. - pub enable_solver: bool, - - /// Reserved: maximum solver time per query in milliseconds. - /// - /// Placeholder for future solver integration. Currently unused. - pub solver_timeout_ms: u64, - /// Minimum confidence score to attempt unflattening (0.0 - 1.0). /// /// Detection assigns confidence based on how strongly the method @@ -224,12 +182,6 @@ pub struct UnflattenConfig { /// produce false positives. pub min_confidence: f64, - /// Maximum depth for constant propagation evaluation. - /// - /// Limits how deeply we trace through SSA definitions when - /// evaluating state values. - pub max_eval_depth: usize, - /// Maximum BFS depth for back-edge transitive reachability check. /// /// Used in confidence scoring to determine if case blocks can @@ -239,36 +191,17 @@ pub struct UnflattenConfig { /// Confidence scoring weights for CFF dispatcher detection. pub confidence_weights: DetectionWeights, - /// Maximum number of blocks to visit during tracing. - /// - /// Prevents infinite loops when tracing through the CFG. If exceeded, - /// tracing stops with a `StopReason::MaxVisitsExceeded`. - pub max_block_visits: usize, - - /// Maximum nesting depth for the trace tree. - /// - /// Limits how deeply nested the tree can become from forking at user - /// branches. Prevents exponential tree growth in methods with many - /// independent conditionals. - pub max_tree_depth: usize, - /// Target pointer size for SSA evaluation. /// - /// Derived from the PE header. Used by the SSA evaluator for - /// pointer-sized arithmetic during tracing. + /// Derived from the PE header. Used when detection evaluates the + /// dispatcher's initial state. pub pointer_size: PointerSize, } impl Default for UnflattenConfig { fn default() -> Self { Self { - max_states: 1000, - enable_solver: true, - solver_timeout_ms: 100, min_confidence: 0.6, - max_eval_depth: 30, - max_block_visits: 50000, - max_tree_depth: 500, pointer_size: PointerSize::Bit32, max_backedge_depth: 10, confidence_weights: DetectionWeights::default(), @@ -279,46 +212,33 @@ impl Default for UnflattenConfig { impl UnflattenConfig { /// Creates a configuration optimized for ConfuserEx samples. /// - /// ConfuserEx uses predictable arithmetic encoding that constant - /// propagation can always resolve, so we can be more aggressive - /// with lower state limits and no solver. + /// ConfuserEx's arithmetic encoding is always resolvable, so detection can + /// accept a lower confidence without risking wasted work. /// /// # Returns /// - /// An `UnflattenConfig` with reduced limits and the solver disabled. + /// An `UnflattenConfig` with a reduced confidence threshold. #[must_use] pub fn confuserex() -> Self { Self { - max_states: 500, - enable_solver: false, - solver_timeout_ms: 50, min_confidence: 0.5, - max_eval_depth: 25, - max_block_visits: 5000, - max_tree_depth: 75, ..Self::default() } } /// Creates a configuration for aggressive analysis. /// - /// Useful for heavily obfuscated code where detection may have - /// lower confidence. Uses higher limits and a lower confidence - /// threshold to catch more CFF patterns at the cost of longer - /// analysis time. + /// Useful for heavily obfuscated code where detection may have lower + /// confidence. Accepts weaker matches and looks further for back edges, at + /// the cost of attempting more methods. /// /// # Returns /// - /// An `UnflattenConfig` with increased limits and the solver enabled. + /// An `UnflattenConfig` with a lower confidence threshold. #[must_use] pub fn aggressive() -> Self { Self { - max_states: 2000, - solver_timeout_ms: 200, min_confidence: 0.4, - max_eval_depth: 50, - max_block_visits: 20000, - max_tree_depth: 150, max_backedge_depth: 15, ..Self::default() } @@ -335,8 +255,13 @@ impl UnflattenConfig { /// dispatchers directly instead of re-running detection. This avoids duplicate /// structural analysis (dominance, SCCs, confidence scoring) for methods /// already analyzed during the detection phase's SSA pass. +/// +/// The pass itself takes no configuration: it consumes dispatchers the detection +/// phase already found and scored, and resolving their edges has no thresholds +/// to tune. [`UnflattenConfig`] belongs to detection, which is where it is now +/// applied. +#[derive(Default)] pub struct CffReconstructionPass { - config: UnflattenConfig, /// Successfully unflattened dispatcher methods (shared with deob engine). unflattened_dispatchers: Arc>, /// All detected dispatcher methods (shared with deob engine). @@ -346,17 +271,6 @@ pub struct CffReconstructionPass { pre_detected: HashMap>, } -impl Default for CffReconstructionPass { - fn default() -> Self { - Self { - config: UnflattenConfig::default(), - unflattened_dispatchers: Arc::new(DashSet::new()), - dispatchers: Arc::new(DashSet::new()), - pre_detected: HashMap::new(), - } - } -} - impl CffReconstructionPass { /// Creates a new CFF reconstruction pass from an analysis context. /// @@ -366,16 +280,13 @@ impl CffReconstructionPass { /// # Arguments /// /// * `ctx` - The analysis context providing shared dispatcher tracking sets. - /// * `config` - Configuration controlling tracing limits, confidence thresholds, - /// and solver usage. /// /// # Returns /// /// A new `CffReconstructionPass` ready for pipeline execution. #[must_use] - pub fn new(ctx: &AnalysisContext, config: UnflattenConfig) -> Self { + pub fn new(ctx: &AnalysisContext) -> Self { Self { - config, unflattened_dispatchers: Arc::clone(&ctx.unflattened_dispatchers), dispatchers: Arc::clone(&ctx.dispatchers), pre_detected: HashMap::new(), @@ -420,6 +331,10 @@ impl SsaPass for CffReconstructionPass { "Recovers original control flow from flattened state machine patterns" } + fn repairs_ssa(&self) -> bool { + true + } + fn provides(&self) -> &[PassCapability] { &[PassCapability::RestoredControlFlow] } @@ -444,14 +359,8 @@ impl SsaPass for CffReconstructionPass { method: &MethodRef, host: &CompilerContext, ) -> analyssa::Result { - let assembly_arc = host - .assembly() - .ok_or_else(|| analyssa::Error::new("CffReconstructionPass requires an assembly"))?; - let assembly: &CilObject = &assembly_arc; let ctx = host; let method_token = method.0; - let mut config = self.config.clone(); - config.pointer_size = PointerSize::from_is_64bit(assembly.file().pe().is_64bit); // Use pre-detected dispatcher block indices from detect_ssa phase, but // refresh variable IDs from the current SSA. Earlier passes (opaque field @@ -463,9 +372,31 @@ impl SsaPass for CffReconstructionPass { .map(|pre| pre.iter().filter_map(|d| d.refresh(ssa)).collect()) .unwrap_or_default(); - match unflatten_with_dispatchers(ssa, &config, Some(assembly), dispatchers) { + let started = std::time::Instant::now(); + match unflatten_with_dispatchers(ssa, dispatchers) { Some(mut patched) => { - patched.rebuild_ssa()?; + // Rewiring can leave a value defined on a path the bypass no + // longer takes. The guards in `resolve` exist to prevent that, + // and `rebuild_ssa` is what proves it: it re-derives the phi + // graph and rejects a definition that no longer reaches its + // uses. Treat a rejection as "this method cannot be unflattened + // safely" and leave it flattened — an unrecovered method is a + // gap in the analysis, whereas propagating the error abandons + // every remaining method in the assembly as well. + if let Err(error) = patched.rebuild_ssa() { + log::warn!( + "CFF {:08x}: rewired form failed SSA validation ({error}); \ + leaving the method flattened", + method_token.value() + ); + return Ok(false); + } + log::debug!( + "CFF {:08x}: {} block(s), {}ms", + method_token.value(), + patched.blocks().len(), + started.elapsed().as_millis() + ); *ssa = patched; diff --git a/dotscope/src/deobfuscation/passes/unflattening/reconstruction.rs b/dotscope/src/deobfuscation/passes/unflattening/reconstruction.rs deleted file mode 100644 index b32fa25b..00000000 --- a/dotscope/src/deobfuscation/passes/unflattening/reconstruction.rs +++ /dev/null @@ -1,1780 +0,0 @@ -//! CFG reconstruction from trace trees. -//! -//! This module handles the reconstruction phase of control flow unflattening: -//! converting a trace tree (built by [`super::tracer`]) back into clean SSA. -//! -//! # Overview -//! -//! The reconstruction process involves: -//! -//! 1. **Patch Plan Extraction**: Analyze the trace tree to determine: -//! - Which blocks should redirect to which targets (bypassing the dispatcher) -//! - Which variables are state-tainted (CFF machinery to remove) -//! - Which blocks need cloning (merge points with different targets) -//! -//! 2. **Patch Application**: Modify the SSA in place: -//! - Redirect block terminators from dispatcher to actual targets -//! - Clone merge blocks when multiple paths converge with different targets -//! - Filter out state-tainted instructions -//! - Clear the dispatcher block -//! -//! 3. **SSA Rebuild**: After patching, `SsaFunction::rebuild_ssa()` is called to -//! reconstruct proper PHI nodes for the new CFG structure. -//! -//! # Algorithm -//! -//! The key insight is that the trace tree captures the *actual* execution flow -//! through the CFF state machine. Each `StateTransition` node represents a -//! state change that we want to eliminate by redirecting control flow directly. -//! -//! For example, if the tree shows: -//! ```text -//! Block 5 -> StateTransition(target=7) -> Block 7 -> ... -//! ``` -//! -//! We patch block 5 to jump directly to block 7, bypassing the dispatcher. -//! -//! ## Merge Point Handling -//! -//! When multiple paths through the tree converge at the same block but with -//! different next targets, we need to clone the block: -//! -//! ```text -//! Path A: Block 5 -> merge_block -> target 7 -//! Path B: Block 6 -> merge_block -> target 9 -//! ``` -//! -//! We keep `merge_block` for path A and clone it as `merge_block'` for path B. - -use std::collections::{BTreeMap, BTreeSet}; - -use analyssa::BitSet; -use rustc_hash::{FxHashMap, FxHashSet}; - -use crate::{ - analysis::{DefSite, PhiOperand, SsaBlock, SsaFunction, SsaInstruction, SsaOp, SsaVarId}, - deobfuscation::passes::unflattening::tracer::{TraceNode, TraceTerminator, TraceTree}, -}; - -/// Maximum dispatcher-phi hops followed when recovering a redirected edge's value. -/// -/// Chains longer than this are pathological; giving up is safe because the caller -/// then leaves the operand absent rather than guessing. -const MAX_PHI_RESOLUTION_HOPS: usize = 8; - -type PhiOperands = Vec<(usize, SsaVarId)>; - -/// Snapshot of every block's phi nodes: `block -> (phi result -> operands)`. -/// -/// Captured before the CFG is rewired, so edges the patch removes can still be -/// consulted when recovering values for the edges it creates. -type PhiSnapshot = BTreeMap>; -type BlockPhiData = Vec<(usize, Vec<(SsaVarId, PhiOperands)>)>; - -/// Result of unflattening a CFG. -#[derive(Debug)] -pub struct ReconstructionResult { - /// Number of state transitions eliminated. - pub state_transitions_removed: usize, - /// Number of user branches preserved. - pub user_branches_preserved: usize, - /// Number of blocks in the function. - pub block_count: usize, - /// Whether the dispatcher block is still needed by unresolved blocks. - /// - /// When `true`, the unflattening was only partial — some blocks still - /// route through the dispatcher because their dispatch values couldn't - /// be resolved. The caller should skip applying the partial result to - /// avoid corrupting SSA phi nodes at the dispatcher merge point. - pub dispatcher_still_needed: bool, -} - -/// A plan for patching the SSA to remove CFF. -#[derive(Debug)] -pub struct PatchPlan { - /// Dispatcher block indices (one per CFF dispatcher in the method). - pub dispatcher_blocks: Vec, - - /// Variables that are state-related (CFF machinery). - pub state_tainted: BitSet, - - /// Block redirects: (source_block, new_target). - /// These blocks currently jump to dispatcher but should jump to new_target. - pub redirects: Vec<(usize, usize)>, - - /// Source blocks whose redirects originated from StateTransition nodes. - /// Redirects from LoopBack nodes (which may be cross-contamination from - /// exploring other dispatchers' switches as user switches) are NOT in - /// this set. Used to filter out cross-contamination in multi-dispatcher - /// methods. - state_transition_sources: BTreeSet, - - /// Blocks that need cloning: merge_block -> [(predecessor, target), ...] - /// These are blocks where multiple paths converge with different targets. - /// We clone the merge block for each path. - pub clone_requests: BTreeMap>, - - /// Blocks in execution order (for debugging/verification). - pub execution_order: Vec, - - /// Membership index for [`execution_order`](Self::execution_order). - /// - /// The order itself has to stay a `Vec`, but it is appended to once per - /// block of every node in the trace tree — millions of times on a - /// flattened method — and each append asked whether the block was already - /// present by scanning the whole list. - execution_order_seen: FxHashSet, - - /// Source block to redirect target, mirroring - /// [`redirects`](Self::redirects). - /// - /// Conflict detection looks up a source on every redirect, which is the - /// same per-node cost as the execution order above. - redirect_targets: FxHashMap, - - /// Branch-collapse requests: `source_block -> new_target`. - /// - /// When a state transition's path enters the dispatcher through a - /// BranchCmp/Branch overflow check (e.g. NETReactor `bcmp state == Const` - /// on the dispatcher's default fall-through chain), the traced arm is the - /// only live path in the unflattened graph. Apply replaces the source's - /// BranchCmp/Branch terminator with a `Jump { target: new_target }` so the - /// dead arm is eliminated while the live arm's user-code blocks still - /// execute sequentially before the separate redirect rewires the final - /// jump-to-dispatcher block. - pub(crate) branch_collapses: BTreeMap, - - /// Number of state transitions removed. - pub state_transitions_removed: usize, - - /// Number of user branches preserved. - pub user_branches_preserved: usize, -} - -impl PatchPlan { - fn new(dispatcher_block: usize, state_tainted: BitSet) -> Self { - Self { - dispatcher_blocks: vec![dispatcher_block], - state_tainted, - redirects: Vec::new(), - state_transition_sources: BTreeSet::new(), - clone_requests: BTreeMap::new(), - execution_order: Vec::new(), - execution_order_seen: FxHashSet::default(), - redirect_targets: FxHashMap::default(), - branch_collapses: BTreeMap::new(), - state_transitions_removed: 0, - user_branches_preserved: 0, - } - } - - fn add_branch_collapse(&mut self, source: usize, target: usize) { - if source == target { - return; - } - // First wins on conflict — multiple traces may visit the same overflow - // check through different state-match arms. Keeping the first avoids - // thrashing the terminator. - self.branch_collapses.entry(source).or_insert(target); - } - - /// Returns true if the given block is one of the dispatcher blocks. - pub fn is_dispatcher_block(&self, block: usize) -> bool { - self.dispatcher_blocks.contains(&block) - } - - fn add_redirect(&mut self, source: usize, target: usize, predecessor: Option) { - // Self-redirects are always wrong — they create infinite loops. - // They occur when the evaluator can't resolve the next CFF state - // (e.g., handler CFF where the second dispatch depends on values - // the evaluator lost track of). Skip them silently. - if source == target { - return; - } - // Check for conflict: same source with different target - if let Some(&existing_target) = self.redirect_targets.get(&source) { - if existing_target != target { - // Conflict! This source is a merge point needing cloning. - // Add both the existing and new paths to clone requests. - let entry = self.clone_requests.entry(source).or_default(); - - // Ensure the original (first) redirect's path is present. - // The first redirect may have been added without a predecessor - // (e.g., from the root-level trace where external_predecessor - // is None), so its entry wasn't recorded in clone_requests. - // The first path keeps the original block during cloning (its - // predecessor isn't used for redirect, only its target matters). - if !entry.iter().any(|(_, t)| *t == existing_target) { - entry.push((usize::MAX, existing_target)); - } - - if let Some(pred) = predecessor { - // Deduplicate by (predecessor, target) pair — not just - // predecessor. A user branch's TRUE and FALSE paths can - // share the same predecessor (the branch block) but route - // the merge block to different CFF targets. Both entries - // are needed for correct cloning. - if !entry.iter().any(|(p, t)| *p == pred && *t == target) { - entry.push((pred, target)); - } - } - } - return; // Don't add duplicate redirect - } - - // First time seeing this source - record it and track predecessor for potential cloning - self.redirect_targets.insert(source, target); - self.redirects.push((source, target)); - - // Also record in clone_requests in case we need it later - if let Some(pred) = predecessor { - let entry = self.clone_requests.entry(source).or_default(); - entry.push((pred, target)); - } - } - - /// Returns redirects that are safe to apply (no conflicts/cloning needed). - pub(crate) fn safe_redirects(&self) -> Vec<(usize, usize)> { - self.redirects - .iter() - .filter(|(source, _)| { - // Safe if this block only has one target (not a merge point) - self.clone_requests.get(source).is_none_or(|v| v.len() <= 1) - }) - .copied() - .collect() - } - - /// Returns blocks that need cloning with their (predecessor, target) pairs. - pub(crate) fn blocks_to_clone(&self) -> Vec<(usize, Vec<(usize, usize)>)> { - self.clone_requests - .iter() - .filter(|(_, paths)| paths.len() > 1) // Only clone if multiple paths - .map(|(&block, paths)| (block, paths.clone())) - .collect() - } - - fn add_to_execution_order(&mut self, block: usize) { - if !self.execution_order_seen.contains(&block) && !self.is_dispatcher_block(block) { - self.execution_order_seen.insert(block); - self.execution_order.push(block); - } - } -} - -/// Extracts a patch plan from a trace tree. -/// -/// This analyzes the trace tree to determine: -/// - Which blocks should redirect to which targets (skip dispatcher) -/// - Which variables are state-tainted (for later instruction filtering) -/// - Which blocks need cloning (merge points with different targets) -/// -/// # Arguments -/// -/// * `tree` - The trace tree built by [`trace_method_tree`] -/// -/// # Returns -/// -/// `Some(PatchPlan)` if the tree contains a dispatcher and can be patched, -/// `None` if no dispatcher was detected. -/// -/// [`trace_method_tree`]: super::tracer::trace_method_tree -pub fn extract_patch_plan(tree: &TraceTree, ssa: &SsaFunction) -> Option { - let dispatcher = tree.dispatcher.as_ref()?; - - let mut plan = PatchPlan::new(dispatcher.block, tree.state_tainted.clone()); - - // Walk the main trace tree and extract redirects - // Start with no external predecessor since we're at the root - extract_redirects_from_node(&tree.root, dispatcher.block, &mut plan, None, ssa); - - // Walk exception handler traces and merge their redirects into the same plan. - // Handler blocks may contain their own CFF dispatchers that need unflattening. - for handler_trace in &tree.handler_traces { - extract_redirects_from_node(&handler_trace.root, dispatcher.block, &mut plan, None, ssa); - } - - // Filter out cross-contamination redirects. When tracing for one dispatcher, - // the tracer explores OTHER dispatchers' switches as user switches. LoopBack - // terminators at those foreign switches generate spurious redirects for blocks - // that don't belong to this dispatcher's CFF region. Only keep redirects that - // originated from StateTransition nodes (the actual CFF state machine redirects). - plan.redirects - .retain(|&(source, _)| plan.state_transition_sources.contains(&source)); - plan.clone_requests - .retain(|source, _| plan.state_transition_sources.contains(source)); - - Some(plan) -} - -/// Merges multiple patch plans into a single combined plan. -/// -/// This is used when a method contains multiple independent CFF dispatchers -/// (e.g., one per exception handler region in ConfuserEx). Each dispatcher -/// is traced independently, producing its own patch plan. This function -/// combines them so all patches are applied in a single pass before -/// `rebuild_ssa()`, avoiding the block renumbering corruption that occurs -/// when dispatchers are processed iteratively. -/// -/// # Arguments -/// -/// * `plans` - Individual patch plans, one per dispatcher. -/// -/// # Returns -/// -/// A single `PatchPlan` with all dispatcher blocks, redirects, taint sets, -/// and clone requests merged. -pub fn merge_patch_plans(plans: Vec) -> PatchPlan { - if plans.is_empty() { - return PatchPlan { - dispatcher_blocks: Vec::new(), - state_tainted: BitSet::new(0), - redirects: Vec::new(), - state_transition_sources: BTreeSet::new(), - clone_requests: BTreeMap::new(), - execution_order: Vec::new(), - execution_order_seen: FxHashSet::default(), - redirect_targets: FxHashMap::default(), - branch_collapses: BTreeMap::new(), - state_transitions_removed: 0, - user_branches_preserved: 0, - }; - } - - if plans.len() == 1 { - if let Some(only) = plans.into_iter().next() { - return only; - } - // Unreachable: we just checked len() == 1, but handle defensively. - return PatchPlan { - dispatcher_blocks: Vec::new(), - state_tainted: BitSet::new(0), - redirects: Vec::new(), - state_transition_sources: BTreeSet::new(), - clone_requests: BTreeMap::new(), - execution_order: Vec::new(), - execution_order_seen: FxHashSet::default(), - redirect_targets: FxHashMap::default(), - branch_collapses: BTreeMap::new(), - state_transitions_removed: 0, - user_branches_preserved: 0, - }; - } - - // Determine the tainted BitSet size (all plans share the same SSA, so same size) - let taint_size = plans - .iter() - .map(|p| p.state_tainted.len()) - .max() - .unwrap_or(0); - let mut merged = PatchPlan { - dispatcher_blocks: Vec::new(), - state_tainted: BitSet::new(taint_size), - redirects: Vec::new(), - state_transition_sources: BTreeSet::new(), - clone_requests: BTreeMap::new(), - execution_order: Vec::new(), - execution_order_seen: FxHashSet::default(), - redirect_targets: FxHashMap::default(), - branch_collapses: BTreeMap::new(), - state_transitions_removed: 0, - user_branches_preserved: 0, - }; - - for plan in plans { - merged.dispatcher_blocks.extend(&plan.dispatcher_blocks); - merged.state_tainted.union_with(&plan.state_tainted); - merged - .state_transition_sources - .extend(&plan.state_transition_sources); - - // Merge redirects, checking for conflicts - for (source, target) in plan.redirects { - if let Some(&existing_target) = merged.redirect_targets.get(&source) { - if existing_target != target { - log::warn!( - "CFF merge: redirect conflict for block {} (target {} vs {}), keeping first", - source, - existing_target, - target - ); - continue; - } - // Duplicate (same source, same target) — skip - continue; - } - merged.redirect_targets.insert(source, target); - merged.redirects.push((source, target)); - } - - // Merge clone requests - for (block, paths) in plan.clone_requests { - merged - .clone_requests - .entry(block) - .or_default() - .extend(paths); - } - - // Merge execution order (deduplicated) - for block in plan.execution_order { - if merged.execution_order_seen.insert(block) { - merged.execution_order.push(block); - } - } - - // Merge branch collapses (first wins on conflict) - for (source, target) in plan.branch_collapses { - merged.branch_collapses.entry(source).or_insert(target); - } - - merged.state_transitions_removed = merged - .state_transitions_removed - .saturating_add(plan.state_transitions_removed); - merged.user_branches_preserved = merged - .user_branches_preserved - .saturating_add(plan.user_branches_preserved); - } - - merged -} - -/// Returns true if an SSA op is safe to consider "pure CFG plumbing" — only -/// data motion (Const/Copy/Phi) or a Jump terminator. Used to identify -/// dispatcher prep-chain blocks that can be bypassed when redirecting case -/// blocks. -fn is_pure_prep_op(op: &SsaOp) -> bool { - matches!( - op, - SsaOp::Const { .. } | SsaOp::Copy { .. } | SsaOp::Jump { .. } | SsaOp::Nop - ) -} - -/// Returns true if `block_idx` is a state-overflow chain block — a block whose -/// terminator is a `BranchCmp` comparing the state variable against a constant -/// (NETReactor's `bcmp state == K, case_K, next_check` pattern on the -/// dispatcher's default arm), and whose non-terminator body is only pure -/// plumbing. -/// -/// Redirecting a case block directly to such a chain block lands the rewritten -/// edge on machinery that will compare a now-undefined state operand and fall -/// through into another case, creating self-loops after state-tainted -/// filtering. The caller should advance the effective redirect target through -/// consecutive chain blocks in the trace's `continues` visit list. -fn is_state_chain_block(ssa: &SsaFunction, block_idx: usize, state_tainted: &BitSet) -> bool { - let Some(block) = ssa.block(block_idx) else { - return false; - }; - let term_is_state_check = matches!(block.terminator_op(), Some(SsaOp::BranchCmp { left, right, .. }) - if state_tainted.contains(left.index()) || state_tainted.contains(right.index())); - if !term_is_state_check { - return false; - } - let instrs = block.instructions(); - let non_term_count = instrs.len().saturating_sub(1); - instrs - .iter() - .take(non_term_count) - .all(|instr| is_pure_prep_op(instr.op())) -} - -/// Extracts redirect information from a trace tree. -/// -/// Walks the tree iteratively via an explicit worklist to avoid blowing the -/// thread stack on methods with hundreds of CFF state transitions (each is -/// its own StateTransition node and the recursive version spent one frame per -/// transition). -/// -/// `external_predecessor` is used when a node is a sub-trace spawned from a -/// UserBranch/UserSwitch. It represents the block that branched to this -/// sub-trace's first block, needed for proper merge point detection when the -/// first block appears in multiple paths. -fn extract_redirects_from_node( - node: &TraceNode, - dispatcher_block: usize, - plan: &mut PatchPlan, - external_predecessor: Option, - ssa: &SsaFunction, -) { - let mut stack: Vec<(&TraceNode, Option)> = Vec::new(); - stack.push((node, external_predecessor)); - - while let Some((node, external_predecessor)) = stack.pop() { - // Add non-dispatcher blocks to execution order - for &block in &node.blocks_visited { - plan.add_to_execution_order(block); - } - - match &node.terminator { - TraceTerminator::StateTransition { - target_block, - continues, - .. - } => { - // When the dispatcher's switch falls through to its default and - // the default is an overflow chain (`bcmp state == K, case : next` - // repeated), `target_block` points at the chain's first block. - // Redirecting there lands on state-check machinery that — after - // state-tainted filtering — folds through an arbitrary chain - // arm and re-enters a case block, producing a self-loop. The - // traced `continues` node already walked the chain with concrete - // state, so the first non-chain block in its visit list is the - // real user-code target. Advance `target_block` to it. - let effective_target = { - let mut t = *target_block; - if is_state_chain_block(ssa, t, &plan.state_tainted) { - for &b in &continues.blocks_visited { - if !is_state_chain_block(ssa, b, &plan.state_tainted) { - t = b; - break; - } - } - } - t - }; - let target_block = &effective_target; - // Find the LAST non-dispatcher block — the one that jumps to dispatcher. - // We redirect it to bypass the dispatcher and go directly to target. - // - // For example, with blocks_visited=[9, 10, 2]: - // - Block 9 jumps to block 10 (keep this edge) - // - Block 10 jumps to dispatcher (redirect to target) - // - Block 2 is the dispatcher (being bypassed) - // - // This preserves block 10's code while bypassing the dispatcher. - // If multiple paths converge at block 10 with different targets, - // conflict detection will prevent the redirect. - let last_pred = node - .blocks_visited - .iter() - .rev() - .find(|&&b| b != dispatcher_block) - .copied(); - let first_pred = node - .blocks_visited - .iter() - .find(|&&b| b != dispatcher_block) - .copied(); - - // When the trace walks through intermediate blocks to reach the - // dispatcher (NETReactor inserts a `stloc state; ldloc state` prep - // chain between every case block and the switch), the "last block - // that jumps to the dispatcher" is a prep block shared across every - // path. Picking it as pred_block causes unresolvable redirect conflicts. - // If the intermediate blocks are pure CFG plumbing (no user-visible - // side effects — only Const/Copy/Phi), use the FIRST non-dispatcher - // block instead so each case block gets its own redirect. - // - // When `first` ends in a BranchCmp/Branch (e.g., NETReactor - // overflow-dispatch `bcmp state == Const ? case : next`), we - // can't simply use it as pred_block — `set_target` would collapse - // both branch arms to the same target, skipping over intervening - // user-code blocks like B103 → B11 (lock body) in the true arm. - // Instead, we keep `last` as pred_block (the block that actually - // jumps to the dispatcher) AND additionally collapse `first`'s - // BranchCmp to a Jump at the NEXT visited block — the arm the - // trace resolved. This preserves the sequential execution of the - // user-code intermediates while still bypassing the dispatcher. - let mut collapse_first_branch: Option<(usize, usize)> = None; - let pred_block = match (first_pred, last_pred) { - (Some(first), Some(last)) if first != last => { - let mut intermediates_are_pure = true; - let start_idx = node - .blocks_visited - .iter() - .position(|&b| b == first) - .unwrap_or(0); - let end_idx = node - .blocks_visited - .iter() - .position(|&b| b == last) - .unwrap_or(node.blocks_visited.len()); - let interior_start = start_idx.saturating_add(1); - if end_idx > interior_start { - let intermediate_blocks: BTreeSet = node - .blocks_visited - .get(interior_start..end_idx) - .map(|s| s.iter().copied().collect()) - .unwrap_or_default(); - // Purity is a property of the blocks themselves, so - // read it from the SSA rather than from a per-node - // instruction log. The tracer walks every - // instruction of a block it passes through, and - // these blocks lie strictly between `first` and - // `last` in the visit order, so they were traversed - // in full — the two formulations see the same - // instructions. - for &block_idx in &intermediate_blocks { - let all_pure = ssa.block(block_idx).is_some_and(|block| { - block.instructions().iter().all(|i| is_pure_prep_op(i.op())) - }); - if !all_pure { - intermediates_are_pure = false; - break; - } - } - } else { - intermediates_are_pure = false; - } - let first_is_jump_terminated = ssa - .block(first) - .and_then(|b| b.terminator_op()) - .is_some_and(|op| { - matches!(op, SsaOp::Jump { .. } | SsaOp::Leave { .. }) - }); - if intermediates_are_pure && first_is_jump_terminated { - Some(first) - } else if intermediates_are_pure { - // first has a Branch/BranchCmp overflow check: collapse - // it to Jump at the next visited block, then use `last` - // for the dispatcher bypass redirect. - let next_in_path = node - .blocks_visited - .get(start_idx.saturating_add(1)) - .copied(); - if let Some(next) = next_in_path { - collapse_first_branch = Some((first, next)); - } - Some(last) - } else { - Some(last) - } - } - _ => last_pred, - }; - - // Find what block leads INTO the pred block (for merge point tracking) - let predecessor_of_pred = if let Some(pred) = pred_block { - node.blocks_visited - .iter() - .position(|&b| b == pred) - .and_then(|pos| { - if pos > 0 { - node.blocks_visited.get(pos.saturating_sub(1)).copied() - } else { - external_predecessor - } - }) - } else { - None - }; - - if let Some(pred) = pred_block { - // This block should redirect to target_block instead of dispatcher - plan.add_redirect(pred, *target_block, predecessor_of_pred); - plan.state_transition_sources.insert(pred); - plan.state_transitions_removed = - plan.state_transitions_removed.saturating_add(1); - } else if let Some(ext_pred) = external_predecessor { - // The sub-trace starts directly at the dispatcher (no preceding user blocks). - // This happens when a user branch at method entry sends one path directly - // to the dispatcher block. Redirect the external predecessor's - // dispatcher-targeting edge to the actual target. - plan.add_redirect(ext_pred, *target_block, None); - plan.state_transition_sources.insert(ext_pred); - plan.state_transitions_removed = - plan.state_transitions_removed.saturating_add(1); - } - - // If the first visited block is a BranchCmp overflow check, add - // a supplementary redirect that collapses its branch to the next - // visited block (the arm the trace resolved). Flagged as a - // branch-collapse redirect so apply_patch_plan can replace the - // BranchCmp terminator with a Jump, preserving the traced arm's - // user-code execution while eliminating the dead state-check. - if let Some((src, next)) = collapse_first_branch { - plan.add_branch_collapse(src, next); - } - - // Continue processing the rest of the trace. - // The continues trace starts at target_block, and its predecessor is - // pred_block (the block that was redirected to go to target_block). - // Fall back to external_predecessor when pred_block is None - // (direct-to-dispatcher path). - stack.push((continues, pred_block.or(external_predecessor))); - } - - TraceTerminator::UserBranch { - block, - true_branch, - false_branch, - .. - } => { - plan.user_branches_preserved = plan.user_branches_preserved.saturating_add(1); - // For user branches, the branch block is the predecessor of both sub-traces. - // This is crucial for proper merge point detection when a block is both - // an entry path target (from the branch) and a CFF case target. - // Push in reverse so pops give true_branch first (preserving recursive order). - stack.push((false_branch, Some(*block))); - stack.push((true_branch, Some(*block))); - } - - TraceTerminator::UserSwitch { - block, - cases, - default, - .. - } => { - plan.user_branches_preserved = plan.user_branches_preserved.saturating_add(1); - // For user switches, the switch block is the predecessor of all - // case sub-traces. The recursive version processed all cases in - // order, then the default — to reproduce that order with a LIFO - // stack, push the default first, then the cases in reverse. - stack.push((default, Some(*block))); - for (_, case_node) in cases.iter().rev() { - stack.push((case_node, Some(*block))); - } - } - - TraceTerminator::Exit { block } => { - plan.add_to_execution_order(*block); - } - - TraceTerminator::LoopBack { target_block, .. } => { - // LoopBack means a loop back-edge through the CFF dispatcher: the - // path "source → dispatcher → target" should become a natural loop - // edge "source → target". - // - // Usually the parent StateTransition already added this redirect. - // However, in CFF patterns where multiple switch cases share the - // same target block (e.g., JIEJIE.NET), the parent's pred_block - // resolution can yield None, causing it to miss the redirect. - // Adding it here as a safety net is always correct — add_redirect - // deduplicates identical (source, target) pairs. - // - // We use external_predecessor (set by the parent call site) as the - // redirect source, since this node's blocks_visited is typically - // just [target_block] (the loop header / destination, not the source). - if let Some(pred) = external_predecessor { - plan.add_redirect(pred, *target_block, external_predecessor); - plan.state_transition_sources.insert(pred); - plan.state_transitions_removed = - plan.state_transitions_removed.saturating_add(1); - } - plan.add_to_execution_order(*target_block); - } - - TraceTerminator::Stopped { .. } | TraceTerminator::PendingStateTransition { .. } => { - // Trace halted due to a limit or unresolvable control flow. - // Nothing to extract. - } - } - } -} - -/// Applies a patch plan to an SSA function, removing CFF machinery. -/// -/// This modifies the SSA in place by: -/// 1. Redirecting block terminators from dispatcher to their actual targets -/// 2. Cloning merge blocks when multiple paths converge with different targets -/// 3. Filtering out state-tainted instructions (CFF machinery) -/// 4. Propagating user variable PHIs to new merge points -/// 5. Cleaning up orphaned references -/// -/// # Arguments -/// -/// * `ssa` - The SSA function to patch (modified in place) -/// * `plan` - The patch plan extracted from the trace tree -/// -/// # Returns -/// -/// A [`ReconstructionResult`] containing statistics about the patching: -/// - Number of state transitions removed -/// - Number of user branches preserved -/// - Final block count -pub fn apply_patch_plan(ssa: &mut SsaFunction, plan: &PatchPlan) -> ReconstructionResult { - // Snapshot the phi graph before any rewiring. Redirects bypass the - // dispatcher, and the dispatcher's phis — which record what each case block - // contributed — are cleared later in this function. Recovering the value for - // a redirected edge means reading those phis, so they must be captured while - // they still exist. - let original_phis: PhiSnapshot = ssa - .blocks() - .iter() - .map(|block| { - let phis = block - .phi_nodes() - .iter() - .map(|phi| { - let operands = phi - .operands() - .iter() - .map(|op| (op.predecessor(), op.value())) - .collect(); - (phi.result(), operands) - }) - .collect(); - (block.id(), phis) - }) - .collect(); - - // Apply only safe redirects (skip conflicting merge points that need cloning) - let safe = plan.safe_redirects(); - let to_clone = plan.blocks_to_clone(); - - // Apply safe redirects (blocks that don't need cloning). - // - // The source may be either the direct CASE→DISP edge block (ConfuserEx-style) - // or the start of a CASE→prep→...→DISP chain (NETReactor-style where the - // dispatcher has intermediate prep blocks for state setup). For Jump - // terminators we overwrite the target outright, which correctly bypasses - // multi-hop prep chains; for Branch/Switch terminators we fall back to - // edge-rewriting relative to any dispatcher block. - for &(source_block, new_target) in &safe { - let is_jump = ssa - .block(source_block) - .and_then(|b| b.terminator_op()) - .is_some_and(|op| matches!(op, SsaOp::Jump { .. } | SsaOp::Leave { .. })); - if let Some(block) = ssa.block_mut(source_block) { - if is_jump { - block.set_target(new_target); - } else { - for &db in &plan.dispatcher_blocks { - block.redirect_target(db, new_target); - } - } - } - } - - // Apply branch collapses: at each source block with a recorded collapse, - // replace the BranchCmp/Branch terminator with a Jump to the traced arm's - // next block. This eliminates NETReactor overflow-dispatch machinery (e.g. - // `bcmp state == 13 ? B_match : B_next_check`) from the unflattened CFG - // while preserving the intervening user-code blocks that the separate - // redirect chain walks to reach the state transition's target. - for (&source_block, &new_target) in &plan.branch_collapses { - if let Some(block) = ssa.block_mut(source_block) { - let current = block - .terminator_op() - .map(|op| matches!(op, SsaOp::Branch { .. } | SsaOp::BranchCmp { .. })) - .unwrap_or(false); - if current { - if let Some(term) = block.instructions_mut().last_mut() { - *term = SsaInstruction::synthetic(SsaOp::Jump { target: new_target }); - } - } - } - } - - // Handle merge points by cloning blocks - // For each merge block, clone it for all predecessors except the first - // Track clone mappings: clone_index -> original_index - let mut clone_map: BTreeMap = BTreeMap::new(); - let mut cloned_blocks = Vec::new(); - for (merge_block, paths) in &to_clone { - if paths.len() < 2 { - continue; - } - - // Get the merge block content - we'll clone it for each path - let merge_content = ssa.block(*merge_block).cloned(); - let Some(original_block) = merge_content else { - continue; - }; - - // First path keeps the original block - update its terminator. - // Only redirect Jump terminators. Branch/BranchCmp blocks are user - // branches that must preserve both targets — cloning them with - // set_target would collapse both branch arms to the same target. - let Some(&(_, first_target)) = paths.first() else { - continue; - }; - let is_user_branch = ssa - .block(*merge_block) - .and_then(|b| b.terminator_op()) - .is_some_and(|op| { - matches!( - op, - SsaOp::Branch { .. } | SsaOp::BranchCmp { .. } | SsaOp::Switch { .. } - ) - }); - if is_user_branch { - // User branch block: don't clone, don't redirect. - // The redirect already points predecessors to this block; - // its branches are user code that should be preserved. - continue; - } - if let Some(block) = ssa.block_mut(*merge_block) { - block.set_target(first_target); - } - - // For remaining paths, create clones - for &(pred, target) in paths.iter().skip(1) { - let new_block_idx = ssa.block_count().saturating_add(cloned_blocks.len()); - - // Track the clone mapping - clone_map.insert(new_block_idx, *merge_block); - - // Clone the block with new ID and updated terminator - let mut cloned = original_block.clone(); - cloned.set_id(new_block_idx); - - // Instead of manually replacing PHI results with operand values (which - // can create cycles when PHI operands have mutual dependencies), we take - // a simpler approach: - // - // 1. Keep only the PHI operand from this predecessor - // 2. Let rebuild_ssa handle the proper conversion - // - // This avoids creating cycles because: - // - PHI nodes are properly handled by the SSA reconstruction algorithm - // - The algorithm correctly handles dominance relationships - // - We don't risk creating references to values defined later in the block - // - // Note: The PHI will be trivial (single operand) after this, and rebuild_ssa - // will either eliminate it or convert it to a proper definition. - for phi in cloned.phi_nodes_mut() { - // Keep only the operand from this predecessor - phi.operands_mut().retain(|op| op.predecessor() == pred); - } - - // Remove PHIs that now have no operands (shouldn't happen, but be safe) - cloned - .phi_nodes_mut() - .retain(|phi| !phi.operands().is_empty()); - - // Set terminator to jump directly to target - cloned.set_target(target); - - // Track: predecessor needs to redirect to this clone - cloned_blocks.push((pred, cloned, *merge_block)); - } - } - - // Add all cloned blocks to SSA - for (pred, mut cloned, original_merge) in cloned_blocks { - let new_block_idx = cloned.id(); - - // Redirect the predecessor to point to the clone instead of the merge block - if let Some(pred_block) = ssa.block_mut(pred) { - pred_block.redirect_target(original_merge, new_block_idx); - } - - // Filter state instructions from the clone. - // - // This must happen before the clone's definitions are renamed below: - // taint is recorded against the original variable ids, so renaming - // first would stop every state instruction from matching. - filter_state_instructions(&mut cloned, &plan.state_tainted, &plan.dispatcher_blocks); - - // Give the clone its own definitions. - // - // Copying a block verbatim defines every one of its variables a second - // time, which breaks SSA's single-definition property. The original code - // relied on `rebuild_ssa()` to sort this out, but rebuild cannot recover - // from duplicate definitions — it produces a function whose uses resolve - // to neither definition, surfacing later as an undefined-variable - // verification failure in an unrelated block. - // - // Uses that refer to definitions *outside* this block are left alone: - // they still resolve to the original definition, which dominates both - // the original block and the clone. - let mut defs_to_rename: Vec = Vec::new(); - for phi in cloned.phi_nodes() { - defs_to_rename.push(phi.result()); - } - for instr in cloned.instructions() { - if let Some(dest) = instr.def() { - defs_to_rename.push(dest); - } - } - - let mut renames: BTreeMap = BTreeMap::new(); - for old_var in defs_to_rename { - let Some(var) = ssa.variable(old_var) else { - continue; - }; - let (origin, var_type) = (var.origin(), var.var_type().clone()); - let new_var = ssa.create_variable(origin, 0, DefSite::phi(new_block_idx), var_type); - renames.insert(old_var, new_var); - } - - if !renames.is_empty() { - for phi in cloned.phi_nodes_mut() { - if let Some(&new_var) = renames.get(&phi.result()) { - phi.set_result(new_var); - } - } - for instr in cloned.instructions_mut() { - instr - .op_mut() - .replace_uses_with(|v| renames.get(&v).copied()); - if let Some(dest) = instr.def() { - if let Some(&new_var) = renames.get(&dest) { - instr.op_mut().replace_def(dest, new_var); - } - } - } - } - - // Add the clone to SSA - ssa.blocks_mut().push(cloned); - } - - // Filter out state-tainted instructions ONLY from blocks that were patched. - // Unresolved blocks (not redirected or cloned) must keep their CFF machinery - // so they remain functional. Stripping state instructions from unresolved - // blocks would remove their Jump-to-dispatcher terminators, orphaning them. - let mut patched_blocks = BitSet::new(ssa.block_count()); - for &(source, _) in &safe { - patched_blocks.insert(source); - } - for (merge_block, _) in &to_clone { - patched_blocks.insert(*merge_block); - } - - for block_idx in patched_blocks.iter() { - // Don't filter state instructions from blocks that still have a - // switch terminator (including handler CFF dispatchers). The switch - // needs its state computation (call + rem.un) to remain functional. - // These blocks will either be fully cleared later (when no longer - // needed) or their state will be folded by constant propagation. - let has_switch = ssa - .block(block_idx) - .and_then(|b| b.terminator_op()) - .is_some_and(|op| matches!(op, SsaOp::Switch { .. })); - if plan.is_dispatcher_block(block_idx) || has_switch { - continue; - } - if let Some(block) = ssa.block_mut(block_idx) { - filter_state_instructions(block, &plan.state_tainted, &plan.dispatcher_blocks); - } - } - - // Materialize dispatcher phi resolutions as explicit copies. - // - // When a case block (e.g., B4) is redirected to bypass the dispatcher and - // jump directly to the next case block (e.g., B5), the phi nodes at the - // dispatcher carried user values between iterations. Without materializing - // these, rebuild_ssa() cannot recover the data flow because the phi - // definitions are destroyed when the dispatcher is cleared. - // - // For each redirected edge (source → new_target), we resolve the phi chain - // through dispatcher blocks and insert Copy instructions at the end of the - // source block. - materialize_dispatcher_phis(ssa, plan, &patched_blocks); - - // Check if the dispatcher is still needed. A block that jumps to the - // dispatcher blocks unflattening only if it's reachable from OUTSIDE the - // dispatcher. Blocks whose only predecessors are dispatcher blocks (dead - // CFF case targets never dispatched during execution) are cleared along - // with the dispatcher. - let unresolved: Vec = (0..ssa.block_count()) - .filter(|&bi| { - !patched_blocks.contains(bi) - && !plan.is_dispatcher_block(bi) - && ssa - .block(bi) - .and_then(|b| b.terminator_op()) - .is_some_and(|op| match op { - SsaOp::Jump { target } => plan.is_dispatcher_block(*target), - SsaOp::BranchCmp { - true_target, - false_target, - .. - } => { - plan.is_dispatcher_block(*true_target) - || plan.is_dispatcher_block(*false_target) - } - _ => false, - }) - }) - .collect(); - - let dispatcher_still_needed = unresolved.iter().any(|&bi| { - ssa.block_predecessors(bi) - .iter() - .any(|&pred| !plan.is_dispatcher_block(pred)) - }); - - if !dispatcher_still_needed { - for &db in &plan.dispatcher_blocks { - if let Some(dispatcher) = ssa.block_mut(db) { - dispatcher.clear(); - } - } - for &bi in &unresolved { - if ssa - .block_predecessors(bi) - .iter() - .all(|&pred| plan.is_dispatcher_block(pred)) - { - if let Some(block) = ssa.block_mut(bi) { - block.clear(); - } - } - } - } - - // Drop phi operands whose value no longer has a definition. - // - // Clearing the dispatcher and dead case blocks, and filtering state - // instructions out of patched blocks, removes definitions that successor - // phis may still reference. `rebuild_ssa()` resolves phi operands into real - // values, so a stale operand becomes an undefined instruction operand and - // fails verification — reported far from here, as an undefined use in - // whichever block the value was propagated into. - // - // Pruning is safe because `rebuild_ssa()` reconstructs phi nodes from the - // patched CFG and the surviving definitions; a stale operand carries no - // information it could use. - { - let mut defined: BTreeSet = BTreeSet::new(); - for block in ssa.blocks() { - for phi in block.phi_nodes() { - defined.insert(phi.result()); - } - for instr in block.instructions() { - if let Some(dest) = instr.def() { - defined.insert(dest); - } - } - } - - // Predecessor sets for the *patched* CFG. Redirects, branch collapses - // and clone edges all rewire terminators without touching phi operands, - // so a phi can still carry an operand for a block that no longer reaches - // it — 53k such operands on a large NETReactor method. `rebuild_ssa()` - // resolves phi operands into concrete values, and a stale operand names - // a value that does not reach the block, so it propagates a definition - // that no longer dominates its uses. - let mut predecessors: BTreeMap> = BTreeMap::new(); - for block in ssa.blocks() { - let id = block.id(); - block.for_each_successor(|succ| { - predecessors.entry(succ).or_default().insert(id); - }); - } - - // Work out, per block, which phi operands survive and which edges need a - // value recovered. Done as a read-only pass first so `resolve_redirected_operand` - // can consult the snapshot without holding a mutable borrow of `ssa`. - let mut repairs: BTreeMap> = BTreeMap::new(); - for block in ssa.blocks() { - let block_id = block.id(); - let Some(preds) = predecessors.get(&block_id) else { - continue; - }; - for phi in block.phi_nodes() { - let result = phi.result(); - for &pred in preds { - if phi.operands().iter().any(|op| op.predecessor() == pred) { - continue; - } - if let Some(value) = resolve_redirected_operand( - &original_phis, - &defined, - block_id, - result, - pred, - preds, - ) { - repairs - .entry(block_id) - .or_default() - .push((result, pred, value)); - } - } - } - } - - for block in ssa.blocks_mut() { - let block_id = block.id(); - let preds = predecessors.remove(&block_id).unwrap_or_default(); - let block_repairs = repairs.remove(&block_id).unwrap_or_default(); - for phi in block.phi_nodes_mut() { - phi.operands_mut().retain(|op| { - preds.contains(&op.predecessor()) && defined.contains(&op.value()) - }); - for &(result, pred, value) in &block_repairs { - if result == phi.result() - && !phi.operands().iter().any(|op| op.predecessor() == pred) - { - phi.add_operand(PhiOperand::new(value, pred)); - } - } - } - block - .phi_nodes_mut() - .retain(|phi| !phi.operands().is_empty()); - } - } - - // NOTE: PHI propagation and variable definition cleanup is NOT done here. - // After apply_patch_plan returns, the caller must call SsaFunction::rebuild_ssa() - // to reconstruct proper SSA form with correct PHI nodes for the new CFG structure. - - ReconstructionResult { - state_transitions_removed: plan.state_transitions_removed, - user_branches_preserved: plan.user_branches_preserved, - block_count: ssa.block_count(), - dispatcher_still_needed, - } -} - -/// Recovers the value a phi should take on an edge created by a redirect. -/// -/// Unflattening rewires `source -> dispatcher -> target` into `source -> target`. -/// The phi at `target` recorded its incoming value against the dispatcher, not -/// against `source`, so after rewiring it has no operand for its new predecessor. -/// The value is still recoverable: the dispatcher's own phi records what each -/// case block contributed, so following `target`'s dispatcher-side operand back -/// through the dispatcher's phi to its operand for `source` yields the value that -/// was live when `source` finished executing. -/// -/// Chains of dispatcher blocks are followed up to [`MAX_PHI_RESOLUTION_HOPS`] -/// times. Returns `None` if the chain cannot be resolved to a value that still -/// has a definition, in which case the caller leaves the operand absent rather -/// than inventing one — a wrong value here would silently corrupt the recovered -/// program, which is worse than a failed unflattening. -fn resolve_redirected_operand( - original_phis: &PhiSnapshot, - defined: &BTreeSet, - block: usize, - phi_result: SsaVarId, - new_pred: usize, - live_preds: &BTreeSet, -) -> Option { - // Operands of this phi that are no longer reachable: these are the edges the - // patch removed, and one of them is the dispatcher path from `new_pred`. - let operands = original_phis.get(&block)?.get(&phi_result)?; - - for &(old_pred, value) in operands { - if live_preds.contains(&old_pred) { - continue; - } - - // Walk back through the removed blocks' phis looking for what `new_pred` - // contributed. - let mut current_block = old_pred; - let mut current_value = value; - for _ in 0..MAX_PHI_RESOLUTION_HOPS { - let Some(block_phis) = original_phis.get(¤t_block) else { - break; - }; - let Some(chain_operands) = block_phis.get(¤t_value) else { - // `current_value` is not defined by a phi here, so it is a plain - // definition that was already live before the dispatcher — usable - // only if it survived the patch. - break; - }; - - if let Some(&(_, from_source)) = - chain_operands.iter().find(|(pred, _)| *pred == new_pred) - { - if defined.contains(&from_source) { - return Some(from_source); - } - current_value = from_source; - current_block = new_pred; - continue; - } - - // Not contributed directly by `new_pred`; step through the single - // unreachable predecessor if there is exactly one candidate. - let mut candidates = chain_operands - .iter() - .filter(|(pred, _)| !live_preds.contains(pred)); - let &(next_block, next_value) = candidates.next()?; - if candidates.next().is_some() { - // Ambiguous — several removed edges could supply this value, and - // guessing risks selecting one that does not dominate `new_pred`. - return None; - } - current_block = next_block; - current_value = next_value; - } - - if defined.contains(¤t_value) { - return Some(current_value); - } - } - - None -} - -/// Filters out state-tainted instructions from a block. -/// -/// An instruction is filtered if: -/// - Its output (def) is state-tainted, OR -/// - Any of its inputs (uses) are state-tainted -/// -/// **Exception**: `Const` instructions are always preserved, even if their -/// def variable is state-tainted. CFF state update constants (e.g., -/// `v_state = Const(42)`) become dead code after the dispatcher redirect -/// and are cleaned up by subsequent dead code elimination passes. Filtering -/// them here risks removing legitimate user constants (strings, integers) -/// whose SSA variables were coincidentally tainted through PHI propagation -/// in combined-protection scenarios (e.g., CFF + string encryption where -/// resolved string constants share PHI merge points with state variables). -/// -/// Terminator instructions are NOT filtered (they're handled by redirect). -fn filter_state_instructions( - block: &mut SsaBlock, - state_tainted: &BitSet, - dispatcher_blocks: &[usize], -) { - block.instructions_mut().retain(|instr| { - // Always keep terminators - they're handled separately - if instr.is_terminator() { - // But skip jumps to any dispatcher (they've been redirected) - if let SsaOp::Jump { target } = instr.op() { - if dispatcher_blocks.contains(target) { - return false; // Remove jump to dispatcher - } - } - return true; - } - - // Preserve Const instructions unconditionally. Dead state constants - // are removed by later DCE; user constants must survive. - if matches!(instr.op(), SsaOp::Const { .. }) { - return true; - } - - // Check if instruction is state-tainted - let def_tainted = instr - .def() - .is_some_and(|d| state_tainted.contains(d.index())); - let uses_tainted = instr - .uses() - .iter() - .any(|u| state_tainted.contains(u.index())); - - // Remove instructions whose output is state infrastructure. - if def_tainted { - return false; - } - - // For uses-only tainted instructions: preserve Call/CallVirt because - // they may compute user-visible values from state-derived constants. - // Example: JIEJIE.NET typeof container — `GetTypeInstance(state_index)` - // returns a System.Type needed by user code, even though the index - // argument is derived from the CFF state machine's Int32ValueContainer. - if uses_tainted { - return matches!(instr.op(), SsaOp::Call { .. } | SsaOp::CallVirt { .. }); - } - - true - }); -} - -/// Materializes dispatcher phi resolutions as explicit Copy instructions. -/// -/// When CFF unflattening redirects a case block to bypass the dispatcher, the -/// dispatcher's phi nodes — which carried user values between loop iterations — -/// are about to be destroyed. This function resolves those phis along each -/// redirected edge and inserts Copy instructions so that `rebuild_ssa()` can -/// see the data flow. -/// -/// For a redirected edge `source → new_target` that originally went through -/// dispatcher blocks `source → D1 → D2 → new_target`, this resolves the phi -/// chain: for each phi at D2, trace back through D1's phis to find the -/// concrete value coming from `source`, and insert `phi_result = value` as a -/// Copy in `source`. -fn materialize_dispatcher_phis(ssa: &mut SsaFunction, plan: &PatchPlan, patched_blocks: &BitSet) { - if plan.dispatcher_blocks.is_empty() { - return; - } - - // Identify all blocks in the dispatcher loop: the dispatcher blocks themselves - // plus any "back-edge" blocks between case blocks and the dispatcher. - // - // The typical CFF pattern has: case_block → back_edge (B1) → dispatcher (B2). - // Both B1 and B2 have phis that carry user values between iterations. - // - // We detect back-edge blocks by two criteria: - // 1. They have phis with operands from patched (case) blocks - // 2. They are predecessors of the dispatcher (directly or transitively) - let dispatcher_set: BTreeSet = plan.dispatcher_blocks.iter().copied().collect(); - let mut loop_blocks: Vec = plan.dispatcher_blocks.clone(); - let mut loop_block_set: BTreeSet = dispatcher_set.clone(); - - // Find back-edge blocks: blocks with phis whose operands come from - // patched case blocks. These blocks sit between case blocks and the - // dispatcher, funneling values through phis. - for block_idx in 0..ssa.block_count() { - if loop_block_set.contains(&block_idx) { - continue; - } - if patched_blocks.contains(block_idx) { - continue; - } - let Some(block) = ssa.block(block_idx) else { - continue; - }; - if block.phi_nodes().is_empty() { - continue; - } - // Check if any phi has an operand from a patched case block - let has_patched_predecessor = block.phi_nodes().iter().any(|phi| { - phi.operands() - .iter() - .any(|op| patched_blocks.contains(op.predecessor())) - }); - if has_patched_predecessor { - loop_blocks.push(block_idx); - loop_block_set.insert(block_idx); - } - } - - // Sort loop blocks: back-edge blocks first (not in dispatcher_set), - // then dispatcher blocks. This ensures we resolve back-edge phis - // before dispatcher phis that depend on them. - loop_blocks.sort_by_key(|&b| if dispatcher_set.contains(&b) { 1 } else { 0 }); - - // Collect the phi data from all loop blocks before mutating. - let dispatcher_phis: BlockPhiData = loop_blocks - .iter() - .filter_map(|&db| { - ssa.block(db).map(|block| { - let phis = block - .phi_nodes() - .iter() - .filter(|phi| !plan.state_tainted.contains(phi.result().index())) - .map(|phi| { - let operands: Vec<(usize, SsaVarId)> = phi - .operands() - .iter() - .map(|op| (op.predecessor(), op.value())) - .collect(); - (phi.result(), operands) - }) - .collect(); - (db, phis) - }) - }) - .collect(); - - if dispatcher_phis.is_empty() { - return; - } - - // Build a lookup: loop_block → (phi_result → operands) - let mut phi_lookup: BTreeMap>> = - BTreeMap::new(); - for (db, phis) in &dispatcher_phis { - let map: BTreeMap> = phis - .iter() - .map(|(result, operands)| (*result, operands.clone())) - .collect(); - phi_lookup.insert(*db, map); - } - - let loop_block_set: BTreeSet = loop_blocks.iter().copied().collect(); - - // For each dispatcher phi, find the concrete value: the operand from a - // patched case block that provides a NON-pass-through value (i.e., a value - // that is NOT itself a dispatcher-phi result). - // - // Then substitute all uses of the phi-result variable in non-provider - // blocks with the concrete value. This eliminates stale references to - // variables defined only by the now-destroyed dispatcher phis. - let phi_result_set: BTreeSet = dispatcher_phis - .iter() - .flat_map(|(_, phis)| phis.iter().map(|(r, _)| *r)) - .collect(); - - // For each phi, find which block(s) provide a concrete (non-phi-result) value - let mut concrete_for_phi: BTreeMap = BTreeMap::new(); - for (_, phis) in &dispatcher_phis { - for (phi_result, operands) in phis { - for &(pred, val) in operands { - if !patched_blocks.contains(pred) || loop_block_set.contains(&pred) { - continue; - } - // Pass-through: val is the phi result itself or another phi result - if val == *phi_result || phi_result_set.contains(&val) { - continue; - } - concrete_for_phi.insert(*phi_result, (pred, val)); - } - } - } - - if concrete_for_phi.is_empty() { - return; - } - - // Resolve phi values along the execution path and substitute uses. - // - // Process blocks in execution order (from the trace tree). For each block, - // resolve the dispatcher phis from the PREVIOUS block's perspective, then - // substitute all uses of phi-result variables in this block. - // - // The key: each block's phi resolution builds on the previous block's. - // When a phi operand from block B is itself a phi result (pass-through), - // we use the accumulated resolution to get the concrete value. - let mut accumulated: BTreeMap = BTreeMap::new(); - - for &block_idx in &plan.execution_order { - if !patched_blocks.contains(block_idx) || loop_block_set.contains(&block_idx) { - continue; - } - - // Apply accumulated resolution to THIS block's instructions - if !accumulated.is_empty() { - if let Some(block) = ssa.block_mut(block_idx) { - for instr in block.instructions_mut() { - for (&phi_var, &concrete_val) in &accumulated { - instr.op_mut().replace_uses(phi_var, concrete_val); - } - } - } - } - - // Compute this block's phi resolution for the NEXT block. - // For each dispatcher phi, find the operand from this block and - // resolve transitively using the accumulated map. - for &db in &loop_blocks { - let Some(phis) = phi_lookup.get(&db) else { - continue; - }; - for (phi_result, operands) in phis { - if let Some(&(_, val)) = operands.iter().find(|&&(pred, _)| pred == block_idx) { - // Resolve transitively - let concrete = accumulated.get(&val).copied().unwrap_or(val); - if *phi_result != concrete { - accumulated.insert(*phi_result, concrete); - } - } - } - } - } -} - -/// Analyzes a trace tree and returns statistics without modifying anything. -/// -/// This is a read-only analysis function that reports what would change -/// if [`apply_patch_plan`] were called. Useful for previewing changes -/// or collecting metrics without side effects. -/// -/// # Arguments -/// -/// * `tree` - The trace tree to analyze -/// * `original` - The original SSA function (for block count) -/// -/// # Returns -/// -/// `Some(ReconstructionResult)` with statistics if the tree has a dispatcher, -/// `None` if no dispatcher was detected. -pub fn reconstruct_from_tree( - tree: &TraceTree, - original: &SsaFunction, -) -> Option { - let plan = extract_patch_plan(tree, original)?; - - Some(ReconstructionResult { - state_transitions_removed: plan.state_transitions_removed, - user_branches_preserved: plan.user_branches_preserved, - block_count: original.block_count(), - dispatcher_still_needed: false, // Preview only — actual value computed during apply - }) -} - -#[cfg(test)] -mod tests { - use crate::{ - analysis::{ - ConstValue, PhiNode, PhiOperand, SsaBlock, SsaFunction, SsaInstruction, SsaOp, - SsaVarId, VariableOrigin, - }, - deobfuscation::passes::unflattening::{ - reconstruction::reconstruct_from_tree, tracer::trace_method_tree, UnflattenConfig, - }, - }; - - /// Creates a simple CFF-like SSA function for testing. - fn create_simple_cff() -> SsaFunction { - let mut ssa = SsaFunction::new(0, 1); - let state_var = SsaVarId::from_index(0); - let const_var = SsaVarId::from_index(1); - - // B0: entry - set initial state and jump to dispatcher - let mut b0 = SsaBlock::new(0); - b0.add_instruction(SsaInstruction::synthetic(SsaOp::Const { - dest: const_var, - value: ConstValue::I32(0), - })); - b0.add_instruction(SsaInstruction::synthetic(SsaOp::Jump { target: 1 })); - ssa.add_block(b0); - - // B1: dispatcher with switch - let mut b1 = SsaBlock::new(1); - let mut phi = PhiNode::new(state_var, VariableOrigin::Local(0)); - phi.add_operand(PhiOperand::new(const_var, 0)); - b1.add_phi(phi); - b1.add_instruction(SsaInstruction::synthetic(SsaOp::Switch { - value: state_var, - targets: vec![2, 3, 4], - default: 5, - })); - ssa.add_block(b1); - - // B2, B3, B4: case blocks that jump back to dispatcher - for i in 2..=4 { - let mut b = SsaBlock::new(i); - b.add_instruction(SsaInstruction::synthetic(SsaOp::Jump { target: 1 })); - ssa.add_block(b); - } - - // B5: exit - let mut b5 = SsaBlock::new(5); - b5.add_instruction(SsaInstruction::synthetic(SsaOp::Return { value: None })); - ssa.add_block(b5); - - ssa - } - - #[test] - fn test_reconstruct_simple_cff() { - let ssa = create_simple_cff(); - let config = UnflattenConfig::default(); - let tree = trace_method_tree(&ssa, &config, None); - - assert!(tree.dispatcher.is_some(), "Should detect dispatcher"); - - let result = reconstruct_from_tree(&tree, &ssa); - assert!(result.is_some(), "Reconstruction should succeed"); - - let result = result.unwrap(); - println!("=== Reconstruction Result ==="); - println!("Blocks: {}", result.block_count); - println!( - "State transitions removed: {}", - result.state_transitions_removed - ); - println!( - "User branches preserved: {}", - result.user_branches_preserved - ); - - // Should have removed state transitions - assert!( - result.state_transitions_removed > 0, - "Should remove at least one state transition" - ); - - // Reconstructed function should have blocks - assert!(result.block_count > 0, "Should have blocks in output"); - } - - /// Creates a CFF with a user branch inside a case block. - fn create_cff_with_user_branch() -> SsaFunction { - let mut ssa = SsaFunction::new(1, 1); // 1 arg, 1 local - let state_var = SsaVarId::from_index(0); - let init_state = SsaVarId::from_index(1); - let const_one = SsaVarId::from_index(2); - let arg0 = SsaVarId::from_index(3); - let user_zero = SsaVarId::from_index(4); - let cmp_result = SsaVarId::from_index(5); - - // B0: entry - set initial state = 0 and jump to dispatcher - let mut b0 = SsaBlock::new(0); - b0.add_instruction(SsaInstruction::synthetic(SsaOp::Const { - dest: init_state, - value: ConstValue::I32(0), - })); - b0.add_instruction(SsaInstruction::synthetic(SsaOp::Const { - dest: arg0, - value: ConstValue::I32(42), // simulate arg > 0 - })); - b0.add_instruction(SsaInstruction::synthetic(SsaOp::Jump { target: 1 })); - ssa.add_block(b0); - - // B1: dispatcher with switch - let mut b1 = SsaBlock::new(1); - let mut phi = PhiNode::new(state_var, VariableOrigin::Local(0)); - phi.add_operand(PhiOperand::new(init_state, 0)); // from entry - phi.add_operand(PhiOperand::new(const_one, 3)); // from B3a - phi.add_operand(PhiOperand::new(const_one, 4)); // from B3b - b1.add_phi(phi); - b1.add_instruction(SsaInstruction::synthetic(SsaOp::Switch { - value: state_var, - targets: vec![2, 5], // case 0 -> B2, case 1 -> B5 - default: 6, - })); - ssa.add_block(b1); - - // B2: case 0 - has USER BRANCH (condition NOT tainted by state) - let mut b2 = SsaBlock::new(2); - b2.add_instruction(SsaInstruction::synthetic(SsaOp::Const { - dest: const_one, - value: ConstValue::I32(1), - })); - b2.add_instruction(SsaInstruction::synthetic(SsaOp::Const { - dest: user_zero, - value: ConstValue::I32(0), - })); - b2.add_instruction(SsaInstruction::synthetic(SsaOp::Cgt { - dest: cmp_result, - left: arg0, - right: user_zero, - unsigned: false, - })); - b2.add_instruction(SsaInstruction::synthetic(SsaOp::Branch { - condition: cmp_result, - true_target: 3, // B3a - false_target: 4, // B3b - })); - ssa.add_block(b2); - - // B3a: true branch of user condition - sets state = 1 - let mut b3a = SsaBlock::new(3); - b3a.add_instruction(SsaInstruction::synthetic(SsaOp::Jump { target: 1 })); - ssa.add_block(b3a); - - // B3b: false branch of user condition - also sets state = 1 - let mut b3b = SsaBlock::new(4); - b3b.add_instruction(SsaInstruction::synthetic(SsaOp::Jump { target: 1 })); - ssa.add_block(b3b); - - // B5: case 1 - exit path - let mut b5 = SsaBlock::new(5); - b5.add_instruction(SsaInstruction::synthetic(SsaOp::Return { value: None })); - ssa.add_block(b5); - - // B6: default - exit - let mut b6 = SsaBlock::new(6); - b6.add_instruction(SsaInstruction::synthetic(SsaOp::Return { value: None })); - ssa.add_block(b6); - - ssa - } - - #[test] - fn test_reconstruct_with_user_branch() { - let ssa = create_cff_with_user_branch(); - let config = UnflattenConfig::default(); - let tree = trace_method_tree(&ssa, &config, None); - - assert!(tree.dispatcher.is_some(), "Should detect dispatcher"); - assert!( - tree.stats.user_branch_count > 0, - "Should have user branches" - ); - - let result = reconstruct_from_tree(&tree, &ssa); - assert!(result.is_some(), "Reconstruction should succeed"); - - let result = result.unwrap(); - println!("=== Reconstruction with User Branch ==="); - println!("Blocks: {}", result.block_count); - println!( - "State transitions removed: {}", - result.state_transitions_removed - ); - println!( - "User branches preserved: {}", - result.user_branches_preserved - ); - - // Should preserve user branches - assert!( - result.user_branches_preserved > 0, - "Should preserve user branches" - ); - - println!("Block count: {}", result.block_count); - } -} diff --git a/dotscope/src/deobfuscation/passes/unflattening/resolve.rs b/dotscope/src/deobfuscation/passes/unflattening/resolve.rs new file mode 100644 index 00000000..17f2699c --- /dev/null +++ b/dotscope/src/deobfuscation/passes/unflattening/resolve.rs @@ -0,0 +1,1778 @@ +//! Direct resolution of dispatcher edges. +//! +//! # Why this exists +//! +//! Control flow flattening replaces every original edge `A -> B` with a pair of +//! edges through a dispatcher: `A` assigns the state value that encodes `B` and +//! jumps to the dispatcher, which switches on that value and lands on `B`. +//! Undoing it means recovering, for each edge that feeds the dispatcher, the +//! state value travelling along it — and then rewiring that edge straight to the +//! block the dispatcher would have chosen. +//! +//! That is a *per-edge* question, and SSA already answers it: the state reaching +//! the dispatcher is a phi whose operands are indexed by predecessor, so the +//! value arriving from `A` is exactly the operand `A` contributes. No execution +//! path has to be explored to read it. +//! +//! The alternative — walking the method from entry and forking at every +//! conditional to see which states show up where — answers the same question by +//! enumerating paths, and there are exponentially many of those. On a flattened +//! NetReactor method that costs millions of trace nodes per dispatcher to +//! recover a few hundred edges, and it re-explores the whole method once per +//! dispatcher because the walk always restarts at entry. Reading phi operands is +//! linear in the number of edges, and each dispatcher only reads its own. +//! +//! # Merge points +//! +//! The state does not always merge at the dispatcher itself. Obfuscators route +//! it through a chain of copies, and several original edges may meet at a phi +//! one or more blocks above the switch. An operand that is not constant is +//! therefore followed to the phi that defines it, and that phi's operands are +//! resolved in turn. Rewiring then targets the edges into *that* block, which is +//! only sound when everything between it and the dispatcher is state plumbing — +//! copies and unconditional jumps — so the blocks skipped carry no program +//! behaviour. [`pure_chain_between`] enforces exactly that. +//! +//! # Partial results are safe +//! +//! An edge whose state cannot be determined is simply left alone: it keeps +//! routing through the dispatcher, which stays correct but flattened. Coverage +//! degrades, never correctness. +//! +//! Resolution itself is a pure read of the function; [`apply_rewires`] performs +//! the mutation. + +use std::collections::{BTreeMap, BTreeSet}; + +use rustc_hash::FxHashMap; + +use crate::{ + analysis::{CmpKind, ConstValue, PhiNode, SsaFunction, SsaInstruction, SsaOp, SsaVarId}, + deobfuscation::passes::unflattening::dispatcher::Dispatcher, +}; + +/// Maximum definition-chain hops followed when folding a state value. +/// +/// State encodings are short arithmetic chains — a constant, or a constant +/// combined with the previous state. A chain longer than this is not a state +/// computation, and giving up leaves the edge routed through the dispatcher. +const MAX_FOLD_DEPTH: usize = 24; + +/// Maximum blocks walked along a dispatcher's overflow-check chain. +/// +/// The chain has one link per state value that falls outside the switch table. +/// The bound only stops a malformed or adversarial CFG from walking forever; +/// links beyond it simply stay unresolved. +const MAX_OVERFLOW_CHAIN: usize = 4096; + +/// Maximum nesting of state merge points followed above the dispatcher. +/// +/// One level covers the usual "cases meet in a preheader" shape; deeper nesting +/// occurs when an obfuscator stacks several merges. Beyond this the remaining +/// edges stay unresolved. +const MAX_MERGE_DEPTH: usize = 8; + +/// Maximum blocks on a state-plumbing chain between a merge point and the +/// dispatcher. +const MAX_CHAIN_LEN: usize = 16; + +/// Maximum distinct states explored by the propagation fixed-point. +/// +/// One state per original basic block is the norm; the bound only stops a +/// mis-detected dispatcher from enumerating an unbounded value space. +const MAX_PROPAGATED_STATES: usize = 4096; + +/// Maximum blocks in one dispatched case's region. +const MAX_REGION: usize = 4096; + +/// One edge to rewire so it bypasses the dispatcher. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Rewire { + /// Block whose terminator is edited. + pub from: usize, + /// Successor currently named by that terminator. + pub old: usize, + /// Block the dispatcher would have transferred control to. + pub new: usize, + /// State value travelling along the edge. + pub state: i64, +} + +/// Summary of one dispatcher's edge resolution, for reporting. +#[derive(Debug, Clone, Default)] +pub struct ResolutionStats { + /// Edges resolved to a dispatch target. + pub resolved: usize, + /// Edges whose state could not be determined. + pub unresolved: usize, + /// Edges dropped because another edge from the same block wanted a + /// different target. + pub conflicts: usize, + /// Entries recovered from the dispatcher's overflow-check chain. + pub overflow_entries: usize, + /// Distinct states in the dispatch table (switch table plus overflow). + pub table_size: usize, + /// Why edges failed to resolve, for diagnosing coverage gaps. + pub reasons: UnresolvedReasons, +} + +/// Counts of why edges could not be resolved. +#[derive(Debug, Clone, Default)] +pub struct UnresolvedReasons { + /// The state folded, but named no block control could be sent to. + pub no_target: usize, + /// The state is not constant and is not defined by a phi either. + pub not_constant: usize, + /// The state merges at a phi, but the blocks between it and the dispatcher + /// carry program behaviour and cannot be skipped. + pub impure_chain: usize, +} + +/// Maps a state value to the block the dispatcher transfers control to. +/// +/// Built once per dispatcher from the switch table plus the chain of equality +/// checks hanging off the default arm, which is where obfuscators put states +/// whose encoded index falls outside the table. +pub struct DispatchTable { + /// Switch targets, indexed by the value the switch operand evaluates to. + cases: Vec, + /// The variable the switch dispatches on. + switch_var: SsaVarId, + /// The phi carrying the raw state into the dispatcher. + state_var: SsaVarId, + /// Raw state value to target, recovered from the overflow-check chain. + overflow: BTreeMap, + /// Block reached when nothing matches — the end of the overflow chain. + fallthrough: Option, + /// Whether each block still holds instructions. + /// + /// Unflattening runs repeatedly as the pipeline iterates, and a previous + /// round empties the machinery it made unreachable. Those husks stay in the + /// block list and remain named by the switch table, so a later round can + /// resolve a state to one of them. Rewiring control into a block with no + /// terminator produces a function that cannot be laid out, and the passes + /// that follow mangle the surrounding branch trying to make sense of it — + /// which is how a live case block loses its arm. A husk is therefore not a + /// valid answer, and the edge stays with the dispatcher instead. + executable: Vec, +} + +impl DispatchTable { + /// Builds the dispatch table for `dispatcher`. + /// + /// `state_var` is the phi whose value a lookup supplies; the switch operand + /// is evaluated from it rather than from a reconstructed transform, so the + /// index is whatever the dispatcher itself would compute. + pub fn build( + ssa: &SsaFunction, + dispatcher: &Dispatcher, + state_var: SsaVarId, + folder: &mut StateFolder<'_>, + ) -> Self { + let mut table = Self { + cases: dispatcher.cases.clone(), + switch_var: dispatcher.switch_var, + state_var, + overflow: BTreeMap::new(), + fallthrough: None, + executable: ssa + .blocks() + .iter() + .map(|block| !block.instructions().is_empty()) + .collect(), + }; + table.walk_overflow_chain(ssa, dispatcher.default, folder); + table + } + + /// Whether control can be sent to `block`. + fn is_executable(&self, block: usize) -> bool { + self.executable.get(block).copied().unwrap_or(false) + } + + /// Walks the equality-check chain on the dispatcher's default arm. + /// + /// Each link compares the state against a constant and branches to that + /// state's real target, falling through to the next check. The walk stops at + /// the first block that is not such a check; that block is where an + /// unmatched state ends up. + fn walk_overflow_chain( + &mut self, + ssa: &SsaFunction, + default: usize, + folder: &mut StateFolder<'_>, + ) { + let mut current = default; + let mut seen: BTreeSet = BTreeSet::new(); + + for _ in 0..MAX_OVERFLOW_CHAIN { + if !seen.insert(current) { + return; + } + let Some(block) = ssa.block(current) else { + return; + }; + let Some(SsaOp::BranchCmp { + left, + right, + cmp: CmpKind::Eq, + true_target, + false_target, + .. + }) = block.terminator_op() + else { + self.fallthrough = Some(current); + return; + }; + + // Exactly one side is the constant the state is tested against; the + // other is the state itself. If both fold, the comparison is already + // decided and is not a dispatch link. + let left_const = folder.fold(*left); + let right_const = folder.fold(*right); + let value = match (left_const, right_const) { + (Some(v), None) | (None, Some(v)) => v, + _ => { + self.fallthrough = Some(current); + return; + } + }; + + self.overflow.entry(value).or_insert(*true_target); + current = *false_target; + } + } + + /// Returns the block the dispatcher sends `state` to. + /// + /// The case index is obtained by evaluating the dispatcher's own switch + /// operand with the state pinned to `state`, so however the obfuscator + /// encodes the index — a modulo, an xor and a modulo, a table lookup folded + /// into arithmetic — the answer is the one the dispatcher would reach. + /// Applying a separately reconstructed transform instead would silently + /// mis-dispatch whenever detection's model of the encoding was incomplete. + pub fn lookup(&self, folder: &mut StateFolder<'_>, state: StateValue) -> Option { + // The overflow chain tests the raw state, so it is consulted first: its + // entries are the states the switch table cannot express. + let target = if let Some(&target) = self.overflow.get(&state.value) { + target + } else { + let index = folder.fold_with(self.switch_var, self.state_var, state)?; + let index = usize::try_from(index.value).ok()?; + self.cases.get(index).copied().or(self.fallthrough)? + }; + + self.is_executable(target).then_some(target) + } + + /// Number of distinct states this table can dispatch. + #[must_use] + pub fn len(&self) -> usize { + self.cases.len().saturating_add(self.overflow.len()) + } + + /// Whether the table can dispatch no states at all. + #[must_use] + pub fn is_empty(&self) -> bool { + self.cases.is_empty() && self.overflow.is_empty() + } + + /// Number of entries recovered from the overflow chain. + #[must_use] + pub fn overflow_len(&self) -> usize { + self.overflow.len() + } +} + +/// What a variable's defining instruction contributes to a fold. +/// +/// Extracted before recursing so the borrow of the SSA ends before the folder +/// needs itself mutably again. +/// A folded state value together with the width its arithmetic wraps at. +/// +/// CIL evaluates `int32` operands at 32 bits and wraps there. State encodings +/// lean on that: `state * 1975223132` is only the intended value once the +/// product is truncated. Folding at 64 bits instead yields a number that +/// matches no case, so the width travels with the value. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct StateValue { + /// The value, sign-extended from its own width. + pub value: i64, + /// Whether arithmetic on it wraps at 64 bits rather than 32. + pub wide: bool, +} + +impl StateValue { + /// A 32-bit value, sign-extended. + fn narrow(value: i64) -> Self { + Self { + value: truncate32(value), + wide: false, + } + } + + /// A 64-bit value. + fn wide(value: i64) -> Self { + Self { value, wide: true } + } + + /// Re-applies this value's width to `value`. + fn rewrap(self, value: i64) -> Self { + if self.wide { + Self::wide(value) + } else { + Self::narrow(value) + } + } +} + +/// Sign-extends the low 32 bits of `value`, as CIL `int32` arithmetic does. +#[allow(clippy::cast_possible_truncation)] +fn truncate32(value: i64) -> i64 { + i64::from(value as i32) +} + +/// Masks a shift amount to the operand width, as CIL does. +/// +/// A shift by the width or more is undefined in CIL and wraps on real hardware; +/// masking reproduces what the obfuscated code actually computes. +fn shift_amount(amount: i64, wide: bool) -> Option { + let mask: u32 = if wide { 63 } else { 31 }; + u32::try_from(amount).ok().map(|a| a & mask) +} + +/// Reinterprets `value` as unsigned at its own width. +/// +/// A negative 32-bit value arrives here sign-extended, so it has to be narrowed +/// before the unsigned reading is taken; otherwise `rem.un` sees a 64-bit +/// quantity and produces a state no case matches. +#[allow(clippy::cast_possible_truncation)] +fn unsigned_bits(value: i64, wide: bool) -> u64 { + if wide { + value.cast_unsigned() + } else { + u64::from(value as u32) + } +} + +enum Folded { + /// The definition is itself the value. + Value(StateValue), + /// The value passes through unchanged from one operand. + Forward(SsaVarId), + /// The value combines two operands. + Binary(SsaVarId, SsaVarId, BinKind), + /// The value negates or inverts one operand. + Unary(SsaVarId, UnKind), + /// Not a foldable definition. + Opaque, +} + +#[derive(Clone, Copy)] +enum BinKind { + Add, + Sub, + Mul, + Div { unsigned: bool }, + Rem { unsigned: bool }, + And, + Or, + Xor, + Shl, + Shr { unsigned: bool }, +} + +#[derive(Clone, Copy)] +enum UnKind { + Neg, + Not, +} + +/// Folds SSA values to constants by walking definition chains. +/// +/// Memoized across queries for one function: a state encoding usually shares its +/// constants between many edges, and the dispatcher's overflow chain re-asks +/// about the same state variable at every link. +pub struct StateFolder<'a> { + ssa: &'a SsaFunction, + /// Cached folds, keyed by the binding they were computed under. + /// + /// A fold is only valid for the pinning it assumed, so the binding is part + /// of the key rather than a reason to discard the cache. Resolution asks + /// about the same variables under hundreds of different states — every + /// dispatch lookup evaluates the switch operand afresh — and clearing + /// between them made each one re-walk its definition chain from scratch. + memo: FxHashMap<(Bindings, SsaVarId), Option>, + /// Variable pinned to a value for the current query, and that value. + /// + /// Encodings that derive each state from the previous one are only constant + /// once the previous state is known. Pinning the dispatcher's state phi to a + /// concrete state makes the whole chain below it fold. + bindings: Bindings, +} + +/// Variables pinned for the current query. +/// +/// Two slots are enough: one for the state a propagation round is exploring, +/// and one for the branch of a conditional state merge being case-split. They +/// are part of the memo key, so folds under different pinnings coexist. +type Bindings = [Option<(SsaVarId, StateValue)>; 2]; + +impl<'a> StateFolder<'a> { + /// Creates a folder over `ssa`. + #[must_use] + pub fn new(ssa: &'a SsaFunction) -> Self { + Self { + ssa, + memo: FxHashMap::default(), + bindings: [None, None], + } + } + + /// Folds `var` to a constant, or returns `None` if it is not constant. + pub fn fold(&mut self, var: SsaVarId) -> Option { + self.fold_value(var).map(|v| v.value) + } + + /// Folds `var`, keeping the width its arithmetic wraps at. + pub fn fold_value(&mut self, var: SsaVarId) -> Option { + self.bindings = [None, None]; + self.fold_at(var, 0) + } + + /// Folds `var` under an explicit set of pinnings. + pub fn fold_bound(&mut self, var: SsaVarId, bindings: Bindings) -> Option { + self.bindings = bindings; + self.fold_at(var, 0) + } + + /// Folds `var` with `state_var` pinned to `state`. + /// + /// Used to evaluate a state-dependent encoding at one concrete state. + pub fn fold_with( + &mut self, + var: SsaVarId, + state_var: SsaVarId, + state: StateValue, + ) -> Option { + self.bindings = [Some((state_var, state)), None]; + self.fold_at(var, 0) + } + + fn fold_at(&mut self, var: SsaVarId, depth: usize) -> Option { + if depth > MAX_FOLD_DEPTH { + return None; + } + let key = (self.bindings, var); + if let Some(&cached) = self.memo.get(&key) { + return cached; + } + // Seed the memo with "not constant" before recursing. A definition chain + // that loops back on itself then terminates instead of recursing to the + // depth bound on every visit. + self.memo.insert(key, None); + + let result = match self.classify(var) { + Folded::Value(v) => Some(v), + Folded::Forward(src) => self.fold_at(src, depth.saturating_add(1)), + Folded::Unary(src, kind) => { + let v = self.fold_at(src, depth.saturating_add(1))?; + Some(v.rewrap(match kind { + UnKind::Neg => v.value.wrapping_neg(), + UnKind::Not => !v.value, + })) + } + Folded::Binary(left, right, kind) => { + let l = self.fold_at(left, depth.saturating_add(1))?; + let r = self.fold_at(right, depth.saturating_add(1))?; + apply_binary(l, r, kind) + } + Folded::Opaque => None, + }; + + self.memo.insert(key, result); + result + } + + /// Follows copies from `var` to the value they ultimately name. + /// + /// State values reach the dispatcher through chains of copies inserted by + /// the obfuscator; the phi that merges them sits at the end of such a chain. + fn copy_root(&self, var: SsaVarId) -> SsaVarId { + let mut current = var; + for _ in 0..MAX_FOLD_DEPTH { + match self.classify(current) { + Folded::Forward(src) => current = src, + _ => break, + } + } + current + } + + /// Reads the defining instruction of `var` into a [`Folded`]. + fn classify(&self, var: SsaVarId) -> Folded { + // A pinned variable stands for its value regardless of how it is + // defined; that is the whole point of pinning the state phi. + for pinned in self.bindings.iter().flatten() { + if pinned.0 == var { + return Folded::Value(pinned.1); + } + } + let Some(variable) = self.ssa.variable(var) else { + return Folded::Opaque; + }; + let site = variable.def_site(); + // No instruction index means the definition is a phi or the function + // entry; neither is a constant on its own. + let Some(index) = site.instruction else { + return Folded::Opaque; + }; + let Some(op) = self + .ssa + .block(site.block) + .and_then(|b| b.instructions().get(index)) + .map(|i| i.op()) + else { + return Folded::Opaque; + }; + + match op { + SsaOp::Const { value, .. } => value.as_i64().map_or(Folded::Opaque, |v| { + // CIL widens every sub-word integer to int32 on the stack, so + // only the genuinely 64-bit constants wrap at 64 bits. + let wide = matches!( + value, + ConstValue::I64(_) + | ConstValue::U64(_) + | ConstValue::NativeInt(_) + | ConstValue::NativeUInt(_) + ); + Folded::Value(if wide { + StateValue::wide(v) + } else { + StateValue::narrow(v) + }) + }), + SsaOp::Copy { src, .. } => Folded::Forward(*src), + // Width conversions preserve the value over the ranges a state + // encoding uses, so treating them as transparent lets a state that + // round-trips through i32/u32 still fold. + SsaOp::IntConv { operand, .. } => Folded::Forward(*operand), + SsaOp::Add { left, right, .. } => Folded::Binary(*left, *right, BinKind::Add), + SsaOp::Sub { left, right, .. } => Folded::Binary(*left, *right, BinKind::Sub), + SsaOp::Mul { left, right, .. } => Folded::Binary(*left, *right, BinKind::Mul), + SsaOp::And { left, right, .. } => Folded::Binary(*left, *right, BinKind::And), + SsaOp::Or { left, right, .. } => Folded::Binary(*left, *right, BinKind::Or), + SsaOp::Xor { left, right, .. } => Folded::Binary(*left, *right, BinKind::Xor), + SsaOp::Shl { value, amount, .. } => Folded::Binary(*value, *amount, BinKind::Shl), + SsaOp::Shr { + value, + amount, + unsigned, + .. + } => Folded::Binary( + *value, + *amount, + BinKind::Shr { + unsigned: *unsigned, + }, + ), + SsaOp::Div { + left, + right, + unsigned, + .. + } => Folded::Binary( + *left, + *right, + BinKind::Div { + unsigned: *unsigned, + }, + ), + SsaOp::Rem { + left, + right, + unsigned, + .. + } => Folded::Binary( + *left, + *right, + BinKind::Rem { + unsigned: *unsigned, + }, + ), + SsaOp::Neg { operand, .. } => Folded::Unary(*operand, UnKind::Neg), + SsaOp::Not { operand, .. } => Folded::Unary(*operand, UnKind::Not), + _ => Folded::Opaque, + } + } +} + +/// Applies a folded binary operation with CIL's wrapping semantics. +/// +/// Division and remainder return `None` on a zero divisor rather than a value: +/// the original would throw, so there is no state to dispatch. +fn apply_binary(left: StateValue, right: StateValue, kind: BinKind) -> Option { + // Mixed widths only occur in malformed input; taking the wider one keeps + // the fold conservative rather than silently truncating a 64-bit value. + let out = StateValue { + value: 0, + wide: left.wide || right.wide, + }; + let (l, r) = (left.value, right.value); + let value = match kind { + BinKind::Add => l.wrapping_add(r), + BinKind::Sub => l.wrapping_sub(r), + BinKind::Mul => l.wrapping_mul(r), + BinKind::And => l & r, + BinKind::Or => l | r, + BinKind::Xor => l ^ r, + // CIL masks the shift amount to the operand width. + BinKind::Shl => { + let amount = shift_amount(r, out.wide)?; + l.wrapping_shl(amount) + } + BinKind::Shr { unsigned } => { + let amount = shift_amount(r, out.wide)?; + if unsigned { + // A logical shift has to start from the value's own width, or + // the sign extension carried in the high bits shifts in. + let bits = if out.wide { + l.cast_unsigned() + } else { + u64::from(l.cast_unsigned() as u32) + }; + bits.wrapping_shr(amount).cast_signed() + } else { + l.wrapping_shr(amount) + } + } + BinKind::Div { unsigned } => { + if r == 0 { + return None; + } + if unsigned { + unsigned_bits(l, out.wide) + .checked_div(unsigned_bits(r, out.wide))? + .cast_signed() + } else { + l.checked_div(r)? + } + } + BinKind::Rem { unsigned } => { + if r == 0 { + return None; + } + if unsigned { + unsigned_bits(l, out.wide) + .checked_rem(unsigned_bits(r, out.wide))? + .cast_signed() + } else { + l.checked_rem(r)? + } + } + }; + Some(out.rewrap(value)) +} + +/// Finds a phi the value depends on, anywhere in its expression. +/// +/// A conditional state is rarely the phi itself. The obfuscator computes +/// `next = phi ^ key`, so the merge sits one or more operations below the value +/// on the edge. Case-splitting needs the phi wherever it is, and following only +/// copies from the top would miss it. +/// +/// The search is breadth-first so the *shallowest* merge is found: that is the +/// one whose predecessors are closest to the value, and therefore the one whose +/// edges are safest to rewire. +fn dependency_phi<'s>( + ssa: &'s SsaFunction, + folder: &StateFolder<'_>, + value: SsaVarId, +) -> Option<(usize, &'s PhiNode)> { + let mut frontier = vec![value]; + let mut seen: BTreeSet = BTreeSet::new(); + + for _ in 0..MAX_FOLD_DEPTH { + let mut next = Vec::new(); + for var in frontier.drain(..) { + if !seen.insert(var) { + continue; + } + if let Some(found) = ssa.find_phi_defining(var) { + return Some(found); + } + match folder.classify(var) { + Folded::Forward(src) | Folded::Unary(src, _) => next.push(src), + Folded::Binary(left, right, _) => { + next.push(left); + next.push(right); + } + Folded::Value(_) | Folded::Opaque => {} + } + } + if next.is_empty() { + break; + } + frontier = next; + } + None +} + +fn pure_chain_between( + ssa: &SsaFunction, + start: usize, + end: usize, + state_only: &BTreeSet, +) -> bool { + let mut chain: Vec = Vec::new(); + let mut current = start; + + for _ in 0..MAX_CHAIN_LEN { + if current == end { + // Nothing on the chain may define a value that outlives it: control + // reaching the dispatch target directly never runs these blocks, so + // a use further on would have no definition. State encoding is + // consumed by the dispatcher itself and dies with it, which is why + // the usual chain passes this. + return chain + .iter() + .all(|&block| defs_stay_within(ssa, block, &chain, end)); + } + let Some(block) = ssa.block(current) else { + return false; + }; + let Some((terminator, body)) = block.instructions().split_last() else { + return false; + }; + // Only value plumbing may be skipped: constants and copies feeding the + // state encoding, `pop` discarding the duplicate an encoding leaves on + // the stack, and nops. Anything else could be program behaviour. + if !body.iter().all(|instr| is_skippable(instr, state_only)) { + return false; + } + match terminator.op() { + SsaOp::Jump { target } => { + chain.push(current); + current = *target; + } + _ => return false, + } + } + false +} + +/// Whether an instruction can be skipped along with its block. +/// +/// An instruction that produces a state-only value does nothing the program can +/// observe once the state machine is gone, whatever its opcode — the arithmetic +/// of an encoding qualifies just as much as a copy. Anything that produces a +/// value the program still uses, or that acts on the world at all, does not. +fn is_skippable(instr: &SsaInstruction, state_only: &BTreeSet) -> bool { + match instr.def() { + Some(def) => instr.is_pure() && state_only.contains(&def), + None => matches!(instr.op(), SsaOp::Nop | SsaOp::Pop { .. }), + } +} + +/// Whether every value defined in `block` is only read inside the chain. +/// +/// `end` is the dispatcher, which reads the state through its phi; those reads +/// disappear along with the dispatcher once its edges are rewired, so they do +/// not stop the chain from being skippable. +fn defs_stay_within(ssa: &SsaFunction, block: usize, chain: &[usize], end: usize) -> bool { + let Some(ssa_block) = ssa.block(block) else { + return false; + }; + ssa_block + .instructions() + .iter() + .filter_map(|instr| instr.def()) + .all(|def| { + ssa.variable(def).is_none_or(|variable| { + variable + .uses() + .iter() + .all(|site| site.block == end || chain.contains(&site.block)) + }) + }) +} + +/// Finds the phi at the dispatcher that carries the state value. +/// +/// Prefers the phi the detector identified, then a phi defining the switch +/// operand itself — the shape when no transform is applied. Failing both, the +/// switch operand's definition chain is walked to find the phi it is computed +/// from, which is what an encoded dispatcher (`(state ^ key) % n`) looks like. +fn state_phi_at<'s>(ssa: &'s SsaFunction, dispatcher: &Dispatcher) -> Option<&'s PhiNode> { + let block = ssa.block(dispatcher.block)?; + if let Some(state) = dispatcher.state_phi { + if let Some(phi) = block.phi_nodes().iter().find(|p| p.result() == state) { + return Some(phi); + } + } + if let Some(phi) = block + .phi_nodes() + .iter() + .find(|p| p.result() == dispatcher.switch_var) + { + return Some(phi); + } + + let folder = StateFolder::new(ssa); + let mut frontier = vec![dispatcher.switch_var]; + let mut seen: BTreeSet = BTreeSet::new(); + for _ in 0..MAX_FOLD_DEPTH { + let mut next = Vec::new(); + for var in frontier.drain(..) { + if !seen.insert(var) { + continue; + } + if let Some(phi) = block.phi_nodes().iter().find(|p| p.result() == var) { + return Some(phi); + } + match folder.classify(var) { + Folded::Forward(src) | Folded::Unary(src, _) => next.push(src), + Folded::Binary(left, right, _) => { + next.push(left); + next.push(right); + } + Folded::Value(_) | Folded::Opaque => {} + } + } + if next.is_empty() { + break; + } + frontier = next; + } + None +} + +/// Values that exist only to drive the state machine. +/// +/// A value is state-only when every instruction that reads it either lives in +/// the dispatcher itself or produces another state-only value. The chain ends +/// at the dispatcher's phi operands, which no instruction reads — they are +/// consumed by the merge, and the merge disappears with the dispatcher. +/// +/// This is a greatest fixed point: everything is assumed state-only, and a +/// value is struck out as soon as some instruction that is *not* machinery +/// reads it. Striking one value out can strike out the values feeding it, so +/// the sweep repeats until nothing changes. +/// +/// The distinction matters because ConfuserEx decodes the state inside the +/// dispatcher and has every case block read the decoded value back to compute +/// its successor. Those reads look like real uses, but they die with the state +/// machine — so the dispatcher can still be bypassed. A case block that used +/// the same value for actual work would not be struck out, and bypassing would +/// then be unsafe. +fn state_only_values(ssa: &SsaFunction, dispatcher_block: usize) -> BTreeSet { + let mut state_only: BTreeSet = ssa + .variables() + .iter() + .map(|variable| variable.id()) + .collect(); + + loop { + let mut struck = false; + for (index, block) in ssa.iter_blocks() { + // Everything the dispatcher itself does is machinery. + if index == dispatcher_block { + continue; + } + for instr in block.instructions() { + if instr.def().is_some_and(|def| state_only.contains(&def)) { + continue; + } + for used in instr.uses() { + if state_only.remove(&used) { + struck = true; + } + } + } + } + if !struck { + break; + } + } + + state_only +} + +/// Whether the dispatcher block computes nothing the blocks after it depend on. +/// +/// A bare `switch` on a merged state is transparent: control that bypasses it +/// misses no computation, so rewiring some edges while others still route +/// through it is safe. +/// +/// A dispatcher that decodes the state in its own body is not. ConfuserEx's +/// `(state ^ key)` is duplicated into a local that each case block reads to +/// derive its successor, so a bypassing edge skips the definition the surviving +/// paths still use. Such a dispatcher may only be bypassed if *every* edge is +/// rewired, which makes the whole block dead and the question moot. +/// +/// Phi nodes do not count: they name a merge rather than compute anything, and +/// `rebuild_ssa` re-derives them from the rewired graph. +fn dispatcher_is_transparent( + ssa: &SsaFunction, + dispatcher_block: usize, + state_only: &BTreeSet, +) -> bool { + let Some(block) = ssa.block(dispatcher_block) else { + return false; + }; + let Some((_terminator, body)) = block.instructions().split_last() else { + return false; + }; + // A definition that leaves the dispatcher is fine as long as everything it + // feeds is state machinery, which dies along with the dispatcher. Requiring + // the definition never to leave at all would refuse every encoded + // dispatcher, and with it every ConfuserEx method. + body.iter() + .filter_map(SsaInstruction::def) + .all(|def| state_only.contains(&def)) +} + +/// Blocks reachable from `start` without entering `stop`. +/// +/// Used to find which edges into the dispatcher belong to one dispatched case, +/// so a state-dependent encoding is evaluated only where that state actually +/// arrives. +fn region_from(ssa: &SsaFunction, start: usize, stop: usize, budget: usize) -> BTreeSet { + let mut seen = BTreeSet::new(); + if start == stop { + return seen; + } + let mut frontier = vec![start]; + while let Some(current) = frontier.pop() { + if current == stop || !seen.insert(current) { + continue; + } + if seen.len() > budget { + break; + } + if let Some(op) = ssa.block(current).and_then(|b| b.terminator_op()) { + frontier.extend(op.successors()); + } + } + seen +} + +/// Resolves every edge feeding `dispatcher` to the block it should reach. +/// +/// Returns the rewires to apply and statistics for reporting. The SSA is only +/// read; use [`apply_rewires`] to perform the change. +pub fn resolve_dispatch_edges( + ssa: &SsaFunction, + dispatcher: &Dispatcher, +) -> (Vec, ResolutionStats) { + let mut folder = StateFolder::new(ssa); + let mut stats = ResolutionStats::default(); + + let Some(phi) = state_phi_at(ssa, dispatcher) else { + // Without a state phi there is nothing per-edge to read: the switch + // operand is computed inside the dispatcher from something that is not + // merged at its entry. + stats.unresolved = ssa.block_predecessors(dispatcher.block).len(); + return (Vec::new(), stats); + }; + + let state_only = state_only_values(ssa, dispatcher.block); + let table = DispatchTable::build(ssa, dispatcher, phi.result(), &mut folder); + stats.overflow_entries = table.overflow_len(); + stats.table_size = table.len(); + + // Pass one: read the states straight out of the phi graph. This resolves + // encodings whose next state is a constant, following merges upward when + // several original edges meet before the jump. + let mut rewires: Vec = Vec::new(); + let mut states: Vec = Vec::new(); + let mut visited: BTreeSet = BTreeSet::new(); + let mut unresolved: BTreeSet = BTreeSet::new(); + resolve_merge( + ssa, + &table, + &mut folder, + dispatcher.block, + phi, + dispatcher.block, + 0, + &mut visited, + &mut rewires, + &mut states, + &mut unresolved, + &mut stats.reasons, + None, + None, + &state_only, + ); + + // Pass two: encodings that derive each state from the previous one leave + // edges no constant can be read from. Propagating concrete states through + // the dispatcher resolves those, seeded with every state pass one proved + // reachable — including states found behind a merge, which is the only way + // a case block reached solely through a conditional becomes visible. + if !unresolved.is_empty() { + let (recovered, covered) = propagate_states( + ssa, + &table, + dispatcher, + phi, + &mut folder, + &states, + &state_only, + ); + for pred in covered { + unresolved.remove(&pred); + } + rewires.extend(recovered); + } + + let rewires = drop_conflicts(rewires, &mut stats); + stats.unresolved = unresolved.len(); + + // A dispatcher that decodes the state in its own body may only be bypassed + // wholesale: leaving one edge routed through it means the surviving path + // still needs the definitions a bypassing edge would skip. + if stats.unresolved > 0 && !dispatcher_is_transparent(ssa, dispatcher.block, &state_only) { + log::debug!( + "CFF resolve b{}: {} edge(s) unresolved and the dispatcher decodes state \ + in its body, so none are rewired", + dispatcher.block, + stats.unresolved + ); + stats.resolved = 0; + return (Vec::new(), stats); + } + + stats.resolved = rewires.len(); + (rewires, stats) +} + +/// Resolves state-dependent encodings by propagating concrete states. +/// +/// When the next state is computed from the current one — ConfuserEx's +/// `next = (state ^ key) * a ^ b` — no edge carries a constant, and the value +/// only becomes concrete once the state that reached the case block is known. +/// +/// Starting from the states that *are* known, this walks the state machine the +/// way it actually runs: pin the state phi to a known state, see which case +/// block the dispatcher selects, evaluate the edges leaving that case block to +/// get the next states, and repeat until no new state appears. It is a +/// fixed-point over states — at most one iteration per original block — not an +/// enumeration of execution paths, so a method with many conditionals costs no +/// more than one with none. +/// +/// An edge is only rewired when exactly one state can reach it. Sharing a tail +/// between cases, or falling through from one case into the next, puts an edge +/// in more than one state's region; the value it carries then depends on how +/// control arrived, which a single successor cannot express. Those edges keep +/// using the dispatcher rather than being wired to whichever state happened to +/// be examined last. +fn propagate_states( + ssa: &SsaFunction, + table: &DispatchTable, + dispatcher: &Dispatcher, + phi: &PhiNode, + folder: &mut StateFolder<'_>, + seeds: &[StateValue], + state_only: &BTreeSet, +) -> (Vec, BTreeSet) { + let state_var = phi.result(); + let operands: Vec<(usize, SsaVarId)> = phi + .operands() + .iter() + .map(|op| (op.predecessor(), op.value())) + .collect(); + + // Seeds: the states already known without any context — the entry edge, and + // whatever the detector recorded as the initial state. + let mut worklist: Vec = seeds.to_vec(); + if let Some(initial) = dispatcher.initial_state { + worklist.push(StateValue::narrow(initial)); + } + for &(_, value) in &operands { + if let Some(state) = folder.fold_value(value) { + worklist.push(state); + } + } + + // Which states can reach each edge, and what the edge yields under each. + let mut merged: Vec = Vec::new(); + // Dispatcher edges this pass accounted for, whether directly or by + // resolving the merge that feeds them. + let mut covered: BTreeSet = BTreeSet::new(); + let mut reaching: BTreeMap> = BTreeMap::new(); + let mut outcome: BTreeMap<(usize, i64), StateValue> = BTreeMap::new(); + let mut seen: BTreeSet = BTreeSet::new(); + + while let Some(state) = worklist.pop() { + if seen.len() >= MAX_PROPAGATED_STATES || !seen.insert(state.value) { + continue; + } + let Some(target) = table.lookup(folder, state) else { + continue; + }; + + let region = region_from(ssa, target, dispatcher.block, MAX_REGION); + + for &(pred, value) in &operands { + if !region.contains(&pred) { + continue; + } + reaching.entry(pred).or_default().insert(state.value); + if let Some(next) = folder.fold_with(value, state_var, state) { + outcome.insert((pred, state.value), next); + worklist.push(next); + continue; + } + + // The edge's value is not constant even with the state pinned: an + // original conditional inside this case picked between two + // successors, and they meet at a phi before the jump. Its operands + // are per-edge, and we are inside the region this state dispatches + // to, so each one can be read under the same pinning. The edges to + // rewire are that phi's, which is only sound when what lies between + // it and the dispatcher is state plumbing. + let Some((inner_block, inner_phi)) = dependency_phi(ssa, folder, value) else { + continue; + }; + if inner_block == dispatcher.block + || !pure_chain_between(ssa, inner_block, dispatcher.block, state_only) + { + continue; + } + let inner: Vec<(usize, SsaVarId)> = inner_phi + .operands() + .iter() + .map(|op| (op.predecessor(), op.value())) + .collect(); + let phi_result = inner_phi.result(); + let pinned = Some((state_var, state)); + for (source, operand) in inner { + // Read what the branch contributes, then the edge value with + // the merge pinned to it — both under the state being explored. + let Some(branch) = folder.fold_bound(operand, [pinned, None]) else { + continue; + }; + let Some(next) = folder.fold_bound(value, [pinned, Some((phi_result, branch))]) + else { + continue; + }; + let Some(target) = table.lookup(folder, next) else { + continue; + }; + merged.push(Rewire { + from: source, + old: inner_block, + new: target, + state: next.value, + }); + // The dispatcher edge is answered by rewiring the merge that + // feeds it, even though no rewire names the edge itself. + covered.insert(pred); + worklist.push(next); + } + } + } + + // Emit what every state reaching an edge implies for it. + // + // An edge reachable under several states is not by itself a problem: the + // question is whether those states disagree about where it should go. They + // usually do not — `region_from` is a forward reachability and over-reports, + // so cases that share a tail all claim the same edge and all compute the + // same successor for it. Where they genuinely disagree the edge would need + // the block duplicated, and `drop_conflicts` removes it. Deciding here on + // the *number* of states instead would discard the agreeing majority along + // with the conflicting few. + let mut rewires: Vec = merged; + for (pred, states) in reaching { + for state in states { + let Some(&next) = outcome.get(&(pred, state)) else { + continue; + }; + let Some(next_target) = table.lookup(folder, next) else { + continue; + }; + covered.insert(pred); + rewires.push(Rewire { + from: pred, + old: dispatcher.block, + new: next_target, + state: next.value, + }); + } + } + + (rewires, covered) +} + +/// Combines the rewires of every dispatcher in a method into one set. +/// +/// Dispatchers are resolved independently against the unmodified function, so +/// two of them can name the same edge. Conflicts are dropped here for the same +/// reason they are dropped within a single dispatcher: an edge can only go one +/// place, and guessing which is worse than leaving it flattened. +#[must_use] +pub fn merge_rewires(per_dispatcher: Vec>) -> (Vec, usize) { + let mut stats = ResolutionStats::default(); + let combined: Vec = per_dispatcher.into_iter().flatten().collect(); + let merged = drop_conflicts(combined, &mut stats); + (merged, stats.conflicts) +} + +/// Resolves the operands of one state merge point. +/// +/// `merge_block` is where the phi lives and therefore which edges get rewired; +/// `dispatcher_block` is the switch those edges ultimately feed, used to check +/// that everything skipped in between is state plumbing. +/// +/// Failures are attributed to `root` — the edge into the dispatcher this +/// resolution ultimately serves — so a merge that only partly resolves is +/// reported against the one dispatcher edge it feeds, not against its own +/// operands. +#[allow(clippy::too_many_arguments)] +fn resolve_merge( + ssa: &SsaFunction, + table: &DispatchTable, + folder: &mut StateFolder<'_>, + merge_block: usize, + phi: &PhiNode, + dispatcher_block: usize, + depth: usize, + visited: &mut BTreeSet, + rewires: &mut Vec, + states: &mut Vec, + unresolved: &mut BTreeSet, + reasons: &mut UnresolvedReasons, + root: Option, + outer: Option<(SsaVarId, StateValue)>, + state_only: &BTreeSet, +) { + if !visited.insert(merge_block) { + return; + } + + // Operands are copied out so the folder can borrow the SSA again. + let operands: Vec<(usize, SsaVarId)> = phi + .operands() + .iter() + .map(|op| (op.predecessor(), op.value())) + .collect(); + + for (pred, value) in operands { + // At the top level each operand answers for itself; inside a merge every + // failure counts against the dispatcher edge the merge feeds. + let blame = root.unwrap_or(pred); + + if let Some(state) = folder.fold_bound(value, [outer, None]) { + if let Some(target) = table.lookup(folder, state) { + states.push(state); + rewires.push(Rewire { + from: pred, + old: merge_block, + new: target, + state: state.value, + }); + } else { + unresolved.insert(blame); + reasons.no_target = reasons.no_target.saturating_add(1); + } + continue; + } + + // Not constant: follow the copies to the phi that merges this value and + // resolve that phi's operands instead. Its block becomes the new set of + // edges to rewire, so the blocks between it and the dispatcher must be + // skippable. + let Some((inner_block, inner_phi)) = dependency_phi(ssa, folder, value) else { + unresolved.insert(blame); + reasons.not_constant = reasons.not_constant.saturating_add(1); + continue; + }; + if depth >= MAX_MERGE_DEPTH + || inner_block == merge_block + || !pure_chain_between(ssa, inner_block, dispatcher_block, state_only) + { + unresolved.insert(blame); + reasons.impure_chain = reasons.impure_chain.saturating_add(1); + continue; + } + + // Case-split: read the value once per branch of the merge, pinning the + // phi to what that branch contributes. When the phi *is* the value this + // is exactly the old recursion; when it sits under an encoding it also + // pushes the split through the arithmetic. + let operands: Vec<(usize, SsaVarId)> = inner_phi + .operands() + .iter() + .map(|op| (op.predecessor(), op.value())) + .collect(); + let phi_result = inner_phi.result(); + let mut split_any = false; + + for (source, operand) in operands { + let Some(branch) = folder.fold_bound(operand, [outer, None]) else { + continue; + }; + let Some(state) = folder.fold_bound(value, [outer, Some((phi_result, branch))]) else { + continue; + }; + let Some(target) = table.lookup(folder, state) else { + continue; + }; + states.push(state); + rewires.push(Rewire { + from: source, + old: inner_block, + new: target, + state: state.value, + }); + split_any = true; + } + + if !split_any { + unresolved.insert(blame); + reasons.not_constant = reasons.not_constant.saturating_add(1); + } + } +} + +/// Drops rewires that disagree about where a block should go. +/// +/// One block can reach a merge point along two edges — both arms of a branch, +/// say — and SSA gives each its own phi operand. If those operands encode +/// different successors the block would need duplicating to express both, so +/// both are dropped and the block keeps using the dispatcher. Rewiring only one +/// of them would silently send the other arm to the wrong place. +fn drop_conflicts(rewires: Vec, stats: &mut ResolutionStats) -> Vec { + let mut chosen: BTreeMap<(usize, usize), Rewire> = BTreeMap::new(); + let mut conflicted: BTreeSet<(usize, usize)> = BTreeSet::new(); + + for rewire in rewires { + let key = (rewire.from, rewire.old); + match chosen.get(&key) { + Some(existing) if existing.new != rewire.new => { + conflicted.insert(key); + } + Some(_) => {} + None => { + chosen.insert(key, rewire); + } + } + } + + for key in &conflicted { + chosen.remove(key); + stats.conflicts = stats.conflicts.saturating_add(1); + stats.unresolved = stats.unresolved.saturating_add(1); + } + + chosen.into_values().collect() +} + +/// Whether an instruction is state machinery rather than program behaviour. +/// +/// State machinery is side-effect free and carries no content an analyst would +/// look for: integer constants encoding states, the copies that move them, and +/// the control flow that dispatches on them. A string constant fails the test +/// even though it is pure — losing it loses evidence. +fn is_state_machinery(instr: &SsaInstruction) -> bool { + if let SsaOp::Const { value, .. } = instr.op() { + return value.as_i64().is_some(); + } + instr.is_pure() + || matches!( + instr.op(), + SsaOp::Jump { .. } + | SsaOp::Leave { .. } + | SsaOp::Switch { .. } + | SsaOp::Branch { .. } + | SsaOp::BranchCmp { .. } + ) +} + +/// Empties the state machinery the rewiring made unreachable. +/// +/// Once all of a dispatcher's edges bypass it, the dispatcher and the constants +/// that encoded the state can no longer execute. They are emptied rather than +/// removed so block indices stay stable for the rewires already applied; +/// `rebuild_ssa` and the dead-code passes drop the remains. +/// +/// Only blocks that are *provably nothing but* state machinery are emptied. +/// Unreachability here is relative to the states resolution managed to +/// discover, and that discovery is deliberately incomplete — an edge it cannot +/// read leaves its case block looking unreachable when it is not. Emptying such +/// a block would turn a gap in coverage into lost program behaviour, so a block +/// holding a call, a store, or a string constant is left alone even when +/// nothing appears to reach it. Dead code costs a little size; deleted code +/// costs the analysis it was kept for. +/// +/// Handler entry blocks are roots alongside the function entry: control reaches +/// them by a runtime exception edge, not from any terminator. +/// +/// Returns the number of blocks emptied. +pub fn clear_unreachable(ssa: &mut SsaFunction) -> usize { + let block_count = ssa.blocks().len(); + if block_count == 0 { + return 0; + } + + let mut roots: Vec = vec![0]; + for handler in ssa.exception_handlers() { + roots.extend(handler.handler_start_block); + roots.extend(handler.filter_start_block); + roots.extend(handler.try_start_block); + } + + let mut reachable = vec![false; block_count]; + let mut frontier = roots; + while let Some(current) = frontier.pop() { + let Some(slot) = reachable.get_mut(current) else { + continue; + }; + if *slot { + continue; + } + *slot = true; + if let Some(op) = ssa.block(current).and_then(|b| b.terminator_op()) { + frontier.extend(op.successors()); + } + } + + let dead: Vec = (0..block_count) + .filter(|&index| !reachable.get(index).copied().unwrap_or(true)) + .filter(|&index| { + ssa.block(index).is_some_and(|block| { + (!block.instructions().is_empty() || !block.phi_nodes().is_empty()) + && block.instructions().iter().all(is_state_machinery) + }) + }) + .collect(); + + for index in &dead { + if let Some(block) = ssa.block_mut(*index) { + block.clear(); + } + } + dead.len() +} + +/// Rewires resolved edges so they bypass the dispatcher. +/// +/// Each edge's source has its terminator's reference to the merge point +/// replaced by the dispatch target. Predecessor lists are derived from +/// terminators, so this is the whole of the CFG change: the stale phi operands +/// and the now-dead constants that encoded the state are cleaned up by the +/// caller's `rebuild_ssa` and the ordinary dead-code passes. +/// +/// Returns the number of edges actually rewired. +pub fn apply_rewires(ssa: &mut SsaFunction, rewires: &[Rewire]) -> usize { + let mut applied: usize = 0; + for rewire in rewires { + let changed = ssa + .block_mut(rewire.from) + .and_then(|block| block.instructions_mut().last_mut()) + .is_some_and(|term| term.op_mut().redirect_target(rewire.old, rewire.new)); + if changed { + applied = applied.saturating_add(1); + } + } + applied +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::analysis::{DefSite, PhiOperand, SsaBlock, SsaInstruction, SsaType, VariableOrigin}; + + /// Appends an instruction to `block`, creating its destination variable with + /// the def site the folder needs to find it again. + fn define( + ssa: &mut SsaFunction, + block: usize, + make: impl FnOnce(SsaVarId) -> SsaOp, + ) -> SsaVarId { + let index = ssa.block(block).map_or(0, |b| b.instructions().len()); + let var = ssa.create_variable( + VariableOrigin::Phi, + 0, + DefSite::instruction(block, index), + SsaType::I32, + ); + let op = make(var); + if let Some(b) = ssa.block_mut(block) { + b.add_instruction(SsaInstruction::synthetic(op)); + } + var + } + + fn constant(ssa: &mut SsaFunction, block: usize, value: i32) -> SsaVarId { + define(ssa, block, |dest| SsaOp::Const { + dest, + value: ConstValue::I32(value), + }) + } + + fn terminate(ssa: &mut SsaFunction, block: usize, op: SsaOp) { + if let Some(b) = ssa.block_mut(block) { + b.add_instruction(SsaInstruction::synthetic(op)); + } + } + + /// A dispatcher whose cases each end in `state = ; jump dispatcher`. + /// + /// Block 0 enters with state 1, block 2 is the dispatcher, blocks 3 and 4 + /// are its cases; case 3 moves to state 0 and case 4 exits. + fn constant_state_cff() -> (SsaFunction, Dispatcher) { + let mut ssa = SsaFunction::new(0, 0); + for index in 0..5 { + ssa.add_block(SsaBlock::new(index)); + } + + let entry_state = constant(&mut ssa, 0, 1); + terminate(&mut ssa, 0, SsaOp::Jump { target: 2 }); + + let case3_state = constant(&mut ssa, 3, 0); + terminate(&mut ssa, 3, SsaOp::Jump { target: 2 }); + + // The dispatcher's state phi: one operand per predecessor edge. + let state = ssa.create_variable(VariableOrigin::Phi, 1, DefSite::phi(2), SsaType::I32); + let mut phi = PhiNode::new(state, VariableOrigin::Phi); + phi.add_operand(PhiOperand::new(entry_state, 0)); + phi.add_operand(PhiOperand::new(case3_state, 3)); + if let Some(block) = ssa.block_mut(2) { + block.add_phi(phi); + block.add_instruction(SsaInstruction::synthetic(SsaOp::Switch { + value: state, + targets: vec![4, 3], + default: 1, + })); + } + + terminate(&mut ssa, 4, SsaOp::Return { value: None }); + terminate(&mut ssa, 1, SsaOp::Return { value: None }); + + let dispatcher = Dispatcher::new(2, state, vec![4, 3], 1).with_state_phi(state); + (ssa, dispatcher) + } + + #[test] + fn folds_arithmetic_at_32_bit_width() { + let mut ssa = SsaFunction::new(0, 0); + ssa.add_block(SsaBlock::new(0)); + + let left = constant(&mut ssa, 0, 1_975_223_132); + let right = constant(&mut ssa, 0, 3); + let product = define(&mut ssa, 0, |dest| SsaOp::Mul { + dest, + left, + right, + flags: None, + }); + + let mut folder = StateFolder::new(&ssa); + // 1975223132 * 3 overflows int32; CIL wraps, and so must the folder — + // at 64 bits the product would be 5925669396 and match no case. + assert_eq!(folder.fold(product), Some(1_630_702_100)); + } + + #[test] + fn folds_through_copy_chains() { + let mut ssa = SsaFunction::new(0, 0); + ssa.add_block(SsaBlock::new(0)); + + let base = constant(&mut ssa, 0, 42); + let first = define(&mut ssa, 0, |dest| SsaOp::Copy { dest, src: base }); + let second = define(&mut ssa, 0, |dest| SsaOp::Copy { dest, src: first }); + + let mut folder = StateFolder::new(&ssa); + assert_eq!(folder.fold(second), Some(42)); + assert_eq!(folder.copy_root(second), base); + } + + #[test] + fn dispatch_table_reads_the_overflow_chain() { + let mut ssa = SsaFunction::new(0, 0); + for index in 0..10 { + ssa.add_block(SsaBlock::new(index)); + } + // Every target must be able to execute, or the table rejects it. + for index in [2, 9] { + terminate(&mut ssa, index, SsaOp::Return { value: None }); + } + + // Default arm: `if state == 700 goto 2` then fall through to block 3. + // The state must be opaque here — a link whose both sides fold is a + // comparison already decided, not a dispatch. + let state = + ssa.create_variable(VariableOrigin::Local(0), 0, DefSite::entry(), SsaType::I32); + let probe = constant(&mut ssa, 0, 700); + terminate( + &mut ssa, + 0, + SsaOp::BranchCmp { + left: state, + right: probe, + cmp: CmpKind::Eq, + unsigned: false, + true_target: 2, + false_target: 3, + }, + ); + terminate(&mut ssa, 3, SsaOp::Return { value: None }); + + let dispatcher = Dispatcher::new(1, state, vec![9], 0); + let mut folder = StateFolder::new(&ssa); + let table = DispatchTable::build(&ssa, &dispatcher, state, &mut folder); + + assert_eq!(table.overflow_len(), 1); + // A state outside the switch table is routed by the chain, not the table. + assert_eq!(table.lookup(&mut folder, StateValue::narrow(700)), Some(2)); + // A state inside the table still uses the table. + assert_eq!(table.lookup(&mut folder, StateValue::narrow(0)), Some(9)); + // Anything else lands where the chain falls through. + assert_eq!(table.lookup(&mut folder, StateValue::narrow(123)), Some(3)); + } + + #[test] + fn resolves_constant_state_edges() { + let (ssa, dispatcher) = constant_state_cff(); + let (rewires, stats) = resolve_dispatch_edges(&ssa, &dispatcher); + + assert_eq!(stats.unresolved, 0, "both edges carry a constant state"); + assert_eq!(stats.resolved, 2); + + // State 1 selects targets[1] = block 3; state 0 selects targets[0] = 4. + let mut targets: Vec<(usize, usize)> = rewires.iter().map(|r| (r.from, r.new)).collect(); + targets.sort_unstable(); + assert_eq!(targets, vec![(0, 3), (3, 4)]); + } + + #[test] + fn applying_rewires_bypasses_the_dispatcher() { + let (mut ssa, dispatcher) = constant_state_cff(); + let (rewires, _) = resolve_dispatch_edges(&ssa, &dispatcher); + + assert_eq!(apply_rewires(&mut ssa, &rewires), 2); + assert!( + ssa.block_predecessors(2).is_empty(), + "no edge should still reach the dispatcher" + ); + + // With every edge rewired the dispatcher is unreachable, and being pure + // state machinery it is emptied, taking its switch with it. + assert_eq!(clear_unreachable(&mut ssa), 1); + assert!(ssa.block(2).is_some_and(|b| b.instructions().is_empty())); + + // The default arm is unreachable too, but it returns — behaviour, not + // machinery — so it is left intact rather than deleted on the strength + // of an analysis that is allowed to be incomplete. + assert!(ssa.block(1).is_some_and(|b| !b.instructions().is_empty())); + } + + #[test] + fn emptied_dispatch_targets_are_not_rewired_into() { + let (mut ssa, dispatcher) = constant_state_cff(); + + // Empty the block state 1 dispatches to, as a previous unflattening + // round does to machinery it made unreachable. The switch table still + // names it. + if let Some(block) = ssa.block_mut(3) { + block.clear(); + } + + let (rewires, stats) = resolve_dispatch_edges(&ssa, &dispatcher); + + assert!( + rewires.iter().all(|r| r.new != 3), + "no edge may be rewired into a block that cannot execute" + ); + // Both edges are lost: one dispatches to the husk, and the other is the + // husk's own edge, whose state the emptying took with it. + assert_eq!(stats.unresolved, 2, "those edges keep using the dispatcher"); + + apply_rewires(&mut ssa, &rewires); + + // The dispatcher's own switch still names the husk — that is the input + // condition, and leaving it is what keeps the edge safe. What must not + // happen is a rewired block acquiring an edge into it. + for rewire in &rewires { + let successors = ssa + .block(rewire.from) + .and_then(|b| b.terminator_op()) + .map(SsaOp::successors) + .unwrap_or_default(); + assert!( + !successors.contains(&3), + "rewired block b{} must not send control into an empty block", + rewire.from + ); + } + } + + #[test] + fn conflicting_edges_are_dropped() { + let mut stats = ResolutionStats::default(); + let kept = drop_conflicts( + vec![ + Rewire { + from: 5, + old: 2, + new: 7, + state: 1, + }, + // Same edge, different destination: unrepresentable without + // duplicating block 5, so neither survives. + Rewire { + from: 5, + old: 2, + new: 9, + state: 2, + }, + Rewire { + from: 6, + old: 2, + new: 7, + state: 1, + }, + ], + &mut stats, + ); + + assert_eq!(kept.len(), 1); + assert_eq!(kept[0].from, 6); + assert_eq!(stats.conflicts, 1); + } + + #[test] + fn unresolved_edges_leave_the_dispatcher_in_place() { + let (mut ssa, dispatcher) = constant_state_cff(); + + // Blank out the entry block's constant. The phi operand still names the + // variable, but its defining instruction no longer produces a value, so + // the edge cannot be folded. + if let Some(instr) = ssa + .block_mut(0) + .and_then(|block| block.instructions_mut().first_mut()) + { + instr.set_op(SsaOp::Nop); + } + + let (rewires, stats) = resolve_dispatch_edges(&ssa, &dispatcher); + assert_eq!(stats.unresolved, 1, "the entry edge no longer folds"); + + apply_rewires(&mut ssa, &rewires); + assert_eq!( + ssa.block_predecessors(2), + vec![0], + "the unresolved edge keeps using the dispatcher" + ); + assert_eq!( + clear_unreachable(&mut ssa), + 0, + "a reachable dispatcher is never emptied" + ); + } +} diff --git a/dotscope/src/deobfuscation/passes/unflattening/tracer/context.rs b/dotscope/src/deobfuscation/passes/unflattening/tracer/context.rs deleted file mode 100644 index 1dad825b..00000000 --- a/dotscope/src/deobfuscation/passes/unflattening/tracer/context.rs +++ /dev/null @@ -1,731 +0,0 @@ -//! Tracing context and state management. -//! -//! [`TreeTraceContext`] owns all mutable state for a single trace pass: -//! -//! - **SSA evaluator**: Concrete evaluation of instructions and PHI nodes -//! - **Taint tracking**: Which variables are derived from the CFF state variable -//! - **Visit tracking**: (block, state) pairs visited to detect loops -//! - **Case dispatch tracking**: How many times each switch case was taken -//! - **Fork snapshotting**: [`ContextSnapshot`] saves/restores state at branch points -//! -//! All fields are private — the [`engine`](super::engine) interacts through -//! semantic methods that encapsulate invariants (visit budgets, case loop -//! detection thresholds, expression switch mode transitions). - -use std::mem; - -use analyssa::BitSet; -use rustc_hash::FxHashSet; - -use crate::{ - analysis::{ - cff_taint_config, ConstValue, EvaluatorMark, SsaEvaluator, SsaFunction, SsaOp, SsaVarId, - SsaVariable, TaintAnalysis, VariableOrigin, - }, - deobfuscation::passes::unflattening::{ - tracer::{helpers, types::TracedDispatcher}, - UnflattenConfig, - }, - CilObject, -}; - -/// Multiplier applied to `max_block_visits` to derive the monotonic global -/// visit cap. -/// -/// Generous enough that methods which trace normally never reach it — only -/// pathological nesting, where the per-arm budget reset would otherwise let -/// tracing run unbounded, is cut off. -const GLOBAL_VISIT_BUDGET_FACTOR: usize = 4; - -/// Per-block properties that depend only on the shape of the SSA and on which -/// blocks are dispatchers — never on the values a particular path carries. -/// -/// The tracer asks these questions once per block *visit*, and a flattened -/// method is visited millions of times, so answering them by re-walking the -/// CFG turns constant-time predicates into a dominant cost. Every field here is -/// a pure function of `(ssa, dispatcher, other_dispatcher_blocks)`, so a single -/// computation is valid for the whole trace. -struct BlockFacts { - /// Whether the block is a direct target of the dispatcher switch. - /// - /// Replaces a linear scan of the switch's target list, which for a - /// flattened method has one entry per case. - dispatch_target: Vec, - - /// Whether the block is a dispatcher of a *different* CFF instance in the - /// same method. - other_dispatcher: Vec, - - /// Jump target of the block if it is a constant-producer, else `None`. - /// - /// See [`helpers::const_producer_target`], which this caches. - const_producer: Vec>, - - /// Memoized [`overflow_dispatch_site`](Self::overflow_dispatch_site) - /// answers: `None` until first queried. - /// - /// Computed lazily rather than eagerly because the predecessor walk is - /// only ever asked about blocks ending in an equality comparison, a small - /// minority of the CFG. - overflow_site: Vec>, -} - -impl BlockFacts { - fn new(ssa: &SsaFunction, dispatcher: Option<&TracedDispatcher>) -> Self { - let count = ssa.block_count(); - let mut dispatch_target = vec![false; count]; - if let Some(d) = dispatcher { - for &target in d.targets.iter().chain(std::iter::once(&d.default)) { - if let Some(slot) = dispatch_target.get_mut(target) { - *slot = true; - } - } - } - - let mut const_producer = vec![None; count]; - for block in ssa.blocks() { - if let Some(slot) = const_producer.get_mut(block.id()) { - *slot = helpers::const_producer_target(block); - } - } - - Self { - dispatch_target, - other_dispatcher: vec![false; count], - const_producer, - overflow_site: vec![None; count], - } - } -} - -/// Context for tree-based tracing. -/// -/// Fields are private — engine and helpers interact through semantic methods -/// that encapsulate the invariants (visit budgets, case dispatch tracking, -/// dispatcher queries, taint propagation). -pub struct TreeTraceContext<'a> { - ssa: &'a SsaFunction, - evaluator: SsaEvaluator<'a>, - assembly: Option<&'a CilObject>, - dispatcher: Option, - state_tainted: BitSet, - next_node_id: usize, - total_visits: usize, - /// Monotonic visit counter that is never reset, unlike `total_visits`. - /// - /// `total_visits` is deliberately reset when entering an expression-switch - /// false arm so each arm gets its own budget. On heavily nested methods that - /// reset fires often enough that the effective budget becomes unbounded — a - /// 348-block ConfuserEx method reached 3.3M block visits against a 50k - /// budget. This counter is the backstop: it bounds total tracing work per - /// context while leaving the per-arm budget semantics untouched. - global_visits: usize, - /// Hard cap on `global_visits`, derived from `max_block_visits`. - max_global_visits: usize, - /// Blocks already entered on the current path, keyed by execution state. - /// - /// Only ever grows while a path is followed — [`mark_visited`](Self::mark_visited) - /// is its sole writer — so a fork records the keys it adds and a restore - /// removes them again, the same trick the evaluator uses. Copying the set - /// per fork was the tracer's largest remaining cost once the evaluator - /// stopped being copied. - visited_states: FxHashSet<(usize, i64)>, - /// Keys added to [`visited_states`](Self::visited_states) since journaling - /// began, oldest first. - visited_journal: Vec<(usize, i64)>, - /// Whether [`visited_journal`](Self::visited_journal) is being recorded. - /// - /// Latched on by the first snapshot, mirroring the evaluator's own journal. - visit_journaling: bool, - last_case_index: usize, - visited_case_counts: Vec, - /// Last state value dispatched to each case, parallel to `visited_case_counts`. - /// - /// Used to distinguish a genuine "stuck" dispatcher (same state value - /// repeated → likely broken state propagation) from legitimate user-loop - /// iteration (state changes correctly between case visits). The A1 - /// overflow-bypass fallback is only appropriate for the former. - last_case_state: Vec>, - max_block_visits: usize, - max_tree_depth: usize, - other_dispatcher_blocks: Vec, - no_fork: bool, - /// Structural per-block answers, shared by every visit. See [`BlockFacts`]. - facts: BlockFacts, - /// SSA variables grouped by the local slot they originate from, in variable - /// table order. - /// - /// The cross-scope bridge in the inner trace loop needs "the first variable - /// for this local that currently has a value". Scanning the whole variable - /// table for that runs inside the per-instruction loop, making the trace - /// quadratic in a method's variable count; this index makes it proportional - /// to the number of versions of the one local involved. - vars_by_local: Vec>, -} - -impl<'a> TreeTraceContext<'a> { - /// Creates a new tracing context without a pre-detected dispatcher. - /// - /// Used by [`trace_method_tree`](super::trace_method_tree) when no CFF - /// dispatcher was detected, or as the base for - /// [`with_dispatcher`](Self::with_dispatcher). - pub fn new( - ssa: &'a SsaFunction, - config: &UnflattenConfig, - assembly: Option<&'a CilObject>, - ) -> Self { - let mut vars_by_local: Vec> = Vec::new(); - for var in ssa.variables() { - if let VariableOrigin::Local(idx) = var.origin() { - let idx = usize::from(idx); - if vars_by_local.len() <= idx { - vars_by_local.resize_with(idx.saturating_add(1), Vec::new); - } - if let Some(slot) = vars_by_local.get_mut(idx) { - slot.push(var.id()); - } - } - } - - Self { - ssa, - evaluator: SsaEvaluator::new(ssa, config.pointer_size), - assembly, - dispatcher: None, - state_tainted: BitSet::new(ssa.var_id_capacity()), - next_node_id: 0, - total_visits: 0, - global_visits: 0, - max_global_visits: config - .max_block_visits - .saturating_mul(GLOBAL_VISIT_BUDGET_FACTOR), - visited_states: FxHashSet::default(), - visited_journal: Vec::new(), - visit_journaling: false, - last_case_index: usize::MAX, - visited_case_counts: Vec::new(), - last_case_state: Vec::new(), - max_block_visits: config.max_block_visits, - max_tree_depth: config.max_tree_depth, - other_dispatcher_blocks: Vec::new(), - no_fork: false, - facts: BlockFacts::new(ssa, None), - vars_by_local, - } - } - - /// Creates a context with a pre-detected dispatcher. - pub fn with_dispatcher( - ssa: &'a SsaFunction, - dispatcher: TracedDispatcher, - config: &UnflattenConfig, - assembly: Option<&'a CilObject>, - ) -> Self { - let mut ctx = Self::new(ssa, config, assembly); - - // Use generic taint analysis for state variable tracking - if let Some(state_var) = dispatcher.state_var { - // Get the state variable's origin to filter PHI chains - let state_origin = ssa.variable(state_var).map(SsaVariable::origin); - - // Create CFF-specific taint configuration - let taint_config = cff_taint_config(ssa, dispatcher.block, state_origin); - - // Initialize taint analysis with the state variable as the seed - let mut taint = TaintAnalysis::new(taint_config); - taint.add_tainted_var(state_var); - - // Also seed the taint with the BACKWARD direction: variables that - // DEFINE the state variable through the dispatcher phi's operands. - // These are the state update values (constants, computed states) from - // each case block. Without tainting these and their definition chains, - // filter_state_instructions can't remove the state update instructions, - // causing CIL stack depth mismatches in codegen. - if let Some(disp_block) = ssa.block(dispatcher.block) { - for phi in disp_block.phi_nodes() { - if phi.result() == state_var { - for op in phi.operands() { - taint.add_tainted_var(op.value()); - } - } - } - } - - // Run propagation through PHI chains and definition chains. - taint.propagate(ssa); - - // Transfer tainted variables to context - for var in taint.tainted_variables() { - ctx.state_tainted.insert(var.index()); - } - } - - // Seed the evaluator with the initial state value when available. - // Optimization passes (copy propagation + DCE) may remove the `ldc.i4 N; stloc` - // that originally set the initial state, leaving the dispatcher PHI's entry operand - // undefined. Pre-seeding the operand variable ensures the first dispatch resolves. - if let (Some(state_var), Some(initial)) = (dispatcher.state_var, dispatcher.initial_state) { - // Find the entry predecessor by walking from block 0 toward the - // dispatcher following Jump terminators. After optimization the entry - // path is linear jumps, but the immediate predecessor of the dispatcher - // may not be block 0 (e.g., B0 → B1 → B_dispatcher). - let entry_pred = { - let mut pred = 0usize; - let mut current = 0usize; - for _ in 0..20 { - if current == dispatcher.block { - break; - } - pred = current; - match ssa.block(current).and_then(|b| b.terminator_op()) { - Some(SsaOp::Jump { target }) => current = *target, - _ => break, - } - } - pred - }; - - if let Some(disp_block) = ssa.block(dispatcher.block) { - for phi in disp_block.phi_nodes() { - if phi.result() == state_var { - // Find the operand from the entry predecessor and seed its value - for op in phi.operands() { - if op.predecessor() == entry_pred { - #[allow(clippy::cast_possible_truncation)] - ctx.evaluator - .set_concrete(op.value(), ConstValue::I32(initial as i32)); - } - } - } - } - } - } - - // Size the case visit counter to fit the dispatcher's switch targets - // (+1 for default, which uses targets.len() as its index). - ctx.visited_case_counts = vec![0u8; dispatcher.targets.len().saturating_add(1)]; - ctx.last_case_state = vec![None; dispatcher.targets.len().saturating_add(1)]; - ctx.facts = BlockFacts::new(ssa, Some(&dispatcher)); - ctx.dispatcher = Some(dispatcher); - ctx - } - - /// Returns a reference to the SSA function being traced. - /// - /// Returns `&'a SsaFunction` (the context's lifetime, not `&self`'s) - /// so the returned reference doesn't borrow `self` and can coexist with - /// `&mut self` calls on the context. - pub fn ssa(&self) -> &'a SsaFunction { - self.ssa - } - - /// Returns a reference to the SSA evaluator. - pub fn evaluator(&self) -> &SsaEvaluator<'a> { - &self.evaluator - } - - /// Returns a mutable reference to the SSA evaluator. - pub fn evaluator_mut(&mut self) -> &mut SsaEvaluator<'a> { - &mut self.evaluator - } - - /// Returns the optional assembly reference for call resolution. - /// - /// Returns the context's `'a` lifetime, not `&self`'s. - pub fn assembly(&self) -> Option<&'a CilObject> { - self.assembly - } - - /// Allocates and returns the next unique node ID. - pub fn next_id(&mut self) -> usize { - let id = self.next_node_id; - self.next_node_id = self.next_node_id.saturating_add(1); - id - } - - /// Returns true if the given block is the CFF dispatcher block. - pub fn is_dispatcher_block(&self, block: usize) -> bool { - self.dispatcher.as_ref().is_some_and(|d| d.block == block) - } - - /// Returns true if the given block is a direct target of the dispatcher - /// (case block or default). - pub fn is_dispatch_target(&self, block: usize) -> bool { - self.facts - .dispatch_target - .get(block) - .copied() - .unwrap_or(false) - } - - /// Returns the state variable (phi at the dispatcher), if detected. - pub fn state_var(&self) -> Option { - self.dispatcher.as_ref().and_then(|d| d.state_var) - } - - /// Returns the dispatcher block index, if detected. - pub fn dispatcher_block(&self) -> Option { - self.dispatcher.as_ref().map(|d| d.block) - } - - /// Returns true if the given block is another CFF dispatcher in the same - /// method (not the one we're tracing for). - pub fn is_other_dispatcher(&self, block: usize) -> bool { - self.facts - .other_dispatcher - .get(block) - .copied() - .unwrap_or(false) - } - - /// Sets the blocks of other CFF dispatchers in this method. - pub fn set_other_dispatcher_blocks(&mut self, blocks: Vec) { - for slot in &mut self.facts.other_dispatcher { - *slot = false; - } - for &block in &blocks { - if let Some(slot) = self.facts.other_dispatcher.get_mut(block) { - *slot = true; - } - } - // `is_overflow_dispatch_site` consults the other-dispatcher set, so any - // answer cached before this point was computed against a stale one. - for slot in &mut self.facts.overflow_site { - *slot = None; - } - self.other_dispatcher_blocks = blocks; - } - - /// Returns the jump target of `block` if it is a constant-producer block. - /// - /// Cached form of [`helpers::const_producer_target`]. - pub fn const_producer_target(&self, block: usize) -> Option { - self.facts.const_producer.get(block).copied().flatten() - } - - /// Returns the cached answer for `block`, computing it with `compute` on - /// first use. - /// - /// See [`BlockFacts::overflow_site`] for why this one is lazy. - pub fn overflow_dispatch_site( - &mut self, - block: usize, - compute: impl FnOnce(&Self) -> bool, - ) -> bool { - if let Some(cached) = self.facts.overflow_site.get(block).copied().flatten() { - return cached; - } - let answer = compute(self); - if let Some(slot) = self.facts.overflow_site.get_mut(block) { - *slot = Some(answer); - } - answer - } - - /// Returns the SSA variables that originate from local slot `local_idx`, in - /// variable table order. - pub fn vars_for_local(&self, local_idx: u16) -> &[SsaVarId] { - self.vars_by_local - .get(usize::from(local_idx)) - .map_or(&[], Vec::as_slice) - } - - /// Checks if a variable is state-tainted. - pub fn is_tainted(&self, var: SsaVarId) -> bool { - self.state_tainted.contains(var.index()) - } - - /// Checks if any of the variables are state-tainted. - pub fn any_tainted(&self, vars: &[SsaVarId]) -> bool { - vars.iter().any(|v| self.is_tainted(*v)) - } - - /// Marks a variable as tainted. - pub fn taint(&mut self, var: SsaVarId) { - self.state_tainted.insert(var.index()); - } - - /// Returns a reference to the state-tainted variable set. - pub fn state_tainted(&self) -> &BitSet { - &self.state_tainted - } - - /// Returns a mutable reference to the state-tainted variable set. - pub fn state_tainted_mut(&mut self) -> &mut BitSet { - &mut self.state_tainted - } - - /// Propagates taint forward through SSA instructions. - /// Encapsulates the borrow of both `ssa` and `state_tainted` within one method - /// to avoid split-borrow issues at call sites. - pub fn propagate_taint_forward(&mut self) { - helpers::propagate_taint_forward(self.ssa, &mut self.state_tainted); - } - - /// Gets the current CFF state value (if we can determine it). - pub fn current_state(&self) -> Option { - self.dispatcher - .as_ref() - .and_then(|d| d.state_var) - .and_then(|v| self.evaluator.get_concrete(v)) - .and_then(ConstValue::as_i64) - } - - /// Computes a visit key for loop detection. - /// - /// Uses the CFF state value when available (after dispatcher evaluation), - /// which allows revisiting blocks with different state machine values - /// (essential for CFF loop iterations). Falls back to a case-index-based - /// key at non-dispatcher blocks to prevent infinite recursion while still - /// allowing re-entry from different CFF case paths. - fn visit_state(&self) -> i64 { - self.current_state().unwrap_or_else(|| { - let count = self - .visited_case_counts - .get(self.last_case_index) - .copied() - .map_or(0, i64::from); - (self.last_case_index as i64) - .wrapping_mul(256) - .wrapping_add(count) - }) - } - - /// Checks if we've visited this block in the current execution context. - pub fn is_visited(&self, block: usize) -> bool { - self.visited_states.contains(&(block, self.visit_state())) - } - - /// Marks a block as visited in the current execution context. - pub fn mark_visited(&mut self, block: usize) { - let key = (block, self.visit_state()); - if self.visited_states.insert(key) && self.visit_journaling { - self.visited_journal.push(key); - } - } - - /// Increments the visit counter and returns true if the budget is exceeded. - pub fn check_visit_budget(&mut self) -> bool { - self.total_visits = self.total_visits.saturating_add(1); - self.global_visits = self.global_visits.saturating_add(1); - self.total_visits > self.max_block_visits || self.global_visits > self.max_global_visits - } - - /// Returns the maximum tree depth allowed. - pub fn max_tree_depth(&self) -> usize { - self.max_tree_depth - } - - /// Records that the dispatcher dispatched to the given case index. - /// Increments the visit count for the case and updates the last case index. - pub fn record_case_dispatch(&mut self, case_idx: usize) { - if let Some(slot) = self.visited_case_counts.get_mut(case_idx) { - *slot = slot.saturating_add(1); - } - self.last_case_index = case_idx; - } - - /// Returns `true` if the current dispatch to `case_idx` repeats the same - /// state value that was last used to dispatch that case. - /// - /// Used by the A1 overflow-bypass fallback to distinguish a truly stuck - /// dispatcher (identical state value, no progress) from a legitimate - /// user-loop iteration (same case but state updated each time). - /// - /// `current_state` is the state value on this incoming dispatch. A return - /// of `true` means "this case was dispatched before with the same state". - pub fn case_state_is_stuck(&self, case_idx: usize, current_state: i64) -> bool { - self.last_case_state - .get(case_idx) - .and_then(|slot| *slot) - .is_some_and(|prev| prev == current_state) - } - - /// Records the state value used for this dispatch of `case_idx`. - pub fn record_case_state(&mut self, case_idx: usize, state: i64) { - if let Some(slot) = self.last_case_state.get_mut(case_idx) { - *slot = Some(state); - } - } - - /// Checks if a case index has been visited enough times to indicate a - /// CFF loop back-edge. The threshold scales with the number of targets - /// to avoid false positives on small dispatchers. - pub fn is_case_loop(&self, case_idx: usize, targets_len: usize) -> bool { - let loop_threshold = (targets_len / 2).max(2) as u8; - self.visited_case_counts - .get(case_idx) - .is_some_and(|count| *count >= loop_threshold) - } - - /// Returns true when the tracer should follow one path instead of forking - /// at user branches/switches. - pub fn no_fork(&self) -> bool { - self.no_fork - } - - /// Takes a snapshot of the mutable context state that must be preserved - /// across branch/switch forks. - /// - /// The evaluator is marked rather than copied — see [`ContextSnapshot`]. - pub fn snapshot(&mut self) -> ContextSnapshot { - self.visit_journaling = true; - ContextSnapshot { - evaluator: self.evaluator.checkpoint(), - visited_mark: self.visited_journal.len(), - last_case_index: self.last_case_index, - visited_case_counts: self.visited_case_counts.clone(), - last_case_state: self.last_case_state.clone(), - } - } - - /// Restores all mutable context fields from a snapshot, consuming it. - pub fn restore(&mut self, snap: ContextSnapshot) { - self.evaluator.rollback(snap.evaluator); - while self.visited_journal.len() > snap.visited_mark { - let Some(key) = self.visited_journal.pop() else { - break; - }; - self.visited_states.remove(&key); - } - self.last_case_index = snap.last_case_index; - self.visited_case_counts = snap.visited_case_counts; - self.last_case_state = snap.last_case_state; - } - - /// Clones the current visited_case_counts for snapshotting. - pub fn case_counts_snapshot(&self) -> Vec { - self.visited_case_counts.clone() - } - - /// Overwrites the visited_case_counts (used when restoring expression - /// switch state across fork arms). - pub fn set_case_counts(&mut self, counts: Vec) { - self.visited_case_counts = counts; - } - - /// Saves the (total_visits, no_fork) pair and sets expression-switch - /// false-arm mode (reset visits, enable no_fork). Returns the saved - /// values for later restoration. - pub fn enter_expr_switch_false_arm(&mut self) -> (usize, bool) { - let saved = (self.total_visits, self.no_fork); - self.total_visits = 0; - self.no_fork = true; - saved - } - - /// Restores total_visits and no_fork from values saved by - /// [`enter_expr_switch_false_arm`]. - pub fn exit_expr_switch_false_arm(&mut self, saved: (usize, bool)) { - self.total_visits = saved.0; - self.no_fork = saved.1; - } - - /// Takes the dispatcher out of the context (for building the final TraceTree). - pub fn take_dispatcher(&mut self) -> Option { - self.dispatcher.take() - } - - /// Takes the state-tainted set out of the context (for building the final TraceTree). - pub fn take_state_tainted(&mut self) -> BitSet { - mem::take(&mut self.state_tainted) - } - - /// Returns the handler start blocks that were not visited by the main trace. - pub fn unvisited_handler_blocks(&self) -> Vec { - self.ssa - .exception_handlers() - .iter() - .filter_map(|h| h.handler_start_block) - .filter(|&block| { - block < self.ssa.block_count() - && !self.visited_states.iter().any(|(b, _)| *b == block) - }) - .collect() - } - - /// Creates an independent context for tracing an exception handler. - /// - /// Handler traces are self-contained: each has its own evaluator, visit - /// state, and visit budget. They share only immutable data from the parent - /// context (`&ssa`, dispatcher info, taint seeds, config limits). - /// This allows handler traces to run in parallel. - pub fn fork_for_handler(&self, node_id_offset: usize) -> Self { - let case_count_len = self.visited_case_counts.len(); - Self { - ssa: self.ssa, - evaluator: SsaEvaluator::new(self.ssa, self.evaluator.pointer_size()), - assembly: self.assembly, - dispatcher: self.dispatcher.clone(), - state_tainted: self.state_tainted.clone(), - next_node_id: node_id_offset, - total_visits: 0, - global_visits: 0, - max_global_visits: self.max_global_visits, - visited_states: FxHashSet::default(), - visited_journal: Vec::new(), - visit_journaling: false, - last_case_index: usize::MAX, - visited_case_counts: vec![0u8; case_count_len], - last_case_state: vec![None; case_count_len], - max_block_visits: self.max_block_visits, - max_tree_depth: self.max_tree_depth, - other_dispatcher_blocks: self.other_dispatcher_blocks.clone(), - no_fork: false, - facts: BlockFacts { - dispatch_target: self.facts.dispatch_target.clone(), - other_dispatcher: self.facts.other_dispatcher.clone(), - const_producer: self.facts.const_producer.clone(), - overflow_site: self.facts.overflow_site.clone(), - }, - vars_by_local: self.vars_by_local.clone(), - } - } - - /// Advances the node ID counter past all handler IDs. - pub fn advance_node_id(&mut self, new_id: usize) { - self.next_node_id = new_id; - } - - /// Returns the max_block_visits budget (used for handler ID stride). - pub fn max_block_visits(&self) -> usize { - self.max_block_visits - } -} - -/// Snapshot of `TreeTraceContext` mutable state saved at branch/switch fork -/// points. Allows the iterative tracer to restore context before tracing -/// each alternative arm. -/// -/// The evaluator is held as a mark into its undo journal, not as a copy. A -/// flattened method forks millions of times and the evaluator's state grows -/// with everything the trace has learned, so copying it per fork costs more -/// than the tracing itself — measured at 99% of tracer time on a NetReactor -/// sample. Rolling back to a mark instead costs one entry per value the -/// abandoned arm actually changed. -/// -/// This makes fork discipline part of the contract: marks are released in -/// reverse order of creation. The work stack in -/// [`engine`](super::engine) guarantees it — a snapshot lives in a work item, -/// and every item pushed after it is popped before it is. -pub struct ContextSnapshot { - evaluator: EvaluatorMark, - visited_mark: usize, - last_case_index: usize, - visited_case_counts: Vec, - last_case_state: Vec>, -} - -impl ContextSnapshot { - /// Clones this snapshot (used when restoring the same snapshot for - /// multiple switch case arms). - pub fn clone_snapshot(&self) -> Self { - Self { - evaluator: self.evaluator.clone(), - visited_mark: self.visited_mark, - last_case_index: self.last_case_index, - visited_case_counts: self.visited_case_counts.clone(), - last_case_state: self.last_case_state.clone(), - } - } -} diff --git a/dotscope/src/deobfuscation/passes/unflattening/tracer/engine.rs b/dotscope/src/deobfuscation/passes/unflattening/tracer/engine.rs deleted file mode 100644 index f75918e9..00000000 --- a/dotscope/src/deobfuscation/passes/unflattening/tracer/engine.rs +++ /dev/null @@ -1,1390 +0,0 @@ -//! Core iterative tracing engine. -//! -//! Builds [`TraceTree`](super::types::TraceTree)s by walking SSA blocks, -//! evaluating instructions, and following CFF state transitions. User -//! branches/switches that don't depend on the state variable are forked -//! to capture all execution paths. -//! -//! # Architecture -//! -//! Tracing uses a three-level architecture to avoid stack overflow while -//! maintaining the tree structure: -//! -//! 1. **[`trace_from_block`]** — Outer driver. Manages an explicit work stack -//! of [`WorkItem`] frames. Processes fork results (branch/switch arms) and -//! assembles completed sub-trees into the parent node. -//! -//! 2. **[`trace_from_block_linear`]** — Middle layer. Handles state transition -//! chains iteratively (dispatcher → case → dispatcher → case → ...) without -//! recursion. When a fork is needed, pushes continuation frames onto the -//! work stack and returns. -//! -//! 3. **[`trace_from_block_inner`]** — Inner loop. Processes blocks linearly -//! until hitting a terminator that requires forking or state transition. -//! Returns either a completed node or a [`ForkRequest`]. - -use std::collections::BTreeSet; - -use crate::{ - analysis::{ - CmpKind, ConstValue, SsaBlock, SsaEvaluator, SsaInstruction, SsaOp, SsaVarId, - VariableOrigin, - }, - deobfuscation::passes::unflattening::tracer::{ - context::{ContextSnapshot, TreeTraceContext}, - helpers::{detect_expression_switch, resolve_call_result}, - types::{StopReason, TraceNode, TraceTerminator}, - }, -}; - -/// Continuation frames for the iterative trace work stack. -/// -/// Each variant represents a pending operation that was deferred when the -/// tracer encountered a branch or switch fork. Instead of recursing, the -/// tracer pushes these frames and processes them one at a time. -enum WorkItem { - /// Start tracing a block. The resulting TraceNode becomes `current_result`. - TraceBlock { block: usize, depth: usize }, - - /// Link a state transition: attach `current_result` as the continuation - /// of `parent_node` via a `StateTransition` terminator. - StateTransitionLink { - parent_node: TraceNode, - from_state: i64, - to_state: i64, - target_block: usize, - }, - - /// After the true arm of a Branch/BranchCmp completes: restore context - /// snapshot, then trace the false arm. Carries the true arm's result. - BranchFalseArm { - parent_node: TraceNode, - block_idx: usize, - condition: SsaVarId, - false_target: usize, - depth: usize, - snapshot: ContextSnapshot, - case_counts_snapshot: Option>, - is_expr_switch: bool, - }, - - /// After the false arm completes: combine true + false into UserBranch. - BranchCombine { - parent_node: TraceNode, - block_idx: usize, - condition: SsaVarId, - true_node: TraceNode, - expr_switch_restore: Option<(usize, bool)>, - }, - - /// After a switch case completes: restore context and trace the next case, - /// or if all cases are done, trace the default arm. - SwitchNextCase { - parent_node: TraceNode, - block_idx: usize, - value: SsaVarId, - targets: Vec, - default_target: usize, - depth: usize, - snapshot: ContextSnapshot, - completed_cases: Vec<(i64, Box)>, - next_case_index: usize, - }, - - /// After the default arm completes: combine all cases into UserSwitch. - SwitchCombine { - parent_node: TraceNode, - block_idx: usize, - value: SsaVarId, - cases: Vec<(i64, Box)>, - }, -} - -/// Traces from a block, building the trace tree iteratively. -/// -/// State transitions (CFF case dispatches) and user branch/switch forks are -/// all handled via an explicit work stack to avoid stack overflow on deeply -/// nested methods. The resulting `TraceNode` tree is identical to what the -/// recursive version would produce. -pub fn trace_from_block( - ctx: &mut TreeTraceContext<'_>, - block_idx: usize, - depth: usize, -) -> TraceNode { - let mut work_stack: Vec = Vec::new(); - let mut current_result: Option = None; - - work_stack.push(WorkItem::TraceBlock { - block: block_idx, - depth, - }); - - loop { - let Some(item) = work_stack.pop() else { - return current_result.unwrap_or_else(|| { - let mut node = TraceNode::new(0, block_idx); - node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { block: block_idx }, - }); - node - }); - }; - - match item { - WorkItem::TraceBlock { block, depth } => { - let result = trace_from_block_linear(ctx, block, depth, &mut work_stack); - current_result = Some(result); - } - - WorkItem::StateTransitionLink { - mut parent_node, - from_state, - to_state, - target_block, - } => { - let Some(child) = current_result.take() else { - parent_node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { - block: target_block, - }, - }); - current_result = Some(parent_node); - continue; - }; - parent_node.set_terminator(TraceTerminator::StateTransition { - from_state, - to_state, - target_block, - continues: Box::new(child), - }); - current_result = Some(parent_node); - } - - WorkItem::BranchFalseArm { - parent_node, - block_idx, - condition, - false_target, - depth, - snapshot, - case_counts_snapshot, - is_expr_switch, - } => { - let Some(true_node) = current_result.take() else { - let mut node = parent_node; - node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { block: block_idx }, - }); - current_result = Some(node); - continue; - }; - - // Restore context to the state before the true arm. - // For expression switches, preserve the visited_case_counts - // accumulated during the true arm (convergence via shared counts). - let saved_case_counts = if case_counts_snapshot.is_none() { - Some(ctx.case_counts_snapshot()) - } else { - None - }; - ctx.restore(snapshot); - if let Some(counts) = saved_case_counts { - ctx.set_case_counts(counts); - } else if let Some(ref counts) = case_counts_snapshot { - ctx.set_case_counts(counts.clone()); - } - - // Expression switch false arms: reset total_visits, set no_fork - let expr_restore = if is_expr_switch { - Some(ctx.enter_expr_switch_false_arm()) - } else { - None - }; - - ctx.evaluator_mut().set_predecessor(Some(block_idx)); - - // Push combine (processed AFTER the false arm completes), - // then the false arm trace (processed FIRST due to stack LIFO). - work_stack.push(WorkItem::BranchCombine { - parent_node, - block_idx, - condition, - true_node, - expr_switch_restore: expr_restore, - }); - work_stack.push(WorkItem::TraceBlock { - block: false_target, - depth: depth.saturating_add(1), - }); - } - - WorkItem::BranchCombine { - mut parent_node, - block_idx, - condition, - true_node, - expr_switch_restore, - } => { - let Some(false_node) = current_result.take() else { - if let Some(saved) = expr_switch_restore { - ctx.exit_expr_switch_false_arm(saved); - } - parent_node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { block: block_idx }, - }); - current_result = Some(parent_node); - continue; - }; - - if let Some(saved) = expr_switch_restore { - ctx.exit_expr_switch_false_arm(saved); - } - - parent_node.set_terminator(TraceTerminator::UserBranch { - block: block_idx, - condition, - true_branch: Box::new(true_node), - false_branch: Box::new(false_node), - }); - current_result = Some(parent_node); - } - - WorkItem::SwitchNextCase { - parent_node, - block_idx, - value, - targets, - default_target, - depth, - snapshot, - mut completed_cases, - next_case_index, - } => { - // Collect the result from the previous case (if any) - if let Some(prev_result) = current_result.take() { - if next_case_index > 0 { - #[allow(clippy::cast_possible_wrap)] - let case_value = next_case_index.saturating_sub(1) as i64; - completed_cases.push((case_value, Box::new(prev_result))); - } - } - - if let Some(&target) = targets.get(next_case_index) { - // More cases to trace — restore and trace the next one - ctx.restore(snapshot.clone_snapshot()); - ctx.evaluator_mut().set_predecessor(Some(block_idx)); - - work_stack.push(WorkItem::SwitchNextCase { - parent_node, - block_idx, - value, - targets, - default_target, - depth, - snapshot, - completed_cases, - next_case_index: next_case_index.saturating_add(1), - }); - work_stack.push(WorkItem::TraceBlock { - block: target, - depth, - }); - } else { - // All cases done — restore and trace the default arm - ctx.restore(snapshot); - ctx.evaluator_mut().set_predecessor(Some(block_idx)); - - work_stack.push(WorkItem::SwitchCombine { - parent_node, - block_idx, - value, - cases: completed_cases, - }); - work_stack.push(WorkItem::TraceBlock { - block: default_target, - depth, - }); - } - } - - WorkItem::SwitchCombine { - mut parent_node, - block_idx, - value, - cases, - } => { - let Some(default_node) = current_result.take() else { - parent_node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { block: block_idx }, - }); - current_result = Some(parent_node); - continue; - }; - parent_node.set_terminator(TraceTerminator::UserSwitch { - block: block_idx, - value, - cases, - default: Box::new(default_node), - }); - current_result = Some(parent_node); - } - } - } -} - -/// Traces linearly from a block, handling state transitions iteratively. -/// -/// Returns a completed `TraceNode` when the trace reaches a leaf (exit, loop, -/// stop). When a user branch/switch fork is needed, pushes continuation frames -/// onto the `work_stack` and returns the node-so-far as `current_result` so -/// the outer `trace_from_block` loop can process the fork. -fn trace_from_block_linear( - ctx: &mut TreeTraceContext<'_>, - block_idx: usize, - depth: usize, - work_stack: &mut Vec, -) -> TraceNode { - // State transition chain — same iterative mechanism as before. - let mut transition_chain: Vec<(TraceNode, i64, usize)> = Vec::new(); - let mut entry_block = block_idx; - - let result = loop { - let (leaf, fork) = trace_from_block_inner(ctx, entry_block, depth); - - if let Some(fork) = fork { - // The inner tracer hit a branch/switch that needs forking. - // Push the transition chain (if any) as StateTransitionLink work - // items, then push the fork work items, and return the leaf node. - - // Push transition chain in forward order: the stack is LIFO, so - // the FIRST transition (closest to method entry) is pushed first and - // popped LAST, making it the outermost node (root) of the tree. - let to_state = ctx.current_state().unwrap_or(0); - for (parent, from_state, target_block) in transition_chain.drain(..) { - work_stack.push(WorkItem::StateTransitionLink { - parent_node: parent, - from_state, - to_state, - target_block, - }); - } - - // Push fork continuation items. Use block_idx from the fork - // (the block where the branch/switch terminator was encountered), - // NOT leaf.start_block (which may be an earlier block). - match fork { - ForkRequest::Branch { - block_idx, - condition, - true_target, - false_target, - snapshot, - case_counts_snapshot, - is_expr_switch, - } => { - ctx.evaluator_mut().set_predecessor(Some(block_idx)); - work_stack.push(WorkItem::BranchFalseArm { - parent_node: leaf, - block_idx, - condition, - false_target, - depth, - snapshot, - case_counts_snapshot, - is_expr_switch, - }); - work_stack.push(WorkItem::TraceBlock { - block: true_target, - depth: depth.saturating_add(1), - }); - } - ForkRequest::Switch { - block_idx, - value, - targets, - default_target, - snapshot, - is_foreign, - } => { - let fork_depth = if is_foreign { - depth - } else { - depth.saturating_add(1) - }; - let Some(&first_target) = targets.first() else { - ctx.evaluator_mut().set_predecessor(Some(block_idx)); - return leaf; - }; - ctx.evaluator_mut().set_predecessor(Some(block_idx)); - work_stack.push(WorkItem::SwitchNextCase { - parent_node: leaf, - block_idx, - value, - targets, - default_target, - depth: fork_depth, - snapshot, - completed_cases: Vec::new(), - next_case_index: 1, - }); - work_stack.push(WorkItem::TraceBlock { - block: first_target, - depth: fork_depth, - }); - } - } - - // Return a dummy node — the actual result will be built by the - // work stack continuations. The first TraceBlock pushed will set - // current_result when it completes. - return TraceNode::new(0, 0); - } - - // Check if the leaf needs a state transition continuation - if let Some((from_state, target_block)) = leaf.pending_state_transition() { - transition_chain.push((leaf, from_state, target_block)); - entry_block = target_block; - continue; - } - - // Leaf is complete - break leaf; - }; - - // Unwind the transition chain - let to_state = ctx.current_state().unwrap_or(0); - let mut leaf = result; - while let Some((mut parent, from_state, target_block)) = transition_chain.pop() { - parent.set_terminator(TraceTerminator::StateTransition { - from_state, - to_state, - target_block, - continues: Box::new(leaf), - }); - leaf = parent; - } - - leaf -} - -/// Fork request returned by the inner tracer when it encounters a user -/// branch or switch that needs to trace multiple arms. -enum ForkRequest { - Branch { - block_idx: usize, - condition: SsaVarId, - true_target: usize, - false_target: usize, - snapshot: ContextSnapshot, - case_counts_snapshot: Option>, - is_expr_switch: bool, - }, - Switch { - block_idx: usize, - value: SsaVarId, - targets: Vec, - default_target: usize, - snapshot: ContextSnapshot, - is_foreign: bool, - }, -} - -/// Inner tracing logic — processes blocks linearly until a decision point. -/// -/// Walks through consecutive blocks, evaluating instructions and propagating -/// taint. Returns when it hits: -/// - A **state transition** (dispatcher switch resolved) → returns -/// `(node, None)` with a pending state transition for the caller to continue -/// - A **user branch/switch** (non-state-dependent fork) → returns -/// `(node, Some(ForkRequest))` for the caller to trace both arms -/// - A **leaf** (return, throw, loop, stop) → returns `(node, None)` as a -/// completed terminal node -fn trace_from_block_inner( - ctx: &mut TreeTraceContext<'_>, - block_idx: usize, - depth: usize, -) -> (TraceNode, Option) { - let mut node = TraceNode::new(ctx.next_id(), block_idx); - - // Safety limits - if depth > ctx.max_tree_depth() { - node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::MaxVisitsExceeded, - }); - return (node, None); - } - - // Check for loop: same block visited in the same execution context. - // The visit key includes the last dispatched case index AND its visit - // count, so re-entering the same case after a loop exit (with an - // incremented count) produces a different key and is allowed. - if ctx.is_visited(block_idx) { - let state = ctx.current_state().unwrap_or(0); - node.set_terminator(TraceTerminator::LoopBack { - target_block: block_idx, - state, - }); - return (node, None); - } - // Dispatcher target blocks are not marked — their revisit detection is - // handled by visited_case_counts at the switch handler. Sub-blocks within - // case chains ARE marked to prevent unbounded expansion. - if !ctx.is_dispatch_target(block_idx) { - ctx.mark_visited(block_idx); - } - - // Bind SSA reference once — it's a shared &SsaFunction that doesn't - // borrow the mutable parts of ctx, avoiding borrow conflicts with - // evaluator_mut()/taint()/etc. - let ssa = ctx.ssa(); - - // Process blocks until we hit a decision point - let mut current_block = block_idx; - - loop { - // Safety: detect cycles in the linear block chain. - // If we revisit a block within the same trace_from_block call, - // we have an unconditional loop (e.g., Jump back-edge). - if ctx.check_visit_budget() { - node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::MaxVisitsExceeded, - }); - return (node, None); - } - // Exempt the dispatcher block — it's intentionally revisited as it dispatches - // to different case blocks based on the state variable. - let is_dispatcher = ctx.is_dispatcher_block(current_block); - let visited_without_last = node - .blocks_visited - .split_last() - .map(|(_, rest)| rest) - .unwrap_or(&[]); - if !is_dispatcher - && current_block != block_idx - && node.blocks_visited.len() > 1 - && visited_without_last.contains(¤t_block) - { - let state = ctx.current_state().unwrap_or(0); - node.set_terminator(TraceTerminator::LoopBack { - target_block: current_block, - state, - }); - return (node, None); - } - - // Handle dispatcher re-entry: clear stale values and fix predecessor - let visited_without_last_for_reentry = node - .blocks_visited - .split_last() - .map(|(_, rest)| rest) - .unwrap_or(&[]); - let is_dispatcher_reentry = is_dispatcher - && node.blocks_visited.len() > 1 - && visited_without_last_for_reentry.contains(¤t_block); - if is_dispatcher_reentry { - if let Some(block) = ssa.block(current_block) { - for instr in block.instructions() { - if let Some(def) = instr.def() { - ctx.evaluator_mut().set_unknown(def); - } - } - - if let Some(state_var) = ctx.state_var() { - for phi in block.phi_nodes() { - if phi.result() == state_var { - for op in phi.operands() { - let op_pred = op.predecessor(); - if node.blocks_visited.contains(&op_pred) - && ctx.evaluator().get_concrete(op.value()).is_some() - { - ctx.evaluator_mut().set_predecessor(Some(op_pred)); - break; - } - } - } - } - } - } - } - - let Some(block) = ssa.block(current_block) else { - node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { - block: current_block, - }, - }); - return (node, None); - }; - - // Set predecessor for phi evaluation - if node.blocks_visited.len() > 1 { - if let Some(&prev) = node - .blocks_visited - .get(node.blocks_visited.len().saturating_sub(2)) - { - ctx.evaluator_mut().set_predecessor(Some(prev)); - } - } - - // Bridge loop-carried phi operand values for dispatcher blocks - if is_dispatcher { - bridge_phi_operands(ctx, current_block); - } - - // Evaluate all phis for this block. - ctx.evaluator_mut().evaluate_phis(current_block); - - // Propagate taint through PHI nodes conservatively: only taint the - // PHI result if ALL operands are already state-tainted. - if let Some(block) = ssa.block(current_block) { - for phi in block.phi_nodes() { - if !phi.operands().is_empty() - && phi.operands().iter().all(|op| ctx.is_tainted(op.value())) - { - ctx.taint(phi.result()); - } - } - } - - // Process instructions with cross-scope variable bridging - for instr in block.instructions() { - trace_instruction(ctx, instr); - - // Bridge unknown local-variable references from known definitions - // of the same local index (cross-scope reaching definitions). - if let SsaOp::Copy { dest, src } = instr.op() { - if ctx.evaluator().get(*dest).is_none() { - if let Some(VariableOrigin::Local(local_idx)) = - ssa.variable(*src).map(|v| v.origin()) - { - // The context indexes variables by local slot, so this - // walks the versions of one local rather than the whole - // variable table — same order, same first match. - let bridged = ctx - .vars_for_local(local_idx) - .iter() - .filter(|&&var| var != *src) - .find_map(|&var| ctx.evaluator().get(var).cloned()); - if let Some(val) = bridged { - ctx.evaluator_mut().set_symbolic_expr(*dest, val); - } - } - } - } - } - - // Handle terminator - match handle_terminator(ctx, block, current_block, &mut node, depth) { - TerminatorResult::Continue(next) => { - node.visit_block(next); - current_block = next; - } - TerminatorResult::Done => return (node, None), - TerminatorResult::StateTransition { - from_state, - target_block, - } => { - node.set_pending_state_transition(from_state, target_block); - return (node, None); - } - TerminatorResult::ForkBranch { - block_idx, - condition, - true_target, - false_target, - snapshot, - case_counts_snapshot, - is_expr_switch, - } => { - return ( - node, - Some(ForkRequest::Branch { - block_idx, - condition, - true_target, - false_target, - snapshot, - case_counts_snapshot, - is_expr_switch, - }), - ); - } - TerminatorResult::ForkSwitch { - block_idx, - value, - targets, - default_target, - snapshot, - is_foreign, - } => { - return ( - node, - Some(ForkRequest::Switch { - block_idx, - value, - targets, - default_target, - snapshot, - is_foreign, - }), - ); - } - } - } -} - -/// Bridges loop-carried phi operand values for dispatcher blocks. -/// -/// In CFF dispatcher loops, the phi operand variable from the case block may -/// differ from the variable the evaluator tracked due to SSA variable renaming -/// at loop boundaries. This bridge ensures evaluate_phis can find the value. -fn bridge_phi_operands(ctx: &mut TreeTraceContext<'_>, block_idx: usize) { - let ssa = ctx.ssa(); - let (Some(sv), Some(block)) = (ctx.state_var(), ssa.block(block_idx)) else { - return; - }; - - let pred = ctx.evaluator().predecessor(); - for phi in block.phi_nodes() { - if phi.result() == sv { - if let Some(op) = phi - .operands() - .iter() - .find(|op| pred.is_some_and(|p| op.predecessor() == p)) - { - let op_var = op.value(); - if ctx.evaluator().get(op_var).is_none() { - if let Some(pred_idx) = pred { - if let Some(pred_block) = ssa.block(pred_idx) { - let bridged = pred_block - .instructions() - .iter() - .rev() - .filter(|i| !i.is_terminator()) - .find_map(|i| { - i.def().and_then(|d| { - ctx.evaluator().get(d).cloned().map(|v| (d, v)) - }) - }); - if let Some((_def_var, val)) = bridged { - ctx.evaluator_mut().set_symbolic_expr(op_var, val); - } - } - } - } - } - } - } -} - -/// Result of handling a block terminator instruction. -/// -/// Tells the inner trace loop how to proceed after processing a block's -/// terminator (jump, branch, switch, return, etc.). -enum TerminatorResult { - /// Follow an unconditional edge to the next block (jump, leave). - /// The caller continues the linear block chain within the same node. - Continue(usize), - /// The node is complete — terminator has been set (exit, stop, loop). - Done, - /// CFF state transition: the dispatcher resolved to a case block. - /// The caller should continue iteratively from `target_block`. - StateTransition { - from_state: i64, - target_block: usize, - }, - /// User branch (non-state-dependent) requires forking into true/false arms. - ForkBranch { - block_idx: usize, - condition: SsaVarId, - true_target: usize, - false_target: usize, - snapshot: ContextSnapshot, - case_counts_snapshot: Option>, - is_expr_switch: bool, - }, - /// User switch (non-state-dependent) requires forking into N case arms + default. - ForkSwitch { - block_idx: usize, - value: SsaVarId, - targets: Vec, - default_target: usize, - snapshot: ContextSnapshot, - is_foreign: bool, - }, -} - -/// Handles a block terminator, potentially forking the trace. -fn handle_terminator( - ctx: &mut TreeTraceContext<'_>, - block: &SsaBlock, - block_idx: usize, - node: &mut TraceNode, - depth: usize, -) -> TerminatorResult { - let Some(terminator) = block.instructions().last() else { - node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { block: block_idx }, - }); - return TerminatorResult::Done; - }; - - match terminator.op() { - SsaOp::Jump { target } => TerminatorResult::Continue(*target), - - SsaOp::Leave { target } => TerminatorResult::Continue(*target), - - SsaOp::Branch { - condition, - true_target, - false_target, - } => { - if ctx.is_tainted(*condition) { - // State-dependent branch - evaluate and follow one path - match ctx - .evaluator() - .get_concrete(*condition) - .and_then(ConstValue::as_i64) - { - Some(v) if v != 0 => TerminatorResult::Continue(*true_target), - Some(_) => TerminatorResult::Continue(*false_target), - None => { - node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { block: block_idx }, - }); - TerminatorResult::Done - } - } - } else { - handle_user_branch_fork( - ctx, - block_idx, - *true_target, - *false_target, - *condition, - depth, - ) - } - } - - SsaOp::BranchCmp { - left, - right, - cmp, - unsigned, - true_target, - false_target, - } => { - let left_val = ctx.evaluator().get_concrete(*left).cloned(); - let right_val = ctx.evaluator().get_concrete(*right).cloned(); - let tainted = ctx.is_tainted(*left) || ctx.is_tainted(*right); - - if tainted { - // State-tainted branch comparison — try to evaluate concretely. - // NETReactor CFF uses beq overflow checks after the switch dispatcher: - // ldloc ; ldc.i4 ; beq - // The state variable is tainted but has a known concrete value. - if let (Some(l), Some(r)) = (&left_val, &right_val) { - if SsaEvaluator::evaluate_comparison(l, r, *cmp, *unsigned) { - return TerminatorResult::Continue(*true_target); - } - return TerminatorResult::Continue(*false_target); - } - } - - // NETReactor overflow dispatch: at blocks on the dispatcher's - // default fall-through chain, `beq state, ` routes each - // overflow state value to its real case target. After LICM - // consolidates per-case Consts to a shared dominator, taint - // attribution is lost and state never becomes concrete here — - // but the check's *structure* (dispatcher-reachable block, - // `BranchCmp(var, Const, Eq)`) is still recognizable. Fork the - // comparison as a CFF continuation: seed state = const on the - // true arm (so the case body traces through to its next state - // update), leave state unknown on the false arm (so it chains to - // subsequent overflow checks or the final fall-through). - if *cmp == CmpKind::Eq && is_overflow_dispatch_site(ctx, block_idx) { - let overflow_seed = match (left_val.clone(), right_val.clone()) { - (None, Some(r)) => Some((*left, r)), - (Some(l), None) => Some((*right, l)), - _ => None, - }; - if let Some((unknown_var, const_val)) = overflow_seed { - let snapshot = ctx.snapshot(); - if let Some(state_var) = ctx.state_var() { - ctx.evaluator_mut() - .set_concrete(state_var, const_val.clone()); - } - ctx.evaluator_mut().set_concrete(unknown_var, const_val); - let case_counts_snapshot = Some(ctx.case_counts_snapshot()); - return TerminatorResult::ForkBranch { - block_idx, - condition: unknown_var, - true_target: *true_target, - false_target: *false_target, - snapshot, - case_counts_snapshot, - is_expr_switch: false, - }; - } - } - - if tainted { - node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { block: block_idx }, - }); - TerminatorResult::Done - } else { - // BranchCmp in no_fork mode has an additional check: allow forking - // for conditional CFF state transitions even when no_fork is set. - let is_expr_switch = - detect_expression_switch(ctx, *true_target, *false_target).is_some(); - - if ctx.no_fork() && !is_expr_switch { - let is_cff_state_transition = - is_conditional_state_transition(ctx, *true_target, *false_target); - if !is_cff_state_transition { - return TerminatorResult::Continue(*true_target); - } - } - - build_fork_branch( - ctx, - block_idx, - *true_target, - *false_target, - *left, - is_expr_switch, - ) - } - } - - SsaOp::Switch { - value, - targets, - default, - } => handle_switch(ctx, node, block_idx, value, targets, default, depth), - - SsaOp::Return { .. } | SsaOp::Throw { .. } => { - node.set_terminator(TraceTerminator::Exit { block: block_idx }); - TerminatorResult::Done - } - - _ => { - node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { block: block_idx }, - }); - TerminatorResult::Done - } - } -} - -/// Handles a user branch fork for both Branch and BranchCmp terminators. -/// -/// Detects expression switches, respects no_fork mode, creates snapshot -/// and returns a ForkBranch result. -fn handle_user_branch_fork( - ctx: &mut TreeTraceContext<'_>, - block_idx: usize, - true_target: usize, - false_target: usize, - condition: SsaVarId, - _: usize, -) -> TerminatorResult { - let is_expr_switch = detect_expression_switch(ctx, true_target, false_target).is_some(); - - // In no_fork mode, only fork for expression switches. - if ctx.no_fork() && !is_expr_switch { - return TerminatorResult::Continue(true_target); - } - - build_fork_branch( - ctx, - block_idx, - true_target, - false_target, - condition, - is_expr_switch, - ) -} - -/// Builds a ForkBranch result with the appropriate snapshot. -fn build_fork_branch( - ctx: &mut TreeTraceContext<'_>, - block_idx: usize, - true_target: usize, - false_target: usize, - condition: SsaVarId, - is_expr_switch: bool, -) -> TerminatorResult { - let snapshot = ctx.snapshot(); - let case_counts_snapshot = if is_expr_switch { - None - } else { - Some(ctx.case_counts_snapshot()) - }; - - TerminatorResult::ForkBranch { - block_idx, - condition, - true_target, - false_target, - snapshot, - case_counts_snapshot, - is_expr_switch, - } -} - -/// Checks if a branch is a conditional CFF state transition. -/// -/// Both arms are const-producers merging at a common block that reaches the -/// dispatcher. This pattern MUST be forked even in no_fork mode to populate -/// merge-point clone requests. -fn is_conditional_state_transition( - ctx: &TreeTraceContext<'_>, - true_target: usize, - false_target: usize, -) -> bool { - let Some(db) = ctx.dispatcher_block() else { - return false; - }; - - let true_merge = ctx.const_producer_target(true_target); - let false_merge = ctx.const_producer_target(false_target); - - match (true_merge, false_merge) { - (Some(tm), Some(fm)) if tm == fm => { - // Verify the merge block reaches the dispatcher - let mut block = tm; - for _ in 0..3 { - if block == db { - return true; - } - match ctx.ssa().block(block).and_then(|b| b.terminator_op()) { - Some(SsaOp::Jump { target }) => block = *target, - _ => break, - } - } - block == db - } - _ => false, - } -} - -/// Returns true if `block_idx` sits on the dispatcher's overflow chain. -/// -/// NETReactor's default dispatcher target chains one or more -/// `BranchCmp(state_var, Const(N), Eq)` checks to route overflow state values -/// (values greater than the switch table size) to their real case targets. -/// These chain blocks contain only Const/Copy/Jump/Nop and state-tainted -/// BranchCmp terminators — no user work. Detection walks predecessors up to a -/// short bound: if any predecessor is the dispatcher (reached directly or -/// through other pure/overflow blocks), the current block is part of the -/// chain and a tainted BranchCmp can safely be forked as a CFF continuation -/// rather than preserved as user code. -/// Memoizing wrapper — the answer depends only on the CFG and on which blocks -/// are dispatchers, both fixed for the lifetime of a trace, while the walk -/// itself allocates and visits up to eight levels of predecessors. On a -/// flattened method the same blocks are asked about millions of times. -fn is_overflow_dispatch_site(ctx: &mut TreeTraceContext<'_>, block_idx: usize) -> bool { - ctx.overflow_dispatch_site(block_idx, |ctx| { - compute_overflow_dispatch_site(ctx, block_idx) - }) -} - -fn compute_overflow_dispatch_site(ctx: &TreeTraceContext<'_>, block_idx: usize) -> bool { - let Some(dispatcher) = ctx.dispatcher_block() else { - return false; - }; - if block_idx == dispatcher { - return false; - } - - const MAX_HOPS: usize = 8; - let mut frontier: Vec = vec![block_idx]; - let mut visited: BTreeSet = BTreeSet::new(); - visited.insert(block_idx); - - for _ in 0..MAX_HOPS { - let mut next_frontier: Vec = Vec::new(); - for &b in &frontier { - for pred in ctx.ssa().block_predecessors(b) { - // Reaching the primary dispatcher (directly or through - // chain blocks) identifies this as an overflow site. A - // predecessor whose terminator is itself a `Switch` — a - // nested CFF dispatcher or any foreign dispatcher in this - // method — also qualifies, since overflow chains hang off - // such switches' default paths. - if pred == dispatcher || ctx.is_other_dispatcher(pred) { - return true; - } - if let Some(pred_block) = ctx.ssa().block(pred) { - if matches!(pred_block.terminator_op(), Some(SsaOp::Switch { .. })) { - return true; - } - } - if !visited.insert(pred) { - continue; - } - let Some(pred_block) = ctx.ssa().block(pred) else { - continue; - }; - // Only chain through blocks that carry no user-visible work: - // pure CFG plumbing (Const/Copy/Jump/Nop) plus overflow - // BranchCmp terminators. This prevents false positives when - // the tainted comparison sits downstream of user code. - let is_chain_block = pred_block.instructions().iter().all(|instr| { - matches!( - instr.op(), - SsaOp::Const { .. } - | SsaOp::Copy { .. } - | SsaOp::Jump { .. } - | SsaOp::Nop - | SsaOp::BranchCmp { .. } - ) - }); - if is_chain_block { - next_frontier.push(pred); - } - } - } - if next_frontier.is_empty() { - break; - } - frontier = next_frontier; - } - false -} - -/// Returns true when `default` leads to an `BranchCmp(_, _, Eq)` overflow -/// check within a short pure-chain. Used to gate the A1 fallback in -/// [`handle_switch`]: we only route a case-loop to the default arm when -/// that arm is structurally an overflow dispatch (not, e.g., a user -/// if/else chain encoded via the switch default). -/// -/// "Pure chain" means blocks whose instructions are only Const/Copy/ -/// Jump/Nop — the same shape `is_overflow_dispatch_site` recognizes on -/// the predecessor side. A BranchCmp-terminated block with one concrete -/// operand is the hallmark of NETReactor's overflow dispatch. -fn default_has_overflow_check(ctx: &TreeTraceContext<'_>, default: usize) -> bool { - const MAX_HOPS: usize = 4; - let mut current = default; - for _ in 0..MAX_HOPS { - let Some(block) = ctx.ssa().block(current) else { - return false; - }; - match block.terminator_op() { - Some(SsaOp::BranchCmp { - cmp: CmpKind::Eq, - left, - right, - .. - }) => { - // Confirm at least one operand is a tracked constant (the - // overflow value). This rules out generic user beq/bne - // structures that happen to sit on a default arm. - return ctx.evaluator().get_concrete(*left).is_some() - || ctx.evaluator().get_concrete(*right).is_some(); - } - Some(SsaOp::Jump { target }) => { - let is_pure = block.instructions().iter().all(|instr| { - matches!( - instr.op(), - SsaOp::Const { .. } | SsaOp::Copy { .. } | SsaOp::Jump { .. } | SsaOp::Nop - ) - }); - if !is_pure { - return false; - } - current = *target; - } - _ => return false, - } - } - false -} - -/// Handles a switch terminator (dispatcher or user switch). -fn handle_switch( - ctx: &mut TreeTraceContext<'_>, - node: &mut TraceNode, - block_idx: usize, - value: &SsaVarId, - targets: &[usize], - default: &usize, - _: usize, -) -> TerminatorResult { - let is_dispatcher = ctx.is_dispatcher_block(block_idx); - - let is_argument = ctx - .ssa() - .variable(*value) - .is_some_and(|v| matches!(v.origin(), VariableOrigin::Argument(_))); - - if !is_argument && (is_dispatcher || ctx.is_tainted(*value)) { - // State-driven switch (dispatcher) - evaluate and follow - let concrete_value = ctx - .evaluator() - .get_concrete(*value) - .and_then(ConstValue::as_u64) - .or_else(|| { - // ECMA-335 §I.12.3.2.2: uninitialized locals are zero-initialized - let var = ctx.ssa().variable(*value)?; - let site = var.def_site(); - let is_entry = site.block == 0 && site.instruction.is_none(); - if is_entry && matches!(var.origin(), VariableOrigin::Local(_)) { - Some(0) - } else { - None - } - }); - - if let Some(idx) = concrete_value { - // Defense-in-depth: verify state variable was resolved - if is_dispatcher { - if let Some(state_var) = ctx.state_var() { - if ctx.evaluator().get(state_var).is_none() { - node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { block: block_idx }, - }); - return TerminatorResult::Done; - } - } - } - - #[allow(clippy::cast_possible_truncation)] - let idx_usize = idx as usize; - let target = targets.get(idx_usize).copied().unwrap_or(*default); - - let from_state = ctx.current_state().unwrap_or(0); - - // Check for CFF loop back-edge - if ctx.is_case_loop(idx_usize, targets.len()) { - // LoopBack on the real case target. This is the correct - // response for both: - // 1. Legitimate user-loop iterations (same case re-entered - // as the loop body executes). - // 2. Tightly-looping CFF paths that the tracer cannot make - // further progress on. - // - // A previous "A1 fallback" variant routed to the dispatcher's - // default (overflow chain) with state cleared, meant to work - // around LICM-induced stuck-state propagation. With the LICM - // hoist guard now preventing per-edge state Consts from being - // hoisted into shared preheaders, that workaround would - // corrupt legitimate user-loop paths (clone the case target - // into the overflow chain and infinite-loop inside it), so we - // always prefer LoopBack here. - - let state = ctx.current_state().unwrap_or(0); - let mut loop_node = TraceNode::new(ctx.next_id(), target); - loop_node.set_terminator(TraceTerminator::LoopBack { - target_block: target, - state, - }); - - node.set_terminator(TraceTerminator::StateTransition { - from_state, - to_state: state, - target_block: target, - continues: Box::new(loop_node), - }); - return TerminatorResult::Done; - } - ctx.record_case_dispatch(idx_usize); - ctx.record_case_state(idx_usize, from_state); - - ctx.evaluator_mut().set_predecessor(Some(block_idx)); - - TerminatorResult::StateTransition { - from_state, - target_block: target, - } - } else if is_dispatcher { - node.set_terminator(TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { block: block_idx }, - }); - TerminatorResult::Done - } else { - handle_user_switch(ctx, block_idx, value, targets, default) - } - } else { - handle_user_switch(ctx, block_idx, value, targets, default) - } -} - -/// Handles a user switch by forking for all cases. -fn handle_user_switch( - ctx: &mut TreeTraceContext<'_>, - block_idx: usize, - value: &SsaVarId, - targets: &[usize], - default: &usize, -) -> TerminatorResult { - // No-fork mode: follow the evaluated path or first target. - if ctx.no_fork() { - let target = ctx - .evaluator() - .get_concrete(*value) - .and_then(|v| v.as_u64()) - .and_then(|idx| targets.get(idx as usize).copied()) - .unwrap_or_else(|| targets.first().copied().unwrap_or(*default)); - return TerminatorResult::Continue(target); - } - - let is_foreign = ctx.is_other_dispatcher(block_idx); - let snapshot = ctx.snapshot(); - - TerminatorResult::ForkSwitch { - block_idx, - value: *value, - targets: targets.to_vec(), - default_target: *default, - snapshot, - is_foreign, - } -} - -/// Evaluates a single instruction, updating the evaluator and taint state. -/// -/// The tracer used to record each instruction together with the concrete values -/// of its operands. Nothing consumed those values — the only reader of the log -/// needed the opcode, which is available from the SSA — while building it cost a -/// `BTreeMap` allocation and a full `SsaInstruction` clone per instruction per -/// visit. On a heavily nested method that is billions of allocations, so the log -/// is no longer produced. -fn trace_instruction(ctx: &mut TreeTraceContext<'_>, instr: &SsaInstruction) { - if ctx.any_tainted(&instr.uses()) { - if let Some(def) = instr.def() { - ctx.taint(def); - } - } - - // Evaluate the instruction - ctx.evaluator_mut().evaluate_op(instr.op()); - - // Resolve calls with concrete arguments (e.g., x86 predicate methods) - if let SsaOp::Call { - dest: Some(dest), - method, - args, - } = instr.op() - { - if let Some(assembly) = ctx.assembly() { - let concrete_args: Option> = args - .iter() - .map(|&a| ctx.evaluator().get_concrete(a).cloned()) - .collect(); - if let Some(concrete_args) = concrete_args { - if let Some(result) = resolve_call_result( - assembly, - method.token(), - &concrete_args, - ctx.evaluator().pointer_size(), - ) { - ctx.evaluator_mut().set_concrete(*dest, result); - } - } - } - } -} diff --git a/dotscope/src/deobfuscation/passes/unflattening/tracer/helpers.rs b/dotscope/src/deobfuscation/passes/unflattening/tracer/helpers.rs deleted file mode 100644 index 2a494f7a..00000000 --- a/dotscope/src/deobfuscation/passes/unflattening/tracer/helpers.rs +++ /dev/null @@ -1,324 +0,0 @@ -//! Standalone helper functions for CFF tracing. -//! -//! These functions are used by the [`super::engine`] and [`super::mod`] modules -//! but are factored out because they represent distinct concerns: -//! -//! - **Exception handler tracing** ([`trace_exception_handlers`]): Traces handler -//! blocks in parallel, each with an independent context -//! - **Taint propagation** ([`propagate_taint_forward`]): Forward-only taint -//! analysis to identify all variables derived from CFF state machinery -//! - **Call resolution** ([`resolve_call_result`]): Evaluates x86 predicate -//! methods (ConfuserEx) by building the callee's SSA and running it -//! - **Expression switch detection** ([`detect_expression_switch`]): Identifies -//! ConfuserEx "expression" mode branches where both arms select a CFF state -//! constant merging at a common block -//! - **Statistics** ([`compute_tree_stats`]): Recursive tree traversal counting -//! nodes, branches, transitions, exits, and max depth - -use analyssa::BitSet; -use rayon::prelude::*; - -use crate::{ - analysis::{ - ConstValue, PhiTaintMode, SsaBlock, SsaEvaluator, SsaFunction, SsaOp, SsaVarId, - TaintAnalysis, TaintConfig, - }, - deobfuscation::passes::unflattening::tracer::{ - context::TreeTraceContext, - engine::trace_from_block, - types::{HandlerTrace, TraceNode, TraceStats, TraceTerminator}, - }, - metadata::{token::Token, typesystem::PointerSize}, - CilObject, -}; - -/// Traces exception handler entry blocks that were not visited by the main trace. -/// -/// Handler blocks (catch, finally, filter) are only reachable via runtime exceptions, -/// not explicit branches, so the main trace from block 0 never reaches them. -/// -/// Each handler trace is fully independent — it uses a fresh evaluator, its own visit -/// state, and its own visit budget. This allows all handlers to be traced in parallel -/// using `fork_for_handler()` to create independent contexts that share only immutable -/// data (SSA, dispatcher info, taint seeds). -pub fn trace_exception_handlers(ctx: &mut TreeTraceContext<'_>) -> Vec { - let handler_blocks = ctx.unvisited_handler_blocks(); - if handler_blocks.is_empty() { - return Vec::new(); - } - - // Give each handler a unique node ID offset so they don't collide. - // Each handler gets a budget of max_block_visits IDs (generous upper bound). - let id_base = ctx.next_id(); - let id_stride = ctx.max_block_visits(); - - // Trace all handlers in parallel — each gets its own independent context - let handler_traces: Vec = handler_blocks - .par_iter() - .enumerate() - .filter_map(|(i, &handler_start)| { - let offset = i.saturating_mul(id_stride); - let mut handler_ctx = ctx.fork_for_handler(id_base.saturating_add(offset)); - let root = trace_from_block(&mut handler_ctx, handler_start, 0); - Some(HandlerTrace { - handler_start_block: handler_start, - root, - }) - }) - .collect(); - - // Advance the parent's node ID counter past all handler IDs - let total_offset = handler_blocks.len().saturating_mul(id_stride); - ctx.advance_node_id(id_base.saturating_add(total_offset)); - - handler_traces -} - -/// Propagates taint forward through instructions using the generic taint analysis. -/// -/// Forward propagation: If an instruction uses a tainted variable, its def becomes tainted. -/// This is used to identify all variables that depend on state machinery. -/// -/// This function uses the generic TaintAnalysis module with forward-only propagation -/// and NoPropagation for PHI nodes (to avoid over-tainting through merge points). -pub fn propagate_taint_forward(ssa: &SsaFunction, tainted: &mut BitSet) { - // Configure for forward-only propagation without PHI propagation. - // PHIs merge values from different paths, and some operands may be user code - // while others are state machinery. Propagating through PHIs would incorrectly - // filter user code that happens to share a merge point with state code. - let config = TaintConfig { - forward: true, - backward: false, - phi_mode: PhiTaintMode::NoPropagation, - max_iterations: 100, - }; - - let mut taint = TaintAnalysis::new(config); - - // Initialize with existing tainted variables - taint.add_tainted_vars(tainted.iter().map(SsaVarId::from_index)); - - // Run propagation - taint.propagate(ssa); - - // Update the tainted set with newly discovered tainted variables - tainted.clear(); - for var in taint.tainted_variables() { - tainted.insert(var.index()); - } - - // Post-pass: taint PHI results where ALL operands are tainted, then - // re-propagate forward from newly tainted PHI results. This identifies - // inner CFF dispatchers (e.g., JIEJIE.NET nested switch patterns) where - // all incoming values are CFF machinery but the PHI result was not - // tainted due to NoPropagation. - let mut changed = true; - while changed { - changed = false; - for block in ssa.blocks() { - for phi in block.phi_nodes() { - if !tainted.contains(phi.result().index()) - && !phi.operands().is_empty() - && phi - .operands() - .iter() - .all(|op| tainted.contains(op.value().index())) - { - tainted.insert(phi.result().index()); - changed = true; - } - } - } - if changed { - // Re-propagate forward from newly tainted PHI results - let config = TaintConfig { - forward: true, - backward: false, - phi_mode: PhiTaintMode::NoPropagation, - max_iterations: 100, - }; - let mut taint = TaintAnalysis::new(config); - taint.add_tainted_vars(tainted.iter().map(SsaVarId::from_index)); - taint.propagate(ssa); - tainted.clear(); - for var in taint.tainted_variables() { - tainted.insert(var.index()); - } - } - } -} - -/// Resolves a method call with concrete arguments by building the callee's SSA and evaluating it. -/// -/// This is used for x86 predicate methods in ConfuserEx CFF, where state computation -/// is done via `call ::predicate(arg1, arg2)` instead of inline arithmetic. -pub fn resolve_call_result( - assembly: &CilObject, - method_token: Token, - concrete_args: &[ConstValue], - pointer_size: PointerSize, -) -> Option { - // Look up the method - let method = assembly.method(&method_token).ok()?; - - // Build SSA for the callee - let callee_ssa = method.ssa(assembly).ok()?; - - // Create evaluator for the callee - let mut eval = SsaEvaluator::new(&callee_ssa, pointer_size); - - // Set concrete argument values - for (var, value) in callee_ssa.argument_variables().zip(concrete_args) { - eval.set_concrete(var.id(), value.clone()); - } - - // Execute with a safety limit of 50 blocks - let trace = eval.execute(0, None, 50); - - // If execution didn't complete, we can't resolve the call - if !trace.is_complete() { - return None; - } - - // Find the return value from the last block - let last_block_idx = trace.last_block()?; - let last_block = callee_ssa.block(last_block_idx)?; - - // Look for a Return instruction with a value - for instr in last_block.instructions() { - if let SsaOp::Return { - value: Some(ret_var), - } = instr.op() - { - return eval.get_concrete(*ret_var).cloned(); - } - } - - None -} - -/// Detects if a Branch is a CFF "expression switch" — both targets are -/// constant-producer blocks that merge into a single block feeding a -/// tainted CFF state computation. Returns the merge block index if matched. -/// -/// ConfuserEx "expression" control flow mode wraps user conditionals so each -/// branch arm selects a different CFF state constant. Without detection, the -/// tracer forks O(2^N) at these branches. With detection, both forks share -/// accumulated tracking state so the false branch stops at the convergence point. -pub fn detect_expression_switch( - ctx: &TreeTraceContext<'_>, - true_target: usize, - false_target: usize, -) -> Option { - // The structural half of the test — are both arms constant-producers that - // meet at the same block — is a property of the CFG, so it is answered from - // the context's precomputed table rather than re-walked. Only the taint - // check below depends on trace state. - let true_merge = ctx.const_producer_target(true_target)?; - let false_merge = ctx.const_producer_target(false_target)?; - - if true_merge != false_merge { - return None; - } - - let merge = ctx.ssa().block(true_merge)?; - let phis = merge.phi_nodes(); - if phis.is_empty() { - return None; - } - - let tainted = ctx.state_tainted(); - let is_phi_result = |var: &SsaVarId| phis.iter().any(|phi| phi.result() == *var); - - let feeds_tainted = merge.instructions().iter().any(|instr| match instr.op() { - SsaOp::Xor { left, right, .. } - | SsaOp::Add { left, right, .. } - | SsaOp::Sub { left, right, .. } - | SsaOp::Mul { left, right, .. } => { - let one_is_phi = is_phi_result(left) || is_phi_result(right); - let one_is_tainted = tainted.contains(left.index()) || tainted.contains(right.index()); - one_is_phi && one_is_tainted - } - _ => false, - }); - - feeds_tainted.then_some(true_merge) -} - -/// Checks if a block is a "constant producer" and returns its jump target. -/// -/// A constant producer block contains at most 2 non-terminator instructions -/// (all `Const`, `Copy`, or `Conv`) and ends with a `Jump`. These blocks -/// appear in ConfuserEx "expression" mode CFF: each branch arm pushes a -/// different state constant and jumps to a merge block that feeds the -/// dispatcher's state computation. -pub fn const_producer_target(block: &SsaBlock) -> Option { - let instrs = block.instructions(); - if instrs.is_empty() { - return None; - } - - let target = match instrs.last()?.op() { - SsaOp::Jump { target } => *target, - _ => return None, - }; - - let non_term: Vec<_> = instrs.iter().filter(|i| !i.is_terminator()).collect(); - if non_term.len() > 2 { - return None; - } - if !non_term.iter().all(|i| { - matches!( - i.op(), - SsaOp::Const { .. } - | SsaOp::Copy { .. } - | SsaOp::IntConv { .. } - | SsaOp::IntToPtr { .. } - | SsaOp::PtrToInt { .. } - | SsaOp::IntToFloat { .. } - | SsaOp::FloatToInt { .. } - | SsaOp::FloatConv { .. } - ) - }) { - return None; - } - - Some(target) -} - -/// Computes statistics for a trace tree. -pub fn compute_tree_stats(node: &TraceNode, stats: &mut TraceStats, depth: usize) { - stats.node_count = stats.node_count.saturating_add(1); - stats.max_depth = stats.max_depth.max(depth); - - let next_depth = depth.saturating_add(1); - match &node.terminator { - TraceTerminator::Exit { .. } => { - stats.exit_count = stats.exit_count.saturating_add(1); - } - TraceTerminator::StateTransition { continues, .. } => { - stats.state_transition_count = stats.state_transition_count.saturating_add(1); - compute_tree_stats(continues, stats, next_depth); - } - TraceTerminator::UserBranch { - true_branch, - false_branch, - .. - } => { - stats.user_branch_count = stats.user_branch_count.saturating_add(1); - compute_tree_stats(true_branch, stats, next_depth); - compute_tree_stats(false_branch, stats, next_depth); - } - TraceTerminator::UserSwitch { cases, default, .. } => { - stats.user_branch_count = stats.user_branch_count.saturating_add(1); - for (_, case_node) in cases { - compute_tree_stats(case_node, stats, next_depth); - } - compute_tree_stats(default, stats, next_depth); - } - TraceTerminator::Stopped { .. } | TraceTerminator::LoopBack { .. } => {} - TraceTerminator::PendingStateTransition { .. } => { - // Internal sentinel — should never appear in the final trace tree - } - } -} diff --git a/dotscope/src/deobfuscation/passes/unflattening/tracer/mod.rs b/dotscope/src/deobfuscation/passes/unflattening/tracer/mod.rs deleted file mode 100644 index ea274e01..00000000 --- a/dotscope/src/deobfuscation/passes/unflattening/tracer/mod.rs +++ /dev/null @@ -1,336 +0,0 @@ -//! Trace-based CFF analysis. -//! -//! This module implements tracing for control flow unflattening. Given a -//! pre-detected CFF dispatcher (from [`super::detection`]), it builds a tree -//! of all execution paths through the method by: -//! -//! 1. **Evaluate from method entry**: Walk through the SSA, evaluating each -//! instruction with concrete values using the [`SsaEvaluator`] -//! 2. **Follow state transitions**: When the dispatcher switch evaluates to a -//! concrete case index, follow that case automatically (CFF machinery) -//! 3. **Fork at user branches**: Branches whose condition does NOT depend on -//! the state variable are forked to capture both paths (original program logic) -//! 4. **Classify via taint analysis**: Variables derived from the state variable -//! are marked as state-tainted — the reconstruction phase removes them -//! 5. **Detect loops**: When the same (block, state) pair is revisited, emit a -//! [`TraceTerminator::LoopBack`] to prevent infinite tree expansion -//! -//! The resulting [`TraceTree`] is consumed by the [`super::reconstruction`] module -//! to extract a [`PatchPlan`](super::reconstruction::PatchPlan) that rewires the -//! CFG, removing dispatcher indirection. -//! -//! [`SsaEvaluator`]: crate::analysis::SsaEvaluator -//! -//! # Module Structure -//! -//! - [`types`]: Public data types ([`TraceTree`], [`TraceNode`], [`TracedDispatcher`], etc.) -//! - [`context`]: Tracing context — owns evaluator, taint state, visit tracking -//! - [`engine`]: Core iterative tracing machine (work stack, terminator dispatch) -//! - [`helpers`]: Standalone helpers (taint propagation, call resolution, statistics) - -mod context; -mod engine; -mod helpers; -mod types; - -pub use types::*; - -use crate::{ - analysis::SsaFunction, - deobfuscation::passes::unflattening::{ - detection::CffDetector, - tracer::{ - context::TreeTraceContext, - engine::trace_from_block, - helpers::{compute_tree_stats, trace_exception_handlers}, - }, - UnflattenConfig, - }, - CilObject, -}; - -/// Traces a method into a tree structure, forking at user branches. -/// -/// This is the main entry point for tree-based tracing. It handles: -/// - Detecting the dispatcher upfront via `CffDetector` (SCCP-based) -/// - Following state transitions automatically -/// - Forking at user branches (non-state-dependent conditions) -/// - Detecting loops to avoid infinite expansion -/// -/// # Arguments -/// -/// * `ssa` - The SSA function to trace -/// * `config` - Configuration controlling tracing limits and behavior -/// * `assembly` - Optional assembly reference for call resolution -/// -/// # Returns -/// -/// A [`TraceTree`] containing all execution paths through the method. -pub fn trace_method_tree( - ssa: &SsaFunction, - config: &UnflattenConfig, - assembly: Option<&CilObject>, -) -> TraceTree { - // Detect dispatcher upfront using CffDetector - let mut detector = CffDetector::new(ssa); - let dispatcher = detector - .detect_best() - .filter(|d| d.confidence >= config.min_confidence) - .map(|d| TracedDispatcher { - block: d.block, - switch_var: d.switch_var, - targets: d.cases.clone(), - default: d.default, - state_var: d.state_phi, - initial_state: d.initial_state, - }); - - let mut ctx = match dispatcher { - Some(d) => TreeTraceContext::with_dispatcher(ssa, d, config, assembly), - None => TreeTraceContext::new(ssa, config, assembly), - }; - - build_trace_tree(&mut ctx) -} - -/// Traces a method for a specific pre-detected dispatcher. -/// -/// Unlike [`trace_method_tree`] which auto-detects the best dispatcher, -/// this function uses a caller-provided dispatcher. This is used when -/// processing multiple independent CFF dispatchers in a single method -/// (e.g., ConfuserEx inserts separate dispatchers per exception handler -/// region). Each dispatcher is traced independently and the resulting -/// patch plans are merged before applying. -/// -/// # Arguments -/// -/// * `ssa` - The SSA function to trace -/// * `config` - Configuration controlling tracing limits and behavior -/// * `assembly` - Optional assembly reference for call resolution -/// * `dispatcher` - The pre-detected dispatcher to trace for -/// * `other_dispatcher_blocks` - Block indices of other dispatchers in this method -/// -/// # Returns -/// -/// A [`TraceTree`] for the given dispatcher. -pub fn trace_for_dispatcher( - ssa: &SsaFunction, - config: &UnflattenConfig, - assembly: Option<&CilObject>, - dispatcher: TracedDispatcher, - other_dispatcher_blocks: &[usize], -) -> TraceTree { - let mut ctx = TreeTraceContext::with_dispatcher(ssa, dispatcher, config, assembly); - ctx.set_other_dispatcher_blocks(other_dispatcher_blocks.to_vec()); - - build_trace_tree(&mut ctx) -} - -/// Shared implementation for building a trace tree from an initialized context. -/// -/// Both [`trace_method_tree`] and [`trace_for_dispatcher`] delegate to this -/// function after context setup. It traces from block 0, traces exception -/// handlers, propagates taint, and computes statistics. -fn build_trace_tree(ctx: &mut TreeTraceContext<'_>) -> TraceTree { - let root = trace_from_block(ctx, 0, 0); - let handler_traces = trace_exception_handlers(ctx); - ctx.propagate_taint_forward(); - - let mut tree = TraceTree::new(root, ctx.ssa().var_id_capacity()); - tree.handler_traces = handler_traces; - tree.dispatcher = ctx.take_dispatcher(); - tree.state_tainted = ctx.take_state_tainted(); - - compute_tree_stats(&tree.root, &mut tree.stats, 0); - for ht in &tree.handler_traces { - compute_tree_stats(&ht.root, &mut tree.stats, 0); - } - tree -} - -#[cfg(test)] -mod tests { - use crate::{ - analysis::{ - ConstValue, PhiNode, PhiOperand, SsaBlock, SsaFunction, SsaInstruction, SsaOp, - SsaVarId, VariableOrigin, - }, - deobfuscation::passes::unflattening::{tracer::trace_method_tree, UnflattenConfig}, - }; - - /// Creates a simple CFF-like SSA function for testing. - fn create_simple_cff() -> SsaFunction { - let mut ssa = SsaFunction::new(0, 1); - let state_var = SsaVarId::from_index(0); - let const_var = SsaVarId::from_index(1); - - // B0: entry - set initial state and jump to dispatcher - let mut b0 = SsaBlock::new(0); - b0.add_instruction(SsaInstruction::synthetic(SsaOp::Const { - dest: const_var, - value: ConstValue::I32(0), - })); - b0.add_instruction(SsaInstruction::synthetic(SsaOp::Jump { target: 1 })); - ssa.add_block(b0); - - // B1: dispatcher with switch - let mut b1 = SsaBlock::new(1); - let mut phi = PhiNode::new(state_var, VariableOrigin::Local(0)); - phi.add_operand(PhiOperand::new(const_var, 0)); - b1.add_phi(phi); - b1.add_instruction(SsaInstruction::synthetic(SsaOp::Switch { - value: state_var, - targets: vec![2, 3, 4], - default: 5, - })); - ssa.add_block(b1); - - // B2, B3, B4: case blocks that jump back to dispatcher - for i in 2..=4 { - let mut b = SsaBlock::new(i); - b.add_instruction(SsaInstruction::synthetic(SsaOp::Jump { target: 1 })); - ssa.add_block(b); - } - - // B5: exit - let mut b5 = SsaBlock::new(5); - b5.add_instruction(SsaInstruction::synthetic(SsaOp::Return { value: None })); - ssa.add_block(b5); - - ssa - } - - #[test] - fn test_tree_trace_simple_cff() { - let ssa = create_simple_cff(); - let config = UnflattenConfig::default(); - let tree = trace_method_tree(&ssa, &config, None); - - // Should find the dispatcher - assert!(tree.dispatcher.is_some(), "Should detect dispatcher"); - let dispatcher = tree.dispatcher.as_ref().unwrap(); - assert_eq!(dispatcher.block, 1); - - // Should have state-tainted variables - assert!(!tree.state_tainted.is_empty(), "Should have tainted vars"); - - // Check stats - println!("Tree stats: {:?}", tree.stats); - assert!(tree.stats.node_count >= 1, "Should have at least one node"); - } - - /// Creates a CFF with a user branch inside a case block. - fn create_cff_with_user_branch() -> SsaFunction { - let mut ssa = SsaFunction::new(1, 1); - let state_var = SsaVarId::from_index(0); - let init_state = SsaVarId::from_index(1); - let const_one = SsaVarId::from_index(2); - let arg0 = SsaVarId::from_index(3); - let user_zero = SsaVarId::from_index(4); - let cmp_result = SsaVarId::from_index(5); - - // B0: entry - set initial state = 0 and jump to dispatcher - let mut b0 = SsaBlock::new(0); - b0.add_instruction(SsaInstruction::synthetic(SsaOp::Const { - dest: init_state, - value: ConstValue::I32(0), - })); - b0.add_instruction(SsaInstruction::synthetic(SsaOp::Const { - dest: arg0, - value: ConstValue::I32(42), - })); - b0.add_instruction(SsaInstruction::synthetic(SsaOp::Jump { target: 1 })); - ssa.add_block(b0); - - // B1: dispatcher with switch - let mut b1 = SsaBlock::new(1); - let mut phi = PhiNode::new(state_var, VariableOrigin::Local(0)); - phi.add_operand(PhiOperand::new(init_state, 0)); - phi.add_operand(PhiOperand::new(const_one, 3)); - phi.add_operand(PhiOperand::new(const_one, 4)); - b1.add_phi(phi); - b1.add_instruction(SsaInstruction::synthetic(SsaOp::Switch { - value: state_var, - targets: vec![2, 5], - default: 6, - })); - ssa.add_block(b1); - - // B2: case 0 - has USER BRANCH - let mut b2 = SsaBlock::new(2); - b2.add_instruction(SsaInstruction::synthetic(SsaOp::Const { - dest: const_one, - value: ConstValue::I32(1), - })); - b2.add_instruction(SsaInstruction::synthetic(SsaOp::Const { - dest: user_zero, - value: ConstValue::I32(0), - })); - b2.add_instruction(SsaInstruction::synthetic(SsaOp::Cgt { - dest: cmp_result, - left: arg0, - right: user_zero, - unsigned: false, - })); - b2.add_instruction(SsaInstruction::synthetic(SsaOp::Branch { - condition: cmp_result, - true_target: 3, - false_target: 4, - })); - ssa.add_block(b2); - - // B3a: true branch - let mut b3a = SsaBlock::new(3); - b3a.add_instruction(SsaInstruction::synthetic(SsaOp::Jump { target: 1 })); - ssa.add_block(b3a); - - // B3b: false branch - let mut b3b = SsaBlock::new(4); - b3b.add_instruction(SsaInstruction::synthetic(SsaOp::Jump { target: 1 })); - ssa.add_block(b3b); - - // B5: case 1 - exit path - let mut b5 = SsaBlock::new(5); - b5.add_instruction(SsaInstruction::synthetic(SsaOp::Return { value: None })); - ssa.add_block(b5); - - // B6: default - exit - let mut b6 = SsaBlock::new(6); - b6.add_instruction(SsaInstruction::synthetic(SsaOp::Return { value: None })); - ssa.add_block(b6); - - ssa - } - - #[test] - fn test_tree_trace_with_user_branch() { - let ssa = create_cff_with_user_branch(); - let config = UnflattenConfig::default(); - let tree = trace_method_tree(&ssa, &config, None); - - println!("=== Tree Trace with User Branch ==="); - println!("Dispatcher: {:?}", tree.dispatcher); - println!("Tainted vars: {:?}", tree.state_tainted); - println!("Stats: {:?}", tree.stats); - - assert!( - tree.stats.user_branch_count > 0, - "Should have forked at branches" - ); - - assert!(tree.stats.exit_count > 0, "Should have exit points"); - - println!("User branch count: {}", tree.stats.user_branch_count); - println!("Exit count: {}", tree.stats.exit_count); - - println!( - "Stats confirm {} user branches were created", - tree.stats.user_branch_count - ); - - assert!( - tree.stats.user_branch_count > 0, - "Stats must show user branches" - ); - } -} diff --git a/dotscope/src/deobfuscation/passes/unflattening/tracer/types.rs b/dotscope/src/deobfuscation/passes/unflattening/tracer/types.rs deleted file mode 100644 index f19ea8b6..00000000 --- a/dotscope/src/deobfuscation/passes/unflattening/tracer/types.rs +++ /dev/null @@ -1,288 +0,0 @@ -//! Data types for trace-based CFF analysis. -//! -//! All types in this module are part of the tracer's public API — they are -//! re-exported from [`super`] and consumed by the -//! [`reconstruction`](crate::deobfuscation::passes::unflattening::reconstruction) -//! module to build patch plans. -//! -//! The central type is [`TraceTree`], which contains a root [`TraceNode`] and -//! optional handler traces. Each node represents a segment of linear execution, -//! terminated by one of the [`TraceTerminator`] variants (state transition, user -//! branch, exit, loop, or stop). - -use analyssa::BitSet; -use smallvec::{smallvec, SmallVec}; - -use crate::analysis::SsaVarId; - -/// Blocks recorded for one trace segment. -/// -/// A segment nearly always covers a single block — a 1200-block NetReactor -/// method produced 4.8 million nodes averaging about one block each — so the -/// common case is stored inline. Heap-allocating a one-element vector per node -/// cost ~92 million allocations across that method's dispatchers, and the -/// allocator churn showed up as the tracer's single largest expense. -pub type VisitedBlocks = SmallVec<[usize; 1]>; - -/// Information about the dispatcher found during tracing. -#[derive(Debug, Clone)] -pub struct TracedDispatcher { - /// Block index of the dispatcher. - pub block: usize, - - /// The switch value variable. - pub switch_var: SsaVarId, - - /// Switch targets (case blocks). - pub targets: Vec, - - /// Default target. - pub default: usize, - - /// The state variable (phi at dispatcher that receives state from case blocks). - pub state_var: Option, - - /// Initial state value captured during detection (before optimization may - /// remove the defining `ldc.i4; stloc` sequence). - pub initial_state: Option, -} - -/// Why tracing stopped. -#[derive(Debug, Clone)] -pub enum StopReason { - /// Hit a return/throw instruction. - Terminator, - - /// Exceeded maximum block visits. - MaxVisitsExceeded, - - /// Couldn't determine next block (unknown branch/switch value). - UnknownControlFlow { block: usize }, - - /// Visited same block too many times (likely infinite loop). - InfiniteLoop { block: usize }, -} - -/// A trace of an exception handler entry block. -/// -/// When CFF exists inside exception handler blocks (catch/finally/filter), -/// the normal trace from block 0 won't reach them because handlers are only -/// reachable via runtime exceptions. This struct holds a separate trace -/// starting from a handler entry block. -#[derive(Debug, Clone)] -pub struct HandlerTrace { - /// The handler's entry block index. - pub handler_start_block: usize, - - /// The root node of the handler's trace tree. - pub root: TraceNode, -} - -/// A trace tree represents all execution paths through a CFF-protected method. -/// -/// Unlike the linear `MethodTrace`, this structure forks at user branches -/// (conditions that don't depend on state) to capture all possible paths. -#[derive(Debug, Clone)] -pub struct TraceTree { - /// The root node of the trace tree. - pub root: TraceNode, - - /// Traces of exception handler entry blocks that were not reached by the main trace. - pub handler_traces: Vec, - - /// Dispatcher information (detected during tracing). - pub dispatcher: Option, - - /// Variables that are tainted by state (CFF machinery). - pub state_tainted: BitSet, - - /// Statistics about the trace. - pub stats: TraceStats, -} - -/// Statistics about a trace tree. -#[derive(Debug, Clone, Default)] -pub struct TraceStats { - /// Total number of nodes in the tree. - pub node_count: usize, - - /// Number of user branches encountered. - pub user_branch_count: usize, - - /// Number of state transitions (dispatcher visits). - pub state_transition_count: usize, - - /// Maximum depth of the tree. - pub max_depth: usize, - - /// Number of exit points (ret/throw). - pub exit_count: usize, -} - -/// A node in the trace tree representing a segment of execution. -#[derive(Debug, Clone)] -pub struct TraceNode { - /// Unique identifier for this node. - pub id: usize, - - /// The block index where this segment starts. - pub start_block: usize, - - /// Blocks visited in this segment (in order). - pub blocks_visited: VisitedBlocks, - - /// How this segment ends. - pub terminator: TraceTerminator, -} - -/// How a trace segment terminates. -#[derive(Debug, Clone)] -pub enum TraceTerminator { - /// Reached method exit (ret/throw). - Exit { - /// The exit block. - block: usize, - }, - - /// State-driven transition through dispatcher (deterministic). - /// We follow this automatically - it's CFF machinery. - StateTransition { - /// The state value that led here. - from_state: i64, - /// The next state value. - to_state: i64, - /// The case block we're transitioning to. - target_block: usize, - /// Continuation of the trace. - continues: Box, - }, - - /// Internal sentinel: state transition that needs iterative continuation. - /// Only used transiently during `trace_from_block` — never appears in - /// the final trace tree. - PendingStateTransition { - from_state: i64, - target_block: usize, - }, - - /// User branch - the condition doesn't depend on state. - /// This represents original program logic that was flattened. - UserBranch { - /// The block containing the branch. - block: usize, - /// The condition variable. - condition: SsaVarId, - /// True branch continuation. - true_branch: Box, - /// False branch continuation. - false_branch: Box, - }, - - /// User switch - value doesn't depend on state. - UserSwitch { - /// The block containing the switch. - block: usize, - /// The switch value variable. - value: SsaVarId, - /// Case continuations: (case_value, node). - cases: Vec<(i64, Box)>, - /// Default continuation. - default: Box, - }, - - /// Trace stopped due to a limit or error. - Stopped { reason: StopReason }, - - /// Loop detected - this path rejoins an earlier point. - /// We don't expand further to avoid infinite trees. - LoopBack { - /// The block we're looping back to. - target_block: usize, - /// The state when reaching this point. - state: i64, - }, -} - -impl TraceTree { - /// Creates a new trace tree with a root node. - /// - /// The `variable_count` parameter is the number of SSA variables, used to - /// size the `state_tainted` bit set. - #[must_use] - pub fn new(root: TraceNode, variable_count: usize) -> Self { - Self { - root, - handler_traces: Vec::new(), - dispatcher: None, - state_tainted: BitSet::new(variable_count), - stats: TraceStats::default(), - } - } - - /// Checks if a variable is state-tainted. - #[must_use] - pub fn is_state_tainted(&self, var: SsaVarId) -> bool { - self.state_tainted.contains(var.index()) - } - - /// Marks a variable as state-tainted. - pub fn mark_tainted(&mut self, var: SsaVarId) { - self.state_tainted.insert(var.index()); - } -} - -impl TraceNode { - /// Creates a new trace node. - pub fn new(id: usize, start_block: usize) -> Self { - Self { - id, - start_block, - blocks_visited: smallvec![start_block], - terminator: TraceTerminator::Stopped { - reason: StopReason::UnknownControlFlow { block: start_block }, - }, - } - } - - /// Records visiting a block. - pub fn visit_block(&mut self, block: usize) { - self.blocks_visited.push(block); - } - - /// Marks this node as needing a state transition continuation. - pub fn set_pending_state_transition(&mut self, from_state: i64, target_block: usize) { - self.terminator = TraceTerminator::PendingStateTransition { - from_state, - target_block, - }; - } - - /// Returns pending state transition info if this node needs continuation. - pub fn pending_state_transition(&self) -> Option<(i64, usize)> { - match &self.terminator { - TraceTerminator::PendingStateTransition { - from_state, - target_block, - } => Some((*from_state, *target_block)), - _ => None, - } - } - - /// Sets the terminator. - pub fn set_terminator(&mut self, terminator: TraceTerminator) { - self.terminator = terminator; - } - - /// Returns true if this node ends at an exit. - pub fn is_exit(&self) -> bool { - matches!(self.terminator, TraceTerminator::Exit { .. }) - } - - /// Returns true if this node has a user branch. - pub fn is_user_branch(&self) -> bool { - matches!( - self.terminator, - TraceTerminator::UserBranch { .. } | TraceTerminator::UserSwitch { .. } - ) - } -} diff --git a/dotscope/src/deobfuscation/renamer/cascade.rs b/dotscope/src/deobfuscation/renamer/cascade.rs index 977cc1d0..d8ba2d44 100644 --- a/dotscope/src/deobfuscation/renamer/cascade.rs +++ b/dotscope/src/deobfuscation/renamer/cascade.rs @@ -341,7 +341,7 @@ impl<'a> CascadeRenamer<'a> { let param_owners = build_param_owner_map(methoddef_table, param_table.row_count); for rid in 1..=param_table.row_count { - let Some(param) = param_table.get(rid) else { + let Some(param) = param_table.get(rid).ok().flatten() else { continue; }; let name_index = param.name; @@ -401,7 +401,7 @@ impl<'a> CascadeRenamer<'a> { let method_order = self.build_method_order(methoddef_table.row_count); for rid in method_order { - let Some(methoddef) = methoddef_table.get(rid) else { + let Some(methoddef) = methoddef_table.get(rid).ok().flatten() else { continue; }; let name_index = methoddef.name; @@ -465,7 +465,7 @@ impl<'a> CascadeRenamer<'a> { .unwrap_or_default(); for rid in 1..=field_table.row_count { - let Some(field) = field_table.get(rid) else { + let Some(field) = field_table.get(rid).ok().flatten() else { continue; }; let name_index = field.name; @@ -524,7 +524,7 @@ impl<'a> CascadeRenamer<'a> { continue; } - let Some(typedef) = typedef_table.get(rid) else { + let Some(typedef) = typedef_table.get(rid).ok().flatten() else { continue; }; let name_index = typedef.type_name; @@ -645,7 +645,7 @@ impl<'a> CascadeRenamer<'a> { } // Base class if let Some(base) = declaring_type.base() { - context.base_class = Some(base.fullname()); + context.base_class = Some(base.fullname().to_string()); } // Siblings: already-renamed methods in the same type for sibling_method in declaring_type.methods() { @@ -782,7 +782,7 @@ impl<'a> CascadeRenamer<'a> { // Base class if let Some(base) = cil_type.base() { - context.base_class = Some(base.fullname()); + context.base_class = Some(base.fullname().to_string()); } // Interfaces @@ -1017,7 +1017,7 @@ fn build_param_owner_map( let mut map = HashMap::new(); for method_rid in 1..=methoddef_table.row_count { - let Some(method) = methoddef_table.get(method_rid) else { + let Some(method) = methoddef_table.get(method_rid).ok().flatten() else { continue; }; let param_start = method.param_list; @@ -1031,6 +1031,8 @@ fn build_param_owner_map( let param_end = if method_rid < methoddef_table.row_count { methoddef_table .get(next_method_rid) + .ok() + .flatten() .map(|next| next.param_list) .unwrap_or(param_end_default) } else { @@ -1057,7 +1059,7 @@ fn build_member_owner_map( let mut map = HashMap::new(); for type_rid in 1..=typedef_table.row_count { - let Some(typedef) = typedef_table.get(type_rid) else { + let Some(typedef) = typedef_table.get(type_rid).ok().flatten() else { continue; }; let start = get_list_start(&typedef); @@ -1070,6 +1072,8 @@ fn build_member_owner_map( let end = if type_rid < typedef_table.row_count { typedef_table .get(next_type_rid) + .ok() + .flatten() .map(|next| get_list_start(&next)) .unwrap_or(end_default) } else { @@ -1150,6 +1154,7 @@ mod tests { validation::ValidationConfig, }, test::helpers::load_sample, + utils::truncate_chars, CilObject, }; @@ -1544,7 +1549,7 @@ mod tests { let methoddef_table = tables.table::().unwrap(); for rid in 1..=methoddef_table.row_count { - let Some(methoddef) = methoddef_table.get(rid) else { + let Some(methoddef) = methoddef_table.get(rid).ok().flatten() else { continue; }; let name_index = methoddef.name; @@ -1616,7 +1621,7 @@ mod tests { let mut obfuscated_with_context = 0u32; for rid in 1..=methoddef_table.row_count { - let Some(methoddef) = methoddef_table.get(rid) else { + let Some(methoddef) = methoddef_table.get(rid).ok().flatten() else { continue; }; let Ok(name) = strings.get(methoddef.name as usize) else { @@ -1675,7 +1680,7 @@ mod tests { // Find the first obfuscated method with call targets let mut found_method_with_calls = false; for rid in 1..=methoddef_table.row_count { - let Some(methoddef) = methoddef_table.get(rid) else { + let Some(methoddef) = methoddef_table.get(rid).ok().flatten() else { continue; }; let Ok(name) = strings.get(methoddef.name as usize) else { @@ -1745,7 +1750,7 @@ mod tests { let mut params_with_parent_or_calls = 0u32; for rid in 1..=param_table.row_count { - let Some(param) = param_table.get(rid) else { + let Some(param) = param_table.get(rid).ok().flatten() else { continue; }; if param.name == 0 { @@ -2026,7 +2031,7 @@ mod tests { if let Some(typedef_table) = tables.table::() { eprintln!("\nTypeDef table: {} rows", typedef_table.row_count); for rid in 1..=typedef_table.row_count { - if let Some(td) = typedef_table.get(rid) { + if let Some(td) = typedef_table.get(rid).ok().flatten() { let name = strings.get(td.type_name as usize).unwrap_or("?"); let ns = strings.get(td.type_namespace as usize).unwrap_or(""); let obf = is_obfuscated_name(name); @@ -2042,7 +2047,7 @@ mod tests { if let Some(methoddef_table) = tables.table::() { eprintln!("\nMethodDef table: {} rows", methoddef_table.row_count); for rid in 1..=methoddef_table.row_count { - if let Some(md) = methoddef_table.get(rid) { + if let Some(md) = methoddef_table.get(rid).ok().flatten() { let name = strings.get(md.name as usize).unwrap_or("?"); let obf = is_obfuscated_name(name); let special = is_special_name(name); @@ -2063,7 +2068,7 @@ mod tests { if let Some(field_table) = tables.table::() { eprintln!("\nField table: {} rows", field_table.row_count); for rid in 1..=field_table.row_count { - if let Some(f) = field_table.get(rid) { + if let Some(f) = field_table.get(rid).ok().flatten() { let name = strings.get(f.name as usize).unwrap_or("?"); let obf = is_obfuscated_name(name); eprintln!( @@ -2077,7 +2082,7 @@ mod tests { if let Some(param_table) = tables.table::() { eprintln!("\nParam table: {} rows", param_table.row_count); for rid in 1..=param_table.row_count { - if let Some(p) = param_table.get(rid) { + if let Some(p) = param_table.get(rid).ok().flatten() { let name = strings.get(p.name as usize).unwrap_or("?"); let obf = is_obfuscated_name(name); eprintln!( @@ -2115,6 +2120,8 @@ mod tests { } let name = methoddef_table .get(rid) + .ok() + .flatten() .and_then(|md| strings.get(md.name as usize).ok()) .unwrap_or("?"); let method = assembly.method(&method_token).ok(); @@ -2152,6 +2159,8 @@ mod tests { }; let name = methoddef_table .get(rid) + .ok() + .flatten() .and_then(|md| strings.get(md.name as usize).ok()) .unwrap_or("?") .to_string(); @@ -2179,8 +2188,8 @@ mod tests { let string_lits = features::collect_string_literals(ssa, &assembly); eprintln!(" String literals ({}):", string_lits.len()); for s in &string_lits { - let display = if s.len() > 60 { - format!("{}...", &s[..57]) + let display = if s.chars().count() > 60 { + format!("{}...", truncate_chars(s, 57)) } else { s.clone() }; @@ -2334,13 +2343,13 @@ mod tests { let param_owners = build_param_owner_map(methoddef_table, param_table.row_count); for rid in 1..=param_table.row_count { - let Some(param) = param_table.get(rid) else { + let Some(param) = param_table.get(rid).ok().flatten() else { continue; }; let name = strings.get(param.name as usize).unwrap_or("?"); let owner_rid = param_owners.get(&rid).copied(); let owner_name = owner_rid - .and_then(|r| methoddef_table.get(r)) + .and_then(|r| methoddef_table.get(r).ok().flatten()) .and_then(|md| strings.get(md.name as usize).ok()) .unwrap_or("?"); @@ -2368,7 +2377,7 @@ mod tests { if let Some(field_table) = tables.table::() { for rid in 1..=field_table.row_count { - let Some(f) = field_table.get(rid) else { + let Some(f) = field_table.get(rid).ok().flatten() else { continue; }; let name = strings.get(f.name as usize).unwrap_or("?"); @@ -2397,7 +2406,7 @@ mod tests { if rid == 1 { continue; // skip } - let Some(td) = typedef_table.get(rid) else { + let Some(td) = typedef_table.get(rid).ok().flatten() else { continue; }; let name = strings.get(td.type_name as usize).unwrap_or("?"); diff --git a/dotscope/src/deobfuscation/renamer/mod.rs b/dotscope/src/deobfuscation/renamer/mod.rs index 4996cc5a..e557af3a 100644 --- a/dotscope/src/deobfuscation/renamer/mod.rs +++ b/dotscope/src/deobfuscation/renamer/mod.rs @@ -218,7 +218,7 @@ fn execute_simple_rename(assembly: &CilObject) -> Result> { // Collect obfuscated names from TypeDef if let Some(typedef_table) = tables.table::() { for rid in 1..=typedef_table.row_count { - if let Some(typedef) = typedef_table.get(rid) { + if let Some(typedef) = typedef_table.get(rid)? { // Skip (RID 1) if rid == 1 { continue; @@ -247,7 +247,7 @@ fn execute_simple_rename(assembly: &CilObject) -> Result> { // Collect obfuscated method names from MethodDef if let Some(methoddef_table) = tables.table::() { for rid in 1..=methoddef_table.row_count { - if let Some(methoddef) = methoddef_table.get(rid) { + if let Some(methoddef) = methoddef_table.get(rid)? { let name_index = methoddef.name; if name_index > 0 { if let Ok(name) = strings.get(name_index as usize) { @@ -271,7 +271,7 @@ fn execute_simple_rename(assembly: &CilObject) -> Result> { // Collect obfuscated field names from Field if let Some(field_table) = tables.table::() { for rid in 1..=field_table.row_count { - if let Some(field) = field_table.get(rid) { + if let Some(field) = field_table.get(rid)? { let name_index = field.name; if name_index > 0 { if let Ok(name) = strings.get(name_index as usize) { @@ -295,7 +295,7 @@ fn execute_simple_rename(assembly: &CilObject) -> Result> { // Collect obfuscated parameter names from Param if let Some(param_table) = tables.table::() { for rid in 1..=param_table.row_count { - if let Some(param) = param_table.get(rid) { + if let Some(param) = param_table.get(rid)? { let name_index = param.name; if name_index > 0 { if let Ok(name) = strings.get(name_index as usize) { @@ -351,7 +351,7 @@ fn update_row_name_field( match table_id { TableId::TypeDef => { let table = tables.table::(); - table.and_then(|t| t.get(rid)).map(|row| { + table.and_then(|t| t.get(rid).ok().flatten()).map(|row| { let mut row = row.clone(); row.type_name = new_string_placeholder; TableDataOwned::TypeDef(row) @@ -359,7 +359,7 @@ fn update_row_name_field( } TableId::MethodDef => { let table = tables.table::(); - table.and_then(|t| t.get(rid)).map(|row| { + table.and_then(|t| t.get(rid).ok().flatten()).map(|row| { let mut row = row.clone(); row.name = new_string_placeholder; TableDataOwned::MethodDef(row) @@ -367,7 +367,7 @@ fn update_row_name_field( } TableId::Field => { let table = tables.table::(); - table.and_then(|t| t.get(rid)).map(|row| { + table.and_then(|t| t.get(rid).ok().flatten()).map(|row| { let mut row = row.clone(); row.name = new_string_placeholder; TableDataOwned::Field(row) @@ -375,7 +375,7 @@ fn update_row_name_field( } TableId::Param => { let table = tables.table::(); - table.and_then(|t| t.get(rid)).map(|row| { + table.and_then(|t| t.get(rid).ok().flatten()).map(|row| { let mut row = row.clone(); row.name = new_string_placeholder; TableDataOwned::Param(row) @@ -616,7 +616,7 @@ mod tests { let mut type_names = Vec::new(); for rid in 1..=typedef_table.row_count { - if let Some(row) = typedef_table.get(rid) { + if let Some(row) = typedef_table.get(rid).ok().flatten() { if let Ok(name) = strings.get(row.type_name as usize) { type_names.push(name.to_string()); } @@ -663,7 +663,7 @@ mod tests { let mut method_names = Vec::new(); for rid in 1..=methoddef_table.row_count { - if let Some(row) = methoddef_table.get(rid) { + if let Some(row) = methoddef_table.get(rid).ok().flatten() { if let Ok(name) = strings.get(row.name as usize) { method_names.push(name.to_string()); } diff --git a/dotscope/src/deobfuscation/renamer/phases.rs b/dotscope/src/deobfuscation/renamer/phases.rs index cc5b8859..ea8c3df5 100644 --- a/dotscope/src/deobfuscation/renamer/phases.rs +++ b/dotscope/src/deobfuscation/renamer/phases.rs @@ -20,6 +20,7 @@ use crate::{ }, utils::is_obfuscated_name, }, + utils::truncate_chars, CilObject, }; @@ -401,8 +402,8 @@ pub fn build_call_site_skeleton(ssa: &SsaFunction, assembly: &CilObject) -> Opti value: ConstValue::DecryptedString(s), .. } => { - let truncated = if s.len() > 30 { - format!("\"{}...\"", &s[..27]) + let truncated = if s.chars().count() > 30 { + format!("\"{}...\"", truncate_chars(s, 27)) } else { format!("\"{s}\"") }; @@ -416,8 +417,8 @@ pub fn build_call_site_skeleton(ssa: &SsaFunction, assembly: &CilObject) -> Opti if let Ok(s) = us.get(*idx as usize) { if let Ok(decoded) = s.to_string() { if !decoded.is_empty() { - let truncated = if decoded.len() > 30 { - format!("\"{}...\"", &decoded[..27]) + let truncated = if decoded.chars().count() > 30 { + format!("\"{}...\"", truncate_chars(&decoded, 27)) } else { format!("\"{decoded}\"") }; @@ -806,14 +807,17 @@ fn classify_op_into_profile(op: &SsaOp, profile: &mut OpcodeProfile) { /// /// The namespace portion of the name. fn extract_namespace(method_name: &str) -> String { - if let Some(idx) = method_name.rfind("::") { - let type_part = &method_name[..idx]; - if let Some(dot_idx) = type_part.rfind('.') { - return type_part[..dot_idx].to_string(); - } - return type_part.to_string(); + // `rsplit_once` rather than `rfind` + slice: the byte offsets from `rfind` are always + // character boundaries so the slicing was sound, but expressing it this way keeps the + // crate free of `str` range indexing and so needs no `clippy::string_slice` escape. + let Some((type_part, _)) = method_name.rsplit_once("::") else { + return method_name.to_string(); + }; + + match type_part.rsplit_once('.') { + Some((namespace, _)) => namespace.to_string(), + None => type_part.to_string(), } - method_name.to_string() } #[cfg(test)] diff --git a/dotscope/src/deobfuscation/renamer/prompt.rs b/dotscope/src/deobfuscation/renamer/prompt.rs index c6553eca..90d9631c 100644 --- a/dotscope/src/deobfuscation/renamer/prompt.rs +++ b/dotscope/src/deobfuscation/renamer/prompt.rs @@ -7,7 +7,10 @@ //! Each [`IdentifierKind`] has a distinct template optimized for the //! information most useful for that kind of rename. -use crate::deobfuscation::renamer::context::{IdentifierKind, ParamInfo, PhaseInfo, RenameContext}; +use crate::{ + deobfuscation::renamer::context::{IdentifierKind, ParamInfo, PhaseInfo, RenameContext}, + utils::truncate_chars, +}; /// Builds a FIM prompt from a rename context. /// @@ -184,8 +187,8 @@ fn render_shared_context(prefix: &mut String, context: &RenameContext) { .iter() .take(5) .map(|s| { - if s.len() > 30 { - format!("\"{}...\"", &s[..27]) + if s.chars().count() > 30 { + format!("\"{}...\"", truncate_chars(s, 27)) } else { format!("\"{s}\"") } @@ -237,8 +240,8 @@ fn render_caller_context(prefix: &mut String, context: &RenameContext) { .iter() .take(3) .map(|s| { - if s.len() > 40 { - format!("\"{}...\"", &s[..37]) + if s.chars().count() > 40 { + format!("\"{}...\"", truncate_chars(s, 37)) } else { format!("\"{s}\"") } diff --git a/dotscope/src/deobfuscation/renamer/validate.rs b/dotscope/src/deobfuscation/renamer/validate.rs index 7efcca05..91daccde 100644 --- a/dotscope/src/deobfuscation/renamer/validate.rs +++ b/dotscope/src/deobfuscation/renamer/validate.rs @@ -5,7 +5,7 @@ use std::collections::HashMap; -use crate::deobfuscation::renamer::context::IdentifierKind; +use crate::{deobfuscation::renamer::context::IdentifierKind, utils::truncate_chars}; /// Validates and normalizes a suggested name. /// @@ -47,11 +47,7 @@ pub fn validate_name(name: &str, kind: IdentifierKind, max_length: usize) -> Opt } // Truncate - let truncated = if cleaned.len() > max_length { - &cleaned[..max_length] - } else { - cleaned - }; + let truncated = truncate_chars(cleaned, max_length); // Apply casing convention let result = match kind { diff --git a/dotscope/src/deobfuscation/techniques/assembly.rs b/dotscope/src/deobfuscation/techniques/assembly.rs index f51abac5..3b66cfef 100644 --- a/dotscope/src/deobfuscation/techniques/assembly.rs +++ b/dotscope/src/deobfuscation/techniques/assembly.rs @@ -167,7 +167,7 @@ impl WorkingAssembly { .view() .tables() .and_then(|t| t.table::()) - .and_then(|table| table.get(fieldrva_rid)) + .and_then(|table| table.get(fieldrva_rid).ok().flatten()) .ok_or_else(|| Error::Other(format!("FieldRVA row {fieldrva_rid} not found")))?; let updated_row = FieldRvaRaw { diff --git a/dotscope/src/deobfuscation/techniques/bitmono/hooks.rs b/dotscope/src/deobfuscation/techniques/bitmono/hooks.rs index d59a5899..5b1fb67c 100644 --- a/dotscope/src/deobfuscation/techniques/bitmono/hooks.rs +++ b/dotscope/src/deobfuscation/techniques/bitmono/hooks.rs @@ -894,7 +894,7 @@ fn is_redirect_stub_memberref( let Some(memberref_table) = tables.table::() else { return false; }; - let Some(memberref) = memberref_table.get(token.row()) else { + let Some(memberref) = memberref_table.get(token.row()).ok().flatten() else { return false; }; diff --git a/dotscope/src/deobfuscation/techniques/bitmono/strings.rs b/dotscope/src/deobfuscation/techniques/bitmono/strings.rs index 43fd095c..d5b9e83b 100644 --- a/dotscope/src/deobfuscation/techniques/bitmono/strings.rs +++ b/dotscope/src/deobfuscation/techniques/bitmono/strings.rs @@ -64,10 +64,10 @@ use crate::{ utils::build_init_array_map, }, metadata::{ - tables::{MemberRefRaw, TableId, TypeDefRaw, TypeRefRaw}, + tables::{MemberRefRaw, MetadataTable, RowReadable, TableId, TypeDefRaw, TypeRefRaw}, token::Token, }, - utils::CryptoParameters, + utils::{CryptoParameters, MAX_DERIVED_KEY_LEN, MAX_PBKDF2_ITERATIONS}, CilObject, }; @@ -476,6 +476,17 @@ fn extract_crypto_parameters(ssa: &SsaFunction, assembly: &CilObject) -> CryptoP params.iv_size = iv_size as usize; } + // Defence in depth. `derive_pbkdf2_key` enforces these ceilings too and is the guard that + // matters, but everything above is lifted verbatim from constants in an attacker-supplied + // method body, so an implausible value here means the detection misfired rather than that a + // real decryptor wants a 2 GB key. Falling back to the .NET defaults keeps the pass working + // on the samples it was written for instead of failing the whole assembly. + if params.iterations > MAX_PBKDF2_ITERATIONS + || params.key_size.saturating_add(params.iv_size) > MAX_DERIVED_KEY_LEN + { + return CryptoParameters::default(); + } + params } @@ -530,54 +541,44 @@ fn resolve_type_name(assembly: &CilObject, token: Token) -> Option { let tables = assembly.tables()?; let strings = assembly.strings()?; + /// Reads one row, treating both "no such row" and "row does not parse" as "no name". + /// + /// This function has no error channel — it answers `Option` — so an unreadable + /// row cannot be propagated. `.ok().flatten()` is that decision written down: the row is + /// given up on here, at the point that can see it, rather than inside the iterator. + fn row(table: &MetadataTable<'_, T>, rid: u32) -> Option { + table.get(rid).ok().flatten() + } + + let type_name = |namespace_idx: u32, name_idx: u32| -> Option { + let name = strings.get(name_idx as usize).ok()?; + let ns = strings.get(namespace_idx as usize).unwrap_or(""); + Some(format!("{ns}.{name}")) + }; + match token.table() { // MemberRef (0x0A) — follow class to get declaring type 0x0A => { - let memberref_table = tables.table::()?; - let memberref = memberref_table.get(token.row())?; + let memberref = row(tables.table::()?, token.row())?; if memberref.class.tag == TableId::TypeRef { - let typeref_table = tables.table::()?; - let typeref = typeref_table.get(memberref.class.row)?; - let name = strings.get(typeref.type_name as usize).ok()?; - let ns = strings - .get(typeref.type_namespace as usize) - .ok() - .unwrap_or(""); - Some(format!("{ns}.{name}")) + let typeref = row(tables.table::()?, memberref.class.row)?; + type_name(typeref.type_namespace, typeref.type_name) } else if memberref.class.tag == TableId::TypeDef { - let typedef_table = tables.table::()?; - let typedef = typedef_table.get(memberref.class.row)?; - let name = strings.get(typedef.type_name as usize).ok()?; - let ns = strings - .get(typedef.type_namespace as usize) - .ok() - .unwrap_or(""); - Some(format!("{ns}.{name}")) + let typedef = row(tables.table::()?, memberref.class.row)?; + type_name(typedef.type_namespace, typedef.type_name) } else { None } } // TypeRef (0x01) 0x01 => { - let typeref_table = tables.table::()?; - let typeref = typeref_table.get(token.row())?; - let name = strings.get(typeref.type_name as usize).ok()?; - let ns = strings - .get(typeref.type_namespace as usize) - .ok() - .unwrap_or(""); - Some(format!("{ns}.{name}")) + let typeref = row(tables.table::()?, token.row())?; + type_name(typeref.type_namespace, typeref.type_name) } // TypeDef (0x02) 0x02 => { - let typedef_table = tables.table::()?; - let typedef = typedef_table.get(token.row())?; - let name = strings.get(typedef.type_name as usize).ok()?; - let ns = strings - .get(typedef.type_namespace as usize) - .ok() - .unwrap_or(""); - Some(format!("{ns}.{name}")) + let typedef = row(tables.table::()?, token.row())?; + type_name(typedef.type_namespace, typedef.type_name) } _ => None, } diff --git a/dotscope/src/deobfuscation/techniques/confuserex/constants.rs b/dotscope/src/deobfuscation/techniques/confuserex/constants.rs index 014e4ade..374c0b4a 100644 --- a/dotscope/src/deobfuscation/techniques/confuserex/constants.rs +++ b/dotscope/src/deobfuscation/techniques/confuserex/constants.rs @@ -313,6 +313,13 @@ impl Technique for ConfuserExConstants { if let Some(fieldrva_table) = tables.table::() { let file = assembly.file(); for row in fieldrva_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if row.rva == 0 { continue; } @@ -363,6 +370,13 @@ impl Technique for ConfuserExConstants { if let Some(tables) = assembly.tables() { if let Some(methodspec_table) = tables.table::() { for spec in methodspec_table { + let spec = match spec { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let references_decryptor = if decryptor_set.contains(&spec.method.token) { true } else if spec.method.token.is_table(TableId::MemberRef) { @@ -663,6 +677,13 @@ fn register_methodspec_mappings( }; for methodspec in methodspec_table { + let methodspec = match methodspec { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let method_token = methodspec.method.token; // Check if this MethodSpec references a known decryptor. diff --git a/dotscope/src/deobfuscation/techniques/confuserex/helpers.rs b/dotscope/src/deobfuscation/techniques/confuserex/helpers.rs index 41692170..873a8154 100644 --- a/dotscope/src/deobfuscation/techniques/confuserex/helpers.rs +++ b/dotscope/src/deobfuscation/techniques/confuserex/helpers.rs @@ -66,7 +66,7 @@ pub(super) fn get_method_rva(assembly: &CilObject, token: Token) -> Option let tables = assembly.tables()?; let method_table = tables.table::()?; let row = token.row(); - let method_row = method_table.get(row)?; + let method_row = method_table.get(row).ok().flatten()?; Some(method_row.rva) } @@ -123,6 +123,13 @@ pub(super) fn extract_decrypted_field_data( }; for row in fieldrva_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let rva = row.rva; if rva == 0 { continue; @@ -204,6 +211,13 @@ pub(super) fn resolve_pinvoke_tokens(assembly: &CilObject, target_name: &str) -> }; for row in implmap_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let Ok(import_name) = strings.get(row.import_name as usize) else { continue; }; diff --git a/dotscope/src/deobfuscation/techniques/confuserex/marker.rs b/dotscope/src/deobfuscation/techniques/confuserex/marker.rs index c0e1fca6..e8da4692 100644 --- a/dotscope/src/deobfuscation/techniques/confuserex/marker.rs +++ b/dotscope/src/deobfuscation/techniques/confuserex/marker.rs @@ -85,6 +85,13 @@ impl Technique for ConfuserExMarker { let mut version = None; for attr in ca_table { + let attr = match attr { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; // Resolve the constructor to find the declaring type name. // ConfuserEx defines marker attributes locally, so the constructor // is typically a MethodDef pointing to a local TypeDef. diff --git a/dotscope/src/deobfuscation/techniques/confuserex/metadata.rs b/dotscope/src/deobfuscation/techniques/confuserex/metadata.rs index 3b0c9567..b4a7bae5 100644 --- a/dotscope/src/deobfuscation/techniques/confuserex/metadata.rs +++ b/dotscope/src/deobfuscation/techniques/confuserex/metadata.rs @@ -123,6 +123,13 @@ impl Technique for ConfuserExMetadata { // Check Module table for invalid name indices (0x7fff7fff marker). if let Some(module_table) = tables.table::() { for row in module_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if row.name == CONFUSEREX_MARKER || row.name as usize >= strings_size { findings.invalid_entries = findings.invalid_entries.saturating_add(1); // Skip rows where the file offset would overflow when adding the @@ -141,6 +148,13 @@ impl Technique for ConfuserExMetadata { // Check Assembly table for invalid name indices. if let Some(assembly_table) = tables.table::() { for row in assembly_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if row.name == CONFUSEREX_MARKER || row.name as usize >= strings_size { findings.invalid_entries = findings.invalid_entries.saturating_add(1); } @@ -150,6 +164,13 @@ impl Technique for ConfuserExMetadata { // Check DeclSecurity for invalid action values. if let Some(declsec_table) = tables.table::() { for row in declsec_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if row.action == CONFUSEREX_MARKER_16 || row.action > 0x000E { findings.invalid_entries = findings.invalid_entries.saturating_add(1); } @@ -159,6 +180,13 @@ impl Technique for ConfuserExMetadata { // Check TypeRef resolution scopes for invalid indices. if let Some(typeref_table) = tables.table::() { for row in typeref_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if row.resolution_scope.tag == TableId::Module && row.resolution_scope.row == 0 { findings.invalid_entries = findings.invalid_entries.saturating_add(1); } diff --git a/dotscope/src/deobfuscation/techniques/confuserex/natives.rs b/dotscope/src/deobfuscation/techniques/confuserex/natives.rs index b6e93438..8d08c66d 100644 --- a/dotscope/src/deobfuscation/techniques/confuserex/natives.rs +++ b/dotscope/src/deobfuscation/techniques/confuserex/natives.rs @@ -229,6 +229,13 @@ fn collect_implmap_methods(tables: &crate::metadata::streams::TablesHeader<'_>) let mut pinvoke_methods = HashSet::new(); if let Some(implmap_table) = tables.table::() { for row in implmap_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; pinvoke_methods.insert(row.member_forwarded.token); } } diff --git a/dotscope/src/deobfuscation/techniques/confuserex/resources.rs b/dotscope/src/deobfuscation/techniques/confuserex/resources.rs index 6e448ab7..6aedb620 100644 --- a/dotscope/src/deobfuscation/techniques/confuserex/resources.rs +++ b/dotscope/src/deobfuscation/techniques/confuserex/resources.rs @@ -47,7 +47,8 @@ use crate::{ error::Error, metadata::{ tables::{ - ManifestResourceAttributes, ManifestResourceBuilder, ManifestResourceRaw, TableId, + skip_unreadable, ManifestResourceAttributes, ManifestResourceBuilder, + ManifestResourceRaw, TableId, }, token::Token, validation::ValidationConfig, @@ -389,6 +390,7 @@ fn find_manifest_resources_by_name( manifest_table .iter() + .filter_map(skip_unreadable) .filter_map(|row| { strings .get(row.name as usize) diff --git a/dotscope/src/deobfuscation/techniques/confuserex/tamper.rs b/dotscope/src/deobfuscation/techniques/confuserex/tamper.rs index 3204fa5a..210bd6be 100644 --- a/dotscope/src/deobfuscation/techniques/confuserex/tamper.rs +++ b/dotscope/src/deobfuscation/techniques/confuserex/tamper.rs @@ -198,6 +198,13 @@ impl Technique for ConfuserExAntiTamper { let text_rva_end = text_rva_start.saturating_add(text.virtual_size as usize); for row in method_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if row.rva == 0 { continue; } @@ -457,7 +464,7 @@ impl Technique for ConfuserExAntiTamper { .view() .tables() .and_then(|t| t.table::()) - .and_then(|table| table.get(rid)) + .and_then(|table| table.get(rid).ok().flatten()) .ok_or_else(|| Error::Deobfuscation(format!("MethodDef row {rid} not found"))) { Ok(v) => v, @@ -495,7 +502,7 @@ impl Technique for ConfuserExAntiTamper { .view() .tables() .and_then(|t| t.table::()) - .and_then(|table| table.get(rid)) + .and_then(|table| table.get(rid).ok().flatten()) .ok_or_else(|| Error::Deobfuscation(format!("FieldRVA row {rid} not found"))) { Ok(v) => v, diff --git a/dotscope/src/deobfuscation/techniques/detection.rs b/dotscope/src/deobfuscation/techniques/detection.rs index c78bc6b4..94338c3a 100644 --- a/dotscope/src/deobfuscation/techniques/detection.rs +++ b/dotscope/src/deobfuscation/techniques/detection.rs @@ -6,7 +6,7 @@ use std::{ any::Any, - collections::{HashMap, HashSet}, + collections::{BTreeSet, HashMap, HashSet}, }; use crate::{cilassembly::CleanupRequest, metadata::token::Token}; @@ -279,9 +279,22 @@ impl Detections { /// Merges all cleanup contributions into a single result. #[must_use] pub fn merged_cleanup(&self) -> CleanupRequest { + self.merged_cleanup_excluding(&BTreeSet::new()) + } + + /// Merges cleanup contributions, skipping the named techniques. + /// + /// A technique fills its cleanup request during detection, before it knows + /// whether the transform those deletions depend on will succeed. When the + /// transform then fails, merging the request removes metadata that nothing + /// replaced — the infrastructure is deleted while the code that needs it is + /// still there in its protected form. Excluding the technique leaves the + /// obfuscation in place, which is the better of the two outcomes. + #[must_use] + pub fn merged_cleanup_excluding(&self, excluded: &BTreeSet<&str>) -> CleanupRequest { let mut request = CleanupRequest::new(); - for detection in self.entries.values() { - if detection.detected { + for (id, detection) in &self.entries { + if detection.detected && !excluded.contains(id.as_str()) { request.merge(&detection.cleanup); } } diff --git a/dotscope/src/deobfuscation/techniques/generic/decompiler.rs b/dotscope/src/deobfuscation/techniques/generic/decompiler.rs index 0997db1b..f39b3ba1 100644 --- a/dotscope/src/deobfuscation/techniques/generic/decompiler.rs +++ b/dotscope/src/deobfuscation/techniques/generic/decompiler.rs @@ -186,7 +186,7 @@ fn detect_antidecompiler_types( continue; }; - let row = typedef_table.get(nested.token.row()); + let row = typedef_table.get(nested.token.row()).ok().flatten(); let Some(row) = row else { continue; }; @@ -231,6 +231,13 @@ fn detect_fake_attributes( }; for attr in custom_attr_table { + let attr = match attr { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let is_module_or_assembly = attr.parent.tag == TableId::Module || attr.parent.tag == TableId::Assembly; if !is_module_or_assembly { diff --git a/dotscope/src/deobfuscation/techniques/generic/flattening.rs b/dotscope/src/deobfuscation/techniques/generic/flattening.rs index 4baf140d..77b3a440 100644 --- a/dotscope/src/deobfuscation/techniques/generic/flattening.rs +++ b/dotscope/src/deobfuscation/techniques/generic/flattening.rs @@ -77,7 +77,15 @@ impl Technique for GenericFlattening { } fn detect_ssa(&self, ctx: &AnalysisContext, _assembly: &CilObject) -> Detection { - let min_confidence = UnflattenConfig::default().min_confidence; + // Detection is the only consumer of these thresholds, so build them from + // the engine configuration here rather than defaulting and discarding + // what the caller asked for. + let config = UnflattenConfig { + max_backedge_depth: ctx.config.unflattening.max_backedge_depth, + confidence_weights: ctx.config.unflattening.confidence_weights.clone(), + ..UnflattenConfig::default() + }; + let min_confidence = config.min_confidence; let mut dispatchers_by_method: HashMap> = HashMap::new(); let mut total_dispatchers = 0usize; @@ -88,7 +96,7 @@ impl Technique for GenericFlattening { // Use the same CffDetector that CffReconstructionPass uses. // This gives us full structural analysis with confidence scoring, // dominance verification, state variable identification, etc. - let mut detector = CffDetector::new(ssa); + let mut detector = CffDetector::with_config(ssa, &config); let all_dispatchers = detector.detect_all_dispatchers(); // Two-tier confidence filtering: high-confidence dispatchers must @@ -143,13 +151,7 @@ impl Technique for GenericFlattening { detection: &Detection, _assembly: &Arc, ) -> Vec>> { - let cff_config = UnflattenConfig { - max_states: ctx.config.unflattening.max_states_per_case, - max_tree_depth: ctx.config.unflattening.max_trace_iterations, - ..UnflattenConfig::default() - }; - - let mut cff_pass = CffReconstructionPass::new(ctx, cff_config); + let mut cff_pass = CffReconstructionPass::new(ctx); if let Some(findings) = detection.findings::() { cff_pass = cff_pass.with_pre_detected(findings.dispatchers.clone()); } diff --git a/dotscope/src/deobfuscation/techniques/generic/handlers.rs b/dotscope/src/deobfuscation/techniques/generic/handlers.rs index 12bb3034..06879f5c 100644 --- a/dotscope/src/deobfuscation/techniques/generic/handlers.rs +++ b/dotscope/src/deobfuscation/techniques/generic/handlers.rs @@ -65,6 +65,13 @@ impl Technique for GenericHandlers { let mut affected = Vec::new(); for row in method_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if row.rva == 0 { continue; } diff --git a/dotscope/src/deobfuscation/techniques/generic/ildasm.rs b/dotscope/src/deobfuscation/techniques/generic/ildasm.rs index 24bb19e8..2e164070 100644 --- a/dotscope/src/deobfuscation/techniques/generic/ildasm.rs +++ b/dotscope/src/deobfuscation/techniques/generic/ildasm.rs @@ -68,6 +68,13 @@ impl Technique for GenericIldasm { }; for attr in custom_attr_table { + let attr = match attr { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let is_module_or_assembly = attr.parent.tag == TableId::Module || attr.parent.tag == TableId::Assembly; if !is_module_or_assembly { diff --git a/dotscope/src/deobfuscation/techniques/generic/metadata.rs b/dotscope/src/deobfuscation/techniques/generic/metadata.rs index f9e93546..9b41b743 100644 --- a/dotscope/src/deobfuscation/techniques/generic/metadata.rs +++ b/dotscope/src/deobfuscation/techniques/generic/metadata.rs @@ -111,6 +111,13 @@ impl Technique for GenericMetadata { let strings_size = strings.as_ref().map(|s| s.data().len()).unwrap_or(0); for row in module_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let name_index = row.name as usize; let is_sentinel = KNOWN_SENTINEL_VALUES.contains(&row.name); if name_index >= strings_size || is_sentinel { @@ -133,6 +140,13 @@ impl Technique for GenericMetadata { let strings_size = strings.as_ref().map(|s| s.data().len()).unwrap_or(0); for row in assembly_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let is_sentinel = KNOWN_SENTINEL_VALUES.contains(&row.name); if row.name as usize >= strings_size || is_sentinel { findings.invalid_assembly_rows = @@ -146,6 +160,13 @@ impl Technique for GenericMetadata { // outside this range is invalid and likely injected by an obfuscator. if let Some(declsec_table) = tables.table::() { for row in declsec_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let is_sentinel = KNOWN_SENTINEL_VALUES_16.contains(&row.action); if row.action > 0x000E || is_sentinel { findings.invalid_declsecurity_rows = @@ -157,6 +178,13 @@ impl Technique for GenericMetadata { // Check TypeRef resolution scopes for invalid indices if let Some(typeref_table) = tables.table::() { for row in typeref_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; // Resolution scope with tag Module but row 0 is suspicious // (valid Module is row 1) if row.resolution_scope.tag == TableId::Module && row.resolution_scope.row == 0 { diff --git a/dotscope/src/deobfuscation/techniques/generic/opaquefields.rs b/dotscope/src/deobfuscation/techniques/generic/opaquefields.rs index 36c6a5ce..e5b79f52 100644 --- a/dotscope/src/deobfuscation/techniques/generic/opaquefields.rs +++ b/dotscope/src/deobfuscation/techniques/generic/opaquefields.rs @@ -102,12 +102,117 @@ fn collect_predicate_static_fields(ssa: &SsaFunction) -> HashSet { static_fields } +/// Collects every field token written by a `StoreField`/`StoreStaticField` anywhere in `ssa`. +/// +/// A field that is assigned after its declaring constructor has run cannot be folded to a +/// constant: the value observed at `.cctor` warm-up time is not the value the program sees. +/// Callers use this across *every* SSA function in the assembly, because the write that +/// invalidates a fold is frequently in a different method from the read. +fn collect_field_stores(ssa: &SsaFunction) -> HashSet { + let mut stored = HashSet::new(); + for block in ssa.blocks() { + for instr in block.instructions() { + match instr.op() { + SsaOp::StoreField { field, .. } | SsaOp::StoreStaticField { field, .. } => { + stored.insert(field.token()); + } + _ => {} + } + } + } + stored +} + +/// Collects methods that only ever execute as part of type initialization. +/// +/// Stores made while a type initializes must not count against a field's immutability — +/// that store is what gives the field its constant value. Matching the `.cctor` name alone +/// is too narrow: an obfuscator can put the stores in a plain static helper and call it from +/// `.cctor`, which is exactly what .NET Reactor does with its key containers. A method whose +/// every caller is itself initialization-only runs only during initialization too, so the set +/// is closed under that rule and grown to a fixed point. +/// +/// A method with no recorded callers is not admitted: unreachable in the SSA call graph is not +/// the same as reachable only from `.cctor`, and assuming otherwise would admit any method the +/// analysis simply failed to see a caller for. +fn collect_initialization_only_methods( + ctx: &AnalysisContext, + assembly: &CilObject, +) -> HashSet { + let mut callers: HashMap> = HashMap::new(); + let mut methods: Vec = Vec::new(); + + for entry in ctx.ssa_functions.iter() { + let caller = *entry.key(); + methods.push(caller); + for block in entry.value().blocks() { + for instr in block.instructions() { + let callee = match instr.op() { + SsaOp::Call { method, .. } + | SsaOp::CallVirt { method, .. } + | SsaOp::LoadFunctionPtr { method, .. } + | SsaOp::LoadVirtFunctionPtr { method, .. } => method.token(), + SsaOp::NewObj { ctor, .. } => ctor.token(), + _ => continue, + }; + callers.entry(callee).or_default().insert(caller); + } + } + } + + let mut init_only: HashSet = methods + .iter() + .copied() + .filter(|token| is_static_constructor(assembly, *token)) + .collect(); + + loop { + let mut grew = false; + for &method in &methods { + if init_only.contains(&method) { + continue; + } + let Some(method_callers) = callers.get(&method) else { + continue; + }; + if !method_callers.is_empty() && method_callers.iter().all(|c| init_only.contains(c)) { + init_only.insert(method); + grew = true; + } + } + if !grew { + break; + } + } + + init_only +} + +/// Returns whether `token` names a static constructor. +/// +/// Used to exclude `.cctor` bodies from the assembly-wide store scan: an `initonly` static is +/// assigned there by definition, so counting those stores would disqualify every field the +/// immutability gate is meant to admit. +fn is_static_constructor(assembly: &CilObject, token: Token) -> bool { + assembly + .resolve_method_name(token) + .is_some_and(|name| name == ".cctor") +} + /// Scans an SSA function for ALL `LoadField(LoadStaticField(..))` patterns /// and collects the static field tokens. /// /// Unlike [`collect_predicate_static_fields`] which only looks at Branch /// terminators, this function scans every instruction. It captures both /// opaque predicate fields AND string encryption XOR key fields. +/// +/// # Why callers must gate the result +/// +/// This matches the plain singleton-access idiom (`Config.Instance.Retries`) just as readily as +/// an opaque predicate — `ldsfld; ldfld` is not an obfuscation signature. Folding an unfiltered +/// match set freezes mutable state at its constructor-time value, turning branches that depend +/// on runtime state into unconditional jumps to the wrong successor. Every caller therefore +/// filters through [`collect_field_stores`]. fn collect_field_load_sources(ssa: &SsaFunction) -> HashSet { let defs = build_def_map(ssa); @@ -291,10 +396,52 @@ impl Technique for GenericOpaquePredicates { let mut affected_fields: HashSet = HashSet::new(); let mut affected_methods: HashSet = HashSet::new(); + // Every field written anywhere in the assembly *except* in a static constructor. A + // store in one method invalidates a fold of the same field in another, so this has to + // be assembly-wide rather than per-method — and it must be collected before any + // folding decision is made. + // + // Initialization is excluded deliberately, and the gate below does not work without + // it: the pass warms those initializers up precisely so their values are known + // constants, so a store there is what makes a field foldable, not what disqualifies + // it. Counting them would leave `all_field_loads` permanently empty and silently + // disable Variant A's field-load detection entirely. + // + // The exclusion covers every initialization-only method, not just `.cctor` itself — + // see [`collect_initialization_only_methods`]. .NET Reactor's `.cctor` calls a plain + // static helper that does the stores, so a name-based test counts them and judges the + // key fields mutable. + let init_only = collect_initialization_only_methods(ctx, assembly); + let mut stored_fields: HashSet = HashSet::new(); + for entry in ctx.ssa_functions.iter() { + if init_only.contains(entry.key()) { + continue; + } + stored_fields.extend(collect_field_stores(entry.value())); + } + for entry in ctx.ssa_functions.iter() { let method_token = *entry.key(); let predicate_fields = collect_predicate_static_fields(entry.value()); - let all_field_loads = collect_field_load_sources(entry.value()); + + // `collect_field_load_sources` matches the ordinary singleton idiom as well as + // opaque predicates, so admit a field only when it is provably immutable: never + // stored outside the `.cctor`s this pass warms up (see the `.cctor` exclusion + // where `stored_fields` is built). Without that condition the fold silently + // freezes runtime-mutated state. + // + // Absence from `stored_fields` is the whole proof. Requiring the `initonly` flag + // as well proves nothing extra -- an `initonly` field is a field the compiler + // already refused to store outside the initializer, so it is a subset of what the + // scan admits -- while excluding every obfuscator that assigns in `.cctor` without + // setting the flag. .NET Reactor is one: its string-decryptor key fields are + // `static` (0x0013) with `initonly` (0x20) clear, so demanding the flag refused + // every fold and left the decryptor's arguments non-constant. + let all_field_loads: HashSet = collect_field_load_sources(entry.value()) + .into_iter() + .filter(|token| !stored_fields.contains(token)) + .collect(); + let combined: HashSet = predicate_fields.union(&all_field_loads).copied().collect(); if !combined.is_empty() { @@ -337,14 +484,27 @@ impl Technique for GenericOpaquePredicates { // Find types that own the Variant A opaque predicate fields. // These types exist solely as opaque predicate infrastructure and can be deleted. + // + // Deletion is whole-type and `build_cleanup_request` merges a type-only request + // unconditionally, so the bar is higher than for folding: require that *every* static + // field the type declares was resolved as predicate infrastructure. A type with even + // one unrelated static field is a real type that something else may still reference, + // and removing its methods and fields would leave surviving call sites dangling. let mut owning_types: HashSet = HashSet::new(); let registry = assembly.types(); for entry in registry.iter() { let type_ref = entry.value(); - let owns_field = type_ref.fields.iter().any(|(_, field)| { - field.flags.is_static() && resolved_fields.contains(&field.token) - }); - if owns_field { + let mut static_fields = type_ref + .fields + .iter() + .filter(|(_, field)| field.flags.is_static()) + .peekable(); + + if static_fields.peek().is_none() { + continue; + } + + if static_fields.all(|(_, field)| resolved_fields.contains(&field.token)) { owning_types.insert(*entry.key()); } } @@ -481,17 +641,25 @@ impl Technique for GenericOpaquePredicates { #[cfg(test)] mod tests { + use super::*; use crate::{ compiler::PassPhase, - deobfuscation::techniques::{ - generic::opaquefields::{GenericOpaquePredicates, OpaquePredicateFindings}, - Technique, TechniqueCategory, - }, + deobfuscation::techniques::{Technique, TechniqueCategory}, test::helpers::load_sample, }; + /// Pins the contract that the non-SSA `detect` entry point reports nothing. + /// + /// This is **not** a negative test for opaque-predicate detection, despite how the previous + /// version of it read. `detect` returns [`Detection::new_empty`] unconditionally — real + /// detection happens in `detect_ssa`, which needs def-use chains — so asserting + /// "nothing detected" here passes for every input, obfuscated or not, and cannot fail. + /// + /// Kept, narrowed, and renamed so it documents the contract instead of implying coverage + /// that does not exist. Genuine negative coverage for `detect_ssa` requires an + /// `AnalysisContext` with built SSA functions and belongs in the integration suite. #[test] - fn test_detect_negative_confuserex_original() { + fn detect_without_ssa_reports_nothing_by_contract() { let asm = load_sample("tests/samples/packers/confuserex/1.6.0/original.exe"); let technique = GenericOpaquePredicates; @@ -499,31 +667,15 @@ mod tests { assert!( !detection.is_detected(), - "GenericOpaquePredicates should not detect anything in a ConfuserEx original sample" - ); - assert!( - detection.evidence().is_empty(), - "No evidence should be present for a non-obfuscated sample" - ); - assert!( - detection.findings::().is_none(), - "No findings should be present for a non-obfuscated sample" + "the non-SSA entry point defers to detect_ssa and must report nothing" ); + assert!(detection.evidence().is_empty()); + assert!(detection.findings::().is_none()); } - #[test] - fn test_detect_negative_obfuscar_sample() { - let asm = load_sample("tests/samples/packers/obfuscar/2.2.50/obfuscar_strings_only.exe"); - - let technique = GenericOpaquePredicates; - let detection = technique.detect(&asm); - - // Obfuscar does not use opaque field predicates - assert!( - !detection.is_detected(), - "GenericOpaquePredicates should not detect anything in an Obfuscar sample" - ); - } + // A second sample-loading copy of the above was removed rather than renamed: it asserted + // the same unconditional-empty contract through a different sample and so added no + // coverage, only the appearance of it plus a sample load. #[test] fn test_technique_metadata() { @@ -543,4 +695,46 @@ mod tests { "GenericOpaquePredicates should run in the Structure SSA phase" ); } + + /// The `.cctor` exclusion that keeps Variant A's immutability gate satisfiable. + /// + /// The gate admits a static field only when it is `initonly` *and* never stored. An + /// `initonly` static can only be assigned in its declaring type's `.cctor`, so unless + /// `.cctor` stores are excluded from the store scan the two halves are mutually exclusive, + /// the candidate set is always empty, and field-load detection is dead code that no + /// assertion in this suite would notice. This pins the predicate that exclusion rests on. + #[test] + fn static_constructors_are_identified_for_the_store_scan() { + let asm = load_sample("tests/samples/packers/confuserex/1.6.0/original.exe"); + + let cctors: Vec<_> = asm + .query_methods() + .static_constructors() + .into_iter() + .collect(); + assert!( + !cctors.is_empty(), + "sample must contain at least one .cctor for this test to mean anything" + ); + for cctor in &cctors { + assert!( + is_static_constructor(&asm, cctor.token), + "a .cctor must be excluded from the assembly-wide store scan" + ); + } + + let non_cctors: Vec<_> = asm + .query_methods() + .filter(|m| !m.is_cctor()) + .into_iter() + .take(8) + .collect(); + assert!(!non_cctors.is_empty()); + for method in &non_cctors { + assert!( + !is_static_constructor(&asm, method.token), + "an ordinary method's stores must still invalidate a fold" + ); + } + } } diff --git a/dotscope/src/deobfuscation/techniques/jiejienet/arrays.rs b/dotscope/src/deobfuscation/techniques/jiejienet/arrays.rs index 36135900..63e51323 100644 --- a/dotscope/src/deobfuscation/techniques/jiejienet/arrays.rs +++ b/dotscope/src/deobfuscation/techniques/jiejienet/arrays.rs @@ -51,8 +51,8 @@ use crate::{ }, }, metadata::{ - signatures::TypeSignature, - tables::{FieldRvaRaw, TableId}, + signatures::{parse_field_signature, TypeSignature}, + tables::{ClassLayoutRaw, FieldRaw, FieldRvaRaw, TableId}, token::Token, typesystem::{wellknown, PointerSize}, }, @@ -805,6 +805,8 @@ fn decrypt_field_rva_data_to_bytes( let field_rid = field_token.row(); let rva_entry = fieldrva_table .iter() + .collect::>>()? + .into_iter() .find(|row| row.field == field_rid) .ok_or_else(|| { Error::Other(format!( @@ -837,11 +839,6 @@ fn decrypt_field_rva_data_to_bytes( /// Uses the ClassLayout table for ValueType fields (the common case for /// array-init backing fields which are ExplicitLayout structs). fn calculate_field_data_size(assembly: &CilObject, field_rid: u32) -> Result { - use crate::metadata::{ - signatures::parse_field_signature, - tables::{ClassLayoutRaw, FieldRaw}, - }; - let tables = assembly .tables() .ok_or_else(|| Error::Other("No metadata tables".to_string()))?; @@ -851,6 +848,8 @@ fn calculate_field_data_size(assembly: &CilObject, field_rid: u32) -> Result>>()? + .into_iter() .find(|r| r.rid == field_rid) .ok_or_else(|| Error::Other(format!("Field {field_rid} not found")))?; @@ -875,6 +874,7 @@ fn calculate_field_data_size(assembly: &CilObject, field_rid: u32) -> Result() { for layout_row in class_layout_table { + let layout_row = layout_row?; if layout_row.parent == row { return Ok(layout_row.class_size as usize); } diff --git a/dotscope/src/deobfuscation/techniques/jiejienet/resources.rs b/dotscope/src/deobfuscation/techniques/jiejienet/resources.rs index d895b7b6..07573233 100644 --- a/dotscope/src/deobfuscation/techniques/jiejienet/resources.rs +++ b/dotscope/src/deobfuscation/techniques/jiejienet/resources.rs @@ -786,6 +786,8 @@ fn extract_and_decrypt_resource( let rva_entry = fieldrva_table .iter() + .collect::>>()? + .into_iter() .find(|row| row.field == field_rid) .ok_or_else(|| { Error::Deobfuscation(format!("No FieldRVA entry for field RID 0x{:X}", field_rid)) diff --git a/dotscope/src/deobfuscation/techniques/mod.rs b/dotscope/src/deobfuscation/techniques/mod.rs index eee81112..290c21b4 100644 --- a/dotscope/src/deobfuscation/techniques/mod.rs +++ b/dotscope/src/deobfuscation/techniques/mod.rs @@ -76,6 +76,7 @@ use crate::{ cilassembly::CleanupRequest, compiler::{CompilerContext, EventLog, PassPhase, SsaPass}, deobfuscation::{config::EngineConfig, context::AnalysisContext}, + metadata::token::Token, CilObject, Result, }; @@ -274,6 +275,22 @@ pub trait Technique: Send + Sync { None } + /// Methods this technique was responsible for restoring, reported when its + /// transform did not succeed. + /// + /// A technique that rewrites method bodies — decrypting them, unpacking + /// them — leaves those methods in their protected form when it fails. Such a + /// method carries no calls, so nothing it references looks reachable and + /// cleanup reads the whole region as dead. It is not dead: it is the + /// original code, and the sample is kept precisely to analyse it later. + /// + /// Only consulted when the technique is detected and its transform did not + /// complete, so an implementation may return its full candidate set without + /// checking whether the run succeeded. + fn unrecovered_methods(&self, _detection: &Detection) -> Vec { + Vec::new() + } + /// Declares the technique's capability patterns. /// /// The engine uses this to understand which lifecycle methods are diff --git a/dotscope/src/deobfuscation/techniques/netreactor/necrobit.rs b/dotscope/src/deobfuscation/techniques/netreactor/necrobit.rs index a3d60d8f..4442cc26 100644 --- a/dotscope/src/deobfuscation/techniques/netreactor/necrobit.rs +++ b/dotscope/src/deobfuscation/techniques/netreactor/necrobit.rs @@ -18,10 +18,19 @@ //! The decryption pipeline: //! 1. Find `::.cctor` and trial check methods via structural analysis //! 2. Register a hook to bypass trial checks (avoids needing DateTime BCL hooks) -//! 3. Emulate `::.cctor` — the protection's own code decrypts method -//! bodies and writes them back to the virtual image via `Marshal.Copy` -//! 4. Extract all method bodies from the decrypted virtual image +//! 3. Emulate `::.cctor` — the protection's own code does the decryption, +//! so the cipher is never reimplemented and version changes to it cost nothing +//! 4. Extract the decrypted bodies from wherever that run left them, which +//! depends on the storage variant (see the format notes below): variant A +//! builds a structured blob on the managed heap, variant B writes complete +//! bodies into the image at each method's RVA //! 5. Rebuild the assembly with restored method bodies +//! +//! Step 3 is only faithful if the emulator models the protection's dynamically +//! resolved `VirtualProtect`. Variant B's writes land in `.text`, which is mapped +//! read-only, and the return value also steers the init's control flow — a call +//! that answers without applying the protection sends it down a path that stores +//! nothing. `docs/research/netreactor/necrobit.md` covers the requirements. use std::{ any::Any, @@ -466,7 +475,7 @@ impl Technique for NetReactorNecroBit { .view() .tables() .and_then(|t| t.table::()) - .and_then(|table| table.get(rid)) + .and_then(|table| table.get(rid).ok().flatten()) .ok_or_else(|| Error::Deobfuscation(format!("MethodDef row {rid} not found"))) { Ok(v) => v, @@ -514,6 +523,17 @@ impl Technique for NetReactorNecroBit { Some(Ok(events)) } + /// The stub methods, which are exactly the ones whose bodies this technique + /// exists to decrypt. When decryption fails they still hold the + /// `nop;nop;X;ret` stub the protection wrote, and the real code is encrypted + /// in the resource rather than gone. + fn unrecovered_methods(&self, detection: &Detection) -> Vec { + detection + .findings::() + .map(|findings| findings.stub_method_tokens.clone()) + .unwrap_or_default() + } + fn requires_regeneration(&self) -> bool { true } @@ -1056,40 +1076,72 @@ mod tests { ); } - #[test] - #[ignore] - fn test_byte_transform() { - let path = "tests/samples/packers/netreactor/7.5.0/reactor_necrobit.exe"; - if !Path::new(path).exists() { - eprintln!("Skipping test: sample not found at {path}"); - return; - } - - let _ = env_logger::builder() - .filter_level(log::LevelFilter::Info) - .is_test(true) - .try_init(); - - let assembly = load_sample("reactor_necrobit.exe"); + /// Decrypts `sample` and returns how many of its stubs were restored. + /// + /// Counting matters: the transform reports success as long as it restored + /// *something*, so asserting only that it returned `Ok` cannot tell a full + /// recovery from a partial one — and for a whole storage variant it could + /// not tell recovery from nothing at all. + fn restore_stub_bodies(sample: &str) -> (usize, usize) { + let assembly = load_sample(sample); let technique = NetReactorNecroBit; let detection = technique.detect(&assembly); - assert!(detection.is_detected(), "Should detect NecroBit"); + assert!(detection.is_detected(), "{sample}: should detect NecroBit"); + + let stub_count = detection + .findings::() + .map(|f| f.stub_method_tokens.len()) + .unwrap_or(0); let mut working = WorkingAssembly::new(assembly); let detections = Detections::new(); - let result = technique.byte_transform(&mut working, &detection, &detections); + let events = match technique.byte_transform(&mut working, &detection, &detections) { + Some(Ok(events)) => events, + Some(Err(e)) => panic!("{sample}: byte_transform returned error: {e}"), + None => panic!("{sample}: byte_transform returned None (skipped)"), + }; - match result { - Some(Ok(events)) => { - eprintln!("byte_transform succeeded with {} events", events.len()); - } - Some(Err(e)) => { - panic!("byte_transform returned error: {e}"); - } - None => { - panic!("byte_transform returned None (skipped)"); - } + let restored = events + .iter() + .filter(|event| event.kind == EventKind::MethodBodyDecrypted && event.method.is_some()) + .count(); + (restored, stub_count) + } + + /// Variant A: bodies live in a structured blob on the managed heap. + #[test] + fn restores_every_stub_in_a_necrobit_only_binary() { + if !Path::new("tests/samples/packers/netreactor/7.5.0/reactor_necrobit.exe").exists() { + eprintln!("Skipping test: sample not found"); + return; } + + let (restored, stubs) = restore_stub_bodies("reactor_necrobit.exe"); + assert!(stubs > 0, "the sample must actually have stubs"); + assert_eq!( + restored, stubs, + "every encrypted body should come back, got {restored} of {stubs}" + ); + } + + /// Variant B: the init writes bodies into the image at each method's RVA, + /// which only works once the dynamically-resolved `VirtualProtect` has made + /// those pages writable. Recovering nothing here is the exact regression + /// this asserts against — it went unnoticed while the only check was that + /// the transform returned `Ok`. + #[test] + fn restores_every_stub_in_a_full_protection_binary() { + if !Path::new("tests/samples/packers/netreactor/7.5.0/reactor_full.exe").exists() { + eprintln!("Skipping test: sample not found"); + return; + } + + let (restored, stubs) = restore_stub_bodies("reactor_full.exe"); + assert!(stubs > 0, "the sample must actually have stubs"); + assert_eq!( + restored, stubs, + "every encrypted body should come back, got {restored} of {stubs}" + ); } #[test] diff --git a/dotscope/src/deobfuscation/techniques/netreactor/resources.rs b/dotscope/src/deobfuscation/techniques/netreactor/resources.rs index 4d3cb0f8..fc7915ee 100644 --- a/dotscope/src/deobfuscation/techniques/netreactor/resources.rs +++ b/dotscope/src/deobfuscation/techniques/netreactor/resources.rs @@ -998,6 +998,13 @@ pub fn find_assembly_typeref(assembly: &CilObject) -> Option { let table = tables.table::()?; let strings = assembly.strings()?; for row in table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let Ok(name) = strings.get(row.type_name as usize) else { continue; }; diff --git a/dotscope/src/deobfuscation/template.rs b/dotscope/src/deobfuscation/template.rs index 3732cd3e..1be47470 100644 --- a/dotscope/src/deobfuscation/template.rs +++ b/dotscope/src/deobfuscation/template.rs @@ -193,7 +193,8 @@ impl EmulationTemplatePool { let tables = assembly.tables()?; let strings = assembly.strings()?; let module_table = tables.table::()?; - let module_row = module_table.iter().next()?; + // Module is RID 1 by definition (ECMA-335 II.22.30). + let module_row = module_table.get(1).ok().flatten()?; strings.get(module_row.name as usize).ok().map(String::from) }); if let Some(ref name) = module_name { @@ -263,7 +264,10 @@ impl EmulationTemplatePool { .map_err(|e| Error::LockError(format!("template pool read lock: {e}")))?; match *guard { - Some(ref template) => template.fork(), + Some(ref template) => { + let limits = self.config.emulation.execution_limits(template.limits()); + template.fork_with_limits(limits) + } None => Err(Error::Emulation(Box::new(EmulationError::InternalError { description: "template pool not warmed up".to_string(), }))), @@ -290,6 +294,8 @@ impl EmulationTemplatePool { pub fn fork_for_targeted_warmup(&self, cctors: &[Token]) -> Option { let guard = self.template.read().ok()?; let template = guard.as_ref()?; + // Targeted warmup runs extra `.cctor`s on the fork, so it gets the warmup budget + // rather than the per-method one — that is what the longer budget is *for*. let mut process = match template.fork() { Ok(p) => p, Err(e) => { diff --git a/dotscope/src/deobfuscation/utils.rs b/dotscope/src/deobfuscation/utils.rs index eddcf237..c83c26aa 100644 --- a/dotscope/src/deobfuscation/utils.rs +++ b/dotscope/src/deobfuscation/utils.rs @@ -21,8 +21,8 @@ use crate::{ signatures::{parse_field_signature, TypeSignature}, streams::Strings, tables::{ - ClassLayoutRaw, FieldRaw, MemberRefRaw, MetadataTable, MethodDefRaw, TableId, - TypeDefRaw, TypeRefRaw, + skip_unreadable, ClassLayoutRaw, FieldRaw, MemberRefRaw, MetadataTable, MethodDefRaw, + TableId, TypeDefRaw, TypeRefRaw, }, token::Token, typesystem::{wellknown, PointerSize}, @@ -41,7 +41,7 @@ pub(crate) fn get_field_data_size(assembly: &CilObject, field_rid: u32) -> Optio let blobs = assembly.blob()?; let field_table = tables.table::()?; - let field_row = field_table.get(field_rid)?; + let field_row = field_table.get(field_rid).ok().flatten()?; let sig_data = blobs.get(field_row.signature as usize).ok()?; let field_sig = parse_field_signature(sig_data).ok()?; @@ -61,6 +61,13 @@ pub(crate) fn get_field_data_size(assembly: &CilObject, field_rid: u32) -> Optio let class_layout_table = tables.table::()?; for layout in class_layout_table { + let layout = match layout { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if layout.parent == type_rid { return Some(layout.class_size as usize); } @@ -407,10 +414,13 @@ pub(crate) fn resolve_methoddef_declaring_type<'a>( ) -> Option> { let methoddef_table = methoddef_table?; let typedef_table = typedef_table?; - let method = methoddef_table.get(method_row)?; + let method = methoddef_table.get(method_row).ok().flatten()?; + // No error channel here (the function answers `Option`), so an unreadable row is + // reported and skipped rather than vanishing inside the iterator. let typedef = typedef_table .iter() + .filter_map(skip_unreadable) .filter(|t| t.method_list <= method.rid) .last()?; @@ -439,12 +449,12 @@ pub(crate) fn resolve_memberref_declaring_type<'a>( strings: &'a Strings<'a>, ) -> Option> { let memberref_table = memberref_table?; - let memberref = memberref_table.get(memberref_row)?; + let memberref = memberref_table.get(memberref_row).ok().flatten()?; match memberref.class.tag { TableId::TypeDef => { let typedef_table = typedef_table?; - let typedef = typedef_table.get(memberref.class.row)?; + let typedef = typedef_table.get(memberref.class.row).ok().flatten()?; let name = strings.get(typedef.type_name as usize).ok()?; let namespace = strings.get(typedef.type_namespace as usize).ok(); Some(ResolvedType { @@ -456,7 +466,7 @@ pub(crate) fn resolve_memberref_declaring_type<'a>( } TableId::TypeRef => { let typeref_table = typeref_table?; - let typeref = typeref_table.get(memberref.class.row)?; + let typeref = typeref_table.get(memberref.class.row).ok().flatten()?; let name = strings.get(typeref.type_name as usize).ok()?; let namespace = strings.get(typeref.type_namespace as usize).ok(); Some(ResolvedType { @@ -947,6 +957,13 @@ mod tests { let mut found_marker = false; for attr in ca_table { + let attr = match attr { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; if let Some(resolved) = resolve_constructor_type( attr.constructor.tag, attr.constructor.row, diff --git a/dotscope/src/emulation/capture/context.rs b/dotscope/src/emulation/capture/context.rs index 7408db70..ade3d8e9 100644 --- a/dotscope/src/emulation/capture/context.rs +++ b/dotscope/src/emulation/capture/context.rs @@ -51,7 +51,12 @@ //! // Strings, file ops, and network ops will not be captured //! ``` -use std::{ops::Range, sync::RwLock}; +use std::{ + collections::{hash_map::DefaultHasher, HashSet as StdHashSet}, + hash::{Hash, Hasher}, + ops::Range, + sync::RwLock, +}; use crate::{ emulation::{ @@ -127,6 +132,19 @@ pub struct CaptureContext { /// operations that may contain decrypted payloads or extracted data. buffers: RwLock>, + /// Total bytes retained across all capture collections. + /// + /// Charged by [`admit`](Self::admit) and compared against + /// [`CaptureConfig::max_total_bytes`]. Captured data lives outside the managed heap, so + /// this is the only thing bounding it. + retained_bytes: RwLock, + + /// Content hashes of captured assemblies, for O(1) duplicate rejection. + /// + /// The previous duplicate check compared the candidate byte-for-byte against every + /// assembly already captured, making capture quadratic in capture count. + assembly_hashes: RwLock>, + /// Return values from monitored methods. /// /// Captures return values from specific methods configured for monitoring, @@ -235,6 +253,9 @@ impl CaptureContext { strings: true, file_operations: true, network_operations: true, + // Explicitly on: this constructor means "capture what is useful". + // `CaptureConfig::default()` captures nothing, which is what that type documents. + buffers: true, ..Default::default() }) } @@ -269,6 +290,8 @@ impl CaptureContext { assemblies: RwLock::new(Vec::new()), strings: RwLock::new(Vec::new()), buffers: RwLock::new(Vec::new()), + retained_bytes: RwLock::new(0), + assembly_hashes: RwLock::new(StdHashSet::new()), method_returns: RwLock::new(Vec::new()), file_operations: RwLock::new(Vec::new()), network_operations: RwLock::new(Vec::new()), @@ -350,13 +373,24 @@ impl CaptureContext { return; } - if let Ok(assemblies) = self.assemblies.read() { - if assemblies.iter().any(|a| a.data == data) { + // Duplicate rejection by content hash. The previous check compared the candidate + // byte-for-byte against every assembly already captured, so capture cost grew + // quadratically in capture count over multi-megabyte payloads. + let mut hasher = DefaultHasher::new(); + data.hash(&mut hasher); + let digest = hasher.finish(); + + if let Ok(mut hashes) = self.assembly_hashes.write() { + if !hashes.insert(digest) { return; // Already captured this exact assembly } } let len = data.len(); + if !self.admit(len) { + return; + } + let assembly = CapturedAssembly { data, source, @@ -365,6 +399,9 @@ impl CaptureContext { }; if let Ok(mut assemblies) = self.assemblies.write() { + if assemblies.len() >= self.config.max_items { + return; + } assemblies.push(assembly); } @@ -438,6 +475,10 @@ impl CaptureContext { return; } + if !self.admit(value.len()) { + return; + } + let string = CapturedString { value, source, @@ -446,6 +487,9 @@ impl CaptureContext { }; if let Ok(mut strings) = self.strings.write() { + if strings.len() >= self.config.max_items { + return; + } strings.push(string); } @@ -518,8 +562,11 @@ impl CaptureContext { /// Captures a raw byte buffer. /// /// Records a byte buffer from memory operations, crypto transforms, or other - /// sources during emulation. Unlike assembly capture, buffer capture is always - /// enabled (no configuration check). + /// sources during emulation. + /// + /// Gated on [`CaptureConfig::buffers`] and bounded by + /// [`CaptureConfig::max_items`]/[`CaptureConfig::max_total_bytes`]. Captures past either + /// ceiling are dropped silently — emulation continues, it simply stops retaining more. /// /// # Arguments /// @@ -534,7 +581,15 @@ impl CaptureContext { buffer_source: BufferSource, label: impl Into, ) { + if !self.config.buffers { + return; + } + let len = data.len(); + if !self.admit(len) { + return; + } + let buffer = CapturedBuffer { data, source, @@ -543,6 +598,9 @@ impl CaptureContext { }; if let Ok(mut buffers) = self.buffers.write() { + if buffers.len() >= self.config.max_items { + return; + } buffers.push(buffer); } @@ -551,6 +609,26 @@ impl CaptureContext { } } + /// Returns whether `len` more captured bytes fit within + /// [`CaptureConfig::max_total_bytes`], charging them if so. + /// + /// Captured data lives outside the managed heap, so the heap budget never sees it; this is + /// the only ceiling on it. The counter is monotonic — captures are never released during a + /// run — so charging on admission is sufficient. + fn admit(&self, len: usize) -> bool { + let Ok(mut retained) = self.retained_bytes.write() else { + return false; + }; + + let next = retained.saturating_add(len); + if next > self.config.max_total_bytes { + return false; + } + + *retained = next; + true + } + /// Returns all captured buffers. /// /// # Returns @@ -950,6 +1028,73 @@ impl Default for CaptureContext { mod tests { use super::*; + /// A default config must capture nothing, which is what `CaptureConfig` documents. + /// + /// Buffer capture is the kind most easily left running unconditionally, which would + /// charge an operator who disabled capture for it anyway. + #[test] + fn default_config_captures_no_buffers() { + let ctx = CaptureContext::with_config(CaptureConfig::default()); + + ctx.capture_buffer( + vec![0u8; 64], + CaptureSource::new(Token::new(0x0600_0001), ThreadId::new(1), 0, 0), + BufferSource::MarshalCopy { address: 0x1000 }, + "test", + ); + + assert_eq!(ctx.buffer_count(), 0, "capture was not enabled"); + } + + /// The item ceiling stops the loop-and-append pattern: emulated code allocates one array, + /// then calls a crypto hook on it repeatedly, appending a full copy each time for a handful + /// of instructions. + #[test] + fn buffer_capture_stops_at_the_item_ceiling() { + let ctx = CaptureContext::with_config(CaptureConfig { + buffers: true, + max_items: 10, + ..Default::default() + }); + + for _ in 0..1000 { + ctx.capture_buffer( + vec![0u8; 8], + CaptureSource::new(Token::new(0x0600_0001), ThreadId::new(1), 0, 0), + BufferSource::MarshalCopy { address: 0x1000 }, + "test", + ); + } + + assert_eq!(ctx.buffer_count(), 10); + } + + /// A few very large buffers must be bounded as well as many small ones. + #[test] + fn buffer_capture_stops_at_the_byte_ceiling() { + let ctx = CaptureContext::with_config(CaptureConfig { + buffers: true, + max_items: usize::MAX, + max_total_bytes: 4096, + ..Default::default() + }); + + for _ in 0..100 { + ctx.capture_buffer( + vec![0u8; 1024], + CaptureSource::new(Token::new(0x0600_0001), ThreadId::new(1), 0, 0), + BufferSource::MarshalCopy { address: 0x1000 }, + "test", + ); + } + + assert_eq!( + ctx.buffer_count(), + 4, + "only what fits in the byte budget is retained" + ); + } + #[test] fn test_capture_context_creation() { let ctx = CaptureContext::new(); diff --git a/dotscope/src/emulation/engine/callresolver.rs b/dotscope/src/emulation/engine/callresolver.rs index 42998abb..48b359a2 100644 --- a/dotscope/src/emulation/engine/callresolver.rs +++ b/dotscope/src/emulation/engine/callresolver.rs @@ -297,8 +297,17 @@ impl CallResolver { is_virtual, pre_push_value: None, is_reflection_invoke: false, - #[allow(clippy::cast_possible_truncation)] - assembly_index: Some(i as u8), + // The index must survive intact: it is what binds the new frame to the + // metadata the token was resolved against. A narrowing conversion here + // wraps a large index onto a different assembly's `EmulationContext`, + // which executes a valid token against the wrong method table. + // Unreachable in practice: `max_loaded_assemblies` bounds the list far + // below `u32::MAX`, so this is an invariant check, not a runtime path. + assembly_index: Some(u32::try_from(i).map_err(|_| { + EmulationError::InternalError { + description: format!("assembly index {i} exceeds u32"), + } + })?), method_type_args: None, })); } @@ -845,8 +854,16 @@ impl CallResolver { // Look up the function name and return an // appropriate value for each native API. if tokens::is_native_function_pointer(target_token) { - let return_value = - resolve_native_delegate_return(target_token, thread, context); + // Element 0 is the delegate itself; the rest are + // the arguments the native function was called + // with, and a hook needs them to do its work. + let native_args: Vec = + arg_values.iter().skip(1).cloned().collect(); + let return_value = self.dispatch_native_delegate( + target_token, + &native_args, + thread, + )?; return Ok(CallResolution::HookedBypass { return_value }); } @@ -1321,6 +1338,72 @@ impl CallResolver { self.execute_hook_with_resolved(context, method_token, thread, &info) } + /// Dispatches a native function reached through a function pointer. + /// + /// A protection that wants its native calls to go unnoticed does not declare + /// them: it resolves the address with `GetProcAddress` and calls it through a + /// delegate. Nothing about the call site says "P/Invoke", so + /// [`Self::try_native_call`] — which starts from a `MethodDef` with an + /// ImplMap — never sees it, and the hooks registered for that function never + /// run. .NET Reactor resolves `VirtualProtect` exactly this way. + /// + /// The identity survives the round trip: `GetProcAddress` records the + /// function against a fake address, `GetDelegateForFunctionPointer` turns + /// that address back into a token, and the token names the function here. So + /// the call is given the same shape a declared P/Invoke would have and run + /// through the ordinary hook path, which is what makes the *effect* happen + /// rather than merely the return value. + /// + /// Falls back to a plausible return value when no hook matches, so a + /// protection calling something unimplemented still makes progress. + fn dispatch_native_delegate( + &self, + target_token: Token, + args: &[EmValue], + thread: &mut EmulationThread, + ) -> Result> { + let function = thread + .runtime_state() + .read() + .ok() + .and_then(|rt| rt.native_functions().lookup_by_token(target_token)); + + let Some(function) = function else { + log::debug!( + "Native delegate invoke: token 0x{:08X} names no known function", + target_token.value() + ); + return Ok(Some(EmValue::I32(0))); + }; + + // A function whose module was never observed can still match a hook that + // does not constrain the DLL, so dispatch is attempted either way. + let dll = function.dll.as_deref().unwrap_or_default(); + let hook_context = + HookContext::native(target_token, dll, &function.name, self.config.pointer_size) + .with_args(args); + + match self.hooks.execute(&hook_context, thread, |_| None)? { + HookOutcome::NoMatch => Ok(native_delegate_fallback( + &function.name, + thread, + target_token, + )), + HookOutcome::Handled(result) + | HookOutcome::ReflectionInvoke { + bypass_value: result, + .. + } => { + log::debug!("Native delegate invoke: {dll}!{} handled", function.name); + Ok(result) + } + HookOutcome::ThrewException { message, .. } => Err(EmulationError::HookError(format!( + "native delegate hook threw CLR exception: {message}" + )) + .into()), + } + } + /// Tries to execute a method call via a native (P/Invoke) stub. /// /// For `MethodDef` tokens that have no IL body (P/Invoke methods), looks up @@ -1732,7 +1815,14 @@ pub fn maybe_run_type_cctor_for_method( return Ok(false); } - // Find the type that declares this method + // Find the type that declares this method. + // + // This runs for every call instruction, including after the type is initialised and this + // function has nothing left to do. The initialisation check cannot be hoisted above it — + // it is keyed by *type* token, which is precisely what this lookup produces. The lookup + // itself is therefore what has to be cheap: `declaring_type` resolves through the method's + // own back-pointer, falling back to a memoised registry index. A scan over every type and + // method would make dispatch O(total methods in assembly). let Some(type_info) = context.assembly().resolver().declaring_type(method) else { return Ok(false); }; @@ -1878,26 +1968,18 @@ fn zero_initialize_static_fields( Ok(()) } -/// Resolves the return value for a native function pointer delegate invocation. +/// Returns a plausible value for a native API that no hook implements. /// -/// Looks up the function name from the native function registry and returns a -/// type-appropriate value. For known Win32 APIs, returns the correct type -/// (e.g., BOOL for VirtualProtect, pointer for VirtualAlloc). For unknown -/// functions, returns I32(0) as a safe default. -fn resolve_native_delegate_return( - target_token: Token, +/// A last resort. It reports success without performing the call's effect, so +/// anything whose effect matters — a memory protection change, a write through a +/// pointer — belongs in a hook, not here. +fn native_delegate_fallback( + name: &str, thread: &EmulationThread, - _context: &EmulationContext, + target_token: Token, ) -> Option { - let func_name = thread - .runtime_state() - .read() - .ok() - .and_then(|rt| rt.native_functions().lookup_by_token(target_token)); - - let name = func_name.as_deref().unwrap_or("unknown"); log::debug!( - "Native delegate invoke: {name} (token 0x{:08X})", + "Native delegate invoke: {name} (token 0x{:08X}) has no hook, returning a default", target_token.value() ); diff --git a/dotscope/src/emulation/engine/context/metadata.rs b/dotscope/src/emulation/engine/context/metadata.rs index e28bb83a..48557823 100644 --- a/dotscope/src/emulation/engine/context/metadata.rs +++ b/dotscope/src/emulation/engine/context/metadata.rs @@ -97,7 +97,7 @@ impl EmulationContext { let blob = self.assembly.blob()?; let table = <_ as TableAccess>::table(tables)?; let row = token.row(); - let raw_sig = table.get(row)?; + let raw_sig = table.get(row).ok().flatten()?; // Parse the raw signature using the blob heap let owned_sig = raw_sig.to_owned(blob).ok()?; Some(owned_sig.parsed_signature.clone()) @@ -126,7 +126,7 @@ impl EmulationContext { let blob = self.assembly.blob()?; let table = <_ as TableAccess>::table(tables)?; let row = token.row(); - let raw_typespec = table.get(row)?; + let raw_typespec = table.get(row).ok().flatten()?; // Parse the raw signature using the blob heap let owned_typespec = raw_typespec.to_owned(blob).ok()?; Some(owned_typespec.signature.clone()) diff --git a/dotscope/src/emulation/engine/context/types.rs b/dotscope/src/emulation/engine/context/types.rs index 6344629d..6f5ea78b 100644 --- a/dotscope/src/emulation/engine/context/types.rs +++ b/dotscope/src/emulation/engine/context/types.rs @@ -134,7 +134,7 @@ impl EmulationContext { if is_bcl_wrapper_token(source_token) && target_token.is_table(TableId::TypeDef) { if let Some(target_type) = self.get_type(target_token) { let name = target_type.fullname(); - if name == "System.Object" || name == "System.ValueType" { + if &*name == "System.Object" || &*name == "System.ValueType" { return true; } if target_type.is_interface() { @@ -373,7 +373,7 @@ impl EmulationContext { #[must_use] pub fn format_type_token(&self, type_token: Token) -> String { if let Some(type_info) = self.get_type(type_token) { - type_info.fullname() + type_info.fullname().to_string() } else { format!("0x{:08X}", type_token.value()) } diff --git a/dotscope/src/emulation/engine/controller.rs b/dotscope/src/emulation/engine/controller.rs index c5f420f0..7e2a896d 100644 --- a/dotscope/src/emulation/engine/controller.rs +++ b/dotscope/src/emulation/engine/controller.rs @@ -138,7 +138,7 @@ pub struct EmulationController { /// /// Keeps each assembly's decoded-method cache alive across the execution /// loop instead of rebuilding it on every instruction. - assembly_contexts: DashMap>, + assembly_contexts: DashMap>, } impl EmulationController { @@ -239,7 +239,7 @@ impl EmulationController { let typespec_row = assembly .tables() .and_then(|t| t.table::()) - .and_then(|table| table.get(token.row()))?; + .and_then(|table| table.get(token.row()).ok().flatten())?; let blob = assembly.blob()?; let parsed = typespec_row.to_owned(blob).ok()?; @@ -312,7 +312,7 @@ impl EmulationController { /// Returns `None` if the assembly index doesn't exist in the `RuntimeState`. /// This is called per-iteration when executing a frame from a loaded assembly, /// but `EmulationContext::new` is trivial (wraps an `Arc`), so the cost is negligible. - fn loaded_assembly_context(&self, index: u8) -> Result>> { + fn loaded_assembly_context(&self, index: u32) -> Result>> { // The execution loop resolves the frame's context on every instruction. // Rebuilding the context each time meant taking the runtime lock and // discarding the assembly's decoded-method cache per step, so memoize it. @@ -327,7 +327,10 @@ impl EmulationController { .map_err(|_| EmulationError::LockPoisoned { description: "runtime state", })?; - let Some(asm) = state.app_domain().get_parsed_assembly(index as usize) else { + let Some(asm) = usize::try_from(index) + .ok() + .and_then(|index| state.app_domain().get_parsed_assembly(index)) + else { return Ok(None); }; @@ -875,51 +878,83 @@ impl EmulationController { } StepResult::EndFilter { value } => { - let should_handle = match value { - EmValue::I32(v) => v != 0, - _ => false, + // ECMA-335 §12.4.2.5: `endfilter` takes an int32. 0 means "continue the + // search", 1 means "this handler runs". Anything else is not a filter + // result — treating a non-int32 as rejection silently mis-routes the + // exception, so it is a program error. + let EmValue::I32(filter_result) = value else { + return Err(EmulationError::InternalError { + description: format!( + "endfilter expects an int32 result, found {}", + value.type_name() + ), + } + .into()); }; + let should_handle = filter_result != 0; thread .exception_state_mut() .set_filter_result(Some(should_handle)); if should_handle { - if let Some(handler_offset) = - thread.exception_state_mut().filter_handler_offset() + // Read the filter state *before* clearing it. `set_in_filter(false)` + // also nulls `filter_handler_offset` (see + // `ThreadExceptionState::set_in_filter`), so reading the offset + // afterwards always yields `None` — and the resulting fall-through + // returns `Continue` without advancing the IP, re-executing this same + // `endfilter` forever. The rejecting path below has the same + // constraint: it needs the offset as its skip key. + let handler_offset = thread.exception_state_mut().filter_handler_offset(); + let origin_offset = self.exception_origin(thread, interpreter); + thread.exception_state_mut().set_in_filter(false); + + let Some(handler_offset) = handler_offset else { + // `enter_filter` records the offset on every path that dispatches a + // filter, so its absence here means the filter body was entered + // without going through exception dispatch. Continuing would spin. + return Err(EmulationError::InternalError { + description: + "endfilter accepted but no filter handler offset was recorded" + .to_string(), + } + .into()); + }; + { - let origin_offset = thread - .exception_state_mut() - .exception_origin_offset() - .unwrap_or(interpreter.ip().offset()); + // `endfilter` consumed the exception object that + // `apply_handler_match` pushed before entering the filter, so the + // handler would otherwise start on whatever the filter left behind + // — typically an underflow on its first `stloc`. Re-push it, as the + // catch path does. + if let Some(exception) = + thread.exception_state_mut().get_exception_value() + { + thread.stack_mut().clear(); + thread.push(exception)?; + } - thread.exception_state_mut().set_in_filter(false); thread.exception_state_mut().enter_catch_handler( current_method, origin_offset, handler_offset, ); interpreter.set_offset(handler_offset); - } else { - thread.exception_state_mut().set_in_filter(false); + LoopAction::Continue } } else { - thread.exception_state_mut().set_in_filter(false); - - if thread.exception_state_mut().has_exception() { - return Ok(EmulationOutcome::UnhandledException { - exception: thread - .exception_state_mut() - .take_exception_as_value() - .unwrap_or_else(|| { - trace!("No pending exception for extraction"); - EmValue::Null - }), - instructions: interpreter.stats().instructions_executed, - }); - } + // Rejection means "keep searching", not "unhandled". Resume the scan + // after the rejected Filter clause and fall into the normal unwind path + // if nothing else matches; returning `UnhandledException` here reports a + // genuinely-caught exception as unhandled and stops the run, which an + // obfuscator can use to hide everything past a rejecting filter. + self.resume_search_after_filter( + interpreter, + thread, + context, + current_method, + )? } - LoopAction::Continue } StepResult::Rethrow => { @@ -1261,6 +1296,29 @@ impl EmulationController { } } + // Run any finally scheduled for the frame we are about to leave, *before* leaving + // it. + // + // `find_exception_handler` queues finallys against the frame it searched, and the + // handler's IL reads locals, arguments and the evaluation stack from + // `thread.current_frame()`. Popping first and then jumping to the handler — without + // pushing a frame — executes that IL against the grandparent's frame instead. It + // also skips the popped frame's own catch clauses, because the handler search below + // only examines the frame that survives the pop. + // + // Every other unwind path in the engine (`unwind_after_error`, `handle_throw`, + // `exhandler::route_clr_exception`) already drains the queue before popping; this + // one was the outlier. + if let Some(pending) = thread.exception_state_mut().pop_finally() { + thread.exception_state_mut().set_in_unwind_finally(true); + interpreter.set_method(pending.method); + interpreter.set_offset(pending.handler_offset); + thread + .exception_state_mut() + .set_leave_target(pending.leave_target); + return Ok(LoopAction::Continue); + } + // Capture the return_offset from the frame being popped if let Some(frame) = thread.current_frame() { call_site_in_caller = frame.return_offset(); @@ -1283,17 +1341,6 @@ impl EmulationController { ))); } - // Check for finally blocks scheduled by find_exception_handler - if let Some(pending) = thread.exception_state_mut().pop_finally() { - thread.exception_state_mut().set_in_unwind_finally(true); - interpreter.set_method(pending.method); - interpreter.set_offset(pending.handler_offset); - thread - .exception_state_mut() - .set_leave_target(pending.leave_target); - return Ok(LoopAction::Continue); - } - // Search for handlers in the caller let caller_frame = thread .current_frame() @@ -1582,6 +1629,91 @@ impl EmulationController { /// Handles a `StepResult::Rethrow` — re-raises the current exception from /// within a catch handler, searching for another handler or unwinding. + /// Resumes the handler search after a filter returned zero. + /// + /// ECMA-335 §12.4.2.5: a filter returning 0 means "continue the search" — the method's + /// remaining clauses are examined, then the caller frames. It does **not** mean the + /// exception is unhandled. + /// + /// Mirrors [`handle_rethrow`](Self::handle_rethrow), differing only in what it skips: the + /// scan resumes past the rejected filter's own handler rather than past an enclosing catch. + /// + /// # Errors + /// + /// Propagates handler-search failures, and returns an error if the call stack is empty + /// while a frame is expected. + /// The IL offset a handler search should resume from. + /// + /// `exception_origin_offset` is only ever written by + /// [`ThreadExceptionState::enter_catch_handler`], so it is `None` for an exception that has + /// not yet entered a catch — which is precisely the case on the filter paths. Reading it + /// alone made the post-filter rescan skip the current method's remaining clauses entirely + /// and jump straight to unwinding the caller. + /// + /// The throw site is already recorded on the in-flight exception itself + /// ([`ExceptionInfo::throw_location`], set by `handle_throw`), so it serves as the fallback. + /// The interpreter's current offset is the last resort, for a filter evaluated with no + /// exception state at all. + fn exception_origin(&self, thread: &EmulationThread, interpreter: &Interpreter) -> u32 { + thread + .exception_state() + .exception_origin_offset() + .or_else(|| { + thread + .exception_state() + .exception() + .map(|info| info.throw_location.offset) + }) + .unwrap_or_else(|| interpreter.ip().offset()) + } + + fn resume_search_after_filter( + &self, + interpreter: &mut Interpreter, + thread: &mut EmulationThread, + context: &EmulationContext, + current_method: Token, + ) -> Result { + let Some(exception) = thread.exception_state_mut().get_exception_value() else { + // No exception in flight: a filter evaluated outside exception dispatch. Nothing to + // route, so just carry on. + return Ok(LoopAction::Continue); + }; + + let exception_type = exhandler::resolve_exception_type(&exception, thread); + // Both reads must precede `set_in_filter(false)`, which nulls the handler offset. + // Without the skip key the rescan below re-matches the filter clause that just + // rejected, re-enters its body, and loops. + let skip_handler = thread.exception_state_mut().filter_handler_offset(); + let origin = self.exception_origin(thread, interpreter); + thread.exception_state_mut().set_in_filter(false); + + // Remaining clauses of the current method, past the rejected filter. + if let Some(handler_match) = exhandler::find_exception_handler( + context, + current_method, + origin, + exception_type, + thread.exception_state_mut(), + skip_handler, + )? { + thread.stack_mut().clear(); + let target_offset = exhandler::apply_handler_match( + &handler_match, + exception, + origin, + current_method, + thread, + )?; + interpreter.set_offset(target_offset); + return Ok(LoopAction::Continue); + } + + // Nothing left in this method — hand over to the normal unwind path, which drains + // pending finallys and searches each caller in turn. + self.unwind_propagating_exception(interpreter, thread, context, false) + } + fn handle_rethrow( &self, interpreter: &mut Interpreter, @@ -1844,7 +1976,7 @@ impl EmulationController { let constraint = initial_constraint; let mut pending_pre_push: Option = None; let mut pending_reflection_invoke = false; - let mut pending_assembly_index: Option = None; + let mut pending_assembly_index: Option = None; let mut pending_method_type_args: Option> = None; for _ in 0..MAX_REDIRECT_DEPTH { diff --git a/dotscope/src/emulation/engine/error.rs b/dotscope/src/emulation/engine/error.rs index 08180ca8..09ada2c5 100644 --- a/dotscope/src/emulation/engine/error.rs +++ b/dotscope/src/emulation/engine/error.rs @@ -42,6 +42,7 @@ use crate::metadata::{token::Token, typesystem::CilFlavor}; /// | 0x7F01_0016 | System.MissingMethodException | /// | 0x7F01_0017 | System.MissingFieldException | /// | 0x7F01_0018 | System.NotImplementedException | +/// | 0x7F01_001B | System.AccessViolationException | pub mod synthetic_exception { use crate::metadata::token::Token; @@ -95,6 +96,25 @@ pub mod synthetic_exception { pub const MISSING_FIELD: Token = Token::new(0x7F01_0017); /// System.NotImplementedException pub const NOT_IMPLEMENTED: Token = Token::new(0x7F01_0018); + /// System.ArgumentOutOfRangeException — subtype of `ARGUMENT_EXCEPTION`. + /// + /// Thrown by BCL hooks that receive a negative or oversized count, index or length from + /// emulated code, so that the emulated program observes the same failure the real runtime + /// would produce and can catch it, instead of the emulator allocating on the bad value. + pub const ARGUMENT_OUT_OF_RANGE: Token = Token::new(0x7F01_0019); + /// System.OutOfMemoryException + /// + /// Thrown when a request would exceed the emulator's heap budget. This is deliberately a + /// catchable CLR exception rather than a hard emulation error: real .NET code that probes + /// allocation limits expects to catch it, and surfacing it as a Rust error instead would + /// abandon the emulation of samples that handle it. + pub const OUT_OF_MEMORY: Token = Token::new(0x7F01_001A); + /// System.AccessViolationException + /// + /// Thrown when emulated code reads or writes memory its protection does not permit, or + /// touches a guard page. Obfuscators probe for an emulator by writing where a real + /// process would fault, so silently permitting the access is itself a detection signal. + pub const ACCESS_VIOLATION: Token = Token::new(0x7F01_001B); } /// Errors that can occur during CIL emulation. @@ -223,6 +243,12 @@ pub enum EmulationError { /// Maximum allowed. limit: u64, }, + /// A configured resource ceiling other than memory, call depth or instruction count was + /// reached — for example the heap object count or a delegate's invocation-list length. + /// + /// Carries a description naming the limit and the value that breached it, because these + /// ceilings are varied enough that a shared structured shape would fit none of them well. + ResourceLimitExceeded(String), /// Execution timeout. Timeout { /// Time elapsed. @@ -354,6 +380,18 @@ pub enum EmulationError { target_type: &'static str, }, + /// Access rejected by the memory protection of the target page. + /// + /// Distinct from [`InvalidAddress`](Self::InvalidAddress): the address *is* mapped, but + /// the access is not permitted — a write to a read-only PE section, a read of a + /// PAGE_NOACCESS page, or any touch of a guard page. + AccessViolation { + /// The address whose page rejected the access. + address: u64, + /// Which protection rejected it. + reason: String, + }, + /// Invalid memory address. /// /// This error occurs when accessing an unmapped or invalid memory address @@ -486,6 +524,9 @@ impl fmt::Display for EmulationError { EmulationError::InstructionLimitExceeded { executed, limit } => { write!(f, "instruction limit exceeded: {executed} (limit: {limit})") } + EmulationError::ResourceLimitExceeded(description) => { + write!(f, "resource limit exceeded: {description}") + } EmulationError::Timeout { elapsed, limit } => { write!( f, @@ -589,6 +630,9 @@ impl fmt::Display for EmulationError { } => { write!(f, "cannot convert {source_type} to {target_type}") } + EmulationError::AccessViolation { address, reason } => { + write!(f, "access violation at 0x{address:08X}: {reason}") + } EmulationError::InvalidAddress { address, reason } => { write!(f, "invalid address 0x{address:08X}: {reason}") } @@ -661,6 +705,7 @@ impl EmulationError { EmulationError::DivisionByZero => synthetic_exception::DIVIDE_BY_ZERO, EmulationError::ArithmeticOverflow => synthetic_exception::OVERFLOW, EmulationError::InvalidCast { .. } => synthetic_exception::INVALID_CAST, + EmulationError::AccessViolation { .. } => synthetic_exception::ACCESS_VIOLATION, _ => synthetic_exception::BASE_EXCEPTION, } } diff --git a/dotscope/src/emulation/engine/exceptions.rs b/dotscope/src/emulation/engine/exceptions.rs index 31214d71..529c1b9b 100644 --- a/dotscope/src/emulation/engine/exceptions.rs +++ b/dotscope/src/emulation/engine/exceptions.rs @@ -76,6 +76,9 @@ pub fn parent(token: Token) -> Option { t if t == synthetic_exception::ARGUMENT_NULL => { Some(synthetic_exception::ARGUMENT_EXCEPTION) } + t if t == synthetic_exception::ARGUMENT_OUT_OF_RANGE => { + Some(synthetic_exception::ARGUMENT_EXCEPTION) + } // Intermediate types → SystemException t if t == synthetic_exception::ARITHMETIC => Some(synthetic_exception::SYSTEM_EXCEPTION), @@ -112,6 +115,10 @@ pub fn parent(token: Token) -> Option { t if t == synthetic_exception::NOT_IMPLEMENTED => { Some(synthetic_exception::SYSTEM_EXCEPTION) } + t if t == synthetic_exception::OUT_OF_MEMORY => Some(synthetic_exception::SYSTEM_EXCEPTION), + t if t == synthetic_exception::ACCESS_VIOLATION => { + Some(synthetic_exception::SYSTEM_EXCEPTION) + } // SystemException → Exception t if t == synthetic_exception::SYSTEM_EXCEPTION => { @@ -195,6 +202,8 @@ pub fn token_from_fullname(fullname: &str) -> Option { "System.FormatException" => Some(synthetic_exception::FORMAT_EXCEPTION), "System.ArgumentException" => Some(synthetic_exception::ARGUMENT_EXCEPTION), "System.ArgumentNullException" => Some(synthetic_exception::ARGUMENT_NULL), + "System.ArgumentOutOfRangeException" => Some(synthetic_exception::ARGUMENT_OUT_OF_RANGE), + "System.OutOfMemoryException" => Some(synthetic_exception::OUT_OF_MEMORY), "System.NotSupportedException" => Some(synthetic_exception::NOT_SUPPORTED), "System.Collections.Generic.KeyNotFoundException" => { Some(synthetic_exception::KEY_NOT_FOUND) @@ -211,6 +220,7 @@ pub fn token_from_fullname(fullname: &str) -> Option { "System.MissingMethodException" => Some(synthetic_exception::MISSING_METHOD), "System.MissingFieldException" => Some(synthetic_exception::MISSING_FIELD), "System.NotImplementedException" => Some(synthetic_exception::NOT_IMPLEMENTED), + "System.AccessViolationException" => Some(synthetic_exception::ACCESS_VIOLATION), _ => None, } } @@ -358,6 +368,7 @@ mod tests { synthetic_exception::MISSING_METHOD, synthetic_exception::MISSING_FIELD, synthetic_exception::NOT_IMPLEMENTED, + synthetic_exception::ACCESS_VIOLATION, ]; for exc in &all_exceptions { diff --git a/dotscope/src/emulation/engine/exhandler.rs b/dotscope/src/emulation/engine/exhandler.rs index e440cd33..69bbceb0 100644 --- a/dotscope/src/emulation/engine/exhandler.rs +++ b/dotscope/src/emulation/engine/exhandler.rs @@ -371,6 +371,21 @@ pub fn find_exception_handler( ExceptionClause::from_metadata_handlers(&body.exception_handlers) }; + // Cleanup clauses encountered while scanning, held locally until the outcome is known. + // + // Committing these to `exception_state` as they are found is wrong: a cleanup clause that + // the exception never unwinds past must not run, and nothing would ever drain it, so it + // would accumulate across every search. But dropping them all when a handler matches is + // equally wrong — that was the defect this shape replaced. For + // `try { try { throw } finally { F } } catch { C }` both clauses cover the throw offset and + // the inner Finally comes first in the clause table, so the exception *does* unwind out of + // the inner try to reach `C`, and `F` has to run on the way. + // + // The rule is containment, not scan order: on a match, commit exactly those candidates + // whose try region is nested inside the matched clause's try region. The full clause is + // kept rather than just its handler offset so that test can be made. + let mut cleanup_candidates: Vec<&ExceptionClause> = Vec::new(); + // Search for a handler that covers the current offset // Handlers are processed in order - innermost handlers first (as per ECMA-335) for clause in &clauses { @@ -397,6 +412,12 @@ pub fn find_exception_handler( }; if is_compatible { + commit_nested_cleanups( + exception_state, + method_token, + clause, + &cleanup_candidates, + ); return Ok(Some(HandlerMatch::Catch { method: method_token, handler_offset: clause.handler_offset(), @@ -409,6 +430,7 @@ pub fn find_exception_handler( // Filter handler - need to execute filter code first // Store the handler offset so EndFilter knows where to jump let handler_offset = clause.handler_offset(); + commit_nested_cleanups(exception_state, method_token, clause, &cleanup_candidates); exception_state.enter_filter(handler_offset); return Ok(Some(HandlerMatch::Filter { method: method_token, @@ -417,23 +439,51 @@ pub fn find_exception_handler( })); } - ExceptionClause::Finally { handler_length, .. } => { - // Finally handler - schedule for execution during unwinding - exception_state.push_finally(method_token, clause.handler_offset(), None); - _ = handler_length; - } - - ExceptionClause::Fault { handler_length, .. } => { - // Fault handler - like finally but only on exception path - exception_state.push_finally(method_token, clause.handler_offset(), None); - _ = handler_length; + ExceptionClause::Finally { .. } | ExceptionClause::Fault { .. } => { + // Runs during unwinding — whether it is reached depends on where the search + // ends up, so it is only a candidate until then. + cleanup_candidates.push(clause); } } } + // No handler matched, so the exception unwinds out of this frame entirely and every + // cleanup clause collected above really does run. + for clause in cleanup_candidates { + exception_state.push_finally(method_token, clause.handler_offset(), None); + } + Ok(None) } +/// Queues the cleanup clauses the exception unwinds past on its way to `matched`. +/// +/// A `finally`/`fault` whose try region is nested inside the matched handler's try region sits +/// between the throw site and that handler, so it must run before the handler is entered — the +/// `try { try { throw } finally { F } } catch { C }` case. One that merely *encloses* the +/// matched handler is not unwound past at all: the exception is caught inside it, and queueing +/// it would run cleanup for a region still being executed and leave an entry nothing drains. +fn commit_nested_cleanups( + exception_state: &mut ThreadExceptionState, + method_token: Token, + matched: &ExceptionClause, + candidates: &[&ExceptionClause], +) { + let (outer_start, outer_end) = (matched.try_offset(), matched.try_end()); + + for clause in candidates { + let nested = clause.try_offset() >= outer_start && clause.try_end() <= outer_end + // An identical region is not nested — `try {} catch {} finally {}` compiles to two + // clauses over the same try, and per ECMA-335 the finally runs *after* the catch + // completes, not before it is entered. Leave it to the normal exit path. + && (clause.try_offset() != outer_start || clause.try_end() != outer_end); + + if nested { + exception_state.push_finally(method_token, clause.handler_offset(), None); + } + } +} + /// Schedules finally blocks to execute when leaving a protected region. /// /// Finds all `finally` handlers between the current offset and the leave @@ -487,18 +537,21 @@ pub fn schedule_finally_blocks( } } - // Sort by try_offset descending (innermost first) - finally_blocks.sort_by_key(|f| std::cmp::Reverse(f.1)); - - // Schedule finally blocks in order (innermost first) - // The last one scheduled will be popped first + // ECMA-335 §12.4.2.5 runs finally handlers innermost-first when a `leave` exits nested try + // regions. `pending_finally` is a LIFO stack, so pop order is the reverse of push order: + // to pop innermost-first we must push **outermost-first**. + // + // For nested regions the inner try starts at or after the outer one, so ascending + // `try_offset` is outermost-first. + finally_blocks.sort_by_key(|f| f.1); + + // The leave target belongs on the entry popped **last** — the outermost, i.e. the first + // pushed. Only once every finally in the chain has run does control transfer to the leave + // target; attaching it to any earlier entry means `handle_end_finally` consumes it and then + // overwrites it with a later entry's `None`, leaving `endfinally` with neither a target nor + // an exception and nothing to advance the instruction pointer. for (i, (handler_offset, _)) in finally_blocks.iter().enumerate() { - // The last finally should have the actual leave target - let target = if i == finally_blocks.len().saturating_sub(1) { - Some(leave_target) - } else { - None - }; + let target = if i == 0 { Some(leave_target) } else { None }; exception_state.push_finally(method_token, *handler_offset, target); } @@ -675,3 +728,86 @@ pub fn track_cctor_failure_if_needed( cctor_tracker.mark_type_failed(type_token, exception_ref)?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn finally_clause(try_offset: u32, try_length: u32, handler_offset: u32) -> ExceptionClause { + ExceptionClause::Finally { + try_offset, + try_length, + handler_offset, + handler_length: 4, + } + } + + fn catch_clause(try_offset: u32, try_length: u32, handler_offset: u32) -> ExceptionClause { + ExceptionClause::Catch { + try_offset, + try_length, + handler_offset, + handler_length: 4, + catch_type: Token::new(0x0100_0001), + } + } + + /// `try { try { throw } finally { F } } catch { C }` must run `F`. + /// + /// Both clauses cover the throw offset, and the inner Finally precedes the outer Catch in + /// the clause table. The exception unwinds out of the inner try to reach `C`, so `F` sits + /// on that path. Dropping every candidate on a match — which is what committing them only + /// on the no-handler path amounts to — silently skipped it. + #[test] + fn a_finally_nested_inside_the_matched_try_is_queued() { + let mut state = ThreadExceptionState::new(); + let inner_finally = finally_clause(10, 10, 30); + let outer_catch = catch_clause(0, 40, 50); + + commit_nested_cleanups( + &mut state, + Token::new(0x0600_0001), + &outer_catch, + &[&inner_finally], + ); + + assert!(state.has_pending_finally(), "the inner finally must run"); + assert_eq!(state.pop_finally().map(|p| p.handler_offset), Some(30)); + } + + /// A `finally` that *encloses* the matched handler is not unwound past. + /// + /// The exception is caught inside that region, so its cleanup has not been reached. + /// Queueing it would run cleanup for a region still executing and leave an entry that + /// nothing drains — the accumulation the local-candidate design exists to prevent. + #[test] + fn a_finally_enclosing_the_matched_try_is_not_queued() { + let mut state = ThreadExceptionState::new(); + let outer_finally = finally_clause(0, 100, 80); + let inner_catch = catch_clause(10, 10, 30); + + commit_nested_cleanups( + &mut state, + Token::new(0x0600_0001), + &inner_catch, + &[&outer_finally], + ); + + assert!(!state.has_pending_finally()); + } + + /// `try { } catch { } finally { }` shares one try region across both clauses. + /// + /// ECMA-335 runs the finally *after* the catch completes, via the normal exit path — not + /// before the catch is entered. An identical region is therefore not "nested". + #[test] + fn a_finally_over_the_same_try_region_is_not_queued() { + let mut state = ThreadExceptionState::new(); + let finally = finally_clause(0, 20, 40); + let catch = catch_clause(0, 20, 20); + + commit_nested_cleanups(&mut state, Token::new(0x0600_0001), &catch, &[&finally]); + + assert!(!state.has_pending_finally()); + } +} diff --git a/dotscope/src/emulation/engine/interpreter/handlers.rs b/dotscope/src/emulation/engine/interpreter/handlers.rs index 0406deba..6a08d4ff 100644 --- a/dotscope/src/emulation/engine/interpreter/handlers.rs +++ b/dotscope/src/emulation/engine/interpreter/handlers.rs @@ -552,9 +552,13 @@ impl Interpreter { length: 0, }) })?, + // `usize::try_from(u64)` is infallible on a 64-bit host, so this arm imposes no + // bound of its own — any value up to u64::MAX reaches the allocator. The heap + // budget in `ManagedHeap::alloc_array` is what actually rejects it; this cast is + // only here to normalise the width. EmValue::NativeUInt(v) => usize::try_from(v).map_err(|_| { Error::from(EmulationError::ArrayIndexOutOfBounds { - index: v as i64, + index: i64::try_from(v).unwrap_or(i64::MAX), length: 0, }) })?, @@ -1547,42 +1551,18 @@ impl Interpreter { _ => unreachable!(), }; - // Helper: read an exact-sized little-endian array from the address space. - // The address_space.read API returns the requested number of bytes on - // success; a length mismatch indicates an internal invariant violation. - let read_exact = |size: usize| -> Result> { - let bytes = address_space.read(ptr_addr, size)?; - if bytes.len() != size { - return Err(EmulationError::InternalError { - description: format!( - "address space read returned {} bytes, expected {}", - bytes.len(), - size - ), - } - .into()); - } - Ok(bytes) - }; - let read_array_2 = || -> Result<[u8; 2]> { - let bytes = read_exact(2)?; - <[u8; 2]>::try_from(bytes.as_slice()).map_err(|_| out_of_bounds_error!()) - }; - let read_array_4 = || -> Result<[u8; 4]> { - let bytes = read_exact(4)?; - <[u8; 4]>::try_from(bytes.as_slice()).map_err(|_| out_of_bounds_error!()) - }; - let read_array_8 = || -> Result<[u8; 8]> { - let bytes = read_exact(8)?; - <[u8; 8]>::try_from(bytes.as_slice()).map_err(|_| out_of_bounds_error!()) - }; + // These fill a stack array. `AddressSpace::read` hands back an owned `Vec`, + // which would put a malloc and a free on every `ldind` — the interpreter's + // innermost loop — for a load of at most eight bytes. + let read_array_2 = || address_space.read_exact::<2>(ptr_addr); + let read_array_4 = || address_space.read_exact::<4>(ptr_addr); + let read_array_8 = || address_space.read_exact::<8>(ptr_addr); // Read from address space based on read_size and expected_type let value = match (expected_type, read_size) { // Small integer reads (1 or 2 bytes) that widen to I32 (&CilFlavor::I4, 1) => { - let bytes = read_exact(1)?; - let b0 = *bytes.first().ok_or(out_of_bounds_error!())?; + let b0 = address_space.read_u8(ptr_addr)?; let val = if signed { // Intentional wrap-around for sign extension from u8 to i8 i32::from(b0.cast_signed()) diff --git a/dotscope/src/emulation/engine/resolution.rs b/dotscope/src/emulation/engine/resolution.rs index 0a8980e5..8d35145a 100644 --- a/dotscope/src/emulation/engine/resolution.rs +++ b/dotscope/src/emulation/engine/resolution.rs @@ -38,7 +38,7 @@ pub enum CallResolution { /// Whether the caller expects a return value on the stack. expects_return: bool, /// Assembly index for dynamically loaded assemblies (`None` = primary). - assembly_index: Option, + assembly_index: Option, /// Method-level generic type arguments (`!!0`, `!!1`, ...) from MethodSpec. method_type_args: Option>, }, @@ -81,7 +81,7 @@ pub enum CallResolution { /// are wrapped in `TargetInvocationException`. is_reflection_invoke: bool, /// Assembly index for cross-assembly redirects (`None` = primary). - assembly_index: Option, + assembly_index: Option, /// Method-level generic type arguments (`!!0`, `!!1`, ...) from MethodSpec. method_type_args: Option>, }, diff --git a/dotscope/src/emulation/engine/typeops/mod.rs b/dotscope/src/emulation/engine/typeops/mod.rs index 03c93ae0..463b69c4 100644 --- a/dotscope/src/emulation/engine/typeops/mod.rs +++ b/dotscope/src/emulation/engine/typeops/mod.rs @@ -178,31 +178,20 @@ pub fn handle_ldsfld( } } - if let Some((type_token, flavor)) = context.get_field_type_info(field) { - if address_space.statics().is_type_initialized(type_token)? { - let default_value = match flavor { - CilFlavor::Boolean - | CilFlavor::Char - | CilFlavor::I1 - | CilFlavor::U1 - | CilFlavor::I2 - | CilFlavor::U2 - | CilFlavor::I4 - | CilFlavor::U4 => EmValue::I32(0), - CilFlavor::I8 | CilFlavor::U8 => EmValue::I64(0), - CilFlavor::R4 => EmValue::F32(0.0), - CilFlavor::R8 => EmValue::F64(0.0), - CilFlavor::I | CilFlavor::U => EmValue::NativeInt(0), - _ => EmValue::Null, - }; - thread.push(default_value)?; - return Ok(()); - } - } - // .NET zero-initializes all static fields before any code runs. // If we have type info, use the type-appropriate default; otherwise use Null // (correct for reference types, the common case for unknown-type fields). + // + // There is deliberately no branch on `is_type_initialized` here. A previous version called + // `get_field_type_info` twice — once inside an initialized-type check and once outside — + // and both arms computed the *same* default from the same flavor, so the check changed + // nothing observable while paying for a second lookup on every `ldsfld` of an + // uninitialized static. Reading a static before its `.cctor` has run yields the same zero + // value it would afterwards, which is exactly what the runtime guarantees. + // + // Note this deliberately differs from `EmValue::default_for_flavor`: on the CIL evaluation + // stack `bool` and the sub-word integers widen to `I32`, whereas that helper preserves + // `Bool`. if let Some((_type_token, flavor)) = context.get_field_type_info(field) { let default_value = match flavor { CilFlavor::Boolean diff --git a/dotscope/src/emulation/exception/handler.rs b/dotscope/src/emulation/exception/handler.rs deleted file mode 100644 index ec911f0b..00000000 --- a/dotscope/src/emulation/exception/handler.rs +++ /dev/null @@ -1,738 +0,0 @@ -//! Exception handler resolution for .NET emulation. -//! -//! This module provides the logic for finding appropriate exception handlers when an exception -//! is thrown during emulation. It implements the handler search algorithm defined by the -//! ECMA-335 specification, which searches exception clauses in innermost-to-outermost order. -//! -//! # Handler Search Algorithm -//! -//! When an exception is thrown, the handler search proceeds as follows: -//! -//! 1. Exception clauses are examined in order (innermost try block first) -//! 2. For each clause whose try block contains the throw point: -//! - **Catch clauses**: Check if the exception type is assignable to the catch type -//! - **Filter clauses**: Return for filter code evaluation -//! - **Finally/Fault clauses**: Queue for execution during unwinding -//! 3. If a matching catch/filter is found, cleanup handlers execute first -//! 4. If no handler is found in the current method, unwinding continues to the caller -//! -//! # Type Matching -//! -//! Exception type matching follows .NET's inheritance rules: a catch clause matches -//! if the thrown exception type is the same as or derives from the catch type. -//! Type compatibility is determined by a caller-provided function, typically using -//! [`EmulationContext::is_type_compatible`](crate::emulation::EmulationContext::is_type_compatible). - -use crate::{ - emulation::exception::{ExceptionClause, HandlerMatch, InstructionLocation}, - metadata::token::Token, -}; - -/// Exception handler resolver for .NET exception handling. -/// -/// This component is responsible for finding appropriate exception handlers when an exception -/// is thrown during emulation. It searches through exception clauses in the current method, -/// respecting the .NET exception handling semantics defined in ECMA-335. -/// -/// # Responsibilities -/// -/// - Finding catch handlers that match the thrown exception type -/// - Identifying filter handlers that need runtime evaluation -/// - Collecting finally and fault handlers for cleanup during unwinding -/// -/// # Type Matching -/// -/// Exception type matching is delegated to the caller via a type checker function. -/// The caller (typically the emulation engine) provides this function using -/// [`EmulationContext::is_type_compatible`](crate::emulation::EmulationContext::is_type_compatible) -/// which has full access to the assembly's type hierarchy. -/// -/// # Usage -/// -/// ```ignore -/// use dotscope::emulation::exception::ExceptionHandler; -/// use dotscope::emulation::EmulationContext; -/// use dotscope::metadata::token::Token; -/// -/// let handler = ExceptionHandler::new(); -/// -/// // Use EmulationContext for type checking -/// let result = handler.find_handler( -/// &clauses, -/// throw_offset, -/// exception_type, -/// method_token, -/// |exc_type, catch_type| ctx.is_type_compatible(exc_type, catch_type), -/// ); -/// ``` -#[derive(Clone, Copy, Debug, Default)] -pub struct ExceptionHandler; - -/// Result of searching for an exception handler within a single method. -/// -/// This enum represents the possible outcomes when searching a method's exception -/// clauses for a handler that matches a thrown exception. -/// -/// # Variants -/// -/// - [`Found`](MethodHandlerResult::Found) - A matching catch or filter handler was found -/// - [`ExecuteCleanup`](MethodHandlerResult::ExecuteCleanup) - Cleanup handlers must run -/// before continuing (finally/fault blocks, or cleanup before catch) -/// - [`NotFound`](MethodHandlerResult::NotFound) - No handler in this method; continue -/// unwinding to the caller -#[derive(Clone, Debug)] -pub enum MethodHandlerResult { - /// Found a suitable catch or filter handler in this method. - /// - /// The handler can be entered after any pending cleanup handlers complete. - Found(HandlerMatch), - - /// Cleanup handlers must be executed before continuing. - /// - /// This variant is returned when: - /// - Finally/fault blocks need to run before entering a catch handler - /// - Finally/fault blocks need to run before continuing the unwind - /// - /// The handlers should be executed in order, with the last one potentially - /// being a catch handler. - ExecuteCleanup { - /// Cleanup handlers to execute in order. - handlers: Vec, - }, - - /// No matching handler was found in this method. - /// - /// The exception handling system should continue unwinding to the caller. - NotFound, -} - -impl ExceptionHandler { - /// Creates a new exception handler resolver. - /// - /// # Returns - /// - /// A new `ExceptionHandler` instance. - #[must_use] - pub fn new() -> Self { - Self - } - - /// Finds an exception handler for an exception thrown at the given location. - /// - /// This method searches through the exception clauses in order (innermost first), - /// looking for handlers that can handle the thrown exception. The search considers: - /// - /// 1. **Catch handlers** - Checked for type compatibility with the exception - /// 2. **Filter handlers** - Returned for runtime evaluation of the filter code - /// 3. **Finally/Fault handlers** - Collected for execution during unwinding - /// - /// # Arguments - /// - /// * `clauses` - Exception clauses for the current method, ordered innermost-first - /// * `throw_offset` - IL offset where the exception was thrown (or current IP during unwind) - /// * `exception_type` - Type token of the thrown exception - /// * `method` - Token of the method being searched - /// * `is_type_compatible` - A function that checks if an exception type is assignable - /// to a catch type. Use [`EmulationContext::is_type_compatible`](crate::emulation::EmulationContext::is_type_compatible). - /// - /// # Returns - /// - /// A [`MethodHandlerResult`] indicating: - /// - `Found` if a matching handler was found with no pending cleanup - /// - `ExecuteCleanup` if cleanup handlers must run (with or without a final catch) - /// - `NotFound` if no handlers apply in this method - /// - /// # Algorithm - /// - /// The search proceeds as follows: - /// 1. Only clauses whose try block contains `throw_offset` are considered - /// 2. Catch clauses are checked for type compatibility via the provided function - /// 3. Filter clauses are returned immediately for evaluation - /// 4. Finally/fault clauses are collected as cleanup handlers - /// 5. If a catch is found, cleanup runs first, then control transfers to the catch - pub fn find_handler( - &self, - clauses: &[ExceptionClause], - throw_offset: u32, - exception_type: Token, - method: Token, - is_type_compatible: F, - ) -> MethodHandlerResult - where - F: Fn(Token, Token) -> bool, - { - let mut cleanup_handlers = Vec::new(); - let mut found_catch = None; - - // Clauses are ordered innermost-first in .NET metadata - for clause in clauses { - // Only consider clauses whose try block contains the throw point - if !clause.is_in_try(throw_offset) { - continue; - } - - match clause { - ExceptionClause::Catch { catch_type, .. } => { - // Check if this catch handler matches the exception type - if is_type_compatible(exception_type, *catch_type) { - found_catch = Some(HandlerMatch::Catch { - method, - handler_offset: clause.handler_offset(), - }); - break; - } - } - - ExceptionClause::Filter { filter_offset, .. } => { - // Filter handlers need to be evaluated at runtime - // Return this so the caller can evaluate the filter - found_catch = Some(HandlerMatch::Filter { - method, - filter_offset: *filter_offset, - handler_offset: clause.handler_offset(), - }); - break; - } - - ExceptionClause::Finally { .. } => { - // Finally blocks must run during unwinding - cleanup_handlers.push(HandlerMatch::Finally { - method, - handler_offset: clause.handler_offset(), - handler_length: clause.handler_length(), - continue_search_after: true, - }); - } - - ExceptionClause::Fault { .. } => { - // Fault blocks only run on exception path - cleanup_handlers.push(HandlerMatch::Fault { - method, - handler_offset: clause.handler_offset(), - handler_length: clause.handler_length(), - }); - } - } - } - - // If we found a catch handler, we still need to run any cleanup handlers first - if let Some(catch) = found_catch { - if cleanup_handlers.is_empty() { - return MethodHandlerResult::Found(catch); - } - // Add the catch as the final handler after cleanup - cleanup_handlers.push(catch); - return MethodHandlerResult::ExecuteCleanup { - handlers: cleanup_handlers, - }; - } - - // No catch found - return cleanup handlers if any - if cleanup_handlers.is_empty() { - MethodHandlerResult::NotFound - } else { - MethodHandlerResult::ExecuteCleanup { - handlers: cleanup_handlers, - } - } - } - - /// Finds finally handlers that must execute for a `leave` instruction. - /// - /// When executing a `leave` instruction to exit a protected region (try block), - /// any finally blocks that protect the current position but not the target must - /// be executed before control transfers to the target. - /// - /// This is distinct from exception handling: `leave` is used for normal control - /// flow out of try blocks (e.g., `return` or `break` inside a try), not for - /// exception propagation. - /// - /// # Arguments - /// - /// * `clauses` - Exception clauses for the method - /// * `leave_offset` - IL offset of the `leave` instruction - /// * `target_offset` - Target IL offset where control will transfer - /// * `method` - Token of the method containing the leave - /// - /// # Returns - /// - /// A vector of [`HandlerMatch::Finally`] handlers to execute in order before - /// transferring control to `target_offset`. Empty if no finally blocks need - /// to run. - /// - /// # Example - /// - /// For code like: - /// ```csharp - /// try { - /// if (condition) return; // leave instruction here - /// } finally { - /// Cleanup(); - /// } - /// ``` - /// This method returns the finally handler so `Cleanup()` runs before returning. - #[must_use] - pub fn find_finally_for_leave( - &self, - clauses: &[ExceptionClause], - leave_offset: u32, - target_offset: u32, - method: Token, - ) -> Vec { - let mut handlers = Vec::new(); - - for clause in clauses { - // Only consider finally clauses - if !clause.is_finally() { - continue; - } - - // Check if we're leaving a try block that this finally protects - let in_try = clause.is_in_try(leave_offset); - let target_in_try = clause.is_in_try(target_offset); - - // If we're in the try block but jumping outside it, run the finally - if in_try && !target_in_try { - handlers.push(HandlerMatch::Finally { - method, - handler_offset: clause.handler_offset(), - handler_length: clause.handler_length(), - continue_search_after: false, // Not searching for catch - }); - } - } - - handlers - } - - /// Checks if an IL offset is within any handler block. - /// - /// This is useful for determining if execution is currently inside a - /// catch, finally, fault, or filter handler block. This information is - /// needed for proper `rethrow` handling and control flow validation. - /// - /// # Arguments - /// - /// * `clauses` - Exception clauses for the method - /// * `offset` - The IL offset to check - /// - /// # Returns - /// - /// `true` if the offset is within any handler block, `false` otherwise. - #[must_use] - pub fn is_in_handler(&self, clauses: &[ExceptionClause], offset: u32) -> bool { - clauses.iter().any(|c| c.is_in_handler(offset)) - } - - /// Gets the exception clause that owns the handler at the given offset. - /// - /// When execution is inside a handler block, this method returns the - /// exception clause that defines that handler. This is useful for - /// determining the type of handler being executed and its properties. - /// - /// # Arguments - /// - /// * `clauses` - Exception clauses for the method - /// * `offset` - The IL offset within a handler block - /// - /// # Returns - /// - /// The exception clause containing the handler, or `None` if the offset - /// is not within any handler block. - #[must_use] - pub fn get_handler_clause<'a>( - &self, - clauses: &'a [ExceptionClause], - offset: u32, - ) -> Option<&'a ExceptionClause> { - clauses.iter().find(|c| c.is_in_handler(offset)) - } - - /// Builds a formatted stack trace string from instruction locations. - /// - /// Creates a human-readable stack trace similar to .NET's exception stack trace - /// format, with each frame on a separate line prefixed with "at". - /// - /// # Arguments - /// - /// * `locations` - Slice of instruction locations representing the call stack - /// - /// # Returns - /// - /// A formatted string with one frame per line. - /// - /// # Example Output - /// - /// ```text - /// at 0x06000001+0x0042 - /// at 0x06000002+0x0010 - /// at 0x06000003+0x0005 - /// ``` - #[must_use] - pub fn build_stack_trace(locations: &[InstructionLocation]) -> String { - locations - .iter() - .map(|loc| format!(" at {loc}")) - .collect::>() - .join("\n") - } -} - -/// State for multi-frame exception handler search across the call stack. -/// -/// When an exception is thrown, the handler search may span multiple stack frames. -/// This structure tracks the state of that search, including which frames have been -/// searched, what cleanup handlers have been found, and whether a catch handler -/// has been located. -/// -/// # Usage -/// -/// ```ignore -/// let mut state = HandlerSearchState::new(exception_type, throw_location); -/// -/// // Add frames from the call stack (deepest first) -/// state.add_frame(method1, offset1, clauses1); -/// state.add_frame(method2, offset2, clauses2); -/// -/// // Search proceeds until a handler is found or all frames are exhausted -/// while !state.is_complete() { -/// // Process each frame... -/// } -/// ``` -#[derive(Clone, Debug)] -pub struct HandlerSearchState { - /// The type token of the exception being handled. - pub exception_type: Token, - - /// The location where the exception was thrown. - pub current_location: InstructionLocation, - - /// Stack frames to search, ordered deepest (innermost) first. - pub frames: Vec, - - /// Index of the current frame being searched (0-based). - pub current_frame: usize, - - /// Cleanup handlers (finally/fault) collected during the search. - /// - /// These handlers must execute before entering a catch handler or - /// before propagating the exception to the next frame. - pub pending_cleanup: Vec, - - /// The catch or filter handler that will handle the exception, if found. - pub handler_found: Option, -} - -/// Information about a single stack frame for exception handler search. -/// -/// This structure contains all the information needed to search for exception -/// handlers within a single method's stack frame during exception propagation. -#[derive(Clone, Debug)] -pub struct FrameSearchInfo { - /// The method token identifying this stack frame. - pub method: Token, - - /// The IL offset where execution was when the exception occurred or propagated. - pub offset: u32, - - /// Exception clauses defined in this method's metadata. - /// - /// These clauses define the protected regions (try blocks) and their - /// associated handlers within this method. - pub clauses: Vec, -} - -impl HandlerSearchState { - /// Creates a new handler search state for an exception. - /// - /// Initializes the search state with the exception type and the location - /// where the exception was thrown. Stack frames should be added using - /// [`add_frame`](Self::add_frame) before searching. - /// - /// # Arguments - /// - /// * `exception_type` - The type token of the thrown exception - /// * `throw_location` - The instruction location where the exception was thrown - /// - /// # Returns - /// - /// A new `HandlerSearchState` ready to accept stack frames for searching. - #[must_use] - pub fn new(exception_type: Token, throw_location: InstructionLocation) -> Self { - Self { - exception_type, - current_location: throw_location, - frames: Vec::new(), - current_frame: 0, - pending_cleanup: Vec::new(), - handler_found: None, - } - } - - /// Adds a stack frame to the search. - /// - /// Frames should be added in order from innermost (current) to outermost (caller). - /// Each frame's exception clauses will be searched for handlers. - /// - /// # Arguments - /// - /// * `method` - The method token for this frame - /// * `offset` - The IL offset within the method - /// * `clauses` - The exception clauses defined in this method - pub fn add_frame(&mut self, method: Token, offset: u32, clauses: Vec) { - self.frames.push(FrameSearchInfo { - method, - offset, - clauses, - }); - } - - /// Checks if the handler search is complete. - /// - /// The search is complete when either: - /// - A matching handler has been found - /// - All frames have been searched without finding a handler - /// - /// # Returns - /// - /// `true` if the search is complete, `false` if more frames need to be searched. - #[must_use] - pub fn is_complete(&self) -> bool { - self.handler_found.is_some() - || (!self.frames.is_empty() && self.current_frame >= self.frames.len()) - } - - /// Checks if there are pending cleanup handlers to execute. - /// - /// Cleanup handlers (finally/fault blocks) must be executed before entering - /// a catch handler or before continuing to unwind to the next frame. - /// - /// # Returns - /// - /// `true` if there are cleanup handlers waiting to be executed. - #[must_use] - pub fn has_pending_cleanup(&self) -> bool { - !self.pending_cleanup.is_empty() - } - - /// Takes the next cleanup handler from the queue. - /// - /// Removes and returns the first pending cleanup handler. Cleanup handlers - /// should be executed in the order they are returned (FIFO). - /// - /// # Returns - /// - /// The next cleanup handler to execute, or `None` if no cleanup handlers remain. - pub fn take_next_cleanup(&mut self) -> Option { - if self.pending_cleanup.is_empty() { - None - } else { - Some(self.pending_cleanup.remove(0)) - } - } -} - -#[cfg(test)] -mod tests { - use crate::{ - emulation::exception::{ - ExceptionClause, ExceptionHandler, HandlerMatch, HandlerSearchState, - InstructionLocation, MethodHandlerResult, - }, - metadata::token::Token, - }; - - fn create_test_clauses() -> Vec { - vec![ - // Inner try/catch - ExceptionClause::Catch { - try_offset: 0x10, - try_length: 0x20, - handler_offset: 0x30, - handler_length: 0x10, - catch_type: Token::new(0x01000010), // Specific exception type - }, - // Outer try/finally - ExceptionClause::Finally { - try_offset: 0x00, - try_length: 0x50, - handler_offset: 0x50, - handler_length: 0x10, - }, - // Another catch for broader type - ExceptionClause::Catch { - try_offset: 0x00, - try_length: 0x50, - handler_offset: 0x60, - handler_length: 0x10, - catch_type: Token::new(0x0100_0001), // System.Exception - }, - ] - } - - /// Simple type checker for tests: exact match or catch-all System.Exception. - fn test_type_checker(exception_type: Token, catch_type: Token) -> bool { - // Exact match - if exception_type == catch_type { - return true; - } - // System.Exception (0x0100_0001) catches everything - if catch_type.value() == 0x0100_0001 { - return true; - } - false - } - - #[test] - fn test_find_matching_catch() { - let handler = ExceptionHandler::new(); - let clauses = create_test_clauses(); - let method = Token::new(0x06000001); - - // Throw in inner try block with matching type - let result = handler.find_handler( - &clauses, - 0x15, - Token::new(0x01000010), - method, - test_type_checker, - ); - - match result { - MethodHandlerResult::Found(HandlerMatch::Catch { handler_offset, .. }) => { - assert_eq!(handler_offset, 0x30); - } - _ => panic!("Expected to find catch handler"), - } - } - - #[test] - fn test_find_base_exception_catch() { - let handler = ExceptionHandler::new(); - let clauses = create_test_clauses(); - let method = Token::new(0x06000001); - - // Throw in outer try (but not inner) with unknown type - // Should match System.Exception catch via the catch-all in test_type_checker - let result = handler.find_handler( - &clauses, - 0x05, - Token::new(0x01000099), - method, - test_type_checker, - ); - - match result { - MethodHandlerResult::ExecuteCleanup { handlers } => { - // Should have finally + catch - assert_eq!(handlers.len(), 2); - assert!(matches!(handlers[0], HandlerMatch::Finally { .. })); - assert!(matches!(handlers[1], HandlerMatch::Catch { .. })); - } - _ => panic!("Expected cleanup handlers"), - } - } - - #[test] - fn test_find_finally_for_leave() { - let handler = ExceptionHandler::new(); - let clauses = create_test_clauses(); - let method = Token::new(0x06000001); - - // Leave from inside try to outside - let handlers = handler.find_finally_for_leave(&clauses, 0x20, 0x70, method); - - assert_eq!(handlers.len(), 1); - match &handlers[0] { - HandlerMatch::Finally { handler_offset, .. } => { - assert_eq!(*handler_offset, 0x50); - } - _ => panic!("Expected finally handler"), - } - } - - #[test] - fn test_no_handler_found() { - let handler = ExceptionHandler::new(); - let clauses = vec![ExceptionClause::Catch { - try_offset: 0x00, - try_length: 0x10, - handler_offset: 0x10, - handler_length: 0x10, - catch_type: Token::new(0x01000010), - }]; - let method = Token::new(0x06000001); - - // Throw outside try block - let result = handler.find_handler( - &clauses, - 0x50, - Token::new(0x01000010), - method, - test_type_checker, - ); - - assert!(matches!(result, MethodHandlerResult::NotFound)); - } - - #[test] - fn test_filter_handler() { - let handler = ExceptionHandler::new(); - let clauses = vec![ExceptionClause::Filter { - try_offset: 0x00, - try_length: 0x20, - handler_offset: 0x30, - handler_length: 0x10, - filter_offset: 0x20, - }]; - let method = Token::new(0x06000001); - - let result = handler.find_handler( - &clauses, - 0x10, - Token::new(0x01000010), - method, - test_type_checker, - ); - - match result { - MethodHandlerResult::Found(HandlerMatch::Filter { - filter_offset, - handler_offset, - .. - }) => { - assert_eq!(filter_offset, 0x20); - assert_eq!(handler_offset, 0x30); - } - _ => panic!("Expected filter handler"), - } - } - - #[test] - fn test_is_in_handler() { - let handler = ExceptionHandler::new(); - let clauses = create_test_clauses(); - - assert!(!handler.is_in_handler(&clauses, 0x15)); // In try block - assert!(handler.is_in_handler(&clauses, 0x35)); // In catch handler - assert!(handler.is_in_handler(&clauses, 0x55)); // In finally handler - } - - #[test] - fn test_handler_search_state() { - let exception_type = Token::new(0x01000010); - let location = InstructionLocation::new(Token::new(0x06000001), 0x20); - - let mut state = HandlerSearchState::new(exception_type, location); - - assert!(!state.is_complete()); - assert!(!state.has_pending_cleanup()); - - state.handler_found = Some(HandlerMatch::Catch { - method: Token::new(0x06000001), - handler_offset: 0x30, - }); - - assert!(state.is_complete()); - } -} diff --git a/dotscope/src/emulation/exception/mod.rs b/dotscope/src/emulation/exception/mod.rs index c9b0af1f..249f5f69 100644 --- a/dotscope/src/emulation/exception/mod.rs +++ b/dotscope/src/emulation/exception/mod.rs @@ -22,48 +22,31 @@ //! finally, fault) //! - [`ThreadExceptionState`] - Manages per-thread exception tracking, including the active //! exception, pending finally blocks, and filter evaluation state -//! - [`ExceptionHandler`] - Provides handler resolution logic for finding appropriate exception -//! handlers based on exception type and protected regions -//! - [`StackUnwinder`] - Manages stack unwinding during exception propagation, ensuring proper -//! execution of cleanup handlers -//! //! # Exception Handling Flow //! -//! When an exception is thrown: -//! -//! 1. The [`ExceptionHandler`] searches for a matching catch or filter handler in the current -//! method's exception clauses -//! 2. If no handler is found, the search continues up the call stack via [`StackUnwinder`] -//! 3. Finally and fault handlers are queued for execution during unwinding -//! 4. Once a handler is found, cleanup handlers execute in order before control transfers -//! to the catch handler +//! The search and unwind logic itself lives in +//! [`engine::exhandler`](crate::emulation::engine), which is the only implementation the +//! execution loop calls. This module provides the data it operates on. //! -//! # Example -//! -//! ```ignore -//! use dotscope::emulation::exception::{ExceptionHandler, ExceptionClause, ThreadExceptionState}; -//! use dotscope::emulation::EmulationContext; -//! -//! // Create exception handler resolver -//! let handler = ExceptionHandler::new(); +//! When an exception is thrown: //! -//! // Find handler for an exception at a given IL offset -//! // Type checking is delegated to EmulationContext -//! let result = handler.find_handler( -//! &clauses, -//! throw_offset, -//! exception_type, -//! method_token, -//! |exc, catch| ctx.is_type_compatible(exc, catch), -//! ); -//! ``` +//! 1. `find_exception_handler` scans the current method's clauses for a `catch` whose type +//! matches, or a `filter` to evaluate +//! 2. Any `finally`/`fault` nested inside the matched clause's try region is queued, since the +//! exception unwinds past it on the way to the handler +//! 3. If no handler matches, every cleanup clause is queued and the search continues up the +//! call stack +//! 4. Queued cleanup handlers execute before control transfers to the handler +//! +//! Two further implementations of this search — an `ExceptionHandler`/`HandlerSearchState` +//! pair here and a `StackUnwinder` — used to sit alongside it. Neither was reachable from the +//! engine, and they had diverged from it on exactly the point above: whether a `finally` found +//! before a matching `catch` still runs. Three copies of one algorithm, two of them dead and +//! silently disagreeing, is how that defect survived, so they were removed rather than +//! resynchronised. -mod handler; mod state; mod types; -mod unwinder; -pub use handler::{ExceptionHandler, FrameSearchInfo, HandlerSearchState, MethodHandlerResult}; pub use state::{ExceptionInfo, PendingFinally, ThreadExceptionState}; pub use types::{ExceptionClause, HandlerMatch, InstructionLocation}; -pub use unwinder::{StackUnwinder, UnwindSequenceBuilder, UnwindStepResult}; diff --git a/dotscope/src/emulation/exception/state.rs b/dotscope/src/emulation/exception/state.rs index 5a47b831..be366c71 100644 --- a/dotscope/src/emulation/exception/state.rs +++ b/dotscope/src/emulation/exception/state.rs @@ -36,6 +36,14 @@ use crate::{ metadata::token::Token, }; +/// Maximum number of cleanup handlers that may be queued for execution at once. +/// +/// A backstop, not a semantic limit: legitimate depth is bounded by try nesting times call +/// depth, both of which are already capped elsewhere, so this is only reachable when entries are +/// queued faster than they are drained. Without it the queue is an unbounded `Vec` free to +/// grow for the lifetime of the emulation. +const MAX_PENDING_FINALLY: usize = 4096; + /// Information about a thrown exception. /// /// This structure captures all the information about an exception at the time @@ -390,7 +398,17 @@ impl ThreadExceptionState { /// * `method` - The method containing the finally block /// * `handler_offset` - The IL offset of the finally handler /// * `leave_target` - Optional target offset for `leave` instructions + /// + /// Silently drops the entry once `MAX_PENDING_FINALLY` entries are queued. The queue is + /// bounded by real control flow — nesting depth times call depth — so reaching the cap means + /// entries are being queued faster than they are drained, which is a defect rather than a + /// deep program. Dropping is the conservative response: a skipped cleanup handler + /// mis-emulates one method, whereas an unbounded queue exhausts host memory. pub fn push_finally(&mut self, method: Token, handler_offset: u32, leave_target: Option) { + if self.pending_finally.len() >= MAX_PENDING_FINALLY { + return; + } + self.pending_finally.push(PendingFinally { method, handler_offset, @@ -398,6 +416,16 @@ impl ThreadExceptionState { }); } + /// Discards queued cleanup handlers belonging to `method`. + /// + /// Called when a frame is popped: entries scheduled against a frame that no longer exists + /// can never run correctly, since the handler IL would execute against whatever frame + /// happens to be current. + pub fn discard_finally_for_method(&mut self, method: Token) { + self.pending_finally + .retain(|pending| pending.method != method); + } + /// Pops and returns the next pending finally block. /// /// Returns the most recently pushed finally block (LIFO order). @@ -791,4 +819,35 @@ mod tests { assert!(!state.has_pending_finally()); assert!(!state.take_rethrow_request()); } + + /// Leaving a filter destroys the handler offset, so every `endfilter` consumer must read + /// it *first*. + /// + /// This coupling is not obvious from either call site and it has already been got wrong + /// once: `set_in_filter(false)` was hoisted above the reads in the controller's + /// `EndFilter` arm, which made `filter_handler_offset` unconditionally `None` there. On + /// the accepting path that skipped the jump and returned without advancing the IP, so + /// `endfilter` re-executed forever; on the rejecting path it emptied the skip key, so the + /// handler search re-matched the filter that had just rejected. Every `catch (E) when + /// (...)` was broken in both directions. + /// + /// If this coupling is ever removed, the ordering comments in + /// `engine::controller`'s `EndFilter` arm and `resume_search_after_filter` become stale. + #[test] + fn leaving_a_filter_clears_the_handler_offset() { + let mut state = ThreadExceptionState::new(); + + state.enter_filter(0x42); + assert!(state.in_filter()); + assert_eq!(state.filter_handler_offset(), Some(0x42)); + + state.set_in_filter(false); + assert!(!state.in_filter()); + assert_eq!( + state.filter_handler_offset(), + None, + "leaving a filter must clear the handler offset; the ordering requirements \ + documented in engine::controller depend on it" + ); + } } diff --git a/dotscope/src/emulation/exception/unwinder.rs b/dotscope/src/emulation/exception/unwinder.rs deleted file mode 100644 index 74ee855e..00000000 --- a/dotscope/src/emulation/exception/unwinder.rs +++ /dev/null @@ -1,906 +0,0 @@ -//! Stack unwinding for .NET exception handling. -//! -//! This module provides the logic for unwinding the call stack during exception -//! propagation, ensuring that finally and fault blocks are executed in the correct -//! order as specified by the ECMA-335 standard. -//! -//! # Overview -//! -//! When an exception is thrown and no handler is found in the current method, -//! the stack must be unwound to propagate the exception to calling methods. -//! During unwinding: -//! -//! 1. Each frame on the call stack is examined for handlers -//! 2. Finally and fault blocks in scope are collected for execution -//! 3. These cleanup handlers execute before continuing the search -//! 4. If a catch/filter handler is found, cleanup runs first, then control transfers -//! 5. If no handler is found, the exception is unhandled -//! -//! # Components -//! -//! - [`StackUnwinder`] - Main component that manages the unwind process -//! - [`UnwindStepResult`] - Indicates the next action after processing a frame -//! - [`UnwindSequenceBuilder`] - Helper for building ordered handler sequences -//! -//! # Two-Pass Exception Handling -//! -//! .NET uses a two-pass model: -//! 1. **First pass**: Search for a handler without running cleanup -//! 2. **Second pass**: Run cleanup handlers, then enter the catch -//! -//! This module primarily implements the second pass, executing handlers -//! in the correct order after a handler has been found. -//! -//! # Type Checking -//! -//! Exception type matching is delegated to the caller via a type checker function, -//! typically using [`EmulationContext::is_type_compatible`](crate::emulation::EmulationContext::is_type_compatible). - -use crate::{ - emulation::{ - exception::{ - ExceptionClause, ExceptionInfo, HandlerMatch, InstructionLocation, ThreadExceptionState, - }, - thread::ThreadCallFrame, - }, - metadata::token::Token, -}; - -/// Result of a single step in the unwind process. -/// -/// After processing a stack frame during exception unwinding, this enum -/// indicates what action the interpreter should take next. The unwind -/// process is driven by repeatedly calling unwind methods and acting -/// on the returned result. -/// -/// # State Machine -/// -/// The typical flow is: -/// 1. `ExecuteHandler` (finally/fault) -> execute cleanup -> `cleanup_complete()` -/// 2. Repeat step 1 for all cleanup handlers -/// 3. `ExecuteHandler` (catch, is_catch=true) or `ContinueUnwind` -/// 4. `HandlerEntered` when entering a catch, or `UnhandledException` when done -#[derive(Clone, Debug)] -pub enum UnwindStepResult { - /// A handler should be executed. - /// - /// The interpreter should transfer control to the handler and execute - /// its code. After the handler completes, call the appropriate completion - /// method (`cleanup_complete()` for finally/fault, or enter catch mode). - ExecuteHandler { - /// The handler to execute. - handler: HandlerMatch, - /// Whether this is the final handler (a catch that will handle the exception). - /// - /// - `true` - This is a catch handler; the exception will be handled - /// - `false` - This is a cleanup handler (finally/fault/filter); - /// more steps may follow - is_catch: bool, - }, - - /// Continue unwinding to the next stack frame. - /// - /// No handler was found in the current frame (or cleanup is complete). - /// Pop the frame and process the next one. - ContinueUnwind, - - /// The exception is unhandled. - /// - /// No handler was found in any stack frame. The exception should be - /// reported as unhandled and execution terminated (or handled by a - /// global exception handler if one exists). - UnhandledException, - - /// Control has been transferred to a catch handler. - /// - /// All cleanup handlers have executed and the exception is now being - /// handled. The interpreter should continue execution at the handler - /// location. - HandlerEntered { - /// The method containing the handler. - method: Token, - /// The IL offset where the handler begins. - offset: u32, - }, -} - -/// Stack unwinder for exception propagation. -/// -/// This component manages the process of unwinding the call stack when an exception -/// is thrown, ensuring proper execution of finally and fault blocks in the correct -/// order. It tracks the unwind state, pending cleanup handlers, and builds the -/// exception's stack trace. -/// -/// # Usage -/// -/// The unwinder is used in a loop, processing each stack frame: -/// -/// ```ignore -/// let mut unwinder = StackUnwinder::new(); -/// unwinder.begin_unwind(&exception_info); -/// -/// loop { -/// let result = unwinder.process_frame(&frame, &clauses, exception_type); -/// match result { -/// UnwindStepResult::ExecuteHandler { handler, is_catch } => { -/// // Execute the handler... -/// if !is_catch { -/// unwinder.cleanup_complete(); -/// } -/// } -/// UnwindStepResult::ContinueUnwind => { -/// // Pop frame and continue -/// } -/// UnwindStepResult::HandlerEntered { method, offset } => { -/// // Transfer control to catch handler -/// break; -/// } -/// UnwindStepResult::UnhandledException => { -/// // No handler found -/// break; -/// } -/// } -/// } -/// ``` -/// -/// # State Management -/// -/// The unwinder maintains internal state across multiple calls. Use [`reset()`](Self::reset) -/// to clear all state and start fresh. -#[derive(Clone, Debug, Default)] -pub struct StackUnwinder { - /// Current unwind state machine state. - state: UnwindState, - - /// Queue of pending handlers to execute during unwinding. - pending_handlers: Vec, - - /// Stack trace being built as frames are unwound. - stack_trace: Vec, -} - -/// Internal state machine for the stack unwinder. -/// -/// This enum tracks where the unwinder is in the unwind process, enabling -/// the state to persist across multiple method calls. -#[derive(Clone, Debug, Default)] -enum UnwindState { - /// Not currently unwinding. - /// - /// The unwinder is idle and ready to begin a new unwind operation. - #[default] - Idle, - - /// Searching for a handler. - /// - /// The unwinder is actively searching stack frames for a matching - /// catch or filter handler. - Searching { - /// The type token of the exception being searched for. - exception_type: Token, - }, - - /// Executing cleanup handlers before entering a catch or continuing. - /// - /// Finally and fault handlers are being executed. Once all cleanup - /// is done, either enter the target handler or continue unwinding. - ExecutingCleanup { - /// The type token of the exception being handled. - exception_type: Token, - /// The catch/filter handler to enter after cleanup, if found. - target: Option, - }, - - /// A handler has been found and is ready to be entered. - /// - /// All cleanup is complete and control should transfer to the handler. - HandlerFound { - /// The handler that will handle the exception. - target: HandlerTarget, - }, -} - -/// A handler queued for execution during unwinding. -/// -/// This internal structure pairs a handler with its containing method, -/// enabling the unwinder to track which handlers need to execute and -/// in which order during the cleanup phase of exception handling. -#[derive(Clone, Debug)] -struct PendingHandler { - /// The handler match information describing the handler type and location. - handler: HandlerMatch, - /// The method token identifying the method containing this handler. - method: Token, -} - -/// Information about the target catch/filter handler. -/// -/// When a catch or filter handler is found during the search phase, this -/// structure stores the information needed to transfer control to it after -/// all cleanup handlers have executed. -#[derive(Clone, Debug)] -struct HandlerTarget { - /// The method token identifying the method containing the handler. - method: Token, - /// The IL offset where the handler code begins. - offset: u32, - /// The type of handler (catch for type-based, filter for condition-based). - handler_type: HandlerType, -} - -/// Type of exception handler (for target tracking). -/// -/// This enum distinguishes between the two types of exception handlers -/// that can actually handle an exception (as opposed to cleanup handlers -/// like finally and fault which only run but don't catch). -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum HandlerType { - /// A catch handler that matches based on exception type compatibility. - Catch, - /// A filter handler that passed its runtime condition evaluation. - Filter, -} - -impl StackUnwinder { - /// Creates a new stack unwinder in the idle state. - /// - /// # Returns - /// - /// A new `StackUnwinder` ready to begin unwinding. - #[must_use] - pub fn new() -> Self { - Self::default() - } - - /// Checks if an unwind operation is in progress. - /// - /// # Returns - /// - /// `true` if the unwinder is actively processing an exception, - /// `false` if it is idle. - #[must_use] - pub fn is_unwinding(&self) -> bool { - !matches!(self.state, UnwindState::Idle) - } - - /// Begins unwinding for a thrown exception. - /// - /// This initializes the unwinder state, clears any previous pending handlers, - /// and starts building a new stack trace from the throw location. - /// - /// # Arguments - /// - /// * `exception_info` - Information about the thrown exception, including - /// its type and throw location - /// - /// # State Change - /// - /// Transitions from `Idle` to `Searching` state. - pub fn begin_unwind(&mut self, exception_info: &ExceptionInfo) { - self.state = UnwindState::Searching { - exception_type: exception_info.type_token, - }; - self.pending_handlers.clear(); - self.stack_trace.clear(); - self.stack_trace.push(exception_info.throw_location); - } - - /// Processes a stack frame during exception unwinding. - /// - /// This method examines the exception clauses of the current frame to find: - /// - Matching catch handlers (by type compatibility) - /// - Filter handlers (for later evaluation) - /// - Finally/fault handlers (queued for cleanup) - /// - /// # Arguments - /// - /// * `frame` - The current call frame being examined - /// * `clauses` - Exception clauses for the method in this frame - /// * `exception_type` - Type token of the exception being handled - /// * `is_type_compatible` - A function that checks if an exception type is assignable - /// to a catch type. Use [`EmulationContext::is_type_compatible`](crate::emulation::EmulationContext::is_type_compatible). - /// - /// # Returns - /// - /// An [`UnwindStepResult`] indicating the next action: - /// - `ExecuteHandler` - Run a handler (cleanup or catch) - /// - `ContinueUnwind` - No handlers in this frame, continue to caller - /// - /// # Stack Trace - /// - /// Each processed frame is automatically added to the exception's stack trace. - pub fn process_frame( - &mut self, - frame: &ThreadCallFrame, - clauses: &[ExceptionClause], - exception_type: Token, - is_type_compatible: F, - ) -> UnwindStepResult - where - F: Fn(Token, Token) -> bool, - { - let current_offset = frame.ip(); - let method = frame.method(); - - // Add to stack trace - self.stack_trace - .push(InstructionLocation::new(method, current_offset)); - - // Find handlers for this frame - let mut finally_handlers = Vec::new(); - let mut fault_handlers = Vec::new(); - let mut catch_handler = None; - let mut filter_handler = None; - - for clause in clauses { - if !clause.is_in_try(current_offset) { - continue; - } - - match clause { - ExceptionClause::Catch { catch_type, .. } => { - // Check type compatibility using the provided checker - if is_type_compatible(exception_type, *catch_type) && catch_handler.is_none() { - catch_handler = Some(HandlerMatch::Catch { - method, - handler_offset: clause.handler_offset(), - }); - } - } - - ExceptionClause::Filter { filter_offset, .. } => { - if filter_handler.is_none() { - filter_handler = Some(HandlerMatch::Filter { - method, - filter_offset: *filter_offset, - handler_offset: clause.handler_offset(), - }); - } - } - - ExceptionClause::Finally { .. } => { - finally_handlers.push(HandlerMatch::Finally { - method, - handler_offset: clause.handler_offset(), - handler_length: clause.handler_length(), - continue_search_after: true, - }); - } - - ExceptionClause::Fault { .. } => { - fault_handlers.push(HandlerMatch::Fault { - method, - handler_offset: clause.handler_offset(), - handler_length: clause.handler_length(), - }); - } - } - } - - // If we have a filter, we need to evaluate it first - if let Some(filter) = filter_handler { - // Queue finally/fault handlers, then filter - for h in finally_handlers { - self.pending_handlers - .push(PendingHandler { handler: h, method }); - } - for h in fault_handlers { - self.pending_handlers - .push(PendingHandler { handler: h, method }); - } - - return UnwindStepResult::ExecuteHandler { - handler: filter, - is_catch: false, // Filter needs evaluation first - }; - } - - // If we have a catch handler - if let Some(catch) = catch_handler { - // Queue finally/fault handlers to run before entering catch - for h in finally_handlers { - self.pending_handlers - .push(PendingHandler { handler: h, method }); - } - for h in fault_handlers { - self.pending_handlers - .push(PendingHandler { handler: h, method }); - } - - // If there are pending handlers, run them first - if !self.pending_handlers.is_empty() { - // Store the catch as the target - self.state = UnwindState::ExecutingCleanup { - exception_type, - target: Some(HandlerTarget { - method, - offset: match &catch { - HandlerMatch::Catch { handler_offset, .. } => *handler_offset, - _ => unreachable!(), - }, - handler_type: HandlerType::Catch, - }), - }; - - let next = self.pending_handlers.remove(0); - return UnwindStepResult::ExecuteHandler { - handler: next.handler, - is_catch: false, - }; - } - - return UnwindStepResult::ExecuteHandler { - handler: catch, - is_catch: true, - }; - } - - // No catch in this frame - queue cleanup and continue unwinding - for h in finally_handlers { - self.pending_handlers - .push(PendingHandler { handler: h, method }); - } - for h in fault_handlers { - self.pending_handlers - .push(PendingHandler { handler: h, method }); - } - - // If we have cleanup handlers, execute them - if !self.pending_handlers.is_empty() { - self.state = UnwindState::ExecutingCleanup { - exception_type, - target: None, - }; - - let next = self.pending_handlers.remove(0); - return UnwindStepResult::ExecuteHandler { - handler: next.handler, - is_catch: false, - }; - } - - // No handlers in this frame, continue unwinding - UnwindStepResult::ContinueUnwind - } - - /// Called when a cleanup handler (finally/fault) completes execution. - /// - /// After executing a finally or fault handler, call this method to - /// continue the unwind process. It will return the next handler to - /// execute, or indicate that the catch handler can be entered. - /// - /// # Returns - /// - /// An [`UnwindStepResult`] indicating the next action: - /// - `ExecuteHandler` - Another cleanup handler to execute - /// - `HandlerEntered` - All cleanup done, entering catch handler - /// - `ContinueUnwind` - All cleanup done, continue to next frame - pub fn cleanup_complete(&mut self) -> UnwindStepResult { - // Check if there are more pending handlers - if !self.pending_handlers.is_empty() { - let next = self.pending_handlers.remove(0); - return UnwindStepResult::ExecuteHandler { - handler: next.handler, - is_catch: false, - }; - } - - // Check if we have a target handler to enter - if let UnwindState::ExecutingCleanup { - target: Some(target), - .. - } = &self.state - { - let result = UnwindStepResult::HandlerEntered { - method: target.method, - offset: target.offset, - }; - self.state = UnwindState::Idle; - return result; - } - - // Continue unwinding to next frame - UnwindStepResult::ContinueUnwind - } - - /// Called when a filter handler completes evaluation. - /// - /// After executing a filter clause's evaluation code, call this method - /// with the result (from the evaluation stack). The filter returns - /// an integer: non-zero means accept, zero means reject. - /// - /// # Arguments - /// - /// * `accepted` - Whether the filter accepted the exception (non-zero result) - /// * `handler_offset` - The IL offset of the handler to enter if accepted - /// * `method` - The method containing the filter and handler - /// - /// # Returns - /// - /// An [`UnwindStepResult`] indicating the next action: - /// - If accepted: `ExecuteHandler` (cleanup) or `HandlerEntered` - /// - If rejected: `ContinueUnwind` to search for other handlers - pub fn filter_complete( - &mut self, - accepted: bool, - handler_offset: u32, - method: Token, - ) -> UnwindStepResult { - if accepted { - // Filter accepted - enter the handler - if !self.pending_handlers.is_empty() { - // Run cleanup first - if let UnwindState::Searching { exception_type } = self.state { - self.state = UnwindState::ExecutingCleanup { - exception_type, - target: Some(HandlerTarget { - method, - offset: handler_offset, - handler_type: HandlerType::Filter, - }), - }; - } - - let next = self.pending_handlers.remove(0); - return UnwindStepResult::ExecuteHandler { - handler: next.handler, - is_catch: false, - }; - } - - UnwindStepResult::HandlerEntered { - method, - offset: handler_offset, - } - } else { - // Filter rejected - continue searching - UnwindStepResult::ContinueUnwind - } - } - - /// Completes the unwind when no handler is found. - /// - /// Called when all stack frames have been searched without finding a - /// matching handler. This resets the unwinder to idle state. - /// - /// # Returns - /// - /// Always returns [`UnwindStepResult::UnhandledException`]. - pub fn complete_unhandled(&mut self) -> UnwindStepResult { - self.state = UnwindState::Idle; - UnwindStepResult::UnhandledException - } - - /// Resets the unwinder to its initial state. - /// - /// Clears all pending handlers, the stack trace, and returns the - /// unwinder to idle state. Call this to abandon an unwind operation - /// or to prepare for a new one. - pub fn reset(&mut self) { - self.state = UnwindState::Idle; - self.pending_handlers.clear(); - self.stack_trace.clear(); - } - - /// Gets the stack trace built during unwinding. - /// - /// This returns all instruction locations added during the unwind - /// process, starting from the throw location. - /// - /// # Returns - /// - /// A slice of instruction locations representing the call stack. - #[must_use] - pub fn stack_trace(&self) -> &[InstructionLocation] { - &self.stack_trace - } - - /// Applies exception state changes when entering a catch handler. - /// - /// This helper method updates the thread's exception state when - /// control transfers to a catch handler. - /// - /// # Arguments - /// - /// * `exception_state` - The thread's exception state to update - pub fn enter_catch_handler(exception_state: &mut ThreadExceptionState) { - exception_state.enter_catch(); - } - - /// Applies exception state changes when entering a finally handler. - /// - /// This helper method updates the thread's exception state when - /// control transfers to a finally handler. - /// - /// # Arguments - /// - /// * `exception_state` - The thread's exception state to update - pub fn enter_finally_handler(exception_state: &mut ThreadExceptionState) { - exception_state.enter_finally(); - } - - /// Applies exception state changes when exiting a finally handler. - /// - /// This helper method updates the thread's exception state when - /// a finally handler completes (via `endfinally` instruction). - /// - /// # Arguments - /// - /// * `exception_state` - The thread's exception state to update - /// * `saved_exception` - Exception to restore if rethrow was requested - pub fn exit_finally_handler( - exception_state: &mut ThreadExceptionState, - saved_exception: Option, - ) { - exception_state.exit_finally(saved_exception); - } -} - -/// Builder for constructing ordered exception handler sequences. -/// -/// This helper is used to build a sequence of handlers that should execute -/// during exception handling. Handlers are added in execution order and -/// can be retrieved as a vector. -/// -/// # Example -/// -/// ```ignore -/// let mut builder = UnwindSequenceBuilder::new(); -/// -/// // Add cleanup handlers first -/// builder.add_handler(HandlerMatch::Finally { ... }); -/// builder.add_handler(HandlerMatch::Finally { ... }); -/// -/// // Then the catch handler -/// builder.add_handler(HandlerMatch::Catch { ... }); -/// -/// // Get the final sequence -/// let handlers = builder.build(); -/// ``` -#[derive(Clone, Debug, Default)] -pub struct UnwindSequenceBuilder { - /// Handlers in execution order. - handlers: Vec, -} - -impl UnwindSequenceBuilder { - /// Creates a new empty sequence builder. - /// - /// # Returns - /// - /// A new `UnwindSequenceBuilder` with no handlers. - #[must_use] - pub fn new() -> Self { - Self::default() - } - - /// Adds a single handler to the sequence. - /// - /// Handlers are executed in the order they are added. - /// - /// # Arguments - /// - /// * `handler` - The handler to add to the sequence - pub fn add_handler(&mut self, handler: HandlerMatch) { - self.handlers.push(handler); - } - - /// Adds multiple handlers to the sequence. - /// - /// The handlers are added in iteration order and will execute - /// after any previously added handlers. - /// - /// # Arguments - /// - /// * `handlers` - An iterator of handlers to add - pub fn add_handlers(&mut self, handlers: impl IntoIterator) { - self.handlers.extend(handlers); - } - - /// Builds the final handler sequence. - /// - /// Consumes the builder and returns the handlers as a vector. - /// - /// # Returns - /// - /// A vector of handlers in execution order. - #[must_use] - pub fn build(self) -> Vec { - self.handlers - } - - /// Checks if the sequence is empty. - /// - /// # Returns - /// - /// `true` if no handlers have been added. - #[must_use] - pub fn is_empty(&self) -> bool { - self.handlers.is_empty() - } - - /// Gets the number of handlers in the sequence. - /// - /// # Returns - /// - /// The count of handlers added to the builder. - #[must_use] - pub fn len(&self) -> usize { - self.handlers.len() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::{emulation::HeapRef, metadata::typesystem::CilFlavor}; - - fn create_test_exception() -> ExceptionInfo { - ExceptionInfo::new( - HeapRef::new(1), - Token::new(0x01000010), - InstructionLocation::new(Token::new(0x06000001), 0x20), - ) - } - - fn create_test_frame(method: Token, offset: u32) -> ThreadCallFrame { - let mut frame = ThreadCallFrame::new(method, None, 0, vec![CilFlavor::I4], vec![], false); - frame.set_ip(offset); - frame - } - - /// Simple type checker for tests: exact match or catch-all System.Exception. - fn test_type_checker(exception_type: Token, catch_type: Token) -> bool { - // Exact match - if exception_type == catch_type { - return true; - } - // System.Exception (0x0100_0001) catches everything - if catch_type.value() == 0x0100_0001 { - return true; - } - false - } - - #[test] - fn test_unwinder_begin() { - let mut unwinder = StackUnwinder::new(); - let exception = create_test_exception(); - - assert!(!unwinder.is_unwinding()); - - unwinder.begin_unwind(&exception); - - assert!(unwinder.is_unwinding()); - assert_eq!(unwinder.stack_trace().len(), 1); - } - - #[test] - fn test_unwinder_reset() { - let mut unwinder = StackUnwinder::new(); - let exception = create_test_exception(); - - unwinder.begin_unwind(&exception); - assert!(unwinder.is_unwinding()); - - unwinder.reset(); - assert!(!unwinder.is_unwinding()); - assert!(unwinder.stack_trace().is_empty()); - } - - #[test] - fn test_process_frame_with_catch() { - let mut unwinder = StackUnwinder::new(); - let exception = create_test_exception(); - - unwinder.begin_unwind(&exception); - - let method = Token::new(0x06000001); - let frame = create_test_frame(method, 0x15); - - let clauses = vec![ExceptionClause::Catch { - try_offset: 0x10, - try_length: 0x20, - handler_offset: 0x30, - handler_length: 0x10, - catch_type: Token::new(0x0100_0001), // Base exception - }]; - - let result = - unwinder.process_frame(&frame, &clauses, exception.type_token, test_type_checker); - - match result { - UnwindStepResult::ExecuteHandler { is_catch, .. } => { - assert!(is_catch); - } - _ => panic!("Expected catch handler"), - } - } - - #[test] - fn test_process_frame_with_finally() { - let mut unwinder = StackUnwinder::new(); - let exception = create_test_exception(); - - unwinder.begin_unwind(&exception); - - let method = Token::new(0x06000001); - let frame = create_test_frame(method, 0x15); - - // Only finally, no catch - let clauses = vec![ExceptionClause::Finally { - try_offset: 0x10, - try_length: 0x20, - handler_offset: 0x30, - handler_length: 0x10, - }]; - - let result = - unwinder.process_frame(&frame, &clauses, exception.type_token, test_type_checker); - - match result { - UnwindStepResult::ExecuteHandler { handler, is_catch } => { - assert!(!is_catch); - assert!(matches!(handler, HandlerMatch::Finally { .. })); - } - _ => panic!("Expected finally handler"), - } - } - - #[test] - fn test_process_frame_continue_unwind() { - let mut unwinder = StackUnwinder::new(); - let exception = create_test_exception(); - - unwinder.begin_unwind(&exception); - - let method = Token::new(0x06000001); - let frame = create_test_frame(method, 0x50); // Outside any try block - - let clauses = vec![ExceptionClause::Catch { - try_offset: 0x10, - try_length: 0x20, - handler_offset: 0x30, - handler_length: 0x10, - catch_type: Token::new(0x01000010), - }]; - - let result = - unwinder.process_frame(&frame, &clauses, exception.type_token, test_type_checker); - - assert!(matches!(result, UnwindStepResult::ContinueUnwind)); - } - - #[test] - fn test_cleanup_complete() { - let mut unwinder = StackUnwinder::new(); - - // No pending handlers - let result = unwinder.cleanup_complete(); - assert!(matches!(result, UnwindStepResult::ContinueUnwind)); - } - - #[test] - fn test_unwind_sequence_builder() { - let mut builder = UnwindSequenceBuilder::new(); - - assert!(builder.is_empty()); - - builder.add_handler(HandlerMatch::Finally { - method: Token::new(0x06000001), - handler_offset: 0x30, - handler_length: 0x10, - continue_search_after: true, - }); - - builder.add_handler(HandlerMatch::Catch { - method: Token::new(0x06000001), - handler_offset: 0x40, - }); - - assert_eq!(builder.len(), 2); - - let sequence = builder.build(); - assert_eq!(sequence.len(), 2); - } -} diff --git a/dotscope/src/emulation/loader/peloader.rs b/dotscope/src/emulation/loader/peloader.rs index b8978f1f..da3886c3 100644 --- a/dotscope/src/emulation/loader/peloader.rs +++ b/dotscope/src/emulation/loader/peloader.rs @@ -117,6 +117,16 @@ pub struct PeLoaderConfig { /// its preferred base, the loader applies relocations to fix up absolute /// addresses in the code and data sections. pub apply_relocations: bool, + + /// Ceiling on the image buffer this loader will materialise, in bytes. + /// + /// `SizeOfImage` is a `u32` taken verbatim from the optional header of an + /// attacker-supplied file, and it alone sizes the zero-initialised image buffer. A few + /// hundred bytes of input can therefore ask for 4 GiB. `alloc_zeroed` makes that cheap in + /// resident memory but not free: a hard cgroup limit, `vm.overcommit_memory=2`, Windows + /// commit charge, a 32-bit build, or enough concurrent loads all turn it into a real + /// failure. + pub max_image_size: u64, } impl Default for PeLoaderConfig { @@ -125,6 +135,7 @@ impl Default for PeLoaderConfig { base_address: None, apply_permissions: true, apply_relocations: true, + max_image_size: 512 * 1024 * 1024, } } } @@ -685,10 +696,48 @@ impl PeLoader { let mut sections = Vec::new(); let mut section_infos = Vec::new(); + // Validate the declared image size before committing an allocation to it. + // + // Two independent bounds, because either alone is too permissive: the configured + // ceiling caps the absolute cost, and the section extents catch a header that claims + // far more than its own sections describe. Sections are the authority on what the + // image actually needs — anything past the highest section end is memory no section + // maps into. + if size_of_image > self.config.max_image_size { + return Err(malformed_error!( + "SizeOfImage {} exceeds the {} byte loader limit", + size_of_image, + self.config.max_image_size + )); + } + + let section_extent = pe + .sections + .iter() + .map(|section| { + u64::from(section.virtual_address).saturating_add(u64::from( + section.virtual_size.max(section.size_of_raw_data), + )) + }) + .max() + .unwrap_or(0); + // Rounded up generously: SizeOfImage is SectionAlignment-aligned and may legitimately + // exceed the last section's end by up to one alignment unit, plus the headers. + let plausible = section_extent.saturating_add(0x100_0000); + if section_extent > 0 && size_of_image > plausible { + return Err(malformed_error!( + "SizeOfImage {} is not plausible for sections ending at {}", + size_of_image, + section_extent + )); + } + // Create the full image data buffer // ToDo: Switch this and the emulator to use mmap for having a disk-backed file if it is very large - #[allow(clippy::cast_possible_truncation)] - let mut image_data = vec![0u8; size_of_image as usize]; + let image_size = usize::try_from(size_of_image).map_err(|_| { + malformed_error!("SizeOfImage {} exceeds the address width", size_of_image) + })?; + let mut image_data = vec![0u8; image_size]; // Copy headers let headers_size = pe diff --git a/dotscope/src/emulation/memory/addressspace.rs b/dotscope/src/emulation/memory/addressspace.rs index 2b4ee5b4..47bdeb6a 100644 --- a/dotscope/src/emulation/memory/addressspace.rs +++ b/dotscope/src/emulation/memory/addressspace.rs @@ -32,9 +32,14 @@ //! assert_eq!(data, vec![0xDE, 0xAD, 0xBE, 0xEF]); //! ``` -use std::sync::{ - atomic::{AtomicU64, Ordering}, - Arc, RwLock, +use std::{ + collections::BTreeMap, + ops::Deref, + result::Result as StdResult, + sync::{ + atomic::{AtomicU64, AtomicUsize, Ordering}, + Arc, RwLock, + }, }; use imbl::HashMap as ImHashMap; @@ -171,7 +176,7 @@ impl Default for SharedHeap { } } -impl std::ops::Deref for SharedHeap { +impl Deref for SharedHeap { type Target = ManagedHeap; fn deref(&self) -> &Self::Target { @@ -179,6 +184,13 @@ impl std::ops::Deref for SharedHeap { } } +/// Default ceiling on unmanaged allocations for a newly created address space. +/// +/// Matches the `max_unmanaged_bytes` default in +/// [`ProcessConfig`](crate::emulation::process::ProcessConfig) so that an address space built +/// directly and one built through the process builder behave the same. +pub const DEFAULT_MAX_UNMANAGED_BYTES: usize = 64 * 1024 * 1024; + /// Metadata for a pinned managed array whose native address aliases /// the managed heap data. Reads and writes through the native address /// are transparently delegated to the managed heap's `Vec`. @@ -254,8 +266,15 @@ pub struct AddressSpace { /// Managed .NET heap (shared across threads). heap: SharedHeap, - /// Memory regions (PE images, mapped data, etc.). - regions: RwLock>, + /// Memory regions (PE images, mapped data, etc.), keyed by base address. + /// + /// Ordered by base so a containing region is found with `range(..=address).next_back()` + /// rather than a scan. Every access an emulated program makes to unmanaged memory goes + /// through that lookup, and the region count is attacker-growable: `Marshal.WriteByte` + /// to an unmapped address maps a fresh region, so a linear structure makes both the + /// per-access lookup and the insert-time overlap check scale with how many regions the + /// program has already created. + regions: RwLock>, /// Static field storage. statics: StaticFieldStorage, @@ -294,6 +313,20 @@ pub struct AddressSpace { /// /// Uses `imbl::HashMap` for O(1) fork via structural sharing. pinned_arrays: RwLock>, + + /// Bytes currently committed to unmanaged allocations. + /// + /// Unmanaged allocations (`localloc`, `Marshal.AllocHGlobal`, `Marshal.AllocCoTaskMem`, + /// `VirtualAlloc`) commit real host pages and are not charged against the managed heap + /// budget, so they need their own accounting. + unmanaged_bytes: AtomicUsize, + + /// Ceiling on [`unmanaged_bytes`](Self::unmanaged_bytes). + /// + /// Atomic so it can be adjusted after construction through + /// [`set_max_unmanaged_bytes`](Self::set_max_unmanaged_bytes), which the process builder + /// uses to apply the configured limit to an address space it does not own exclusively. + max_unmanaged_bytes: AtomicUsize, } impl AddressSpace { @@ -321,13 +354,15 @@ impl AddressSpace { pub fn with_config(heap_size: usize, address_space_size: u64) -> Self { Self { heap: SharedHeap::new(heap_size), - regions: RwLock::new(Vec::new()), + regions: RwLock::new(BTreeMap::new()), statics: StaticFieldStorage::new(), next_address: AtomicU64::new(0x1000_0000), // Start at 256MB size: address_space_size, protection_overrides: RwLock::new(ImHashMap::new()), monitor_locks: RwLock::new(ImHashMap::new()), pinned_arrays: RwLock::new(ImHashMap::new()), + unmanaged_bytes: AtomicUsize::new(0), + max_unmanaged_bytes: AtomicUsize::new(DEFAULT_MAX_UNMANAGED_BYTES), } } @@ -344,13 +379,15 @@ impl AddressSpace { pub fn with_heap(heap: SharedHeap) -> Self { Self { heap, - regions: RwLock::new(Vec::new()), + regions: RwLock::new(BTreeMap::new()), statics: StaticFieldStorage::new(), next_address: AtomicU64::new(0x1000_0000), size: 0x1_0000_0000, protection_overrides: RwLock::new(ImHashMap::new()), monitor_locks: RwLock::new(ImHashMap::new()), pinned_arrays: RwLock::new(ImHashMap::new()), + unmanaged_bytes: AtomicUsize::new(0), + max_unmanaged_bytes: AtomicUsize::new(DEFAULT_MAX_UNMANAGED_BYTES), } } @@ -441,8 +478,17 @@ impl AddressSpace { }) })?; - // Check for overlaps - for existing in regions.iter() { + // Only two regions can overlap a new one in a base-ordered map: the last starting at + // or before it, and the first starting after it. Everything else starts earlier and + // ends before that predecessor does, or starts later than the successor. + let base = region.base(); + let predecessor = regions.range(..=base).next_back().map(|(_, r)| r); + let successor = regions + .range(base.saturating_add(1)..) + .next() + .map(|(_, r)| r); + + for existing in [predecessor, successor].into_iter().flatten() { if Self::regions_overlap(existing, ®ion) { return Err(EmulationError::InvalidAddress { address, @@ -452,7 +498,7 @@ impl AddressSpace { } } - regions.push(region); + regions.insert(base, region); Ok(()) } @@ -513,8 +559,7 @@ impl AddressSpace { }) })?; - if let Some(pos) = regions.iter().position(|r| r.base() == base) { - regions.remove(pos); + if regions.remove(&base).is_some() { Ok(()) } else { Err(EmulationError::InvalidAddress { @@ -541,29 +586,97 @@ impl AddressSpace { return result; } + // Validate before allocating. `read_mapped` below checks the range too, but only + // after it has been handed a `len`-sized buffer — so an unmapped address with an + // attacker-chosen length would commit that much host memory before being refused. + // `cpblk` reaches here with a size straight off the evaluation stack, so "refused" + // has to cost nothing. Same rule `init_block` follows: validate first, fill second. + if !self.covers_range(address, len)? { + return Err(EmulationError::InvalidAddress { + address, + reason: format!("read of {len} bytes is not fully mapped"), + } + .into()); + } + let regions = self.regions.read().map_err(|_| { Error::from(EmulationError::InternalError { description: "region lock poisoned".to_string(), }) })?; - for region in regions.iter() { - if region.contains_range(address, len) { - return region.read(address, len).ok_or_else(|| { - EmulationError::InvalidAddress { - address, - reason: "read failed".to_string(), - } - .into() - }); + let mut buffer = vec![0u8; len]; + self.read_mapped(®ions, address, &mut buffer)?; + Ok(buffer) + } + + /// Reads into a caller-supplied buffer, allocating nothing. + /// + /// `read` hands back an owned `Vec`, which makes a 1-, 2-, 4- or 8-byte `ldind` — the + /// interpreter's inner loop — pay a malloc and a free per load. The fixed-size accessors + /// below fill a stack array through this instead. + /// + /// # Errors + /// + /// Returns an error if the range is unmapped, unreadable under its protection, or the + /// region lock is poisoned. + pub fn read_into(&self, address: u64, dest: &mut [u8]) -> Result<()> { + if let Some(result) = self.read_pinned(address, dest.len()) { + let bytes = result?; + if bytes.len() != dest.len() { + return Err(EmulationError::InvalidAddress { + address, + reason: "pinned read length mismatch".to_string(), + } + .into()); } + dest.copy_from_slice(&bytes); + return Ok(()); } - Err(EmulationError::InvalidAddress { - address, - reason: "address not mapped".to_string(), + let regions = self.regions.read().map_err(|_| { + Error::from(EmulationError::InternalError { + description: "region lock poisoned".to_string(), + }) + })?; + + self.read_mapped(®ions, address, dest) + } + + /// Shared body of [`Self::read`] and [`Self::read_into`], with the regions lock held. + fn read_mapped( + &self, + regions: &BTreeMap, + address: u64, + dest: &mut [u8], + ) -> Result<()> { + let Some(region) = Self::region_containing(regions, address) else { + return Err(EmulationError::InvalidAddress { + address, + reason: "address not mapped".to_string(), + } + .into()); + }; + + if !region.contains_range(address, dest.len()) { + return Err(EmulationError::InvalidAddress { + address, + reason: "address not mapped".to_string(), + } + .into()); + } + + self.check_access(region, address, dest.len(), MemoryProtection::READ)?; + + if region.read_into(address, dest) { + Ok(()) + } else { + Err(EmulationError::InvalidAddress { + address, + reason: "read failed".to_string(), + } + .into()) } - .into()) } /// Writes bytes to any mapped region. @@ -589,24 +702,149 @@ impl AddressSpace { }) })?; - for region in regions.iter() { - if region.contains_range(address, data.len()) { - if region.write(address, data) { - return Ok(()); + let Some(region) = Self::region_containing(®ions, address) else { + return Err(EmulationError::InvalidAddress { + address, + reason: "address not mapped".to_string(), + } + .into()); + }; + + if !region.contains_range(address, data.len()) { + return Err(EmulationError::InvalidAddress { + address, + reason: "address not mapped".to_string(), + } + .into()); + } + + self.check_access(region, address, data.len(), MemoryProtection::WRITE)?; + + if region.write(address, data) { + Ok(()) + } else { + // Protection is rejected above, so reaching here means the paged write itself + // failed — a range or page-level fault, not a permission one. + Err(EmulationError::InvalidAddress { + address, + reason: "write failed".to_string(), + } + .into()) + } + } + + /// Rejects an access whose protection does not permit it. + /// + /// Protection is per page, so an access is checked at the first and last page it touches; + /// a region's pages carry uniform protection unless `VirtualProtect` has overridden some + /// of them, and an override is page-aligned, so those two are what can differ. + /// + /// `GUARD` faults on any access regardless of `required`, matching a Windows guard page. + /// + /// Takes the region rather than looking it up, because both callers already hold the + /// regions read lock — re-acquiring it here would be a recursive read acquisition, which + /// `std::sync::RwLock` may deadlock on when a writer is queued between the two. + /// + /// # Errors + /// + /// Returns [`EmulationError::AccessViolation`] when the access is not permitted. + fn check_access( + &self, + region: &MemoryRegion, + address: u64, + len: usize, + required: MemoryProtection, + ) -> Result<()> { + if len == 0 { + return Ok(()); + } + + let last = address.saturating_add(len.saturating_sub(1) as u64); + for probe in [address, last] { + let Some(protection) = self.protection_of(region, probe) else { + continue; + }; + + if protection.contains(MemoryProtection::GUARD) { + return Err(EmulationError::AccessViolation { + address: probe, + reason: "guard page".to_string(), } - return Err(EmulationError::InvalidAddress { - address, - reason: "write failed (possibly read-only)".to_string(), + .into()); + } + + if !protection.contains(required) { + return Err(EmulationError::AccessViolation { + address: probe, + reason: format!("protection {protection:?} does not permit {required:?}"), } .into()); } } - Err(EmulationError::InvalidAddress { - address, - reason: "address not mapped".to_string(), + Ok(()) + } + + /// Resolves the protection of one address inside a region already in hand. + /// + /// Same resolution order as [`Self::get_protection`] — `VirtualProtect` override first, + /// then the region's own (per-PE-section) protection — without the regions lookup. + fn protection_of(&self, region: &MemoryRegion, address: u64) -> Option { + let page_addr = address & !(Self::PAGE_SIZE - 1); + if let Ok(overrides) = self.protection_overrides.read() { + if let Some(&protection) = overrides.get(&page_addr) { + return Some(protection); + } } - .into()) + + region.protection_at(address).ok() + } + + /// Reads a fixed-size little-endian value without allocating. + /// + /// # Errors + /// + /// Returns an error if the range is unmapped or unreadable under its protection. + pub fn read_exact(&self, address: u64) -> Result<[u8; N]> { + let mut buffer = [0u8; N]; + self.read_into(address, &mut buffer)?; + Ok(buffer) + } + + /// Reads one byte. + /// + /// # Errors + /// + /// See [`Self::read_exact`]. + pub fn read_u8(&self, address: u64) -> Result { + Ok(u8::from_le_bytes(self.read_exact::<1>(address)?)) + } + + /// Reads a little-endian `u16`. + /// + /// # Errors + /// + /// See [`Self::read_exact`]. + pub fn read_u16(&self, address: u64) -> Result { + Ok(u16::from_le_bytes(self.read_exact::<2>(address)?)) + } + + /// Reads a little-endian `u32`. + /// + /// # Errors + /// + /// See [`Self::read_exact`]. + pub fn read_u32(&self, address: u64) -> Result { + Ok(u32::from_le_bytes(self.read_exact::<4>(address)?)) + } + + /// Reads a little-endian `u64`. + /// + /// # Errors + /// + /// See [`Self::read_exact`]. + pub fn read_u64(&self, address: u64) -> Result { + Ok(u64::from_le_bytes(self.read_exact::<8>(address)?)) } /// Returns `true` if the address is within a mapped region. @@ -629,7 +867,7 @@ impl AddressSpace { let Ok(regions) = self.regions.read() else { return false; }; - regions.iter().any(|r| r.contains(address)) + Self::region_containing(®ions, address).is_some() } /// Returns the region containing the given address, if any. @@ -640,7 +878,7 @@ impl AddressSpace { #[must_use] pub fn get_region(&self, address: u64) -> Option { let regions = self.regions.read().ok()?; - regions.iter().find(|r| r.contains(address)).cloned() + Self::region_containing(®ions, address).cloned() } /// Returns the memory protection flags for an address. @@ -665,10 +903,7 @@ impl AddressSpace { // Fall back to region's inherent protection let regions = self.regions.read().ok()?; - regions - .iter() - .find(|r| r.contains(address)) - .and_then(|r| r.protection_at(address).ok()) + Self::region_containing(®ions, address).and_then(|r| r.protection_at(address).ok()) } /// Sets the memory protection for a range of addresses. @@ -792,8 +1027,78 @@ impl AddressSpace { /// /// Returns an error if the mapping fails. pub fn alloc_unmanaged(&self, size: usize) -> Result { + // Unmanaged allocations commit real host pages and are invisible to the managed heap + // budget, so they are charged here — before the region is constructed, since + // `MemoryRegion::unmanaged_alloc` allocates the backing store eagerly. + let max = self.max_unmanaged_bytes.load(Ordering::Relaxed); + let current = self.unmanaged_bytes.load(Ordering::Relaxed); + if current.saturating_add(size) > max { + return Err(EmulationError::HeapMemoryLimitExceeded { + current, + limit: max, + } + .into()); + } + let region = MemoryRegion::unmanaged_alloc(0, size); - self.map(region) + let address = self.map(region)?; + self.unmanaged_bytes.fetch_add(size, Ordering::Relaxed); + Ok(address) + } + + /// Maps a zero-filled region at a fixed address, charged against the unmanaged budget. + /// + /// This backs the `Marshal.Write*` auto-allocation path, where emulated code writes to an + /// address that is not mapped and the emulator materialises memory there rather than + /// aborting. The address is attacker-chosen and the loop is driven by ordinary CIL, so + /// without a ceiling each unmapped page written to commits real host pages that are never + /// reclaimed. Charging the same counter as [`Self::alloc_unmanaged`] also bounds the + /// region *count*, since every auto-allocation is the same fixed size. + /// + /// # Errors + /// + /// Returns [`EmulationError::HeapMemoryLimitExceeded`] once the unmanaged budget is + /// exhausted, or an error if the range overlaps an existing mapping. + pub fn alloc_unmanaged_at(&self, address: u64, size: usize, label: &str) -> Result<()> { + let max = self.max_unmanaged_bytes.load(Ordering::Relaxed); + let current = self.unmanaged_bytes.load(Ordering::Relaxed); + if current.saturating_add(size) > max { + return Err(EmulationError::HeapMemoryLimitExceeded { + current, + limit: max, + } + .into()); + } + + let region = MemoryRegion::mapped_data( + address, + &vec![0u8; size], + label, + MemoryProtection::READ_WRITE, + ); + self.map_at(address, region)?; + self.unmanaged_bytes.fetch_add(size, Ordering::Relaxed); + Ok(()) + } + + /// Sets the ceiling on total unmanaged allocation, in bytes. + /// + /// Applied after construction because the process builder configures limits on an address + /// space it shares rather than owns. + pub fn set_max_unmanaged_bytes(&self, max: usize) { + self.max_unmanaged_bytes.store(max, Ordering::Relaxed); + } + + /// Returns the number of bytes currently committed to unmanaged allocations. + #[must_use] + pub fn unmanaged_bytes(&self) -> usize { + self.unmanaged_bytes.load(Ordering::Relaxed) + } + + /// Returns the ceiling on unmanaged allocation, in bytes. + #[must_use] + pub fn max_unmanaged_bytes(&self) -> usize { + self.max_unmanaged_bytes.load(Ordering::Relaxed) } /// Frees unmanaged memory previously allocated with [`alloc_unmanaged`](Self::alloc_unmanaged). @@ -813,14 +1118,28 @@ impl AddressSpace { }) })?; - let is_unmanaged = regions - .iter() - .any(|r| r.base() == address && r.is_unmanaged_alloc()); + let freed_size = regions + .get(&address) + .filter(|r| r.is_unmanaged_alloc()) + .map(MemoryRegion::size); drop(regions); - if is_unmanaged { - self.unmap(address) + if let Some(size) = freed_size { + self.unmap(address)?; + // Return the budget so an alloc/free cycle cannot ratchet the accounted total + // upwards and starve later allocations. + // + // Saturating rather than a plain `fetch_sub`: a region can be flagged as an + // unmanaged allocation without having been charged here (a forked address space + // carries regions over, for instance), and a wrapping subtract would underflow the + // counter to near `usize::MAX` and reject every later allocation. + let _ = + self.unmanaged_bytes + .try_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + Some(current.saturating_sub(size)) + }); + Ok(()) } else { Err(EmulationError::InvalidAddress { address, @@ -833,11 +1152,71 @@ impl AddressSpace { /// Reserves an address range without creating a backing memory region. /// /// Used for pinned arrays where the backing store is the managed heap. - /// The returned address is guaranteed not to conflict with other allocations. - pub fn reserve_address_range(&self, size: usize) -> u64 { - let aligned_size = size.saturating_add(0xFFF) & !0xFFF; // Align to 4KB - self.next_address - .fetch_add(aligned_size as u64, Ordering::SeqCst) + /// + /// The cursor starts at `0x1000_0000` — the default `ImageBase` the C# compiler emits — + /// so it walks straight into a mapped PE image on any host that opted into + /// `ProcessBuilder::map_pe_image`. Bumping it blindly there hands out a base that aliases + /// the image: `read`/`write` consult pins before regions, so the pin would silently + /// intercept accesses to the image and redirect them into a managed array. The same + /// cursor also feeds `map`, so once it is inside an image every `alloc_unmanaged` fails + /// on the overlap check. This therefore skips past any range that overlaps a mapping. + /// + /// # Returns + /// + /// The reserved base address, or `None` if no free range of this size remains below the + /// address space limit. + pub fn reserve_address_range(&self, size: usize) -> Option { + let aligned_size = u64::try_from(size.saturating_add(0xFFF) & !0xFFF).ok()?; + let regions = self.regions.read().ok()?; + let region_end = |region: &MemoryRegion| { + region + .base() + .saturating_add(u64::try_from(region.size()).unwrap_or(u64::MAX)) + }; + + loop { + let base = self.next_address.load(Ordering::SeqCst); + let end = base.checked_add(aligned_size)?; + if end > self.size { + return None; + } + + // Regions never overlap, so the only ones that can intersect [base, end) are the + // last starting before it — if it reaches past base — and the first starting + // inside it. + let blocker = regions + .range(..base) + .next_back() + .map(|(_, region)| region) + .filter(|region| region_end(region) > base) + .or_else(|| regions.range(base..end).next().map(|(_, region)| region)); + + let next = match blocker { + // Resume at the page after the blocking region and try again. + Some(region) => (region_end(region).checked_add(0xFFF)?) & !0xFFF, + None => { + if self + .next_address + .compare_exchange(base, end, Ordering::SeqCst, Ordering::SeqCst) + .is_ok() + { + return Some(base); + } + continue; + } + }; + + // Another thread may have moved the cursor further along already; never rewind it. + let _ = self + .next_address + .try_update(Ordering::SeqCst, Ordering::SeqCst, |current| { + if current >= next { + None + } else { + Some(next) + } + }); + } } /// Registers a pinned managed array so that native pointer access @@ -860,14 +1239,28 @@ impl AddressSpace { element_size: usize, element_count: usize, ) -> Result<()> { + let byte_length = element_size.checked_mul(element_count).ok_or_else(|| { + EmulationError::InternalError { + description: "pinned array byte length overflow".to_string(), + } + })?; + + // Pins are consulted before regions on every access, so one registered over a mapped + // range would shadow it: reads and writes to the mapping would be answered from the + // managed array instead. `reserve_address_range` picks bases that avoid this, but the + // base is a parameter here, so the invariant is enforced rather than assumed. + if byte_length > 0 && self.region_overlaps_range(base_addr, byte_length)? { + return Err(EmulationError::InvalidAddress { + address: base_addr, + reason: "pinned array would shadow a mapped region".to_string(), + } + .into()); + } + let entry = PinnedArrayEntry { array_ref, base_addr, - byte_length: element_size.checked_mul(element_count).ok_or_else(|| { - EmulationError::InternalError { - description: "pinned array byte length overflow".to_string(), - } - })?, + byte_length, element_size, }; let mut pins = self.pinned_arrays.write().map_err(|_| { @@ -879,6 +1272,21 @@ impl AddressSpace { Ok(()) } + /// Returns the native base address an array is already pinned at, if any. + /// + /// Lets a re-pin reuse the existing reservation instead of adding another entry for the + /// same array — the pin table is scanned on every unmanaged access, so duplicates cost + /// every later access, not just the memory. + #[must_use] + pub fn pinned_base_of(&self, array_ref: HeapRef) -> Option { + let pins = self.pinned_arrays.read().ok()?; + let base = pins + .values() + .find(|entry| entry.array_ref == array_ref) + .map(|entry| entry.base_addr); + base + } + /// Attempts to read from a pinned array region. /// /// Returns `None` if the address is not within a pinned array range. @@ -1119,8 +1527,111 @@ impl AddressSpace { return Ok(()); } - let data = vec![value; size]; - self.write(address, &data) + // `size` comes from the emulated stack via `initblk`. Building the whole fill buffer + // first would commit it to the host allocator before the destination is checked, so an + // out-of-range address with a huge size costs the memory anyway and only then fails. + // Validate first, fill second. + if !self.covers_range(address, size)? { + return Err(EmulationError::InvalidAddress { + address, + reason: format!("initblk destination range of {size} bytes is not mapped"), + } + .into()); + } + + // Fill through a fixed-size buffer so peak overhead is constant in `size`. + const CHUNK: usize = 64 * 1024; + let chunk_len = size.min(CHUNK); + let chunk = vec![value; chunk_len]; + + let mut written: usize = 0; + while written < size { + let remaining = size.saturating_sub(written); + let n = remaining.min(chunk_len); + let slice = chunk + .get(..n) + .ok_or_else(|| EmulationError::InternalError { + description: "initblk chunk slice out of range".to_string(), + })?; + let offset = u64::try_from(written).map_err(|_| EmulationError::InvalidAddress { + address, + reason: "initblk offset exceeds address width".to_string(), + })?; + let target = + address + .checked_add(offset) + .ok_or_else(|| EmulationError::InvalidAddress { + address, + reason: "initblk range overflows the address space".to_string(), + })?; + self.write(target, slice)?; + written = written.saturating_add(n); + } + + Ok(()) + } + + /// Reports whether any mapped region intersects `[address, address + size)`. + /// + /// Unlike [`Self::covers_range`] this asks about *intersection*, not containment, and + /// ignores pins: it exists to keep a new pin from being laid over a mapping. + fn region_overlaps_range(&self, address: u64, size: usize) -> Result { + let len = u64::try_from(size).map_err(|_| EmulationError::InvalidAddress { + address, + reason: "range length exceeds the address width".to_string(), + })?; + let Some(end) = address.checked_add(len) else { + return Ok(false); + }; + + let regions = self.regions.read().map_err(|_| { + Error::from(EmulationError::InternalError { + description: "region lock poisoned".to_string(), + }) + })?; + + let predecessor_overlaps = + regions + .range(..address) + .next_back() + .is_some_and(|(_, region)| { + region + .base() + .saturating_add(u64::try_from(region.size()).unwrap_or(u64::MAX)) + > address + }); + + Ok(predecessor_overlaps || regions.range(address..end).next().is_some()) + } + + /// Reports whether a contiguous range is backed by a pinned array or a mapped region. + /// + /// Used to validate a destination before committing memory proportional to its size. + fn covers_range(&self, address: u64, size: usize) -> Result { + let len = u64::try_from(size).map_err(|_| EmulationError::InvalidAddress { + address, + reason: "range length exceeds the address width".to_string(), + })?; + let Some(end) = address.checked_add(len) else { + return Ok(false); + }; + + if let Ok(pins) = self.pinned_arrays.read() { + for entry in pins.values() { + let pin_end = entry.base_addr.saturating_add(entry.byte_length as u64); + if address >= entry.base_addr && end <= pin_end { + return Ok(true); + } + } + } + + let regions = self.regions.read().map_err(|_| { + Error::from(EmulationError::InternalError { + description: "region lock poisoned".to_string(), + }) + })?; + Ok(Self::region_containing(®ions, address) + .is_some_and(|r| r.contains_range(address, size))) } /// Maps a PE image at its preferred base address. @@ -1174,7 +1685,7 @@ impl AddressSpace { pub fn regions(&self) -> Vec<(u64, usize, String)> { match self.regions.read() { Ok(regions) => regions - .iter() + .values() .map(|r| (r.base(), r.size(), r.label().to_string())) .collect(), Err(_) => Vec::new(), @@ -1185,11 +1696,26 @@ impl AddressSpace { #[must_use] pub fn mapped_size(&self) -> usize { match self.regions.read() { - Ok(regions) => regions.iter().map(MemoryRegion::size).sum(), + Ok(regions) => regions.values().map(MemoryRegion::size).sum(), Err(_) => 0, } } + /// Finds the region containing `address`, in `O(log regions)`. + /// + /// Regions never overlap — [`Self::map_at`] rejects a mapping that would create one — so + /// the only candidate is the last region starting at or before the address. + fn region_containing( + regions: &BTreeMap, + address: u64, + ) -> Option<&MemoryRegion> { + regions + .range(..=address) + .next_back() + .map(|(_, region)| region) + .filter(|region| region.contains(address)) + } + /// Checks if two regions overlap in the address space. /// /// Uses the standard interval overlap test: two intervals [a_start, a_end) @@ -1316,7 +1842,7 @@ impl AddressSpace { // Clone regions - this is cheap because pages use CoW internally let regions = match self.regions.read() { Ok(r) => r.clone(), - Err(_) => Vec::new(), + Err(_) => BTreeMap::new(), }; Self { @@ -1336,6 +1862,8 @@ impl AddressSpace { monitor_locks: RwLock::new(ImHashMap::new()), // Fresh pinned array mappings pinned_arrays: RwLock::new(ImHashMap::new()), + unmanaged_bytes: AtomicUsize::new(0), + max_unmanaged_bytes: AtomicUsize::new(DEFAULT_MAX_UNMANAGED_BYTES), } } @@ -1396,8 +1924,8 @@ impl AddressSpace { description: "address space regions", })? .iter() - .map(|region| region.fork()) - .collect::, _>>()?; + .map(|(&base, region)| region.fork().map(|forked| (base, forked))) + .collect::, _>>()?; // Fork protection overrides (O(1) due to imbl) let protection_overrides = self @@ -1443,6 +1971,10 @@ impl AddressSpace { monitor_locks: RwLock::new(monitor_locks), // Fork pinned array mappings - O(1) due to imbl pinned_arrays: RwLock::new(pinned_arrays), + // The forked regions are carried over, so their bytes stay charged; the fork + // inherits the parent's ceiling. + unmanaged_bytes: AtomicUsize::new(self.unmanaged_bytes.load(Ordering::Relaxed)), + max_unmanaged_bytes: AtomicUsize::new(self.max_unmanaged_bytes.load(Ordering::Relaxed)), }) } } @@ -1457,7 +1989,7 @@ mod tests { }, EmValue, }, - metadata::token::Token, + metadata::{token::Token, typesystem::CilFlavor}, }; #[test] @@ -1488,6 +2020,90 @@ mod tests { assert_eq!(read, vec![0xCA, 0xFE]); } + /// Protection is tracked per page; an access that its page forbids must fault rather + /// than succeed silently, or `VirtualProtect` emulation is pure bookkeeping and an + /// obfuscator probing by writing where a real process faults detects the emulator. + #[test] + fn write_to_read_only_memory_faults() { + let space = AddressSpace::new(); + space.map_data(0x1000, &[0u8; 16], "test").unwrap(); + space + .set_protection(0x1000, 16, MemoryProtection::READ) + .unwrap(); + + assert!(space.read(0x1000, 4).is_ok(), "reads must still work"); + assert!(space.write(0x1000, &[0xFF]).is_err()); + } + + /// A guard page faults on *any* access, read included. + #[test] + fn guard_page_faults_on_read_and_write() { + let space = AddressSpace::new(); + space.map_data(0x2000, &[0u8; 16], "test").unwrap(); + space + .set_protection( + 0x2000, + 16, + MemoryProtection::READ_WRITE | MemoryProtection::GUARD, + ) + .unwrap(); + + assert!(space.read(0x2000, 1).is_err()); + assert!(space.write(0x2000, &[0x01]).is_err()); + } + + /// Fixed-size reads must agree with the allocating path byte for byte. + #[test] + fn fixed_size_reads_match_the_allocating_read() { + let space = AddressSpace::new(); + space + .map_data( + 0x3000, + &[0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88], + "test", + ) + .unwrap(); + + assert_eq!(space.read_u8(0x3000).unwrap(), 0x11); + assert_eq!(space.read_u16(0x3000).unwrap(), 0x2211); + assert_eq!(space.read_u32(0x3000).unwrap(), 0x4433_2211); + assert_eq!(space.read_u64(0x3000).unwrap(), 0x8877_6655_4433_2211); + assert_eq!( + space.read_exact::<4>(0x3000).unwrap().to_vec(), + space.read(0x3000, 4).unwrap() + ); + assert!(space.read_u32(0x9000).is_err(), "unmapped must still fault"); + } + + /// The reservation cursor starts at 0x1000_0000 — the default managed `ImageBase` — so + /// it must step over a mapping rather than hand out a base that aliases it. + #[test] + fn reserved_ranges_skip_mapped_regions() { + let space = AddressSpace::new(); + let cursor = space.reserve_address_range(0x1000).unwrap(); + + // Map exactly where the next reservation would otherwise land. + let blocked = cursor + 0x1000; + space.map_data(blocked, &[0u8; 0x2000], "image").unwrap(); + + let next = space.reserve_address_range(0x1000).unwrap(); + assert!( + next >= blocked + 0x2000, + "reservation {next:#x} overlaps the region at {blocked:#x}" + ); + } + + /// Pins are consulted before regions on every access, so one laid over a mapping would + /// shadow it. + #[test] + fn pinned_array_cannot_shadow_a_mapped_region() { + let space = AddressSpace::new(); + space.map_data(0x4000, &[0u8; 0x1000], "image").unwrap(); + let array = space.managed_heap().alloc_array(CilFlavor::I4, 4).unwrap(); + + assert!(space.register_pinned_array(0x4000, array, 4, 4).is_err()); + } + #[test] fn test_static_fields() { let space = AddressSpace::new(); @@ -1692,4 +2308,36 @@ mod tests { // Original doesn't see it assert!(heap.get_string(new_ref).is_err()); } + + /// An unmapped read must be refused without first allocating a buffer for it. + /// + /// `read` built its destination `Vec` before `read_mapped` validated the range, so the + /// length — which arrives from the emulated evaluation stack via `cpblk` — was committed + /// to the host allocator on the reject path. At these sizes that is `handle_alloc_error` + /// and an uncatchable abort of the analysis host, not an `Err`. The value below is far + /// beyond any plausible mapping, so this test asserts the refusal is decided from the + /// region index rather than by trying and failing to allocate. + #[test] + fn unmapped_read_is_refused_without_allocating() { + let space = AddressSpace::new(); + space.map_data(0x1000, &[0u8; 16], "test").unwrap(); + + // Wholly unmapped address. + assert!(space.read(0xDEAD_0000, 1 << 40).is_err()); + // Mapped base, but the range runs off the end of the region. + assert!(space.read(0x1000, 1 << 40).is_err()); + // The in-bounds case still works. + assert_eq!(space.read(0x1000, 4).unwrap(), vec![0u8; 4]); + } + + /// `cpblk` reaches `copy_block` with a size taken straight off the evaluation stack, so + /// the same refusal has to hold there — this is the path that made the ordering bug + /// reachable from three emulated instructions. + #[test] + fn copy_block_with_an_unmapped_source_is_refused_without_allocating() { + let space = AddressSpace::new(); + space.map_data(0x1000, &[0u8; 16], "test").unwrap(); + + assert!(space.copy_block(0x1000, 0xDEAD_0000, 1 << 40).is_err()); + } } diff --git a/dotscope/src/emulation/memory/heap/arrays.rs b/dotscope/src/emulation/memory/heap/arrays.rs index 38d9e9b6..eb0fe5eb 100644 --- a/dotscope/src/emulation/memory/heap/arrays.rs +++ b/dotscope/src/emulation/memory/heap/arrays.rs @@ -23,9 +23,20 @@ impl ManagedHeap { /// /// # Errors /// - /// Returns [`EmulationError::HeapMemoryLimitExceeded`] if heap is out of memory. + /// Returns [`EmulationError::HeapMemoryLimitExceeded`] if the array would exceed the heap + /// budget. The check happens **before** any host memory is committed — `length` comes + /// straight off the emulated evaluation stack via `newarr`, so building the backing store + /// first and asking afterwards means the limit can never fire. pub fn alloc_array(&self, element_type: CilFlavor, length: usize) -> Result { - let elements = vec![EmValue::default_for_flavor(&element_type); length]; + self.reserve_elements(length)?; + let mut elements = Vec::new(); + elements.try_reserve_exact(length).map_err(|_| { + EmulationError::HeapMemoryLimitExceeded { + current: self.current_size(), + limit: self.max_size(), + } + })?; + elements.resize(length, EmValue::default_for_flavor(&element_type)); self.alloc_object_internal( HeapObject::Array { element_type, @@ -64,8 +75,25 @@ impl ManagedHeap { element_type: CilFlavor, dimensions: Vec, ) -> Result { - let total_elements: usize = dimensions.iter().product(); - let elements = vec![EmValue::default_for_flavor(&element_type); total_elements]; + // `iter().product()` wraps on overflow in release builds and panics in debug, so the + // element count is computed with checked arithmetic before it is used for anything. + let total_elements = dimensions + .iter() + .try_fold(1usize, |acc, &d| acc.checked_mul(d)) + .ok_or(EmulationError::HeapMemoryLimitExceeded { + current: self.current_size(), + limit: self.max_size(), + })?; + + self.reserve_elements(total_elements)?; + let mut elements = Vec::new(); + elements.try_reserve_exact(total_elements).map_err(|_| { + EmulationError::HeapMemoryLimitExceeded { + current: self.current_size(), + limit: self.max_size(), + } + })?; + elements.resize(total_elements, EmValue::default_for_flavor(&element_type)); self.alloc_object_internal( HeapObject::MultiArray { element_type, @@ -118,6 +146,19 @@ impl ManagedHeap { /// Sets an array element. /// + /// # Heap accounting + /// + /// Deliberately none. This overwrites an existing slot rather than extending the array, and + /// [`HeapObject::estimated_size`] charges arrays a flat `EMVALUE_SIZE` per element, so the + /// accounted footprint is identical before and after — a delta calculation here would always + /// be zero. + /// + /// The residual gap is that the estimate is *shallow*: an `EmValue::ValueType` carries a + /// `Vec` whose contents are not walked, so a deeply nested value costs more host + /// memory than it is charged. Building such values is bounded by the instruction budget + /// rather than by the heap budget. Closing that properly means making `estimated_size` + /// recursive, which costs a walk on every size query; it is not fixed here. + /// /// # Panics /// /// Panics if the internal `RwLock` is poisoned. @@ -228,7 +269,17 @@ impl ManagedHeap { /// /// Returns [`EmulationError::HeapMemoryLimitExceeded`] if heap is out of memory. pub fn alloc_byte_array(&self, data: &[u8]) -> Result { - let elements: Vec = data.iter().map(|&b| EmValue::I32(i32::from(b))).collect(); + // Each byte becomes a full `EmValue`, so a byte array costs an order of magnitude more + // host memory than its logical size. Charge for that before expanding. + self.reserve_elements(data.len())?; + let mut elements = Vec::new(); + elements.try_reserve_exact(data.len()).map_err(|_| { + EmulationError::HeapMemoryLimitExceeded { + current: self.current_size(), + limit: self.max_size(), + } + })?; + elements.extend(data.iter().map(|&b| EmValue::I32(i32::from(b)))); self.alloc_array_with_values(CilFlavor::U1, elements) } @@ -374,4 +425,85 @@ mod tests { .set_array_element(array_ref, 10, EmValue::I32(0)) .is_err()); } + + // The allocators below take their element count from emulated, attacker-controlled values, + // so the heap budget has to be consulted before any host memory is committed. These cases + // assert that: a failure shows up as an OOM-killed or aborted test process rather than an + // assertion failure, because the rejected length would otherwise be allocated for real. + + /// 64 MiB budget: large enough for ordinary test allocations, small enough that the + /// hostile lengths below are rejected rather than attempted. + const HEAP_BUDGET: usize = 64 * 1024 * 1024; + + #[test] + fn alloc_array_rejects_absurd_length_without_allocating() { + let heap = ManagedHeap::new(HEAP_BUDGET); + // The length must be refused before materialisation: constructing the object first + // would evaluate `vec![EmValue; 2^40]` before any guard could fire. + assert!(heap.alloc_array(CilFlavor::I4, 1 << 40).is_err()); + assert_eq!(heap.current_size(), 0, "rejected array must not be charged"); + } + + #[test] + fn alloc_array_rejects_usize_max() { + let heap = ManagedHeap::new(HEAP_BUDGET); + // The value a bare `-1 as usize` produces. + assert!(heap.alloc_array(CilFlavor::I4, usize::MAX).is_err()); + } + + #[test] + fn alloc_array_still_serves_reasonable_requests() { + let heap = ManagedHeap::new(HEAP_BUDGET); + let r = heap + .alloc_array(CilFlavor::I4, 1000) + .expect("1000 elements"); + assert_eq!(heap.get_array_length(r).unwrap(), 1000); + } + + #[test] + fn alloc_multi_array_rejects_overflowing_dimensions() { + let heap = ManagedHeap::new(HEAP_BUDGET); + // The product of these wraps `usize`; `iter().product()` would have silently produced + // a small number and allocated an array of the wrong size. + let dims = vec![usize::MAX, 2, 2]; + assert!(heap.alloc_multi_array(CilFlavor::I4, dims).is_err()); + } + + #[test] + fn alloc_multi_array_rejects_oversized_product() { + let heap = ManagedHeap::new(HEAP_BUDGET); + assert!(heap + .alloc_multi_array(CilFlavor::I4, vec![1 << 20, 1 << 20]) + .is_err()); + } + + #[test] + fn alloc_byte_array_is_charged_per_emvalue_not_per_byte() { + let heap = ManagedHeap::new(HEAP_BUDGET); + // Each byte expands to a full EmValue, so a buffer that looks like it fits by byte + // count must still be rejected when its real cost exceeds the budget. + let too_big = vec![0u8; HEAP_BUDGET / 2]; + assert!(heap.alloc_byte_array(&too_big).is_err()); + } + + #[test] + fn estimated_size_charges_the_real_element_width() { + let heap = ManagedHeap::new(HEAP_BUDGET); + heap.alloc_array(CilFlavor::I4, 100).expect("alloc"); + // A hard-coded 8 bytes/element would under-charge by roughly 30x. + assert!( + heap.current_size() >= 100 * std::mem::size_of::(), + "charged {} for 100 elements of {} bytes each", + heap.current_size(), + std::mem::size_of::() + ); + } + + #[test] + fn reserve_gates_without_charging() { + let heap = ManagedHeap::new(HEAP_BUDGET); + assert!(heap.reserve(HEAP_BUDGET * 2).is_err()); + assert!(heap.reserve(1024).is_ok()); + assert_eq!(heap.current_size(), 0, "reserve must not itself account"); + } } diff --git a/dotscope/src/emulation/memory/heap/collections.rs b/dotscope/src/emulation/memory/heap/collections.rs index 07b67ae6..330eda99 100644 --- a/dotscope/src/emulation/memory/heap/collections.rs +++ b/dotscope/src/emulation/memory/heap/collections.rs @@ -4,12 +4,15 @@ //! on [`ManagedHeap`] objects. Each collection type is stored as a variant of //! [`HeapObject`] and supports the standard .NET collection operations. -use std::collections::{HashMap, HashSet as StdHashSet, VecDeque}; +use std::{ + collections::{HashMap, HashSet as StdHashSet, VecDeque}, + sync::atomic::Ordering, +}; use crate::{ emulation::{ engine::EmulationError, - memory::heap::{DictionaryKey, HeapObject, ManagedHeap}, + memory::heap::{DictionaryKey, HeapObject, ManagedHeap, EMVALUE_SIZE}, EmValue, HeapRef, }, Result, @@ -354,10 +357,19 @@ impl ManagedHeap { /// Appends an element to a List. /// + /// The appended element is charged against the heap budget. A `List` grows in place + /// without ever going back through the allocator, so nothing else on this path would + /// observe the growth: the budget is otherwise consulted only when an object is created. + /// /// # Errors /// - /// Returns [`EmulationError::LockPoisoned`] if the internal `RwLock` is poisoned. + /// Returns [`EmulationError::LockPoisoned`] if the internal `RwLock` is poisoned, or + /// [`EmulationError::HeapMemoryLimitExceeded`] if the append would exceed the budget. pub fn list_add(&self, heap_ref: HeapRef, value: EmValue) -> Result<()> { + // Gate before taking the lock and before the push, so a refused append does not grow + // the host `Vec` first and ask permission afterwards. + self.check_allocation(EMVALUE_SIZE)?; + let mut state = self .state .write() @@ -366,6 +378,9 @@ impl ManagedHeap { })?; if let Some(HeapObject::List { elements }) = state.objects.get_mut(&heap_ref.id()) { elements.push(value); + // Only charge when the push actually happened — a mismatched or missing reference + // leaves the footprint unchanged. + self.current_size.fetch_add(EMVALUE_SIZE, Ordering::Relaxed); } Ok(()) } @@ -959,3 +974,71 @@ impl ManagedHeap { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + /// Appending to a `List` must be charged against the heap budget. + /// + /// A list grows in place without going back through the allocator, so nothing else on this + /// path observes the growth — the budget is otherwise consulted only at object creation. + #[test] + fn list_add_charges_the_appended_element() { + let heap = ManagedHeap::new(1024 * 1024); + let list = heap.alloc_list_with_elements(Vec::new()).unwrap(); + let before = heap.current_size.load(Ordering::Relaxed); + + heap.list_add(list, EmValue::I32(1)).unwrap(); + + assert_eq!( + heap.current_size.load(Ordering::Relaxed), + before.saturating_add(EMVALUE_SIZE), + "the appended element must be charged exactly once" + ); + } + + /// Once the budget is exhausted, further appends are refused rather than growing the host + /// `Vec` and asking permission afterwards. + #[test] + fn list_add_is_refused_past_the_budget() { + // Enough for the list object itself, but only a handful of elements. + let heap = ManagedHeap::new(EMVALUE_SIZE.saturating_mul(8).saturating_add(256)); + let list = heap.alloc_list_with_elements(Vec::new()).unwrap(); + + let mut appended = 0usize; + for _ in 0..1000 { + if heap.list_add(list, EmValue::I32(1)).is_err() { + break; + } + appended = appended.saturating_add(1); + } + + assert!(appended > 0, "some appends should succeed"); + assert!( + appended < 1000, + "the budget must eventually refuse an append" + ); + assert_eq!( + heap.list_count(list).unwrap(), + appended, + "a refused append must not have pushed" + ); + } + + /// A mismatched reference changes nothing, so it must not be charged either. + #[test] + fn list_add_on_non_list_does_not_charge() { + let heap = ManagedHeap::new(1024 * 1024); + let not_a_list = heap.alloc_string("x").unwrap(); + let before = heap.current_size.load(Ordering::Relaxed); + + heap.list_add(not_a_list, EmValue::I32(1)).unwrap(); + + assert_eq!( + heap.current_size.load(Ordering::Relaxed), + before, + "nothing was appended, so nothing should be charged" + ); + } +} diff --git a/dotscope/src/emulation/memory/heap/mod.rs b/dotscope/src/emulation/memory/heap/mod.rs index dfd62468..55b446a9 100644 --- a/dotscope/src/emulation/memory/heap/mod.rs +++ b/dotscope/src/emulation/memory/heap/mod.rs @@ -55,9 +55,47 @@ use crate::{ assembly::InstructionAssembler, emulation::{engine::EmulationError, tokens, EmValue, HeapRef}, metadata::{signatures::TypeSignature, token::Token, typesystem::CilFlavor}, + utils::truncate_chars, Result, }; +/// Host cost of one element in an emulated array or list. +/// +/// Derived from the type rather than hard-coded, so the budget stays honest if [`EmValue`] +/// changes width. A fixed per-element figure is easy to get badly wrong: `EmValue` is +/// currently on the order of a couple of hundred bytes, so under-charging it would let an +/// array sit nominally inside a 256 MB cap while really costing gigabytes. +pub(crate) const EMVALUE_SIZE: usize = std::mem::size_of::(); + +/// Object-header overhead charged for a single-dimensional array. +const ARRAY_HEADER_BYTES: usize = 24; + +/// Object-header overhead charged for a multi-dimensional array. +const MULTI_ARRAY_HEADER_BYTES: usize = 32; + +/// Object-header overhead charged for a `List`. +const LIST_HEADER_BYTES: usize = 32; + +/// Object-header overhead charged for a delegate. +const DELEGATE_HEADER_BYTES: usize = 48; + +/// Host cost of one entry in a delegate's invocation list. +/// +/// Derived from the type so it tracks [`DelegateEntry`] rather than drifting from it. +const DELEGATE_ENTRY_BYTES: usize = std::mem::size_of::(); + +/// Maximum number of entries in a single delegate's invocation list. +/// +/// `Delegate.Combine` concatenates two invocation lists and nothing stops both arguments from +/// being the same delegate, so each call can double the length — exponential growth in a linear +/// number of emulated instructions. Charging the list against the heap budget (see +/// [`HeapObject::estimated_size`]) bounds the total, but a dedicated cap fails the operation at +/// the point of the defect with a message that names it, instead of surfacing as a generic +/// out-of-memory once the doubling happens to cross the budget. +/// +/// Real multicast delegates hold a handful of entries; .NET itself has no fixed limit. +const MAX_DELEGATE_INVOCATION_LIST: usize = 65_536; + /// Info about a symmetric algorithm: (algorithm_type, key, iv, mode, padding). pub type SymmetricAlgorithmInfo = (Arc, Option>, Option>, u8, u8); @@ -628,12 +666,14 @@ impl HeapObject { // saturated `usize::MAX` correctly signals "huge" for limit checks. match self { HeapObject::String(s) => s.len().saturating_mul(2).saturating_add(24), // Object header + UTF-16 - HeapObject::Array { elements, .. } => { - elements.len().saturating_mul(8).saturating_add(24) - } - HeapObject::MultiArray { elements, .. } => { - elements.len().saturating_mul(8).saturating_add(32) - } + HeapObject::Array { elements, .. } => elements + .len() + .saturating_mul(EMVALUE_SIZE) + .saturating_add(ARRAY_HEADER_BYTES), + HeapObject::MultiArray { elements, .. } => elements + .len() + .saturating_mul(EMVALUE_SIZE) + .saturating_add(MULTI_ARRAY_HEADER_BYTES), HeapObject::Object { fields, .. } => fields.len().saturating_mul(16).saturating_add(24), HeapObject::TypedReference { .. } | HeapObject::BoxedValue { .. } @@ -648,7 +688,16 @@ impl HeapObject { HeapObject::CryptoTransform { key, iv, .. } => { 48usize.saturating_add(key.len()).saturating_add(iv.len()) } - HeapObject::Delegate { .. } => 48, + // Charged by content, like every other variadic variant. A flat constant here was a + // hole in the memory limit rather than an inaccuracy: `Delegate.Combine(d, d)` is + // legal and doubles the invocation list, so ~30 emulated calls reach a billion + // entries while the heap accounts each result at the same 48 bytes. + HeapObject::Delegate { + invocation_list, .. + } => invocation_list + .len() + .saturating_mul(DELEGATE_ENTRY_BYTES) + .saturating_add(DELEGATE_HEADER_BYTES), HeapObject::Encoding { .. } => 24, HeapObject::SymmetricAlgorithm { key, iv, .. } => 32usize .saturating_add(key.as_ref().map_or(0, Vec::len)) @@ -656,7 +705,10 @@ impl HeapObject { HeapObject::Dictionary { entries } => { entries.len().saturating_mul(32).saturating_add(48) } - HeapObject::List { elements } => elements.len().saturating_mul(8).saturating_add(32), + HeapObject::List { elements } => elements + .len() + .saturating_mul(EMVALUE_SIZE) + .saturating_add(LIST_HEADER_BYTES), HeapObject::StringBuilder { buffer, .. } => 32usize.saturating_add(buffer.len()), HeapObject::Stack { elements } => elements.len().saturating_mul(8).saturating_add(32), HeapObject::Queue { elements } => elements.len().saturating_mul(8).saturating_add(32), @@ -685,8 +737,8 @@ impl fmt::Display for HeapObject { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { HeapObject::String(s) => { - if s.len() > 50 { - write!(f, "\"{}...\"", &s[..47]) + if s.chars().count() > 50 { + write!(f, "\"{}...\"", truncate_chars(s, 47)) } else { write!(f, "\"{s}\"") } @@ -815,11 +867,11 @@ impl fmt::Display for HeapObject { write!(f, "list({} elements)", elements.len()) } HeapObject::StringBuilder { buffer, .. } => { - if buffer.len() > 50 { + if buffer.chars().count() > 50 { write!( f, "stringbuilder({}... len={})", - &buffer[..47], + truncate_chars(buffer, 47), buffer.len() ) } else { @@ -960,6 +1012,13 @@ pub struct ManagedHeap { pub(crate) current_size: AtomicUsize, /// Maximum allowed heap size in bytes. pub(crate) max_size: usize, + /// Maximum number of live heap objects, independent of their total size. + /// + /// A byte ceiling alone does not bound object *count*: a flood of tiny objects is charged + /// only its header estimate each while really costing an `imbl::HashMap` node plus an + /// `original_types` entry per object. Set from `EmulationLimits::max_heap_objects`; + /// `usize::MAX` means unlimited. + pub(crate) max_objects: AtomicUsize, } impl ManagedHeap { @@ -978,9 +1037,19 @@ impl ManagedHeap { next_id: AtomicU64::new(1), current_size: AtomicUsize::new(0), max_size, + // Unlimited until a builder applies the configured ceiling, so constructing a heap + // directly keeps its previous behaviour. + max_objects: AtomicUsize::new(usize::MAX), } } + /// Sets the maximum number of live heap objects. + /// + /// See `max_objects` for why a byte ceiling alone is insufficient. + pub fn set_max_objects(&self, max: usize) { + self.max_objects.store(max, Ordering::Relaxed); + } + /// Creates a managed heap with default size (64MB). #[must_use] pub fn default_size() -> Self { @@ -1004,6 +1073,41 @@ impl ManagedHeap { Ok(()) } + /// Reserves budget for `size` bytes *before* the caller commits any host memory. + /// + /// Hook and allocator code that builds a payload whose size is derived from emulated, + /// attacker-controlled values must call this first. Checking after the payload exists is + /// useless: the host allocator has already been asked for the memory, so the budget cannot + /// prevent the very exhaustion it exists to prevent. + /// + /// This performs no accounting of its own — the eventual + /// `alloc_object_internal` records the real size. It is a + /// gate, not a reservation ledger, so a caller that reserves and then does not allocate + /// leaks nothing. + /// + /// # Errors + /// + /// Returns [`EmulationError::HeapMemoryLimitExceeded`] if `size` would exceed the budget. + pub fn reserve(&self, size: usize) -> Result<()> { + self.check_allocation(size) + } + + /// Reserves budget for an array of `count` [`EmValue`] elements. + /// + /// Charges the same per-element cost that [`HeapObject::estimated_size`] charges, so the + /// pre-flight gate and the post-allocation accounting cannot disagree. + /// + /// # Errors + /// + /// Returns [`EmulationError::HeapMemoryLimitExceeded`] if the array would exceed the budget. + pub(crate) fn reserve_elements(&self, count: usize) -> Result<()> { + self.check_allocation( + count + .saturating_mul(EMVALUE_SIZE) + .saturating_add(ARRAY_HEADER_BYTES), + ) + } + /// Internal helper to allocate an object on the heap. /// /// This consolidates the common allocation pattern: check size limits, @@ -1038,6 +1142,19 @@ impl ManagedHeap { .map_err(|_| EmulationError::LockPoisoned { description: "managed heap", })?; + + // Object-count ceiling, checked under the same lock that performs the insert so the + // count cannot drift. The byte budget above does not bound count: many small objects + // pay only their header estimate while each costs a map node plus an entry in + // `original_types`. + let max_objects = self.max_objects.load(Ordering::Relaxed); + if state.objects.len() >= max_objects { + return Err(EmulationError::ResourceLimitExceeded(format!( + "heap object count limit reached ({max_objects} objects)" + )) + .into()); + } + state.objects.insert(heap_ref.id(), obj); if let Some(token) = original_type { state.original_types.insert(heap_ref.id(), token); @@ -1115,6 +1232,9 @@ impl ManagedHeap { next_id: AtomicU64::new(self.next_id.load(Ordering::SeqCst)), current_size: AtomicUsize::new(self.current_size.load(Ordering::Relaxed)), max_size: self.max_size, + // The fork must inherit the ceiling; otherwise forking would itself be an escape + // from it. + max_objects: AtomicUsize::new(self.max_objects.load(Ordering::Relaxed)), }) } @@ -1308,14 +1428,27 @@ impl ManagedHeap { /// Allocates a multicast delegate on the heap. /// + /// The invocation list is capped at `MAX_DELEGATE_INVOCATION_LIST`; see that constant for + /// why an unbounded list is reachable in a few dozen emulated instructions. + /// /// # Errors /// - /// Returns [`EmulationError::HeapMemoryLimitExceeded`] if heap is out of memory. + /// Returns [`EmulationError::HeapMemoryLimitExceeded`] if heap is out of memory, or + /// [`EmulationError::ResourceLimitExceeded`] if the invocation list is too long. pub fn alloc_multicast_delegate( &self, type_token: Token, entries: Vec, ) -> Result { + if entries.len() > MAX_DELEGATE_INVOCATION_LIST { + return Err(EmulationError::ResourceLimitExceeded(format!( + "delegate invocation list of {} entries exceeds maximum {}", + entries.len(), + MAX_DELEGATE_INVOCATION_LIST + )) + .into()); + } + self.alloc_object_internal( HeapObject::Delegate { type_token, @@ -1766,9 +1899,19 @@ impl ManagedHeap { /// /// Used by collection constructors and StringBuilder mutations. /// + /// The size difference between the outgoing and incoming object is charged against the heap + /// budget. This is the write-back path for `StringBuilder`, `List`, `Dictionary` and + /// `HashSet`, so without it every one of those could grow without bound after allocation: + /// the budget was consulted only when an object was first created, leaving in-place growth + /// both unchecked and invisible to `current_size`. + /// + /// A replacement that would exceed the budget is refused and the existing object is left + /// untouched. + /// /// # Errors /// - /// Returns [`EmulationError::LockPoisoned`] if the internal `RwLock` is poisoned. + /// Returns [`EmulationError::LockPoisoned`] if the internal `RwLock` is poisoned, or + /// [`EmulationError::HeapMemoryLimitExceeded`] if the growth would exceed the budget. pub fn replace_object(&self, heap_ref: HeapRef, obj: HeapObject) -> Result<()> { let mut state = self .state @@ -1779,9 +1922,24 @@ impl ManagedHeap { let id = heap_ref.id(); Self::preserve_original_type(&mut state, id); - if state.objects.contains_key(&id) { - state.objects.insert(id, obj); + let Some(existing) = state.objects.get(&id) else { + return Ok(()); + }; + + let old_size = existing.estimated_size(); + let new_size = obj.estimated_size(); + + // Charge before mutating, so a refused replacement leaves the heap as it was. + let growth = new_size.saturating_sub(old_size); + if growth > 0 { + self.check_allocation(growth)?; + self.current_size.fetch_add(growth, Ordering::Relaxed); + } else { + self.current_size + .fetch_sub(old_size.saturating_sub(new_size), Ordering::Relaxed); } + + state.objects.insert(id, obj); Ok(()) } } @@ -1810,6 +1968,133 @@ mod tests { use super::*; use crate::Error; + /// A delegate's accounted size must scale with its invocation list. + /// + /// While this was a flat constant, `Delegate.Combine(d, d)` — which is legal, since nothing + /// stops both arguments being the same handle — doubled the list on every call while the + /// heap charged the same 48 bytes each time. Roughly thirty calls reach a billion entries + /// with the memory limit offering no resistance at all. + #[test] + fn delegate_size_scales_with_invocation_list() { + let type_token = Token::new(0x0200_0001); + let entry = || DelegateEntry { + target: None, + method_token: Token::new(0x0600_0001), + }; + + let one = HeapObject::Delegate { + type_token, + invocation_list: vec![entry()], + }; + let many = HeapObject::Delegate { + type_token, + invocation_list: (0..1000).map(|_| entry()).collect(), + }; + + assert!( + many.estimated_size() > one.estimated_size(), + "a longer invocation list must cost more" + ); + assert!( + many.estimated_size() >= 1000usize.saturating_mul(DELEGATE_ENTRY_BYTES), + "every entry must be charged" + ); + } + + /// The invocation-list cap fails the doubling at its source. + #[test] + fn multicast_delegate_invocation_list_is_capped() { + let heap = ManagedHeap::new(usize::MAX); + let entries: Vec = (0..MAX_DELEGATE_INVOCATION_LIST.saturating_add(1)) + .map(|_| DelegateEntry { + target: None, + method_token: Token::new(0x0600_0001), + }) + .collect(); + + assert!(heap + .alloc_multicast_delegate(Token::new(0x0200_0001), entries) + .is_err()); + } + + /// Growing an object in place must be charged, and refused when it does not fit. + /// + /// `replace_object` is the write-back path for `StringBuilder`, `List`, `Dictionary` and + /// `HashSet`. A budget consulted only at creation cannot see growth after allocation. + #[test] + fn replace_object_charges_growth_and_refuses_overflow() { + let heap = ManagedHeap::new(4096); + let heap_ref = heap.alloc_string("small").unwrap(); + let before = heap.current_size.load(Ordering::Relaxed); + + // A modest growth is charged. + heap.replace_object(heap_ref, HeapObject::String("a".repeat(100).into())) + .unwrap(); + let after = heap.current_size.load(Ordering::Relaxed); + assert!( + after > before, + "in-place growth must increase the accounted size" + ); + + // A growth past the budget is refused, and leaves the object untouched. + let err = heap.replace_object(heap_ref, HeapObject::String("b".repeat(100_000).into())); + assert!(err.is_err(), "growth past the budget must be refused"); + assert_eq!( + heap.get_string(heap_ref).unwrap().len(), + 100, + "a refused replacement must not mutate the object" + ); + } + + /// Shrinking must give the budget back, or a long-lived object would ratchet the heap up. + #[test] + fn replace_object_credits_shrinkage() { + let heap = ManagedHeap::new(1024 * 1024); + let heap_ref = heap.alloc_string(&"a".repeat(1000)).unwrap(); + let before = heap.current_size.load(Ordering::Relaxed); + + heap.replace_object(heap_ref, HeapObject::String("x".into())) + .unwrap(); + + assert!( + heap.current_size.load(Ordering::Relaxed) < before, + "shrinking must return budget" + ); + } + + /// The object-count ceiling bounds a flood of tiny objects, which the byte budget does not: + /// each pays only its header estimate while costing a map node plus an `original_types` entry. + #[test] + fn heap_object_count_is_capped() { + let heap = ManagedHeap::new(usize::MAX); + heap.set_max_objects(8); + + for _ in 0..8 { + heap.alloc_string("x").unwrap(); + } + + assert!( + heap.alloc_string("x").is_err(), + "allocation past the object ceiling must fail" + ); + } + + /// A fork must inherit the object ceiling, or forking would itself escape it. + #[test] + fn fork_inherits_object_ceiling() { + let heap = ManagedHeap::new(usize::MAX); + heap.set_max_objects(4); + heap.alloc_string("x").unwrap(); + + let forked = heap.fork().unwrap(); + assert_eq!(forked.max_objects.load(Ordering::Relaxed), 4); + + for _ in 0..3 { + forked.alloc_string("y").unwrap(); + } + assert!(forked.alloc_string("z").is_err()); + } + #[test] fn test_heap_alloc_string() { let heap = ManagedHeap::new(1024 * 1024); diff --git a/dotscope/src/emulation/memory/heap/streams.rs b/dotscope/src/emulation/memory/heap/streams.rs index fbf77104..767c7cd4 100644 --- a/dotscope/src/emulation/memory/heap/streams.rs +++ b/dotscope/src/emulation/memory/heap/streams.rs @@ -307,7 +307,13 @@ impl ManagedHeap { /// /// # Errors /// - /// Returns [`EmulationError::LockPoisoned`] if the internal `RwLock` is poisoned. + /// Returns [`EmulationError::LockPoisoned`] if the internal `RwLock` is poisoned, or + /// [`EmulationError::HeapMemoryLimitExceeded`] if growing the stream to `new_length` would + /// exceed the heap budget. + /// + /// Growth is charged against the heap budget and reserved fallibly before the buffer is + /// resized. This method is `pub`, and `new_length` reaches it from emulated code via + /// `Stream.SetLength`, so the bound has to live here rather than in any one caller. pub fn truncate_stream(&self, heap_ref: HeapRef, new_length: usize) -> Result<()> { let mut state = self .state @@ -316,6 +322,39 @@ impl ManagedHeap { description: "managed heap", })?; if let Some(HeapObject::Stream { data, position }) = state.objects.get_mut(&heap_ref.id()) { + let old_length = data.len(); + if new_length > old_length { + let growth = new_length.saturating_sub(old_length); + let current = self.current_size.load(Ordering::Relaxed); + if current.saturating_add(growth) > self.max_size { + return Err(EmulationError::HeapMemoryLimitExceeded { + current, + limit: self.max_size, + } + .into()); + } + data.try_reserve(growth) + .map_err(|_| EmulationError::HeapMemoryLimitExceeded { + current, + limit: self.max_size, + })?; + self.current_size.fetch_add(growth, Ordering::Relaxed); + } else { + // Shrinking returns budget so a grow/shrink cycle cannot ratchet the + // accounted size upwards. + // + // Saturating, not `fetch_sub`: in-place stream writes do not currently update + // `current_size`, so a stream can hold more bytes than were ever charged. + // Crediting those back with a wrapping subtract would underflow the counter to + // near `usize::MAX` and make every later allocation fail. + let refund = old_length.saturating_sub(new_length); + let _ = + self.current_size + .try_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + Some(current.saturating_sub(refund)) + }); + } + data.resize(new_length, 0); if *position > new_length { *position = new_length; diff --git a/dotscope/src/emulation/memory/region.rs b/dotscope/src/emulation/memory/region.rs index bd38e4e4..7c1ea25c 100644 --- a/dotscope/src/emulation/memory/region.rs +++ b/dotscope/src/emulation/memory/region.rs @@ -503,35 +503,74 @@ impl MemoryRegion { return Some(Vec::new()); } + let mut result = vec![0u8; len]; + if self.read_into(address, &mut result) { + Some(result) + } else { + None + } + } + + /// Reads into a caller-supplied buffer, allocating nothing. + /// + /// This is the paging loop; [`Self::read`] is this plus a `Vec`. Fixed-size loads + /// (`ldind.i4` and friends) go through here with a stack array, which is why the loop + /// lives on this side of the allocation. + /// + /// # Arguments + /// + /// * `address` - The address to read from + /// * `dest` - Buffer to fill; its length is the number of bytes read + /// + /// # Returns + /// + /// `true` if the whole buffer was filled, `false` if the range is outside this region or + /// a page read failed. + #[must_use] + pub fn read_into(&self, address: u64, dest: &mut [u8]) -> bool { + let len = dest.len(); + if len == 0 { + return true; + } + if !self.contains_range(address, len) { - return None; + return false; } // Safe: offset within a memory region always fits in usize #[allow(clippy::cast_possible_truncation)] let offset = address.saturating_sub(self.base) as usize; - let mut result = vec![0u8; len]; let mut bytes_read = 0; while bytes_read < len { - let current_offset = offset.checked_add(bytes_read)?; + let Some(current_offset) = offset.checked_add(bytes_read) else { + return false; + }; let page_index = current_offset / PAGE_SIZE; let page_offset = current_offset % PAGE_SIZE; - let page = self.pages.get(page_index)?; + let Some(page) = self.pages.get(page_index) else { + return false; + }; - let remaining = len.checked_sub(bytes_read)?; + let Some(remaining) = len.checked_sub(bytes_read) else { + return false; + }; let bytes_in_page = PAGE_SIZE.saturating_sub(page_offset).min(remaining); - let read_end = bytes_read.checked_add(bytes_in_page)?; - let dest = result.get_mut(bytes_read..read_end)?; + let Some(read_end) = bytes_read.checked_add(bytes_in_page) else { + return false; + }; + let Some(slot) = dest.get_mut(bytes_read..read_end) else { + return false; + }; - if page.read(page_offset, dest).is_err() { - return None; + if page.read(page_offset, slot).is_err() { + return false; } bytes_read = bytes_read.saturating_add(bytes_in_page); } - Some(result) + true } /// Writes bytes to this region. diff --git a/dotscope/src/emulation/mod.rs b/dotscope/src/emulation/mod.rs index 7b5536d2..1cef64be 100644 --- a/dotscope/src/emulation/mod.rs +++ b/dotscope/src/emulation/mod.rs @@ -149,9 +149,8 @@ pub use engine::{ }; // Re-export primary types from exception module pub use exception::{ - ExceptionClause, ExceptionHandler, ExceptionInfo, FrameSearchInfo, HandlerMatch, - HandlerSearchState, InstructionLocation, MethodHandlerResult, PendingFinally, StackUnwinder, - ThreadExceptionState, UnwindSequenceBuilder, UnwindStepResult, + ExceptionClause, ExceptionInfo, HandlerMatch, InstructionLocation, PendingFinally, + ThreadExceptionState, }; // Re-export primary types from fakeobjects module pub use fakeobjects::{FakeObjects, SharedFakeObjects}; @@ -163,10 +162,10 @@ pub use loader::{ }; // Re-export primary types from memory module pub use memory::{ - AddressSpace, ArgumentStorage, DictionaryKey, EncodingType, EvaluationStack, HeapObject, - LocalVariables, ManagedHeap, MemoryProtection, MemoryRegion, Page, SectionInfo, SharedHeap, - StaticFieldStorage, ThreadId, TypeInitState, TypeWrapper, UnmanagedMemory, UnmanagedRef, - PAGE_SIZE, + AddressSpace, ArgumentStorage, DelegateEntry, DictionaryKey, EncodingType, EvaluationStack, + HeapObject, LocalVariables, ManagedHeap, MemoryProtection, MemoryRegion, Page, SectionInfo, + SharedHeap, StaticFieldStorage, ThreadId, TypeInitState, TypeWrapper, UnmanagedMemory, + UnmanagedRef, PAGE_SIZE, }; // Re-export primary types from process module pub use process::{ diff --git a/dotscope/src/emulation/process/builder.rs b/dotscope/src/emulation/process/builder.rs index 81996503..001fcfb3 100644 --- a/dotscope/src/emulation/process/builder.rs +++ b/dotscope/src/emulation/process/builder.rs @@ -137,6 +137,7 @@ fn populate_fieldrva_statics(assembly: &CilObject, address_space: &AddressSpace) let ptr_size = PointerSize::from_is_64bit(file.pe().is_64bit); for row in fieldrva_table { + let row = row?; if row.rva == 0 { continue; } @@ -1122,13 +1123,28 @@ impl ProcessBuilder { "Creating emulation process: instruction_limit={}, call_depth={}", self.config.limits.max_instructions, self.config.limits.max_call_depth ); - let heap_size = self.config.memory.max_heap_size; + // Two fields describe the same ceiling: `limits.max_heap_bytes` and + // `memory.max_heap_size`. Only `ProcessBuilder::with_max_heap_bytes` keeps them in sync, + // so a caller who builds `EmulationConfig` as a struct literal — the pattern the config + // module's own docs demonstrate — sets `max_heap_bytes` and silently gets the default. + // Honour the stricter of the two rather than picking one and ignoring the other. + let heap_size = self + .config + .memory + .max_heap_size + .min(self.config.limits.max_heap_bytes); let heap = SharedHeap::new(heap_size); + heap.heap() + .set_max_objects(self.config.limits.max_heap_objects); // Initialize fake BCL objects before anything else uses the heap let fake_objects = SharedFakeObjects::new(heap.heap()); let address_space = Arc::new(AddressSpace::with_heap(heap)); + // Apply the configured unmanaged ceiling. Unmanaged allocations bypass the managed + // heap budget entirely, so without this the `max_unmanaged_bytes` setting has no + // effect on anything. + address_space.set_max_unmanaged_bytes(self.config.limits.max_unmanaged_bytes); let mut config = self.config.clone(); if !self.register_defaults { @@ -1280,7 +1296,8 @@ impl ProcessBuilder { let tables = assembly.tables()?; let strings = assembly.strings()?; let module_table = tables.table::()?; - let module_row = module_table.iter().next()?; + // Module is RID 1 by definition (ECMA-335 II.22.30). + let module_row = module_table.get(1).ok().flatten()?; strings.get(module_row.name as usize).ok().map(String::from) }); diff --git a/dotscope/src/emulation/process/config.rs b/dotscope/src/emulation/process/config.rs index e5708ff0..70c1247a 100644 --- a/dotscope/src/emulation/process/config.rs +++ b/dotscope/src/emulation/process/config.rs @@ -195,6 +195,8 @@ pub struct EmulationConfig { /// | `max_heap_objects` | 100,000 | /// | `max_heap_bytes` | 256 MB | /// | `max_unmanaged_bytes` | 64 MB | +/// | `max_loaded_assemblies` | 64 | +/// | `max_loaded_assembly_bytes` | 32 MB | /// | `timeout_ms` | 60,000 (1 minute) | #[derive(Clone, Debug)] pub struct EmulationLimits { @@ -228,6 +230,25 @@ pub struct EmulationLimits { /// similar unmanaged allocation methods. pub max_unmanaged_bytes: usize, + /// Maximum number of assemblies emulated code may load at runtime. + /// + /// `Assembly.Load(byte[])` parses its payload into a `CilObject` that the AppDomain + /// retains for the lifetime of the emulation, for cross-assembly resolution. That + /// memory is host memory, so `max_heap_bytes` does not see it, and mutating a single + /// byte between calls defeats any content dedup while leaving the emulated heap + /// footprint unchanged. This bounds how many are kept. + /// + /// The value must stay below `u32::MAX` for the assembly index to remain + /// representable in a call frame. + pub max_loaded_assemblies: usize, + + /// Maximum size in bytes of a single assembly payload accepted for parsing. + /// + /// Checked before `Assembly.Load(byte[])` hands the bytes to the metadata parser, + /// which is a full re-entry into attacker-controlled parsing from inside a hook and + /// so runs outside the instruction and timeout budgets. + pub max_loaded_assembly_bytes: usize, + /// Timeout in milliseconds. /// /// Set to 0 for no timeout. When exceeded, emulation stops @@ -560,7 +581,7 @@ pub enum UnknownMethodBehavior { /// # Default Values /// /// By default, no capture is enabled to minimize overhead. -#[derive(Clone, Debug, Default)] +#[derive(Clone, Debug)] pub struct CaptureConfig { /// Capture assemblies loaded via Assembly.Load. /// @@ -591,6 +612,45 @@ pub struct CaptureConfig { /// Records socket operations, HTTP requests, and other network /// activity during emulation. pub network_operations: bool, + + /// Capture raw byte buffers from memory operations and crypto transforms. + /// + /// Gated like every other capture kind, so a default config records no buffers and + /// `captured_buffers()` stays empty until this is set to `true`. + pub buffers: bool, + + /// Maximum number of items retained in each capture collection. + /// + /// Capture is driven by hook invocations rather than heap allocations, and captured data + /// lives outside the managed heap, so `max_heap_bytes` does not see it. Emulated code can + /// allocate one array and then loop calling `set_Key`/`TransformFinalBlock` on it, appending + /// a full copy per iteration for a few instructions each. Without a ceiling, emulation + /// reaches a clean `Completed` outcome having consumed tens of gigabytes. + /// + /// Once reached, further captures of that kind are dropped. `usize::MAX` means unlimited. + pub max_items: usize, + + /// Maximum total bytes retained across all captured buffers, strings and assemblies. + /// + /// Complements [`max_items`](Self::max_items): a small number of very large buffers is as + /// damaging as a large number of small ones. `usize::MAX` means unlimited. + pub max_total_bytes: usize, +} + +impl Default for CaptureConfig { + /// No capture enabled, with ceilings applied to whatever the caller does enable. + fn default() -> Self { + Self { + assemblies: false, + memory_regions: Vec::new(), + strings: false, + file_operations: false, + network_operations: false, + buffers: false, + max_items: 10_000, + max_total_bytes: 256 * 1024 * 1024, + } + } } /// Environment simulation configuration. @@ -728,7 +788,9 @@ impl Default for EmulationLimits { max_heap_objects: 100_000, max_heap_bytes: 256 * 1024 * 1024, // 256 MB max_unmanaged_bytes: 64 * 1024 * 1024, // 64 MB - timeout_ms: 60_000, // 1 minute + max_loaded_assemblies: 64, + max_loaded_assembly_bytes: 32 * 1024 * 1024, // 32 MB + timeout_ms: 60_000, // 1 minute } } } @@ -1139,6 +1201,36 @@ impl EmulationLimits { self } + /// Sets the maximum number of assemblies emulated code may load at runtime. + /// + /// # Arguments + /// + /// * `max` - Maximum retained runtime-loaded assemblies + /// + /// # Returns + /// + /// Returns `self` for method chaining. + #[must_use] + pub fn with_max_loaded_assemblies(mut self, max: usize) -> Self { + self.max_loaded_assemblies = max; + self + } + + /// Sets the maximum size of a single runtime-loaded assembly payload. + /// + /// # Arguments + /// + /// * `max` - Maximum payload size in bytes + /// + /// # Returns + /// + /// Returns `self` for method chaining. + #[must_use] + pub fn with_max_loaded_assembly_bytes(mut self, max: usize) -> Self { + self.max_loaded_assembly_bytes = max; + self + } + /// Sets the execution timeout in milliseconds. /// /// # Arguments diff --git a/dotscope/src/emulation/process/execution.rs b/dotscope/src/emulation/process/execution.rs index 20435f98..d63d2082 100644 --- a/dotscope/src/emulation/process/execution.rs +++ b/dotscope/src/emulation/process/execution.rs @@ -49,9 +49,12 @@ //! } //! ``` -use std::sync::{ - atomic::{AtomicU64, Ordering}, - Arc, RwLock, +use std::{ + fmt, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, RwLock, + }, }; use log::{debug, warn}; @@ -66,7 +69,7 @@ use crate::{ filesystem::VirtualFs, loader::{LoadedImage, MappedRegionInfo}, memory::AddressSpace, - process::EmulationConfig, + process::{EmulationConfig, EmulationLimits}, runtime::RuntimeState, thread::ThreadContext, tracer::TraceWriter, @@ -590,7 +593,7 @@ impl EmulationProcess { .filter(move |t| { let fqn = t.fullname(); fqn.ends_with(type_name) - || fqn == type_name + || &*fqn == type_name || fqn.split('.').next_back() == Some(type_name) }) .methods() @@ -757,19 +760,30 @@ impl EmulationProcess { /// /// - **Address space**: Memory regions, managed heap, and static fields /// are forked with per-page/per-object CoW semantics + /// - **Runtime state**: The fork gets its own `AppDomainState`, so assemblies it loads + /// via `Assembly.Load(byte[])`, strings it interns and resolve handlers it registers + /// are invisible to siblings. The `Arc` handles present at fork time are + /// copied, not the metadata behind them + /// - **Virtual filesystem**: Forked, so writes stay local /// /// # What Gets Shared (Immutable) /// /// - **Assembly**: Metadata is read-only, shared via `Arc` /// - **Configuration**: Immutable after creation, shared via `Arc` - /// - **Runtime state**: Method stubs and type info, shared via `Arc` - /// - **Loaded image metadata**: List of loaded images (shallow copy) - /// - **Mapped region metadata**: List of mapped regions (shallow copy) + /// - **Hook manager**: Registered after construction and not mutated during execution + /// - **Native function registry**: Synthetic `GetProcAddress` addresses have to denote + /// the same function in every fork + /// - **Synthetic method counter**: Shared so two forks cannot mint the same synthetic + /// token for different method bodies /// /// # What Gets Fresh /// /// - **Capture context**: Each fork gets fresh captures so results don't mix /// - **Instruction count**: Reset to 0 for independent tracking + /// - **Synthetic methods**: Seeded from the parent, then independent + /// + /// Forks are run in real parallel by callers such as the constant decryption pass, so + /// "independent" here is a concurrency claim, not just a bookkeeping one. /// /// # Performance /// @@ -823,6 +837,48 @@ impl EmulationProcess { }) } + /// The limits this process is currently running under. + /// + /// Callers building a derived budget for [`fork_with_limits`](Self::fork_with_limits) + /// should start from this rather than from [`EmulationLimits::default`]: that replacement + /// is wholesale, so every field the caller does not set explicitly comes from whatever it + /// started from. Deriving from the parent keeps a fork bounded by the same ceilings the + /// parent was warmed under. + #[must_use] + pub fn limits(&self) -> &EmulationLimits { + &self.context.config.limits + } + + /// Forks this process, giving the fork its own execution limits. + /// + /// A template process is built once under whatever budget its *warmup* needs, and every + /// per-method execution is then forked from it. Those are different jobs with different + /// budgets, and [`fork`](Self::fork) cannot express that: it shares the configuration + /// `Arc`, so the fork silently inherits the warmup budget. Overriding the limits here is + /// preferable to mutating the shared configuration, which would retroactively change the + /// template's own budget and race with any sibling fork. + /// + /// `limits` replaces the fork's budget **wholesale**, so build it from + /// [`limits`](Self::limits) rather than from a default, or every field the caller does + /// not name silently reverts. + /// + /// # Errors + /// + /// Returns an error if the process state cannot be forked. + pub fn fork_with_limits(&self, limits: EmulationLimits) -> Result { + let mut config = (*self.context.config).clone(); + config.limits = limits; + + Ok(Self { + name: self.name.clone(), + context: Arc::new(self.context.fork_with_config(Arc::new(config))?), + loaded_images: self.loaded_images.clone(), + mapped_regions: self.mapped_regions.clone(), + instruction_count: AtomicU64::new(0), + trace_writer: self.trace_writer.clone(), + }) + } + /// Forks this process, preserving captured data. /// /// Like [`fork`](Self::fork), but the capture context is also forked, @@ -923,7 +979,7 @@ impl EmulationProcess { } } -impl std::fmt::Debug for EmulationProcess { +impl fmt::Debug for EmulationProcess { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("EmulationProcess") .field("name", &self.name) diff --git a/dotscope/src/emulation/runtime/appdomain.rs b/dotscope/src/emulation/runtime/appdomain.rs index c80f7b18..9e7d6de6 100644 --- a/dotscope/src/emulation/runtime/appdomain.rs +++ b/dotscope/src/emulation/runtime/appdomain.rs @@ -107,7 +107,7 @@ use crate::{emulation::HeapRef, metadata::token::Token, CilObject}; /// // The domain is ready for emulation /// assert!(domain.executing_assembly().is_some()); /// ``` -#[derive(Default)] +#[derive(Clone, Default)] pub struct AppDomainState { /// Loaded assemblies indexed by their simple name. /// @@ -444,13 +444,25 @@ impl AppDomainState { /// to associate frames with their /// originating assembly. /// + /// Returns `None` when the domain already holds `max` assemblies. Emulated code drives + /// this list through `Assembly.Load(byte[])`, and each entry retains a whole parsed + /// metadata graph in host memory that no heap budget accounts for, so the list has to be + /// bounded by something. The caller decides what to do on refusal — the .NET-visible + /// behaviour is a failed load, not a silently dropped registration. + /// /// # Arguments /// /// * `asm` - The parsed [`CilObject`] wrapped in `Arc` - pub fn register_parsed_assembly(&mut self, asm: Arc) -> usize { + /// * `max` - Maximum assemblies this domain may hold, from + /// [`EmulationLimits::max_loaded_assemblies`](crate::emulation::process::EmulationLimits::max_loaded_assemblies) + pub fn register_parsed_assembly(&mut self, asm: Arc, max: usize) -> Option { + if self.loaded_cilobjects.len() >= max { + return None; + } + let index = self.loaded_cilobjects.len(); self.loaded_cilobjects.push(asm); - index + Some(index) } /// Retrieves a previously registered parsed assembly by index. @@ -505,6 +517,60 @@ mod tests { assert!(domain.loaded_assemblies().next().is_none()); } + /// Loads the smallest sample available; only its identity as a parsed assembly matters. + fn sample_assembly() -> Arc { + let path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/samples/WindowsBase.dll"); + Arc::new(CilObject::from_path(&path).expect("sample assembly must load")) + } + + /// Emulated code drives this list through `Assembly.Load(byte[])`, so it must refuse to + /// grow past the configured limit rather than retain a parsed assembly per call. + #[test] + fn register_parsed_assembly_refuses_past_the_cap() { + let mut domain = AppDomainState::new(); + let asm = sample_assembly(); + + assert_eq!( + domain.register_parsed_assembly(Arc::clone(&asm), 2), + Some(0) + ); + assert_eq!( + domain.register_parsed_assembly(Arc::clone(&asm), 2), + Some(1) + ); + assert_eq!(domain.register_parsed_assembly(Arc::clone(&asm), 2), None); + + assert_eq!(domain.parsed_assembly_count(), 2); + } + + /// A cap of zero must refuse everything rather than wrap into "unlimited". + #[test] + fn register_parsed_assembly_honours_a_zero_cap() { + let mut domain = AppDomainState::new(); + + assert_eq!(domain.register_parsed_assembly(sample_assembly(), 0), None); + assert_eq!(domain.parsed_assembly_count(), 0); + } + + /// Forks execute in parallel, so an assembly one fork loads must not appear in another — + /// otherwise a planted type steers a sibling's method resolution. + #[test] + fn cloning_the_domain_keeps_later_registrations_local() { + let mut parent = AppDomainState::new(); + parent.register_parsed_assembly(sample_assembly(), 8); + + let mut fork = parent.clone(); + fork.register_parsed_assembly(sample_assembly(), 8); + + assert_eq!(fork.parsed_assembly_count(), 2); + assert_eq!( + parent.parsed_assembly_count(), + 1, + "a fork's load must not be visible to the parent" + ); + } + #[test] fn test_register_assembly() { let mut domain = AppDomainState::new(); diff --git a/dotscope/src/emulation/runtime/bcl/appdomain.rs b/dotscope/src/emulation/runtime/bcl/appdomain.rs index d150c2f5..effaae20 100644 --- a/dotscope/src/emulation/runtime/bcl/appdomain.rs +++ b/dotscope/src/emulation/runtime/bcl/appdomain.rs @@ -62,7 +62,7 @@ use crate::{ thread::EmulationThread, EmValue, HeapObject, }, - metadata::{token::Token, typesystem::CilFlavor}, + metadata::{token::Token, typesystem::CilFlavor, validation::ValidationConfig}, CilObject, Result, }; @@ -322,15 +322,48 @@ fn assembly_load_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> Pre None, ); - // Parse the loaded assembly for cross-assembly resolution - if let Ok(loaded_asm) = CilObject::from_mem(bytes) { + // Parse the loaded assembly for cross-assembly resolution. + // + // This is a full re-entry into the metadata parser on attacker-chosen bytes from + // inside a hook, so it runs outside the instruction and timeout budgets that bound + // ordinary execution. Both the payload size and the number of assemblies retained + // are therefore checked here rather than left to those budgets. + let limits = &thread.config().limits; + if bytes.len() > limits.max_loaded_assembly_bytes { + debug!( + "Assembly.Load(byte[]): refusing {} byte payload, limit is {}", + bytes.len(), + limits.max_loaded_assembly_bytes + ); + } else if thread.runtime_state().read().is_ok_and(|state| { + state.app_domain().parsed_assembly_count() >= limits.max_loaded_assemblies + }) { + debug!( + "Assembly.Load(byte[]): refusing load, already at the {} assembly limit", + limits.max_loaded_assemblies + ); + } else if let Ok(loaded_asm) = + CilObject::from_mem_with_validation(bytes, ValidationConfig::minimal()) + { + // `minimal()` rather than the production preset: this payload is being parsed + // for cross-assembly *resolution*, not accepted as trustworthy, and the full + // validation pipeline is a second unbudgeted walk over the same hostile bytes. let asm_arc = Arc::new(loaded_asm); if let Ok(mut state) = thread.runtime_state().write() { - let index = state.app_domain_mut().register_parsed_assembly(asm_arc); - debug!( - "Assembly.Load(byte[]): parsed and registered as index {}", - index - ); + match state + .app_domain_mut() + .register_parsed_assembly(asm_arc, limits.max_loaded_assemblies) + { + Some(index) => debug!( + "Assembly.Load(byte[]): parsed and registered as index {index}" + ), + // The count is re-checked under the write lock, so a concurrent fork + // cannot push the domain past the cap between the check above and here. + None => debug!( + "Assembly.Load(byte[]): parsed but not retained, at the {} assembly limit", + limits.max_loaded_assemblies + ), + } } } } diff --git a/dotscope/src/emulation/runtime/bcl/crypto/derivation.rs b/dotscope/src/emulation/runtime/bcl/crypto/derivation.rs index a7b8acef..88bcc10e 100644 --- a/dotscope/src/emulation/runtime/bcl/crypto/derivation.rs +++ b/dotscope/src/emulation/runtime/bcl/crypto/derivation.rs @@ -29,7 +29,12 @@ use crate::utils::derive_pbkdf1_key; use crate::{ emulation::{ - runtime::hook::{Hook, HookContext, HookManager, PreHookResult}, + runtime::{ + bcl::limits::{ + negative_argument, oversized_argument, MAX_DERIVED_KEY_BYTES, MAX_KDF_ITERATIONS, + }, + hook::{Hook, HookContext, HookManager, PreHookResult}, + }, thread::EmulationThread, EmValue, }, @@ -161,14 +166,23 @@ fn password_derive_bytes_get_bytes_pre( ctx: &HookContext<'_>, thread: &mut EmulationThread, ) -> PreHookResult { - let size = ctx + // `cb` is attacker-controlled and sizes both the derivation work and the output buffer. + // Derived keys are inherently small, so a tight ceiling costs nothing and stops a single + // call from allocating gigabytes inside one hook. + let size = match ctx .args .first() .map(usize::try_from) .transpose() .ok() .flatten() - .unwrap_or(16); + { + Some(n) if n > MAX_DERIVED_KEY_BYTES => { + return oversized_argument("DeriveBytes.GetBytes", "cb", n, MAX_DERIVED_KEY_BYTES) + } + Some(n) => n, + None => 16, + }; let heap_ref = if let Some(EmValue::ObjectRef(hr)) = ctx.this { *hr @@ -242,10 +256,26 @@ fn rfc2898_derive_bytes_ctor_pre( _ => Vec::new(), }; - // Intentional cast: iteration count is always positive in crypto operations - #[allow(clippy::cast_sign_loss)] + // The comment this replaced asserted the iteration count "is always positive in crypto + // operations". It is not: the value comes from emulated code, and `as u32` turned a + // negative Int32 into roughly 4.3 billion rounds executed inside a single hook call, with + // no emulation budget evaluated in between. Reject negatives and clamp the magnitude — + // real obfuscator key schedules use values in the low thousands. let iterations: u32 = match ctx.args.get(2) { - Some(val) => i32::try_from(val).unwrap_or(1000) as u32, + Some(val) => match i32::try_from(val).ok().and_then(|v| u32::try_from(v).ok()) { + Some(n) if n > MAX_KDF_ITERATIONS => { + return oversized_argument( + "Rfc2898DeriveBytes..ctor", + "iterations", + n as usize, + MAX_KDF_ITERATIONS as usize, + ) + } + Some(n) => n, + None => { + return negative_argument("Rfc2898DeriveBytes..ctor", "iterations"); + } + }, None => 1000, }; @@ -277,14 +307,23 @@ fn rfc2898_derive_bytes_get_bytes_pre( ctx: &HookContext<'_>, thread: &mut EmulationThread, ) -> PreHookResult { - let size = ctx + // `cb` is attacker-controlled and sizes both the derivation work and the output buffer. + // Derived keys are inherently small, so a tight ceiling costs nothing and stops a single + // call from allocating gigabytes inside one hook. + let size = match ctx .args .first() .map(usize::try_from) .transpose() .ok() .flatten() - .unwrap_or(16); + { + Some(n) if n > MAX_DERIVED_KEY_BYTES => { + return oversized_argument("DeriveBytes.GetBytes", "cb", n, MAX_DERIVED_KEY_BYTES) + } + Some(n) => n, + None => 16, + }; let heap_ref = if let Some(EmValue::ObjectRef(hr)) = ctx.this { *hr @@ -298,11 +337,25 @@ fn rfc2898_derive_bytes_get_bytes_pre( let params = thread.heap().get_key_derivation_params(heap_ref); + // A derivation that cannot be performed is reported, not papered over. Returning + // `vec![0u8; size]` here would hand back an all-zero key that decrypts to plausible + // garbage, and the analyst has no way to tell that apart from a real result. let derived_key = match params { Ok(Some((password, salt, iterations, hash_algorithm))) => { - derive_pbkdf2_key(&password, &salt, iterations, size, &hash_algorithm) + match derive_pbkdf2_key(&password, &salt, iterations, size, &hash_algorithm) { + Ok(key) => key, + Err(e) => return PreHookResult::Error(format!("PBKDF2 derivation failed: {e}")), + } + } + Ok(None) => { + return PreHookResult::Error( + "Rfc2898DeriveBytes.GetBytes called before the key derivation parameters were set" + .to_string(), + ) + } + Err(e) => { + return PreHookResult::Error(format!("failed to read key derivation parameters: {e}")) } - _ => vec![0u8; size], }; match thread.heap().alloc_byte_array(&derived_key) { diff --git a/dotscope/src/emulation/runtime/bcl/interop/marshal.rs b/dotscope/src/emulation/runtime/bcl/interop/marshal.rs index c92cb49b..99374a2b 100644 --- a/dotscope/src/emulation/runtime/bcl/interop/marshal.rs +++ b/dotscope/src/emulation/runtime/bcl/interop/marshal.rs @@ -64,12 +64,18 @@ use crate::{ emulation::{ - runtime::hook::{Hook, HookContext, HookManager, PreHookResult}, + engine::synthetic_exception, + runtime::{ + bcl::limits::{checked_len, MAX_HOOK_BUFFER}, + hook::{Hook, HookContext, HookManager, PreHookResult}, + }, thread::EmulationThread, - tokens, EmValue, HeapObject, + tokens, + value::{ManagedPointer, PointerTarget}, + EmValue, EmulationError, HeapObject, }, metadata::token::Token, - Result, + Error, Result, }; /// Extracts a memory address from a hook argument. @@ -117,11 +123,7 @@ fn resolve_address(arg: &EmValue, thread: &EmulationThread) -> Option { /// Reads the value that the pointer points to (local variable, argument, or field) /// and converts it to i64. Used by IntPtr hooks when the `this` argument is a /// `ldloca`-produced managed pointer instead of a direct value. -fn resolve_managed_ptr_as_i64( - ptr: &crate::emulation::value::ManagedPointer, - thread: &EmulationThread, -) -> Option { - use crate::emulation::value::PointerTarget; +fn resolve_managed_ptr_as_i64(ptr: &ManagedPointer, thread: &EmulationThread) -> Option { let value = match &ptr.target { PointerTarget::Local(idx) => thread.get_local(*idx as usize).ok().cloned(), PointerTarget::StaticField(token) => { @@ -342,26 +344,55 @@ pub fn register(manager: &HookManager) -> Result<()> { Ok(()) } +/// Size of the window materialised for a write to an unmapped address. +const AUTO_ALLOC_SIZE: usize = 0x1_0000; // 64KB + /// Writes bytes to a possibly-unmapped address, auto-allocating if needed. /// -/// Some obfuscated code writes to addresses that don't exist in emulation -/// (e.g., CLR method table addresses). Rather than aborting emulation, we -/// silently allocate a page at the target address and proceed. -fn write_with_auto_alloc(thread: &EmulationThread, addr: u64, data: &[u8]) { - if thread.address_space().write(addr, data).is_err() { - let page_base = addr & !0xFFFF; - let page_size = 0x1_0000usize; // 64KB - log::debug!( - "Auto-allocating 0x{page_size:X} bytes at 0x{page_base:X} for write to 0x{addr:X}" - ); - if thread - .address_space() - .map_data(page_base, &vec![0u8; page_size], "auto-alloc") - .is_ok() - { - let _ = thread.address_space().write(addr, data); - } +/// Some obfuscated code writes to addresses that don't exist in emulation (e.g. CLR method +/// table addresses). Rather than aborting emulation, a window is materialised at the target +/// and the write proceeds. +/// +/// The address is attacker-chosen and this is reached from ordinary CIL, so the allocation is +/// charged against the unmanaged memory budget. Past it the write is refused: a wild pointer +/// then faults, as it would on a real process, instead of quietly becoming valid memory — +/// which is both the containment property and what stops the analyst's picture of the +/// program's memory from including regions the program never legitimately had. +/// +/// # Errors +/// +/// Returns the address space's own error when the window cannot be mapped — the unmanaged +/// budget is exhausted, or the range collides with an existing mapping — or when the write +/// itself fails. +fn write_with_auto_alloc(thread: &EmulationThread, addr: u64, data: &[u8]) -> Result<()> { + let failure = match thread.address_space().write(addr, data) { + Ok(()) => return Ok(()), + Err(failure) => failure, + }; + + // A refusal is not an absence. `AccessViolation` means the page is mapped + // and rejected the write — a PE section carrying no `IMAGE_SCN_MEM_WRITE` + // is the common case, and a protection writing decrypted method bodies back + // into `.text` lands there. Materialising a window would try to map over the + // image the address already belongs to, fail on the overlap, and report that + // instead of the permission that actually stopped the write. + if matches!( + &failure, + Error::Emulation(inner) if matches!(**inner, EmulationError::AccessViolation { .. }) + ) { + return Err(failure); } + + let page_base = addr & !0xFFFF; + log::debug!( + "Auto-allocating 0x{AUTO_ALLOC_SIZE:X} bytes at 0x{page_base:X} for write to 0x{addr:X}" + ); + + thread + .address_space() + .alloc_unmanaged_at(page_base, AUTO_ALLOC_SIZE, "auto-alloc")?; + + thread.address_space().write(addr, data) } /// Hook for `System.Runtime.InteropServices.Marshal.GetHINSTANCE` method. @@ -435,15 +466,45 @@ fn marshal_copy_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> PreH EmValue::ObjectRef(r) => *r, _ => return PreHookResult::Bypass(None), }; + // `cast_unsigned() as usize` turned a negative Int32 into a value near 4 GiB, which + // then drove both an address-space read and a per-byte loop taking the heap lock. + // .NET throws ArgumentOutOfRangeException for either argument being negative. let start_idx = match arg2 { - EmValue::I32(v) => (*v).cast_unsigned() as usize, + EmValue::I32(v) => match checked_len(*v, MAX_HOOK_BUFFER, "Marshal.Copy", "startIndex") + { + Ok(n) => n, + Err(result) => return result, + }, _ => return PreHookResult::Bypass(None), }; let length = match arg3 { - EmValue::I32(v) => (*v).cast_unsigned() as usize, + EmValue::I32(v) => match checked_len(*v, MAX_HOOK_BUFFER, "Marshal.Copy", "length") { + Ok(n) => n, + Err(result) => return result, + }, _ => return PreHookResult::Bypass(None), }; + // Validate the destination window before reading anything, so an overrun is reported + // as the managed exception rather than discovered part-way through the copy. + match thread.heap().get_array_length(dst_ref) { + Ok(dst_len) => { + if start_idx + .checked_add(length) + .is_none_or(|end| end > dst_len) + { + return PreHookResult::Throw { + exception_type: synthetic_exception::ARGUMENT_EXCEPTION, + message: format!( + "Marshal.Copy: destination range {start_idx}..+{length} exceeds \ + array length {dst_len}" + ), + }; + } + } + Err(e) => return PreHookResult::Error(format!("Marshal.Copy: {e}")), + } + if let Ok(bytes) = thread.address_space().read(src_addr, length) { for (i, &byte) in bytes.iter().enumerate() { let Some(idx) = start_idx.checked_add(i) else { @@ -464,7 +525,10 @@ fn marshal_copy_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> PreH return PreHookResult::Bypass(None); }; let start_idx = match arg1 { - EmValue::I32(v) => (*v).cast_unsigned() as usize, + EmValue::I32(v) => match checked_len(*v, MAX_HOOK_BUFFER, "Marshal.Copy", "startIndex") { + Ok(n) => n, + Err(result) => return result, + }, _ => return PreHookResult::Bypass(None), }; let dest_addr = match arg2 { @@ -473,10 +537,33 @@ fn marshal_copy_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> PreH _ => return PreHookResult::Bypass(None), }; let length = match arg3 { - EmValue::I32(v) => (*v).cast_unsigned() as usize, + EmValue::I32(v) => match checked_len(*v, MAX_HOOK_BUFFER, "Marshal.Copy", "length") { + Ok(n) => n, + Err(result) => return result, + }, _ => return PreHookResult::Bypass(None), }; + // Validate the source window up front rather than letting per-element read failures + // degrade into a silent run of zero bytes. + match thread.heap().get_array_length(*src_ref) { + Ok(src_len) => { + if start_idx + .checked_add(length) + .is_none_or(|end| end > src_len) + { + return PreHookResult::Throw { + exception_type: synthetic_exception::ARGUMENT_EXCEPTION, + message: format!( + "Marshal.Copy: source range {start_idx}..+{length} exceeds array \ + length {src_len}" + ), + }; + } + } + Err(e) => return PreHookResult::Error(format!("Marshal.Copy: {e}")), + } + let mut bytes = Vec::with_capacity(length); for i in 0..length { let Some(idx) = start_idx.checked_add(i) else { @@ -494,7 +581,9 @@ fn marshal_copy_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> PreH bytes.push(byte_val); } - write_with_auto_alloc(thread, dest_addr, &bytes); + if let Err(e) = write_with_auto_alloc(thread, dest_addr, &bytes) { + return PreHookResult::throw_access_violation(&e.to_string()); + } PreHookResult::Bypass(None) } @@ -592,7 +681,9 @@ fn marshal_write_byte_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) - _ => return PreHookResult::Bypass(None), }; - write_with_auto_alloc(thread, addr, &[value]); + if let Err(e) = write_with_auto_alloc(thread, addr, &[value]) { + return PreHookResult::throw_access_violation(&e.to_string()); + } PreHookResult::Bypass(None) } @@ -626,7 +717,9 @@ fn marshal_write_int32_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) _ => return PreHookResult::Bypass(None), }; - write_with_auto_alloc(thread, addr, &value.to_le_bytes()); + if let Err(e) = write_with_auto_alloc(thread, addr, &value.to_le_bytes()) { + return PreHookResult::throw_access_violation(&e.to_string()); + } PreHookResult::Bypass(None) } @@ -773,7 +866,9 @@ fn marshal_write_int64_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) }; let final_addr = (addr as i64).wrapping_add(offset).cast_unsigned(); - write_with_auto_alloc(thread, final_addr, &value.to_le_bytes()); + if let Err(e) = write_with_auto_alloc(thread, final_addr, &value.to_le_bytes()) { + return PreHookResult::throw_access_violation(&e.to_string()); + } PreHookResult::Bypass(None) } @@ -813,7 +908,9 @@ fn marshal_write_int16_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) }; let final_addr = (addr as i64).wrapping_add(offset).cast_unsigned(); - write_with_auto_alloc(thread, final_addr, &value.to_le_bytes()); + if let Err(e) = write_with_auto_alloc(thread, final_addr, &value.to_le_bytes()) { + return PreHookResult::throw_access_violation(&e.to_string()); + } PreHookResult::Bypass(None) } @@ -858,11 +955,15 @@ fn marshal_write_intptr_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) let final_addr = (addr as i64).wrapping_add(offset).cast_unsigned(); let ptr_size = ctx.pointer_size.bytes(); if ptr_size == 8 { - write_with_auto_alloc(thread, final_addr, &value.to_le_bytes()); + if let Err(e) = write_with_auto_alloc(thread, final_addr, &value.to_le_bytes()) { + return PreHookResult::throw_access_violation(&e.to_string()); + } } else { #[allow(clippy::cast_possible_truncation)] let val32 = value as i32; - write_with_auto_alloc(thread, final_addr, &val32.to_le_bytes()); + if let Err(e) = write_with_auto_alloc(thread, final_addr, &val32.to_le_bytes()) { + return PreHookResult::throw_access_violation(&e.to_string()); + } } PreHookResult::Bypass(None) } @@ -878,31 +979,57 @@ fn marshal_alloc_cotaskmem_pre( ctx: &HookContext<'_>, thread: &mut EmulationThread, ) -> PreHookResult { - let size = ctx - .args - .first() - .and_then(EmValue::as_i32) - .unwrap_or(0) - .max(0) as usize; + let size = match checked_len( + ctx.args.first().and_then(EmValue::as_i32).unwrap_or(0), + MAX_HOOK_BUFFER, + "Marshal.AllocCoTaskMem", + "cb", + ) { + Ok(n) => n.max(1), // Minimum 1 byte allocation + Err(result) => return result, + }; - let size = size.max(1); // Minimum 1 byte allocation match thread.address_space().alloc_unmanaged(size) { Ok(addr) => PreHookResult::Bypass(Some(EmValue::NativeInt(addr as i64))), - Err(_) => PreHookResult::Bypass(Some(EmValue::NativeInt(0))), + // .NET signals allocation failure by throwing, and returning a null IntPtr instead + // lets emulated code carry on writing through a pointer that was never allocated. + Err(e) => PreHookResult::Throw { + exception_type: synthetic_exception::OUT_OF_MEMORY, + message: format!("Marshal.AllocCoTaskMem: {e}"), + }, } } /// Hook for `System.Runtime.InteropServices.Marshal.FreeCoTaskMem` method. /// -/// No-op in emulation — memory is not individually freed. +/// Releases the region and returns its bytes to the unmanaged allocation budget. Emulated +/// code that allocates and frees in a loop would otherwise exhaust that budget, since the +/// allocations are real but the frees were not. /// /// # Handled Overloads /// /// - `Marshal.FreeCoTaskMem(IntPtr) -> void` fn marshal_free_cotaskmem_pre( - _ctx: &HookContext<'_>, - _thread: &mut EmulationThread, + ctx: &HookContext<'_>, + thread: &mut EmulationThread, ) -> PreHookResult { + free_unmanaged_arg(ctx, thread) +} + +/// Shared body for `FreeHGlobal` and `FreeCoTaskMem`. +/// +/// A free of a null or unrecognised pointer is ignored rather than reported: .NET treats +/// `Free*(IntPtr.Zero)` as a no-op, and emulated code frequently frees pointers this emulator +/// never handed out (for instance when a paired allocation was bypassed by another hook). +fn free_unmanaged_arg(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> PreHookResult { + let address = match ctx.args.first() { + Some(EmValue::UnmanagedPtr(a)) => *a, + Some(EmValue::NativeInt(a)) => (*a).cast_unsigned(), + _ => return PreHookResult::Bypass(None), + }; + if address != 0 { + let _ = thread.address_space().free_unmanaged(address); + } PreHookResult::Bypass(None) } @@ -916,30 +1043,35 @@ fn marshal_free_cotaskmem_pre( /// - `Marshal.AllocHGlobal(IntPtr) -> IntPtr` fn marshal_alloc_hglobal_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> PreHookResult { let size = match ctx.args.first() { - Some(EmValue::I32(v)) => (*v).max(0) as usize, - Some(EmValue::NativeInt(v)) => (*v).max(0) as usize, - _ => 0, + Some(EmValue::I32(v)) => checked_len(*v, MAX_HOOK_BUFFER, "Marshal.AllocHGlobal", "cb"), + Some(EmValue::NativeInt(v)) => { + checked_len(*v, MAX_HOOK_BUFFER, "Marshal.AllocHGlobal", "cb") + } + _ => Ok(0), + }; + let size = match size { + Ok(n) => n.max(1), + Err(result) => return result, }; - let size = size.max(1); match thread.address_space().alloc_unmanaged(size) { Ok(addr) => PreHookResult::Bypass(Some(EmValue::NativeInt(addr as i64))), - Err(_) => PreHookResult::Bypass(Some(EmValue::NativeInt(0))), + Err(e) => PreHookResult::Throw { + exception_type: synthetic_exception::OUT_OF_MEMORY, + message: format!("Marshal.AllocHGlobal: {e}"), + }, } } /// Hook for `System.Runtime.InteropServices.Marshal.FreeHGlobal` method. /// -/// No-op in emulation — memory is not individually freed. +/// Releases the region and returns its bytes to the unmanaged allocation budget. /// /// # Handled Overloads /// /// - `Marshal.FreeHGlobal(IntPtr) -> void` -fn marshal_free_hglobal_pre( - _ctx: &HookContext<'_>, - _thread: &mut EmulationThread, -) -> PreHookResult { - PreHookResult::Bypass(None) +fn marshal_free_hglobal_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> PreHookResult { + free_unmanaged_arg(ctx, thread) } /// Hook for `System.Runtime.InteropServices.Marshal.SizeOf` method. @@ -1028,10 +1160,12 @@ fn marshal_get_delegate_for_function_pointer_pre( runtime .as_ref() .and_then(|rt| { - let name = rt.native_functions().lookup_by_address(addr)?; - let token = rt.native_functions().allocate_token(&name); + let function = rt.native_functions().lookup_by_address(addr)?; + let token = rt.native_functions().allocate_token(&function); log::debug!( - "GetDelegateForFunctionPointer(0x{addr:X}) → {name} → token 0x{:08X}", + "GetDelegateForFunctionPointer(0x{addr:X}) → {}!{} → token 0x{:08X}", + function.dll.as_deref().unwrap_or("?"), + function.name, token.value() ); Some(token) @@ -1080,9 +1214,12 @@ fn marshal_get_function_pointer_for_delegate_pre( let method_token = entry.method_token; if tokens::is_native_function_pointer(method_token) { if let Ok(rt) = thread.runtime_state().read() { - if let Some(name) = rt.native_functions().lookup_by_token(method_token) { - if let Some(addr) = rt.native_functions().lookup_address_by_name(&name) + if let Some(function) = rt.native_functions().lookup_by_token(method_token) + { + if let Some(addr) = + rt.native_functions().lookup_address_by_name(&function.name) { + let name = &function.name; log::debug!("GetFunctionPointerForDelegate → {name} at 0x{addr:X}"); return PreHookResult::Bypass(Some(EmValue::NativeInt( addr as i64, diff --git a/dotscope/src/emulation/runtime/bcl/io/binaryreader.rs b/dotscope/src/emulation/runtime/bcl/io/binaryreader.rs index 289ae49d..1c9f5952 100644 --- a/dotscope/src/emulation/runtime/bcl/io/binaryreader.rs +++ b/dotscope/src/emulation/runtime/bcl/io/binaryreader.rs @@ -36,7 +36,10 @@ use crate::{ emulation::{ runtime::{ - bcl::io::stream::{stream_close_pre, stream_dispose_pre}, + bcl::{ + io::stream::{stream_close_pre, stream_dispose_pre}, + limits::{checked_len, MAX_HOOK_BUFFER}, + }, hook::{Hook, HookContext, HookManager, PreHookResult}, }, thread::EmulationThread, @@ -311,10 +314,13 @@ fn binary_reader_read_bytes_pre( ctx: &HookContext<'_>, thread: &mut EmulationThread, ) -> PreHookResult { - // Safe: value validated as non-negative - #[allow(clippy::cast_sign_loss)] let count = match ctx.args.first() { - Some(EmValue::I32(v)) => *v as usize, + Some(EmValue::I32(v)) => { + match checked_len(*v, MAX_HOOK_BUFFER, "BinaryReader.ReadBytes", "count") { + Ok(n) => n, + Err(result) => return result, + } + } _ => 0, }; @@ -716,9 +722,13 @@ fn binary_reader_read_chars_pre( ctx: &HookContext<'_>, thread: &mut EmulationThread, ) -> PreHookResult { - #[allow(clippy::cast_sign_loss)] let count = match ctx.args.first() { - Some(EmValue::I32(v)) => *v as usize, + Some(EmValue::I32(v)) => { + match checked_len(*v, MAX_HOOK_BUFFER, "BinaryReader.ReadChars", "count") { + Ok(n) => n, + Err(result) => return result, + } + } _ => 0, }; @@ -732,8 +742,12 @@ fn binary_reader_read_chars_pre( }; let Some(chars) = try_hook!(thread.heap().with_stream(stream_ref, |data, position| { - // Decode `count` UTF-8 characters - let mut chars = Vec::with_capacity(count); + // Reserve for what the stream can actually supply, not for what was asked. One char + // needs at least one byte, so the remaining byte count is a hard upper bound on the + // number of chars this call can return. + let available = data.len().saturating_sub(*position); + let mut chars = Vec::with_capacity(count.min(available)); + for _ in 0..count { let Some(remaining) = data.get(*position..) else { break; @@ -741,12 +755,36 @@ fn binary_reader_read_chars_pre( if remaining.is_empty() { break; } - let s = String::from_utf8_lossy(remaining); - if let Some(ch) = s.chars().next() { - chars.push(ch); - *position = position.saturating_add(ch.len_utf8()); - } else { - break; + + // Decode a single scalar from the front of the slice. Running + // `String::from_utf8_lossy` over the whole remainder on each iteration would + // re-scan and re-allocate the entire stream tail per character, making this loop + // quadratic in stream length. + let mut decoded = None; + let window = remaining.len().min(4); + for take in 1..=window { + let Some(prefix) = remaining.get(..take) else { + break; + }; + if let Ok(valid) = std::str::from_utf8(prefix) { + if let Some(ch) = valid.chars().next() { + decoded = Some((ch, take)); + break; + } + } + } + + match decoded { + Some((ch, len)) => { + chars.push(ch); + *position = position.saturating_add(len); + } + // Not a valid scalar in any 1..=4 byte prefix: substitute and advance one + // byte, matching `from_utf8_lossy`'s behaviour without its cost. + None => { + chars.push(char::REPLACEMENT_CHARACTER); + *position = position.saturating_add(1); + } } } chars diff --git a/dotscope/src/emulation/runtime/bcl/io/filestream.rs b/dotscope/src/emulation/runtime/bcl/io/filestream.rs index 9b14f936..e147448f 100644 --- a/dotscope/src/emulation/runtime/bcl/io/filestream.rs +++ b/dotscope/src/emulation/runtime/bcl/io/filestream.rs @@ -538,11 +538,9 @@ fn path_get_directory_name_pre( None => return PreHookResult::Bypass(Some(EmValue::Null)), }; - let dir = if let Some(pos) = path.rfind(['\\', '/']) { - &path[..pos] - } else { - "" - }; + let dir = path + .rsplit_once(['\\', '/']) + .map_or("", |(parent, _)| parent); alloc_string_result(thread, dir) } @@ -593,11 +591,7 @@ fn path_get_file_name_without_extension_pre( None => return PreHookResult::Bypass(Some(EmValue::Null)), }; let filename = path_filename(&path); - let without_ext = if let Some(pos) = filename.rfind('.') { - &filename[..pos] - } else { - filename - }; + let without_ext = filename.rsplit_once('.').map_or(filename, |(stem, _)| stem); alloc_string_result(thread, without_ext) } @@ -608,11 +602,12 @@ fn path_get_extension_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) - None => return PreHookResult::Bypass(Some(EmValue::Null)), }; let filename = path_filename(&path); - let ext = if let Some(pos) = filename.rfind('.') { - &filename[pos..] - } else { - "" - }; + // `split_at` rather than `&filename[pos..]`: `rfind` always returns a character boundary, so + // both are sound, but this keeps the crate free of `str` range indexing (`clippy::string_slice`). + // The extension keeps its leading dot, which `rsplit_once` would consume. + let ext = filename + .rfind('.') + .map_or("", |pos| filename.split_at(pos).1); alloc_string_result(thread, ext) } @@ -656,16 +651,10 @@ fn path_change_extension_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread let new_ext = extract_nth_string_arg(ctx, thread, 1).unwrap_or_default(); // Strip existing extension - let base = if let Some(pos) = path.rfind('.') { - // Only strip if the dot is in the filename portion - let last_sep = path.rfind(['\\', '/']).unwrap_or(0); - if pos > last_sep { - &path[..pos] - } else { - &path - } - } else { - &path + let base = match path.rfind('.') { + // Only strip if the dot is in the filename portion, not in a parent directory name. + Some(pos) if pos > path.rfind(['\\', '/']).unwrap_or(0) => path.split_at(pos).0, + _ => &path, }; // Append new extension (ensure it starts with '.') @@ -785,7 +774,10 @@ fn fileinfo_get_extension_pre( None => return PreHookResult::Bypass(Some(EmValue::Null)), }; let filename = path_filename(&path); - let ext = filename.rfind('.').map_or("", |pos| &filename[pos..]); + // Extension keeps its leading dot; see the note in `path_get_extension_pre`. + let ext = filename + .rfind('.') + .map_or("", |pos| filename.split_at(pos).1); alloc_string_result(thread, ext) } @@ -798,11 +790,9 @@ fn fileinfo_get_directory_name_pre( Some(p) => p, None => return PreHookResult::Bypass(Some(EmValue::Null)), }; - let dir = if let Some(pos) = path.rfind(['\\', '/']) { - &path[..pos] - } else { - "" - }; + let dir = path + .rsplit_once(['\\', '/']) + .map_or("", |(parent, _)| parent); alloc_string_result(thread, dir) } diff --git a/dotscope/src/emulation/runtime/bcl/io/stream.rs b/dotscope/src/emulation/runtime/bcl/io/stream.rs index 0a29b53c..2047ff09 100644 --- a/dotscope/src/emulation/runtime/bcl/io/stream.rs +++ b/dotscope/src/emulation/runtime/bcl/io/stream.rs @@ -66,7 +66,10 @@ use crate::{ emulation::{ - runtime::hook::{Hook, HookContext, HookManager, PreHookResult}, + runtime::{ + bcl::limits::{checked_len, MAX_HOOK_BUFFER}, + hook::{Hook, HookContext, HookManager, PreHookResult}, + }, thread::EmulationThread, EmValue, }, @@ -725,10 +728,19 @@ fn stream_set_length_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> _ => return PreHookResult::Bypass(None), }; - #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] let new_length = match ctx.args.first() { - Some(EmValue::I64(v)) => *v as usize, - Some(EmValue::I32(v)) => *v as usize, + Some(EmValue::I64(v)) => { + match checked_len(*v, MAX_HOOK_BUFFER, "Stream.SetLength", "value") { + Ok(n) => n, + Err(result) => return result, + } + } + Some(EmValue::I32(v)) => { + match checked_len(*v, MAX_HOOK_BUFFER, "Stream.SetLength", "value") { + Ok(n) => n, + Err(result) => return result, + } + } _ => return PreHookResult::Bypass(None), }; diff --git a/dotscope/src/emulation/runtime/bcl/limits.rs b/dotscope/src/emulation/runtime/bcl/limits.rs new file mode 100644 index 00000000..0608ec4c --- /dev/null +++ b/dotscope/src/emulation/runtime/bcl/limits.rs @@ -0,0 +1,193 @@ +//! Argument validation and allocation ceilings shared by BCL method hooks. +//! +//! # Why hooks need their own ceilings +//! +//! The emulator's instruction budget is evaluated between CIL instructions, in the dispatch +//! loop. A BCL hook runs entirely *inside* one such instruction: once dispatch enters +//! `String.PadLeft` or `Rfc2898DeriveBytes.GetBytes`, no budget is consulted again until the +//! hook returns. A hook that loops or allocates in proportion to an argument it took from the +//! emulated evaluation stack is therefore unbounded by construction, no matter how small the +//! configured instruction limit is. +//! +//! Every hook that sizes work from an emulated value must bound that value itself, before +//! doing the work. The helpers here provide the two checks that need to happen together: +//! +//! 1. **Sign.** Emulated counts arrive as `int32`/`int64`. A bare `as usize` turns `-1` into +//! `usize::MAX`, which downstream code then tries to honour. Use [`checked_len`], which +//! converts with `TryFrom` and reports a negative argument as the managed exception .NET +//! itself raises. +//! 2. **Magnitude.** A positive but absurd count is just as effective — `int.MaxValue` is a +//! perfectly valid `usize`. [`checked_len`] also enforces a ceiling. +//! +//! Failures are returned as [`PreHookResult::Throw`] rather than +//! [`PreHookResult::Error`](crate::emulation::runtime::hook::PreHookResult::Error), so they +//! flow through CIL exception handling. Emulated code that guards a large allocation with +//! `try`/`catch` — which obfuscator runtimes and packers routinely do — then behaves as it +//! would on a real runtime, instead of the whole emulation being abandoned. + +use crate::{ + emulation::{engine::synthetic_exception, runtime::hook::PreHookResult}, + utils::MAX_DERIVED_KEY_LEN, +}; + +/// Ceiling for a general-purpose buffer materialised inside a single hook. +/// +/// Sized to be far above anything real code asks for in one call while staying small enough +/// that a rejected request costs nothing. Hooks whose output has a naturally tighter bound +/// should use a more specific constant rather than this one. +pub(crate) const MAX_HOOK_BUFFER: usize = 16 * 1024 * 1024; + +/// Ceiling for a single string built by a hook, in `char`s. +/// +/// Padding and repeat operations take a target width from emulated code; this bounds the +/// resulting string independently of [`MAX_HOOK_BUFFER`] because strings are measured in +/// scalars rather than bytes. +pub(crate) const MAX_HOOK_STRING_CHARS: usize = 4 * 1024 * 1024; + +/// Ceiling for key material produced by a key-derivation hook. +/// +/// Derived keys are small by nature — the largest symmetric key in common use is 32 bytes, +/// and obfuscator key schedules do not exceed a few hundred. +/// +/// Defined as the limit [`derive_pbkdf2_key`] actually enforces rather than as its own value. +/// The two were independent (4096 here, 1024 there), so a `GetBytes(n)` with +/// `1024 < n <= 4096` passed this check and then hard-failed inside the derivation — an +/// internal error where the hook should have thrown a managed exception. Whichever bound is +/// tighter has to be the one the hook applies, or the hook is not the gate it appears to be. +/// +/// [`derive_pbkdf2_key`]: crate::utils::derive_pbkdf2_key +pub(crate) const MAX_DERIVED_KEY_BYTES: usize = MAX_DERIVED_KEY_LEN; + +/// Ceiling for a key-derivation iteration count. +/// +/// PBKDF2 work is linear in this value and runs to completion inside one hook call, so a +/// large count is a wall-clock denial of service even though it allocates nothing. Real +/// obfuscators use values in the low thousands. +pub(crate) const MAX_KDF_ITERATIONS: u32 = 100_000; + +/// Builds the exception thrown when an emulated count argument is negative. +pub(crate) fn negative_argument(method: &str, param: &str) -> PreHookResult { + PreHookResult::Throw { + exception_type: synthetic_exception::ARGUMENT_OUT_OF_RANGE, + message: format!("{method}: '{param}' must be non-negative"), + } +} + +/// Builds the exception thrown when an emulated count argument exceeds its ceiling. +pub(crate) fn oversized_argument( + method: &str, + param: &str, + requested: usize, + max: usize, +) -> PreHookResult { + PreHookResult::Throw { + exception_type: synthetic_exception::OUT_OF_MEMORY, + message: format!( + "{method}: '{param}' of {requested} exceeds the emulator's per-call limit of {max}" + ), + } +} + +/// Converts an emulated count to `usize`, rejecting negative and oversized values. +/// +/// This is the single entry point hooks should use for any argument that will size an +/// allocation or bound a loop. +/// +/// # Arguments +/// +/// * `value` — the count as it came off the emulated evaluation stack. +/// * `max` — the ceiling to enforce; see the constants in this module. +/// * `method` — the .NET method name, for the exception message (e.g. `"String.PadLeft"`). +/// * `param` — the .NET parameter name, for the exception message (e.g. `"totalWidth"`). +/// +/// # Errors +/// +/// Returns the [`PreHookResult`] the caller should return directly: an +/// `ArgumentOutOfRangeException` for a negative value, or an `OutOfMemoryException` for one +/// above `max`. +// The `Err` variant is a `PreHookResult`, which is large because it can carry an `EmValue`. +// That is inherent rather than incidental: the error here *is* the value the calling hook +// returns, and every hook in the tree already returns `PreHookResult` by value. Boxing it +// would add a deref at each of these call sites without removing the cost anywhere. +#[allow(clippy::result_large_err)] +pub(crate) fn checked_len( + value: T, + max: usize, + method: &str, + param: &str, +) -> Result +where + usize: TryFrom, +{ + let Ok(len) = usize::try_from(value) else { + return Err(negative_argument(method, param)); + }; + if len > max { + return Err(oversized_argument(method, param, len, max)); + } + Ok(len) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Asserts that `result` is a throw of `expected` exception type. + fn assert_throws(result: &PreHookResult, expected: crate::metadata::token::Token) { + match result { + PreHookResult::Throw { exception_type, .. } => assert_eq!(*exception_type, expected), + other => panic!("expected a throw, got {other:?}"), + } + } + + #[test] + fn accepts_a_reasonable_count() { + assert_eq!( + checked_len(1024_i32, MAX_HOOK_BUFFER, "M", "count").unwrap(), + 1024 + ); + } + + #[test] + fn accepts_zero() { + assert_eq!( + checked_len(0_i32, MAX_HOOK_BUFFER, "M", "count").unwrap(), + 0 + ); + } + + #[test] + fn rejects_negative_i32_as_argument_out_of_range() { + let err = checked_len(-1_i32, MAX_HOOK_BUFFER, "M", "count").unwrap_err(); + assert_throws(&err, synthetic_exception::ARGUMENT_OUT_OF_RANGE); + } + + #[test] + fn rejects_negative_i64_as_argument_out_of_range() { + let err = checked_len(i64::MIN, MAX_HOOK_BUFFER, "M", "count").unwrap_err(); + assert_throws(&err, synthetic_exception::ARGUMENT_OUT_OF_RANGE); + } + + #[test] + fn rejects_oversized_count_as_out_of_memory() { + let err = checked_len(i32::MAX, 1024, "M", "count").unwrap_err(); + assert_throws(&err, synthetic_exception::OUT_OF_MEMORY); + } + + #[test] + fn accepts_exactly_the_ceiling() { + assert_eq!(checked_len(1024_i32, 1024, "M", "count").unwrap(), 1024); + } + + #[test] + fn message_names_the_method_and_parameter() { + let err = checked_len(-5_i32, MAX_HOOK_BUFFER, "String.PadLeft", "totalWidth").unwrap_err(); + match err { + PreHookResult::Throw { message, .. } => { + assert!(message.contains("String.PadLeft"), "got: {message}"); + assert!(message.contains("totalWidth"), "got: {message}"); + } + other => panic!("expected a throw, got {other:?}"), + } + } +} diff --git a/dotscope/src/emulation/runtime/bcl/mod.rs b/dotscope/src/emulation/runtime/bcl/mod.rs index 0fd6fdb9..65b1d78c 100644 --- a/dotscope/src/emulation/runtime/bcl/mod.rs +++ b/dotscope/src/emulation/runtime/bcl/mod.rs @@ -72,6 +72,7 @@ mod collections; mod crypto; mod interop; mod io; +mod limits; mod reflection; mod runtime; mod statics; diff --git a/dotscope/src/emulation/runtime/bcl/reflection/helpers.rs b/dotscope/src/emulation/runtime/bcl/reflection/helpers.rs index bd599d82..5cb2a463 100644 --- a/dotscope/src/emulation/runtime/bcl/reflection/helpers.rs +++ b/dotscope/src/emulation/runtime/bcl/reflection/helpers.rs @@ -9,7 +9,7 @@ use crate::{ metadata::{ method::MethodRc, token::Token, - typesystem::{CilFlavor, CilPrimitiveKind, CilTypeReference}, + typesystem::{CilFlavor, CilPrimitiveKind, CilType, CilTypeReference}, }, CilObject, }; @@ -81,35 +81,40 @@ pub(crate) fn bcl_name_to_primitive_token(name: &str) -> Option { Some(kind.token()) } -/// Finds a method by name on a type, searching the inheritance chain. +/// Finds the best-matching method declared directly on `cil_type`. /// /// If multiple overloads match, prefers the one with fewer parameters -/// (common obfuscator pattern). Walks `cil_type.base()` if not found on -/// the immediate type. -pub(crate) fn find_method_by_name(asm: &CilObject, type_token: Token, name: &str) -> Option { - if let Some(cil_type) = asm.types().resolve(&type_token) { - // Search the type's own methods - let mut best: Option<(Token, usize)> = None; - for (_, method_weak) in cil_type.methods.iter() { - if let Some(method) = method_weak.upgrade() { - if method.name == name { - let param_count = method.signature.params.len(); - if best.is_none_or(|(_, n)| param_count < n) { - best = Some((method.token, param_count)); - } +/// (common obfuscator pattern). +fn find_method_on_type(cil_type: &CilType, name: &str) -> Option { + let mut best: Option<(Token, usize)> = None; + for (_, method_weak) in cil_type.methods.iter() { + if let Some(method) = method_weak.upgrade() { + if method.name == name { + let param_count = method.signature.params.len(); + if best.is_none_or(|(_, n)| param_count < n) { + best = Some((method.token, param_count)); } } } - if let Some((token, _)) = best { - return Some(token); - } + } + best.map(|(token, _)| token) +} - // Walk the inheritance chain - if let Some(base_rc) = cil_type.base() { - return find_method_by_name(asm, base_rc.token, name); - } +/// Finds a method by name on a type, searching the inheritance chain. +/// +/// If multiple overloads match, prefers the one with fewer parameters +/// (common obfuscator pattern). Ancestors are visited via +/// [`CilType::base_chain`](crate::metadata::typesystem::CilType::base_chain), which bounds the +/// walk — a hostile assembly can describe an inheritance cycle, and searching for a name that +/// appears nowhere in it would otherwise recurse until the native stack is exhausted. +pub(crate) fn find_method_by_name(asm: &CilObject, type_token: Token, name: &str) -> Option { + let cil_type = asm.types().resolve(&type_token)?; + if let Some(token) = find_method_on_type(&cil_type, name) { + return Some(token); } - None + cil_type + .base_chain() + .find_map(|ancestor| find_method_on_type(&ancestor, name)) } /// Resolves a method token (MethodDef, MemberRef, or MethodSpec) to a [`MethodRc`]. diff --git a/dotscope/src/emulation/runtime/bcl/reflection/modules.rs b/dotscope/src/emulation/runtime/bcl/reflection/modules.rs index 330e559b..3ffed5bc 100644 --- a/dotscope/src/emulation/runtime/bcl/reflection/modules.rs +++ b/dotscope/src/emulation/runtime/bcl/reflection/modules.rs @@ -257,7 +257,10 @@ fn module_get_fully_qualified_name_pre( let tables = asm.tables()?; let strings = asm.strings()?; let module_table = tables.table::()?; - let module_row = module_table.iter().next()?; + // The Module table has exactly one row and it is RID 1 (ECMA-335 II.22.30), so + // fetch it by RID. Taking the first row that *parses* silently promotes row 2 + // when row 1 is malformed. + let module_row = module_table.get(1).ok().flatten()?; strings.get(module_row.name as usize).ok().map(String::from) }) .unwrap_or_else(|| "module.exe".to_string()); @@ -444,7 +447,10 @@ fn assembly_get_location_pre( let tables = asm.tables()?; let strings = asm.strings()?; let module_table = tables.table::()?; - let module_row = module_table.iter().next()?; + // The Module table has exactly one row and it is RID 1 (ECMA-335 II.22.30), so + // fetch it by RID. Taking the first row that *parses* silently promotes row 2 + // when row 1 is malformed. + let module_row = module_table.get(1).ok().flatten()?; strings.get(module_row.name as usize).ok().map(String::from) }) .unwrap_or_else(|| "module.exe".to_string()); diff --git a/dotscope/src/emulation/runtime/bcl/reflection/types.rs b/dotscope/src/emulation/runtime/bcl/reflection/types.rs index f4d026cf..247ee2cb 100644 --- a/dotscope/src/emulation/runtime/bcl/reflection/types.rs +++ b/dotscope/src/emulation/runtime/bcl/reflection/types.rs @@ -1000,17 +1000,16 @@ fn type_get_is_value_type_pre( } // Matches Mono: IsSubclassOf(typeof(ValueType)) - // Walk the base type chain looking for System.ValueType - let mut current = cil_type.base(); - while let Some(ancestor) = current { + // Walk the base type chain looking for System.ValueType. + // `base_chain` bounds the walk against cyclic `extends` graphs. + for ancestor in cil_type.base_chain() { let name = ancestor.fullname(); - if name == "System.ValueType" { + if &*name == "System.ValueType" { return PreHookResult::Bypass(Some(EmValue::I32(1))); } - if name == "System.Object" { + if &*name == "System.Object" { break; } - current = ancestor.base(); } // Fallback: check CilFlavor for value types whose base chain @@ -1863,17 +1862,17 @@ fn type_is_assignable_from_pre( if this_token == other_token { return PreHookResult::Bypass(Some(EmValue::I32(1))); } - // Walk inheritance chain of 'other' looking for 'this' + // Walk inheritance chain of 'other' looking for 'this'. + // `base_chain` bounds the walk: a hostile assembly can describe a cyclic + // `extends` graph, on which a hand-rolled token loop never terminates. if let Some(asm) = thread.assembly().cloned() { - let mut current = Some(other_token); - while let Some(tok) = current { - if tok == this_token { + if let Some(other_type) = asm.types().resolve(&other_token) { + if other_type + .base_chain() + .any(|ancestor| ancestor.token == this_token) + { return PreHookResult::Bypass(Some(EmValue::I32(1))); } - current = asm - .types() - .get(&tok) - .and_then(|t| t.base().map(|b| b.token)); } // Check interfaces if let Some(other_type) = asm.types().resolve(&other_token) { @@ -1905,19 +1904,15 @@ fn type_is_subclass_of_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) try_hook!(thread.heap().get_reflection_type_token(*other_ref)) { if let Some(asm) = thread.assembly().cloned() { - // Walk inheritance chain of 'this' looking for 'other' - let mut current = asm - .types() - .get(&this_token) - .and_then(|t| t.base().map(|b| b.token)); - while let Some(tok) = current { - if tok == other_token { + // Walk inheritance chain of 'this' looking for 'other'. + // `base_chain` bounds the walk against cyclic `extends` graphs. + if let Some(this_type) = asm.types().resolve(&this_token) { + if this_type + .base_chain() + .any(|ancestor| ancestor.token == other_token) + { return PreHookResult::Bypass(Some(EmValue::I32(1))); } - current = asm - .types() - .get(&tok) - .and_then(|t| t.base().map(|b| b.token)); } } } @@ -2553,10 +2548,27 @@ fn type_get_interface_map_pre( /// Finds the implementation of an interface method on a concrete type. /// Used by GetInterfaceMap to map interface methods → concrete implementations. +/// +/// Ancestors are visited via [`CilType::base_chain`], which bounds the walk against the cyclic +/// `extends` graphs a hostile assembly can describe. fn find_interface_impl_for_map( type_info: &CilType, interface_method: Token, base_method: &Method, +) -> Option { + if let Some(token) = find_interface_impl_on_type(type_info, interface_method, base_method) { + return Some(token); + } + type_info + .base_chain() + .find_map(|ancestor| find_interface_impl_on_type(&ancestor, interface_method, base_method)) +} + +/// Finds an interface-method implementation declared directly on `type_info`. +fn find_interface_impl_on_type( + type_info: &CilType, + interface_method: Token, + base_method: &Method, ) -> Option { // Step 1: Explicit MethodImpl overrides for (_, method_ref) in type_info.methods.iter() { @@ -2588,11 +2600,7 @@ fn find_interface_impl_for_map( } } - // Step 3: Walk base type - if let Some(base) = type_info.base() { - return find_interface_impl_for_map(&base, interface_method, base_method); - } - + // Step 3: base types are handled by the caller's `base_chain` walk. None } diff --git a/dotscope/src/emulation/runtime/bcl/runtime.rs b/dotscope/src/emulation/runtime/bcl/runtime.rs index e67a2085..38b5849a 100644 --- a/dotscope/src/emulation/runtime/bcl/runtime.rs +++ b/dotscope/src/emulation/runtime/bcl/runtime.rs @@ -374,6 +374,13 @@ fn runtime_helpers_initialize_array_pre( // Find the RVA for this field token let mut rva: Option = None; for row in fieldrva_table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; // Convert field index to full token (table 0x04 = Field) let row_token = row.field | 0x0400_0000; if row_token == field_token && row.rva > 0 { diff --git a/dotscope/src/emulation/runtime/bcl/system/array.rs b/dotscope/src/emulation/runtime/bcl/system/array.rs index 93c429ca..f530ac7c 100644 --- a/dotscope/src/emulation/runtime/bcl/system/array.rs +++ b/dotscope/src/emulation/runtime/bcl/system/array.rs @@ -58,6 +58,7 @@ use crate::{ emulation::{ + engine::synthetic_exception, memory::HeapObject, runtime::hook::{Hook, HookContext, HookManager, PreHookResult}, thread::EmulationThread, @@ -882,12 +883,21 @@ fn array_create_instance_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread }; // arg[1] = int length (or int[] for multi-dimensional, which we only handle for 1D) + // + // A bare `as usize` turns `Array.CreateInstance(typeof(int), -1)` into `usize::MAX`, which + // the allocator then tries to honour. .NET throws ArgumentOutOfRangeException for a + // negative length, so do the same and let emulated code catch it. let length = match length_arg { - EmValue::I32(n) => *n as usize, - EmValue::I64(n) => *n as usize, - EmValue::NativeInt(n) => *n as usize, + EmValue::I32(n) => usize::try_from(*n).ok(), + EmValue::I64(n) | EmValue::NativeInt(n) => usize::try_from(*n).ok(), _ => return PreHookResult::Continue, }; + let Some(length) = length else { + return PreHookResult::Throw { + exception_type: synthetic_exception::ARGUMENT_OUT_OF_RANGE, + message: "Array.CreateInstance: length must be non-negative".to_string(), + }; + }; // arg[0] = Type (ObjectRef to ReflectionType on the heap) // Resolve the element type through the proper metadata path @@ -914,7 +924,13 @@ fn array_create_instance_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread match thread.heap_mut().alloc_array(element_flavor, length) { Ok(array_ref) => PreHookResult::Bypass(Some(EmValue::ObjectRef(array_ref))), - Err(_) => PreHookResult::Continue, + // Falling through to `Continue` on a budget rejection would let the emulated program + // proceed as though the allocation had succeeded. Surface it as the CLR exception the + // real runtime raises so `catch` blocks in the sample still behave correctly. + Err(e) => PreHookResult::Throw { + exception_type: synthetic_exception::OUT_OF_MEMORY, + message: format!("Array.CreateInstance: {e}"), + }, } } diff --git a/dotscope/src/emulation/runtime/bcl/system/diagnostics.rs b/dotscope/src/emulation/runtime/bcl/system/diagnostics.rs index d92585d5..d352f17c 100644 --- a/dotscope/src/emulation/runtime/bcl/system/diagnostics.rs +++ b/dotscope/src/emulation/runtime/bcl/system/diagnostics.rs @@ -99,7 +99,8 @@ fn process_get_current_process_pre( let tables = asm.tables()?; let strings = asm.strings()?; let module_table = tables.table::()?; - let module_row = module_table.iter().next()?; + // Module is RID 1 by definition (ECMA-335 II.22.30); see `modules.rs`. + let module_row = module_table.get(1).ok().flatten()?; strings.get(module_row.name as usize).ok().map(String::from) }) .unwrap_or_else(|| "module.exe".to_string()); diff --git a/dotscope/src/emulation/runtime/bcl/system/exception.rs b/dotscope/src/emulation/runtime/bcl/system/exception.rs index 2ab07a18..3d4828e5 100644 --- a/dotscope/src/emulation/runtime/bcl/system/exception.rs +++ b/dotscope/src/emulation/runtime/bcl/system/exception.rs @@ -59,10 +59,12 @@ pub fn register(manager: &HookManager) -> Result<()> { register_exception_type(manager, "System", "FormatException")?; register_exception_type(manager, "System", "NotSupportedException")?; register_exception_type(manager, "System", "NotImplementedException")?; + register_exception_type(manager, "System", "AccessViolationException")?; register_exception_type(manager, "System", "NullReferenceException")?; register_exception_type(manager, "System", "IndexOutOfRangeException")?; register_exception_type(manager, "System", "InvalidCastException")?; register_exception_type(manager, "System", "OverflowException")?; + register_exception_type(manager, "System", "OutOfMemoryException")?; register_exception_type(manager, "System", "ArithmeticException")?; register_exception_type(manager, "System", "TypeInitializationException")?; register_exception_type(manager, "System", "ObjectDisposedException")?; diff --git a/dotscope/src/emulation/runtime/bcl/system/string.rs b/dotscope/src/emulation/runtime/bcl/system/string.rs index abf15f95..81eff286 100644 --- a/dotscope/src/emulation/runtime/bcl/system/string.rs +++ b/dotscope/src/emulation/runtime/bcl/system/string.rs @@ -6,7 +6,10 @@ use crate::{ emulation::{ memory::HeapObject, - runtime::hook::{Hook, HookContext, HookManager, PreHookResult}, + runtime::{ + bcl::limits::{oversized_argument, MAX_HOOK_STRING_CHARS}, + hook::{Hook, HookContext, HookManager, PreHookResult}, + }, thread::EmulationThread, EmValue, }, @@ -713,12 +716,16 @@ fn string_pad_left_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> P Err(e) => return PreHookResult::Error(format!("heap allocation failed: {e}")), }; - let total_width = ctx - .args - .first() - .map(usize::try_from) - .and_then(|r| r.ok()) - .unwrap_or(0); + // `totalWidth` comes off the emulated stack; without a ceiling an `int.MaxValue` argument + // builds a multi-gigabyte string inside a single hook call, where no emulation budget is + // evaluated. + let total_width = match ctx.args.first().map(usize::try_from).and_then(|r| r.ok()) { + Some(w) if w > MAX_HOOK_STRING_CHARS => { + return oversized_argument("String.Pad", "totalWidth", w, MAX_HOOK_STRING_CHARS) + } + Some(w) => w, + None => 0, + }; let pad_char = ctx .args .get(1) @@ -731,12 +738,20 @@ fn string_pad_left_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> P }) .unwrap_or(' '); - let result = if s.len() >= total_width { + // .NET measures `totalWidth` in characters, not UTF-8 bytes: comparing against `s.len()` + // under-pads any string containing a multi-byte scalar. + let char_count = s.chars().count(); + let result = if char_count >= total_width { s } else { - let padding: String = - std::iter::repeat_n(pad_char, total_width.saturating_sub(s.len())).collect(); - format!("{padding}{s}") + let pad_count = total_width.saturating_sub(char_count); + let mut padded = String::with_capacity( + s.len() + .saturating_add(pad_count.saturating_mul(pad_char.len_utf8())), + ); + padded.extend(std::iter::repeat_n(pad_char, pad_count)); + padded.push_str(&s); + padded }; match thread.heap_mut().alloc_string(&result) { @@ -770,12 +785,16 @@ fn string_pad_right_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> Err(e) => return PreHookResult::Error(format!("heap allocation failed: {e}")), }; - let total_width = ctx - .args - .first() - .map(usize::try_from) - .and_then(|r| r.ok()) - .unwrap_or(0); + // `totalWidth` comes off the emulated stack; without a ceiling an `int.MaxValue` argument + // builds a multi-gigabyte string inside a single hook call, where no emulation budget is + // evaluated. + let total_width = match ctx.args.first().map(usize::try_from).and_then(|r| r.ok()) { + Some(w) if w > MAX_HOOK_STRING_CHARS => { + return oversized_argument("String.Pad", "totalWidth", w, MAX_HOOK_STRING_CHARS) + } + Some(w) => w, + None => 0, + }; let pad_char = ctx .args .get(1) @@ -788,12 +807,20 @@ fn string_pad_right_pre(ctx: &HookContext<'_>, thread: &mut EmulationThread) -> }) .unwrap_or(' '); - let result = if s.len() >= total_width { + // .NET measures `totalWidth` in characters, not UTF-8 bytes: comparing against `s.len()` + // under-pads any string containing a multi-byte scalar. + let char_count = s.chars().count(); + let result = if char_count >= total_width { s } else { - let padding: String = - std::iter::repeat_n(pad_char, total_width.saturating_sub(s.len())).collect(); - format!("{s}{padding}") + let pad_count = total_width.saturating_sub(char_count); + let mut padded = String::with_capacity( + s.len() + .saturating_add(pad_count.saturating_mul(pad_char.len_utf8())), + ); + padded.push_str(&s); + padded.extend(std::iter::repeat_n(pad_char, pad_count)); + padded }; match thread.heap_mut().alloc_string(&result) { diff --git a/dotscope/src/emulation/runtime/bcl/text/stringbuilder.rs b/dotscope/src/emulation/runtime/bcl/text/stringbuilder.rs index 3a69be27..d61f7acd 100644 --- a/dotscope/src/emulation/runtime/bcl/text/stringbuilder.rs +++ b/dotscope/src/emulation/runtime/bcl/text/stringbuilder.rs @@ -29,7 +29,10 @@ use crate::{ emulation::{ memory::HeapObject, - runtime::hook::{Hook, HookContext, HookManager, PreHookResult}, + runtime::{ + bcl::limits::{checked_len, MAX_HOOK_STRING_CHARS}, + hook::{Hook, HookContext, HookManager, PreHookResult}, + }, thread::EmulationThread, EmValue, }, @@ -393,15 +396,25 @@ fn stringbuilder_set_length_pre( if let Some(EmValue::ObjectRef(sb_ref)) = ctx.this { if let Some((buffer, capacity)) = read_sb(thread, *sb_ref) { if let Some(EmValue::I32(new_len)) = ctx.args.first() { - let target = (*new_len).max(0) as usize; + // The growth branch pads one character at a time, bounded only by the + // caller's Int32, so the target width needs a ceiling before it is used. + let target = match checked_len( + *new_len, + MAX_HOOK_STRING_CHARS, + "StringBuilder.set_Length", + "value", + ) { + Ok(n) => n, + Err(result) => return result, + }; let current = buffer.chars().count(); let new_buffer = if target <= current { buffer.chars().take(target).collect() } else { + let pad = target.saturating_sub(current); let mut s = buffer; - for _ in 0..target.saturating_sub(current) { - s.push('\0'); - } + s.reserve(pad); + s.extend(std::iter::repeat_n('\0', pad)); s }; try_hook!(write_sb(thread, *sb_ref, new_buffer, capacity)); diff --git a/dotscope/src/emulation/runtime/hook/types.rs b/dotscope/src/emulation/runtime/hook/types.rs index f8bc1600..b3cb1f29 100644 --- a/dotscope/src/emulation/runtime/hook/types.rs +++ b/dotscope/src/emulation/runtime/hook/types.rs @@ -380,6 +380,15 @@ impl PreHookResult { } } + /// Creates a `Throw` for `System.AccessViolationException`. + #[must_use] + pub fn throw_access_violation(msg: &str) -> Self { + Self::Throw { + exception_type: synthetic_exception::ACCESS_VIOLATION, + message: format!("AccessViolationException: {msg}"), + } + } + /// Creates a `Throw` for `System.InvalidOperationException`. #[must_use] pub fn throw_invalid_operation(msg: &str) -> Self { diff --git a/dotscope/src/emulation/runtime/native.rs b/dotscope/src/emulation/runtime/native.rs index b67d0ba7..685168c6 100644 --- a/dotscope/src/emulation/runtime/native.rs +++ b/dotscope/src/emulation/runtime/native.rs @@ -109,16 +109,36 @@ use crate::{ Result, }; +/// A native function resolved at runtime, and the module it came from. +/// +/// The module matters because hooks are registered against a `(dll, function)` +/// pair. A function resolved through `GetProcAddress` only reaches its hook if +/// the library that produced the handle is known, so the name is carried +/// alongside the function rather than reconstructed later. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct NativeFunction { + /// Owning module, as passed to `LoadLibrary` (e.g. `"kernel32.dll"`). + /// + /// `None` when the handle did not come from an observed `LoadLibrary` — the + /// function is still dispatchable, but only against hooks that do not + /// constrain the module. + pub dll: Option>, + /// Exported function name, as passed to `GetProcAddress`. + pub name: Arc, +} + /// Shared registry that maps fake function pointer addresses to their names /// and native function tokens to their names. This allows the full chain -/// (GetProcAddress → GetDelegateForFunctionPointer → delegate dispatch) to -/// preserve function identity. +/// (`LoadLibrary` → `GetProcAddress` → `GetDelegateForFunctionPointer` → +/// delegate dispatch) to preserve function identity. #[derive(Clone, Debug)] pub struct NativeFunctionRegistry { - address_to_name: Arc>>, - token_to_name: Arc>>, + address_to_name: Arc>, + token_to_name: Arc>, + module_handles: Arc>>, next_address: Arc, next_token_id: Arc, + next_module: Arc, } impl NativeFunctionRegistry { @@ -126,63 +146,93 @@ impl NativeFunctionRegistry { Self { address_to_name: Arc::new(DashMap::new()), token_to_name: Arc::new(DashMap::new()), + module_handles: Arc::new(DashMap::new()), next_address: Arc::new(AtomicU64::new(native_addresses::PROC_ADDRESS_BASE)), next_token_id: Arc::new(AtomicU64::new(1)), + next_module: Arc::new(AtomicU64::new( + native_addresses::LOADED_LIBRARY.cast_unsigned(), + )), } } - /// Registers a native function name and returns a unique fake address for it. - pub fn register_proc(&self, name: &str) -> u64 { + /// Registers a loaded module and returns the handle standing in for it. + /// + /// Handles are distinct per module so a later `GetProcAddress` can name the + /// library its function came from. Loading the same module twice returns the + /// same handle, as the real loader does. + pub fn register_module(&self, name: &str) -> u64 { let name: Arc = Arc::from(name); + for entry in self.module_handles.iter() { + if *entry.value() == name { + return *entry.key(); + } + } + let handle = self + .next_module + .fetch_add(native_addresses::PROC_ADDRESS_PAGE, Ordering::Relaxed); + self.module_handles.insert(handle, name); + handle + } + + /// Looks up the module registered for a handle. + pub fn lookup_module(&self, handle: u64) -> Option> { + self.module_handles + .get(&handle) + .map(|v| Arc::clone(v.value())) + } + + /// Registers a native function and returns a unique fake address for it. + pub fn register_proc(&self, dll: Option<&str>, name: &str) -> u64 { + let function = NativeFunction { + dll: dll.map(Arc::from), + name: Arc::from(name), + }; // Check if already registered for entry in self.address_to_name.iter() { - if *entry.value() == name { + if *entry.value() == function { return *entry.key(); } } let addr = self .next_address .fetch_add(native_addresses::PROC_ADDRESS_PAGE, Ordering::Relaxed); - self.address_to_name.insert(addr, name); + self.address_to_name.insert(addr, function); addr } - /// Looks up the function name for a given fake address. - pub fn lookup_by_address(&self, addr: u64) -> Option> { - self.address_to_name - .get(&addr) - .map(|v| Arc::clone(v.value())) + /// Looks up the function registered at a given fake address. + pub fn lookup_by_address(&self, addr: u64) -> Option { + self.address_to_name.get(&addr).map(|v| v.value().clone()) } - /// Allocates a unique native function pointer token for a function name and - /// returns the token. If a token was already allocated for this name, returns + /// Allocates a unique native function pointer token for a function and + /// returns the token. If a token was already allocated for it, returns /// the existing one. - pub fn allocate_token(&self, name: &str) -> Token { - let name: Arc = Arc::from(name); + pub fn allocate_token(&self, function: &NativeFunction) -> Token { for entry in self.token_to_name.iter() { - if *entry.value() == name { + if entry.value() == function { return tokens::native::token_for_id(*entry.key()); } } let id = self.next_token_id.fetch_add(1, Ordering::Relaxed); - self.token_to_name.insert(id as u32, name); + self.token_to_name.insert(id as u32, function.clone()); tokens::native::token_for_id(id as u32) } /// Finds the fake address registered for the given function name. pub fn lookup_address_by_name(&self, name: &str) -> Option { for entry in self.address_to_name.iter() { - if entry.value().as_ref() == name { + if entry.value().name.as_ref() == name { return Some(*entry.key()); } } None } - /// Looks up the function name for a native function pointer token. - pub fn lookup_by_token(&self, token: Token) -> Option> { + /// Looks up the function for a native function pointer token. + pub fn lookup_by_token(&self, token: Token) -> Option { let id = token.value() & 0x0000_FFFF; - self.token_to_name.get(&id).map(|v| Arc::clone(v.value())) + self.token_to_name.get(&id).map(|v| v.value().clone()) } } @@ -225,7 +275,7 @@ pub fn register(manager: &HookManager, registry: &NativeFunctionRegistry) -> Res // Module loading hooks register_get_module_handle(manager)?; register_get_proc_address(manager, registry)?; - register_load_library(manager)?; + register_load_library(manager, registry)?; // Anti-debug bypass hooks register_is_debugger_present(manager)?; @@ -535,19 +585,38 @@ fn register_get_proc_address( .match_native("kernel32", "GetProcAddress") .pre(move |ctx, thread| { // GetProcAddress(hModule, lpProcName) - let func_name = ctx + let func_name = ctx.args.get(1).and_then(|arg| match arg { + EmValue::ObjectRef(href) => { + thread.heap().get_string(*href).ok().map(|s| s.to_string()) + } + _ => None, + }); + + // An unreadable name cannot be told apart from any other + // unreadable name. Registering them all under one label would + // collapse them onto a single address and dispatch the wrong + // function; leave the resolution to fail instead. + let Some(func_name) = func_name else { + log::debug!("GetProcAddress: could not read the requested name"); + return PreHookResult::Bypass(Some(EmValue::NativeInt(0))); + }; + + let dll = ctx .args - .get(1) + .first() .and_then(|arg| match arg { - EmValue::ObjectRef(href) => { - thread.heap().get_string(*href).ok().map(|s| s.to_string()) - } + EmValue::NativeInt(v) => Some((*v).cast_unsigned()), + EmValue::NativeUInt(v) => Some(*v), + EmValue::UnmanagedPtr(v) => Some(*v), _ => None, }) - .unwrap_or_else(|| "unknown".to_string()); + .and_then(|handle| registry.lookup_module(handle)); - let addr = registry.register_proc(&func_name); - log::debug!("GetProcAddress({func_name}) → 0x{addr:X}"); + let addr = registry.register_proc(dll.as_deref(), &func_name); + log::debug!( + "GetProcAddress({}!{func_name}) → 0x{addr:X}", + dll.as_deref().unwrap_or("?") + ); PreHookResult::Bypass(Some(EmValue::NativeInt(addr as i64))) }), )?; @@ -560,31 +629,48 @@ fn register_get_proc_address( /// Some obfuscators (e.g., .NET Reactor) use the unsuffixed `LoadLibrary` /// import name in their P/Invoke declarations, which doesn't match the /// standard `LoadLibraryA`/`LoadLibraryW` exports. -fn register_load_library(manager: &HookManager) -> Result<()> { - let load_library_handler = - |_ctx: &HookContext<'_>, _thread: &mut EmulationThread| -> PreHookResult { - PreHookResult::Bypass(Some(EmValue::NativeInt(native_addresses::LOADED_LIBRARY))) - }; +fn register_load_library(manager: &HookManager, registry: &NativeFunctionRegistry) -> Result<()> { + // Each library gets its own handle so a later `GetProcAddress` can name the + // module its function came from — which is what lets the resolved function + // reach a hook registered against a `(dll, function)` pair. + let make_handler = |registry: NativeFunctionRegistry| { + move |ctx: &HookContext<'_>, thread: &mut EmulationThread| -> PreHookResult { + let module = ctx.args.first().and_then(|arg| match arg { + EmValue::ObjectRef(href) => { + thread.heap().get_string(*href).ok().map(|s| s.to_string()) + } + _ => None, + }); + let Some(module) = module else { + return PreHookResult::Bypass(Some(EmValue::NativeInt( + native_addresses::LOADED_LIBRARY, + ))); + }; + let handle = registry.register_module(&module); + log::debug!("LoadLibrary({module}) → 0x{handle:X}"); + PreHookResult::Bypass(Some(EmValue::NativeInt(handle as i64))) + } + }; manager.register( Hook::new("native-load-library-a") .with_priority(HookPriority::HIGH) .match_native("kernel32", "LoadLibraryA") - .pre(load_library_handler), + .pre(make_handler(registry.clone())), )?; manager.register( Hook::new("native-load-library-w") .with_priority(HookPriority::HIGH) .match_native("kernel32", "LoadLibraryW") - .pre(load_library_handler), + .pre(make_handler(registry.clone())), )?; manager.register( Hook::new("native-load-library") .with_priority(HookPriority::HIGH) .match_native("kernel32", "LoadLibrary") - .pre(load_library_handler), + .pre(make_handler(registry.clone())), )?; Ok(()) @@ -712,6 +798,68 @@ mod tests { assert!(manager.len() >= 12); } + /// A function resolved through `GetProcAddress` only reaches its hook if the + /// module that produced the handle is known, because hooks are registered + /// against a `(dll, function)` pair. + #[test] + fn resolved_function_remembers_its_module() { + let registry = NativeFunctionRegistry::new(); + + let handle = registry.register_module("kernel32.dll"); + assert_eq!( + registry.lookup_module(handle).as_deref(), + Some("kernel32.dll") + ); + + let dll = registry.lookup_module(handle); + let addr = registry.register_proc(dll.as_deref(), "VirtualProtect"); + + let resolved = registry + .lookup_by_address(addr) + .expect("the address just registered must resolve"); + assert_eq!(resolved.name.as_ref(), "VirtualProtect"); + assert_eq!(resolved.dll.as_deref(), Some("kernel32.dll")); + + // The token carried through GetDelegateForFunctionPointer must round-trip + // to the same pair, which is what the delegate dispatch matches on. + let token = registry.allocate_token(&resolved); + assert_eq!(registry.lookup_by_token(token), Some(resolved)); + } + + /// Loading the same module twice yields one handle, as the real loader does; + /// distinct modules must not collide, or a function would be attributed to + /// whichever library was loaded last. + #[test] + fn module_handles_are_stable_and_distinct() { + let registry = NativeFunctionRegistry::new(); + + let kernel32 = registry.register_module("kernel32.dll"); + let user32 = registry.register_module("user32.dll"); + + assert_eq!(registry.register_module("kernel32.dll"), kernel32); + assert_ne!(kernel32, user32); + assert_eq!( + registry.lookup_module(user32).as_deref(), + Some("user32.dll") + ); + } + + /// The same exported name in two libraries is two different functions, and + /// must not share one fake address. + #[test] + fn same_name_in_two_modules_stays_distinct() { + let registry = NativeFunctionRegistry::new(); + + let a = registry.register_proc(Some("kernel32.dll"), "Sleep"); + let b = registry.register_proc(Some("winmm.dll"), "Sleep"); + + assert_ne!(a, b); + assert_eq!( + registry.lookup_by_address(a).and_then(|f| f.dll), + Some("kernel32.dll".into()) + ); + } + #[test] fn test_is_debugger_present_hook() { let manager = HookManager::new(); diff --git a/dotscope/src/emulation/runtime/state.rs b/dotscope/src/emulation/runtime/state.rs index b6436956..f2d0d282 100644 --- a/dotscope/src/emulation/runtime/state.rs +++ b/dotscope/src/emulation/runtime/state.rs @@ -267,6 +267,36 @@ impl RuntimeState { self.unknown_method_behavior } + /// Creates an independent runtime state for a forked execution. + /// + /// # What separates + /// + /// The [`AppDomainState`] is cloned by value. It is the mutable half of this type: + /// emulated code appends to it through `Assembly.Load(byte[])`, interns strings into it + /// and registers resolve handlers on it. Forks execute in real parallel — the constant + /// decryption pass drives them from a rayon iterator — so a shared domain lets a type + /// planted by one fork steer method resolution in another, and lets every fork's loads + /// accumulate against one budget. Cloning `Vec>` copies the handles, not + /// the parsed metadata, so the assemblies present at fork time stay shared while + /// additions stay local. + /// + /// # What stays shared + /// + /// `hooks` and `config` are immutable after construction. `native_functions` is a + /// handle onto `Arc`-shared maps and is deliberately kept shared: it assigns the + /// synthetic addresses returned by `GetProcAddress`, and those must mean the same + /// function in every fork for a pointer captured in one to be interpretable in another. + #[must_use] + pub fn fork(&self) -> Self { + Self { + hooks: Arc::clone(&self.hooks), + app_domain: self.app_domain.clone(), + unknown_method_behavior: self.unknown_method_behavior, + config: Arc::clone(&self.config), + native_functions: self.native_functions.clone(), + } + } + /// Sets the behavior for unknown method calls. /// /// This controls what happens when a method call is encountered diff --git a/dotscope/src/emulation/thread/context.rs b/dotscope/src/emulation/thread/context.rs index 45a0ae7c..74841e12 100644 --- a/dotscope/src/emulation/thread/context.rs +++ b/dotscope/src/emulation/thread/context.rs @@ -16,7 +16,7 @@ use dashmap::DashMap; use crate::{ emulation::{ capture::CaptureContext, - engine::SyntheticMethodBody, + engine::{EmulationError, SyntheticMethodBody}, fakeobjects::SharedFakeObjects, filesystem::VirtualFs, memory::{AddressSpace, ManagedHeap, StaticFieldStorage}, @@ -104,27 +104,80 @@ impl ThreadContext { token } - /// Forks this context with a forked address space. + /// Forks this context into an independent execution environment. + /// + /// Everything emulated code can mutate is separated, so two forks running concurrently + /// cannot observe each other: /// - /// The forked context shares all `Arc`-wrapped resources except: /// - `address_space` — forked with CoW semantics - /// - `capture` — fresh context (same config, empty captures) + /// - `runtime` — forked: the fork gets its own `AppDomainState`, so assemblies it loads + /// and strings it interns stay local (see [`RuntimeState::fork`]) + /// - `capture` — fresh context, same config, empty captures /// - `virtual_fs` — forked (falls back to shared on error) + /// - `synthetic_methods` — seeded from the parent, then independent, so a `DynamicMethod` + /// emitted in one fork is not callable from another + /// + /// Genuinely immutable state is shared: `config`, `assembly`, `fake_objects` and the + /// hook manager inside the runtime. + /// + /// `synthetic_method_counter` stays shared on purpose. It is the one mutable thing that + /// must *not* be forked: two forks allocating from private counters would mint the same + /// synthetic token for different bodies, and those tokens outlive the fork in captured + /// output. + /// + /// # Errors + /// + /// Returns an error if the address space cannot be forked, or if the parent's runtime + /// lock is poisoned. pub fn fork(&self) -> crate::Result { + self.fork_with_config(Arc::clone(&self.config)) + } + + /// Forks this context, giving the fork its own configuration. + /// + /// The configuration is the one thing a fork legitimately needs to *differ* on. A + /// template process is warmed up under a long budget because warmup genuinely takes it; + /// the per-method executions forked from it are supposed to run under the much smaller + /// per-method budget. Sharing the `Arc` makes that impossible to express, and mutating it + /// in place would retroactively change the template's own budget. + /// + /// See [`Self::fork`] for what else separates and what stays shared. + /// + /// # Errors + /// + /// Returns an error if the address space cannot be forked, or if the parent's runtime + /// lock is poisoned. + pub fn fork_with_config(&self, config: Arc) -> crate::Result { let virtual_fs = match self.virtual_fs.fork() { Ok(forked) => Arc::new(forked), Err(_) => Arc::clone(&self.virtual_fs), }; + let runtime = self + .runtime + .read() + .map_err(|_| { + crate::Error::Emulation(Box::new(EmulationError::LockPoisoned { + description: "runtime state", + })) + })? + .fork(); + + let synthetic_methods: DashMap = self + .synthetic_methods + .iter() + .map(|entry| (*entry.key(), entry.value().clone())) + .collect(); + Ok(Self { address_space: Arc::new(self.address_space.fork()?), - runtime: Arc::clone(&self.runtime), + runtime: Arc::new(RwLock::new(runtime)), capture: Arc::new(CaptureContext::with_config(self.capture.config().clone())), - config: Arc::clone(&self.config), + config, assembly: self.assembly.clone(), fake_objects: self.fake_objects.clone(), virtual_fs, - synthetic_methods: Arc::clone(&self.synthetic_methods), + synthetic_methods: Arc::new(synthetic_methods), synthetic_method_counter: Arc::clone(&self.synthetic_method_counter), }) } diff --git a/dotscope/src/emulation/thread/scheduler.rs b/dotscope/src/emulation/thread/scheduler.rs index 2eebce7f..39c41081 100644 --- a/dotscope/src/emulation/thread/scheduler.rs +++ b/dotscope/src/emulation/thread/scheduler.rs @@ -33,6 +33,7 @@ use std::collections::{BinaryHeap, HashMap}; use crate::{ emulation::{ + engine::EmulationError, thread::{EmulationThread, ThreadPriority, ThreadState, WaitReason}, EmValue, HeapRef, ThreadId, }, @@ -192,6 +193,12 @@ pub struct ThreadScheduler { /// Next thread ID to assign. next_thread_id: u32, + + /// Maximum number of threads this scheduler will hold. + /// + /// Set from `EmulationLimits::max_threads`. `usize::MAX` means unlimited, which is the + /// default so that constructing a scheduler directly is unchanged. + max_threads: usize, } impl ThreadScheduler { @@ -223,9 +230,17 @@ impl ThreadScheduler { total_instructions: 0, next_sequence: 0, next_thread_id: 2, // 1 is reserved for main thread + max_threads: usize::MAX, } } + /// Sets the maximum number of threads this scheduler will hold. + /// + /// [`spawn`](Self::spawn) fails once the scheduler is at this many threads. + pub fn set_max_threads(&mut self, max: usize) { + self.max_threads = max; + } + /// Creates a scheduler with the default quantum of 1000 instructions. /// /// This is a convenience constructor for typical use cases where the @@ -288,12 +303,27 @@ impl ThreadScheduler { /// # Returns /// /// The thread ID of the spawned thread. - pub fn spawn(&mut self, thread: EmulationThread) -> ThreadId { + /// + /// # Errors + /// + /// Returns [`EmulationError::ResourceLimitExceeded`] if the scheduler already holds + /// [`max_threads`](Self::set_max_threads) threads. Emulated code controls how many threads + /// it asks for, so this ceiling is what stops it from spawning without bound; note that + /// threads are counted in all states, including completed ones. + pub fn spawn(&mut self, thread: EmulationThread) -> Result { + if self.threads.len() >= self.max_threads { + return Err(EmulationError::ResourceLimitExceeded(format!( + "thread limit reached ({} threads)", + self.max_threads + )) + .into()); + } + let id = thread.id(); let priority = thread.priority(); self.threads.insert(id, thread); self.enqueue_ready(id, priority); - id + Ok(id) } /// Allocates a new unique thread ID. @@ -883,7 +913,7 @@ mod tests { let mut scheduler = ThreadScheduler::new(100); let thread = create_test_thread(1); - let id = scheduler.spawn(thread); + let id = scheduler.spawn(thread).unwrap(); assert_eq!(id, ThreadId::new(1)); assert_eq!(scheduler.thread_count(), 1); @@ -896,8 +926,8 @@ mod tests { let thread1 = create_test_thread(1); let thread2 = create_test_thread(2); - scheduler.spawn(thread1); - scheduler.spawn(thread2); + scheduler.spawn(thread1).unwrap(); + scheduler.spawn(thread2).unwrap(); // First selection should get a thread let selected = scheduler.select_next(); @@ -913,7 +943,7 @@ mod tests { let mut scheduler = ThreadScheduler::new(3); let thread = create_test_thread(1); - scheduler.spawn(thread); + scheduler.spawn(thread).unwrap(); scheduler.select_next(); // Execute 3 instructions to exhaust quantum @@ -937,9 +967,9 @@ mod tests { high_thread.set_priority(ThreadPriority::Highest); // Add low priority first - scheduler.spawn(low_thread); + scheduler.spawn(low_thread).unwrap(); // Add high priority second - scheduler.spawn(high_thread); + scheduler.spawn(high_thread).unwrap(); // High priority should be selected first let selected = scheduler.select_next(); @@ -951,7 +981,7 @@ mod tests { let mut scheduler = ThreadScheduler::new(100); let thread = create_test_thread(1); - scheduler.spawn(thread); + scheduler.spawn(thread).unwrap(); scheduler.select_next(); scheduler.complete_current(Some(EmValue::I32(42))); @@ -966,7 +996,7 @@ mod tests { let mut scheduler = ThreadScheduler::new(100); let thread = create_test_thread(1); - scheduler.spawn(thread); + scheduler.spawn(thread).unwrap(); assert!(!scheduler.all_completed()); scheduler.select_next(); @@ -982,8 +1012,8 @@ mod tests { let thread1 = create_test_thread(1); let thread2 = create_test_thread(2); - scheduler.spawn(thread1); - scheduler.spawn(thread2); + scheduler.spawn(thread1).unwrap(); + scheduler.spawn(thread2).unwrap(); let first = scheduler.select_next(); scheduler.yield_current(); @@ -998,7 +1028,7 @@ mod tests { let mut scheduler = ThreadScheduler::new(100); let thread = create_test_thread(1); - scheduler.spawn(thread); + scheduler.spawn(thread).unwrap(); scheduler.select_next(); // Block thread with sleep @@ -1023,7 +1053,7 @@ mod tests { let mut scheduler = ThreadScheduler::new(100); let thread = create_test_thread(1); - scheduler.spawn(thread); + scheduler.spawn(thread).unwrap(); scheduler.select_next(); // Block thread (use a dummy HeapRef for the monitor object) diff --git a/dotscope/src/emulation/thread/state.rs b/dotscope/src/emulation/thread/state.rs index 6a9656e1..21c21d9f 100644 --- a/dotscope/src/emulation/thread/state.rs +++ b/dotscope/src/emulation/thread/state.rs @@ -282,7 +282,7 @@ pub struct ThreadCallFrame { /// This enables cross-assembly execution: when a method from a dynamically /// loaded assembly calls another method, the controller uses this index to /// fetch instructions and resolve metadata from the correct assembly. - assembly_index: Option, + assembly_index: Option, } impl ThreadCallFrame { @@ -520,7 +520,7 @@ impl ThreadCallFrame { /// Sets the assembly index for this frame. /// /// `None` = primary assembly, `Some(i)` = i-th dynamically loaded assembly. - pub fn set_assembly_index(&mut self, index: Option) { + pub fn set_assembly_index(&mut self, index: Option) { self.assembly_index = index; } @@ -528,7 +528,7 @@ impl ThreadCallFrame { /// /// `None` = primary assembly, `Some(i)` = i-th dynamically loaded assembly. #[must_use] - pub fn assembly_index(&self) -> Option { + pub fn assembly_index(&self) -> Option { self.assembly_index } @@ -937,8 +937,24 @@ impl EmulationThread { /// Pops and returns the current call frame. /// /// Used when returning from a method. + /// + /// Also discards any cleanup handlers still queued against the popped method. Such an entry + /// can never run correctly — the handler IL reads locals, arguments and the evaluation stack + /// from whatever frame happens to be current — and left in place it accumulates for the + /// lifetime of the emulation. + /// + /// The check is against the *remaining* call stack rather than the popped method alone, + /// because a recursive method occupies several frames at once and the outer frames' entries + /// are still live. pub fn pop_frame(&mut self) -> Option { - self.call_stack.pop() + let popped = self.call_stack.pop()?; + + let method = popped.method(); + if !self.call_stack.iter().any(|frame| frame.method() == method) { + self.exception_state.discard_finally_for_method(method); + } + + Some(popped) } /// Returns a reference to the evaluation stack. @@ -1086,6 +1102,16 @@ impl EmulationThread { /// address space. Reads and writes through the native address are /// transparently delegated to the managed heap's `Vec`, /// ensuring a single source of truth. + /// + /// Re-pinning an array that is already pinned returns its existing base rather than + /// reserving a second range. A loop that walks an array with `ldelema` + `conv.u` reaches + /// here on every iteration, so without this the pin table grows once per element access — + /// and every memory access scans it. + /// + /// # Errors + /// + /// Returns an error if the array's size overflows, if no free address range remains, or + /// if pin registration fails. pub fn pin_array_element( &self, array: HeapRef, @@ -1097,7 +1123,16 @@ impl EmulationThread { let length = heap.get_array_length(array).unwrap_or(0); if length == 0 { - return Ok(self.context.address_space.reserve_address_range(1)); + return self + .context + .address_space + .reserve_address_range(1) + .ok_or_else(|| { + Error::from(EmulationError::InvalidAddress { + address: 0, + reason: "no free address range for pinned array".to_string(), + }) + }); } // Determine element size from the element type or first element @@ -1117,15 +1152,28 @@ impl EmulationThread { let total_size = length .checked_mul(elem_size) .ok_or(EmulationError::ArithmeticOverflow)?; - let base_addr = self - .context - .address_space - .reserve_address_range(total_size.max(1)); - // Register the pinned mapping for transparent read/write delegation - self.context - .address_space - .register_pinned_array(base_addr, array, elem_size, length)?; + let base_addr = match self.context.address_space.pinned_base_of(array) { + Some(existing) => existing, + None => { + let base_addr = self + .context + .address_space + .reserve_address_range(total_size.max(1)) + .ok_or_else(|| { + Error::from(EmulationError::InvalidAddress { + address: 0, + reason: "no free address range for pinned array".to_string(), + }) + })?; + + // Register the pinned mapping for transparent read/write delegation + self.context + .address_space + .register_pinned_array(base_addr, array, elem_size, length)?; + base_addr + } + }; let index_offset = index .checked_mul(elem_size) @@ -1489,6 +1537,27 @@ mod tests { use super::*; use crate::test::emulation::{create_test_context, create_test_thread}; + /// The index selects which loaded assembly's metadata a frame executes against, and + /// emulated code grows that list. Narrowing it binds the frame to a different assembly + /// than the one the method token was resolved in — a valid token run against the wrong + /// method table. + #[test] + fn assembly_index_survives_beyond_a_byte() { + let mut frame = ThreadCallFrame::new( + Token::new(0x0600_0001), + None, + 0, + Vec::new(), + Vec::new(), + false, + ); + + for index in [0_u32, 255, 256, 4096, u32::from(u16::MAX) + 1] { + frame.set_assembly_index(Some(index)); + assert_eq!(frame.assembly_index(), Some(index)); + } + } + #[test] fn test_thread_creation() { let thread = create_test_thread(); diff --git a/dotscope/src/emulation/value/emvalue.rs b/dotscope/src/emulation/value/emvalue.rs index 402ad251..1a1d5398 100644 --- a/dotscope/src/emulation/value/emvalue.rs +++ b/dotscope/src/emulation/value/emvalue.rs @@ -157,6 +157,20 @@ pub enum EmValue { Symbolic(SymbolicValue), } +/// `EmValue` is the unit of currency in the interpreter: every `push`/`pop`, every local slot, +/// argument slot, array element and inline field is one of these, and `EvaluationStack::new` +/// preallocates up to 256 per thread. Its width is therefore a whole-interpreter cost, and it +/// is set by whichever variant is widest — `Symbolic`, through the `CilFlavor` it carries. +/// `CilFlavor::FnPtr` boxes its `SignatureMethod` for exactly this reason: embedding one by +/// value costs every `EmValue` in the process 96 further bytes. +/// +/// The bound below is the measured width, not a round number with slack in it: nothing else in +/// the build would catch a newly added fat variant, and a loose bound would let one land +/// silently. +/// +/// If this assertion fails, box the offending variant rather than raising the bound. +const _: () = assert!(std::mem::size_of::() <= 104); + impl EmValue { /// Returns the type token for this value, if available. /// diff --git a/dotscope/src/error.rs b/dotscope/src/error.rs index 0bf92017..40fc39a0 100644 --- a/dotscope/src/error.rs +++ b/dotscope/src/error.rs @@ -93,7 +93,7 @@ //! propagation in concurrent parsing and analysis operations. //! -use std::io; +use std::{io, sync::Arc}; use thiserror::Error; @@ -117,9 +117,9 @@ impl std::fmt::Display for EmulationError { /// Helper macro for creating malformed data errors with source location information. /// -/// This macro simplifies the creation of [`crate::Error::Malformed`] errors by automatically -/// capturing the current file and line number. It supports both simple string messages -/// and format string patterns with arguments. +/// This macro simplifies the creation of malformed-data errors by automatically capturing the +/// current file and line number. It supports both simple string messages and format string +/// patterns with arguments. /// /// # Arguments /// @@ -128,8 +128,9 @@ impl std::fmt::Display for EmulationError { /// /// # Returns /// -/// Returns a [`crate::Error::Malformed`] variant with the provided message and -/// automatically captured source location information. +/// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] at +/// [`crate::ParseStage::Generic`], with the provided message and automatically captured +/// source location information. /// /// # Examples /// @@ -203,8 +204,9 @@ macro_rules! out_of_bounds_error { /// # Error Categories /// /// ## File Parsing Errors -/// - [`crate::Error::Malformed`] - Corrupted or invalid file structure -/// - [`crate::Error::OutOfBounds`] - Attempted to read beyond file boundaries +/// - [`crate::Error::Parse`] - Corrupted structure, truncation, or a read past the end of the +/// input; the [`crate::ParseFailure`] it carries says which, and its [`crate::ParseStage`] +/// says where /// - [`crate::Error::NotSupported`] - Unsupported file format or feature /// /// ## I/O and External Errors @@ -229,7 +231,8 @@ macro_rules! out_of_bounds_error { /// This error enum is [`std::marker::Send`] and [`std::marker::Sync`] as all variants contain thread-safe types. /// This includes owned strings, primitive values, and errors from external crates that are themselves /// thread-safe. Errors can be safely passed between threads and shared across thread boundaries. -#[derive(Error, Debug)] +#[derive(Error, Debug, Clone)] +#[non_exhaustive] pub enum Error { // File parsing Errors /// This file type is not supported. @@ -243,8 +246,13 @@ pub enum Error { /// /// Wraps standard I/O errors that can occur during file operations /// such as reading from disk, permission issues, or filesystem errors. + /// + /// Held behind an `Arc` because [`io::Error`] is not [`Clone`], and this enum is. Without + /// it `Clone` cannot be derived, and a hand-written one has to decide what to do with + /// this variant — the answer being to flatten it into a string, which loses the + /// [`io::ErrorKind`] a caller matches on. #[error("{0}")] - Io(#[from] io::Error), + Io(Arc), /// Other errors that don't fit specific categories. /// @@ -259,8 +267,10 @@ pub enum Error { /// /// The goblin crate is used for low-level PE format parsing. /// This error wraps any failures from that parsing layer. + /// + /// `Arc`-wrapped for the same reason as [`Io`](Self::Io). #[error("{0}")] - Goblin(#[from] goblin::error::Error), + Goblin(Arc), /// Failed to find type in `TypeSystem`. /// @@ -285,8 +295,7 @@ pub enum Error { /// Wraps a [`ParseFailure`] so consumers can categorize parse failures /// (truncated headers, bad magic, unsupported schemas, heap corruption, /// invalid fields) without parsing string messages. Returned by every - /// parse-pipeline error site in [`crate::file`], [`crate::metadata::root`], - /// and [`crate::metadata::streams`]. + /// parse-pipeline error site in the file, metadata-root and stream parsers. /// /// Match on `Error::Parse(_)` to recover the structured failure, or on a /// specific variant of [`ParseFailure`] to react to a particular failure @@ -621,7 +630,7 @@ pub enum Error { /// Failure modes for method-by-token lookups. /// /// Returned by [`crate::CilObject::method`] and -/// [`crate::CilObject::method_spec`]. Propagates into [`Error`] via the +/// [`crate::CilObject::method_spec`]. Propagates into [`enum@Error`] via the /// [`Error::LookupMethod`] variant — call sites that already use /// `Result<_, Error>` can propagate with `?` without manual conversion. /// @@ -640,7 +649,7 @@ pub enum MethodLookupError { /// inspect [`crate::metadata::method::Method::rva_kind`] to see why no IL /// is present. /// - /// [`MethodDef`]: crate::metadata::tables::MethodDef + /// [`MethodDef`]: crate::metadata::tables::MethodDefRaw #[error("MethodDef token {0} not found")] NotFound(Token), @@ -823,13 +832,11 @@ impl std::fmt::Display for StreamKind { /// Structured parse-pipeline failure. /// -/// Reported through [`Error::Parse`] for every error site in -/// [`crate::file`], [`crate::metadata::root`], and -/// [`crate::metadata::streams`], plus per-table parse paths that read raw -/// PE/metadata bytes. Replaces the stringly-typed [`Error::Malformed`] / -/// [`Error::OutOfBounds`] / [`Error::HeapBoundsError`] variants for parse -/// sites — those remain valid for non-parse code (validation, lookups, -/// emulation), but new parse code must use [`ParseFailure`]. +/// Reported through [`Error::Parse`] for every error site in the file, metadata-root and +/// stream parsers, plus per-table parse paths that read raw +/// PE/metadata bytes. This replaced the stringly-typed `Malformed` / `OutOfBounds` / +/// `HeapBoundsError` variants, which no longer exist; every parse site reports through +/// [`ParseFailure`]. /// /// # Stability /// @@ -933,47 +940,22 @@ pub enum ParseFailure { }, } -impl Clone for Error { - fn clone(&self) -> Self { - match self { - // Handle non-cloneable variants by converting to string representation - Error::Io(io_err) => Error::Other(io_err.to_string()), - Error::Goblin(goblin_err) => Error::Other(goblin_err.to_string()), - // For validation errors that have Box sources, clone them recursively - Error::ValidationStage1Failed { source, message } => Error::ValidationStage1Failed { - source: source.clone(), - message: message.clone(), - }, - Error::ValidationRawFailed { validator, message } => Error::ValidationRawFailed { - validator: validator.clone(), - message: message.clone(), - }, - Error::ValidationOwnedFailed { validator, message } => Error::ValidationOwnedFailed { - validator: validator.clone(), - message: message.clone(), - }, - // Emulation errors are cloneable (boxed) - Error::Emulation(e) => Error::Emulation(e.clone()), - // Deobfuscation errors are cloneable - Error::Deobfuscation(s) => Error::Deobfuscation(s.clone()), - // X86 errors are cloneable - Error::X86Error(s) => Error::X86Error(s.clone()), - // Tracing errors are cloneable - Error::TracingError(s) => Error::TracingError(s.clone()), - // Method-lookup errors are pure data and Clone-derived. - Error::LookupMethod(e) => Error::LookupMethod(e.clone()), - // Parse failures are pure data and Clone-derived. - Error::Parse(e) => Error::Parse(e.clone()), - // For all other variants, convert to their string representation and use Other - other => Error::Other(other.to_string()), - } +impl From for Error { + fn from(err: io::Error) -> Self { + Error::Io(Arc::new(err)) + } +} + +impl From for Error { + fn from(err: goblin::error::Error) -> Self { + Error::Goblin(Arc::new(err)) } } impl From for Error { fn from(err: cowfile::Error) -> Self { match err { - cowfile::Error::Io(io_err) => Error::Io(io_err), + cowfile::Error::Io(io_err) => Error::Io(Arc::new(io_err)), cowfile::Error::OutOfBounds { .. } => Error::Parse(ParseFailure::OutOfBounds { stage: ParseStage::Generic, }), @@ -1000,3 +982,61 @@ impl From for analyssa::Error { analyssa::Error::new(err.to_string()) } } + +#[cfg(test)] +mod tests { + use super::*; + + /// `Clone` must preserve the variant. + /// + /// A hand-written `Clone` with a wildcard arm flattens most variants into + /// `Error::Other(String)`, so `matches!(e, Error::NotSupported)` is true before the clone + /// and false after — a contract violation no compiler catches. Deriving `Clone` is what + /// makes that impossible; this pins the property rather than the derive. + #[test] + fn clone_preserves_the_variant() { + let cases = [ + Error::NotSupported, + Error::TypeNotFound(Token::new(0x0200_0001)), + Error::RecursionLimit(32), + Error::Other("plain".to_string()), + ]; + + for error in cases { + let cloned = error.clone(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&cloned), + "clone changed the variant of {error:?}" + ); + assert_eq!(error.to_string(), cloned.to_string()); + } + } + + /// The nested error list of a stage-2 failure must survive a clone; a wildcard arm keeps + /// only the one-line summary and discards every nested `Error`. + #[test] + fn clone_preserves_nested_validation_errors() { + let nested = Error::ValidationStage2Failed { + errors: vec![Error::NotSupported, Error::RecursionLimit(32)], + error_count: 2, + summary: "two failures".to_string(), + }; + + let Error::ValidationStage2Failed { errors, .. } = nested.clone() else { + panic!("clone changed the variant"); + }; + assert_eq!(errors.len(), 2); + } + + /// An `io::Error`'s kind must survive a clone; stringifying it loses what callers match on. + #[test] + fn clone_preserves_the_io_error_kind() { + let error = Error::from(io::Error::new(io::ErrorKind::PermissionDenied, "denied")); + + let Error::Io(cloned) = error.clone() else { + panic!("clone changed the variant"); + }; + assert_eq!(cloned.kind(), io::ErrorKind::PermissionDenied); + } +} diff --git a/dotscope/src/file/mod.rs b/dotscope/src/file/mod.rs index 5564ebc4..c0d5cf26 100644 --- a/dotscope/src/file/mod.rs +++ b/dotscope/src/file/mod.rs @@ -134,7 +134,7 @@ use crate::{ repair::{repair_pe_cow, RepairAction}, }, utils::align_to, - Error::{self, Goblin, LayoutFailed, Other}, + Error::{self, LayoutFailed, Other}, ParseFailure, ParseStage, Result, }; @@ -395,7 +395,7 @@ impl File { cowfile.commit()?; } - let goblin_pe = PE::parse(cowfile.data()).map_err(Goblin)?; + let goblin_pe = PE::parse(cowfile.data()).map_err(Error::from)?; let pe = Pe::from_goblin_pe(&goblin_pe)?; Ok(File { diff --git a/dotscope/src/file/parser.rs b/dotscope/src/file/parser.rs index 04fbf643..c5b57f85 100644 --- a/dotscope/src/file/parser.rs +++ b/dotscope/src/file/parser.rs @@ -253,7 +253,7 @@ impl<'a> Parser<'a> { /// * `pos` - The position to move the cursor to /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if position is beyond the data length. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if position is beyond the data length. /// /// # Examples /// @@ -280,7 +280,7 @@ impl<'a> Parser<'a> { /// Move the position forward by one byte. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if advancing would exceed the data length. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if advancing would exceed the data length. /// /// # Examples /// @@ -304,7 +304,7 @@ impl<'a> Parser<'a> { /// * `step` - Amount of bytes to advance /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if advancing by step would exceed the data length. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if advancing by step would exceed the data length. /// /// # Examples /// @@ -364,7 +364,7 @@ impl<'a> Parser<'a> { /// Peek at the next byte without advancing the position. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if position is at or beyond the data length. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if position is at or beyond the data length. /// /// # Examples /// @@ -390,7 +390,7 @@ impl<'a> Parser<'a> { /// parser state, allowing inspection of upcoming data before deciding how to proceed. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading `T` would exceed the data length. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading `T` would exceed the data length. /// /// # Examples /// @@ -475,7 +475,7 @@ impl<'a> Parser<'a> { /// * `alignment` - The boundary to align to (must be a power of 2) /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if aligning would exceed the data length. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if aligning would exceed the data length. /// /// # Examples /// @@ -506,7 +506,7 @@ impl<'a> Parser<'a> { /// Read a type `T` from the current position in little-endian format and advance the position. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading would exceed the data length. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading would exceed the data length. /// /// # Examples /// @@ -527,7 +527,7 @@ impl<'a> Parser<'a> { /// Read a type `T` from the current position in big-endian format and advance the position. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading would exceed the data length. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading would exceed the data length. /// /// # Examples /// @@ -553,8 +553,8 @@ impl<'a> Parser<'a> { /// - Values 16384-536870911: 4 bytes (11xxxxxx xxxxxxxx xxxxxxxx xxxxxxxx) /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading would exceed the data length or - /// [`crate::Error::Malformed`] for invalid compressed uint format. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading would exceed the data length or + /// [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] for invalid compressed uint format. /// /// # Examples /// @@ -608,8 +608,8 @@ impl<'a> Parser<'a> { /// but with the least significant bit indicating the sign and the remaining bits shifted right. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading would exceed the data length or - /// [`crate::Error::Malformed`] for invalid encoding. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading would exceed the data length or + /// [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] for invalid encoding. /// /// # Examples /// @@ -663,8 +663,8 @@ impl<'a> Parser<'a> { /// Encountering this tag value indicates a malformed compressed token. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading would exceed the data length or - /// [`crate::Error::Malformed`] if tag 0x3 is encountered (invalid encoding). + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading would exceed the data length or + /// [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] if tag 0x3 is encountered (invalid encoding). /// /// # Examples /// @@ -711,8 +711,8 @@ impl<'a> Parser<'a> { /// concatenating the lower 7 bits of each byte in little-endian order. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading would exceed the data length or - /// [`crate::Error::Malformed`] for invalid encoding (overflow). + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading would exceed the data length or + /// [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] for invalid encoding (overflow). /// /// # Examples /// @@ -768,8 +768,8 @@ impl<'a> Parser<'a> { /// then decodes the bytes as UTF-8. The position is advanced past the null terminator. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading would exceed the data length or - /// [`crate::Error::Malformed`] for invalid UTF-8 encoding. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading would exceed the data length or + /// [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] for invalid UTF-8 encoding. /// /// # Examples /// @@ -825,8 +825,8 @@ impl<'a> Parser<'a> { /// UTF-8 bytes. This format is commonly used in .NET metadata streams. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading would exceed the data length or - /// [`crate::Error::Malformed`] for invalid UTF-8 encoding. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading would exceed the data length or + /// [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] for invalid UTF-8 encoding. /// /// # Examples /// @@ -877,8 +877,8 @@ impl<'a> Parser<'a> { /// /// # Errors /// - /// Returns [`crate::Error::OutOfBounds`] if reading would exceed the data length or - /// [`crate::Error::Malformed`] for invalid UTF-8 encoding. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading would exceed the data length or + /// [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] for invalid UTF-8 encoding. /// /// # Examples /// @@ -917,8 +917,8 @@ impl<'a> Parser<'a> { /// security permissions, and other metadata structures that follow ECMA-335 blob format. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading would exceed the data length or - /// [`crate::Error::Malformed`] for invalid UTF-8 encoding. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading would exceed the data length or + /// [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] for invalid UTF-8 encoding. /// /// # Examples /// @@ -984,7 +984,7 @@ impl<'a> Parser<'a> { /// * `needed` - The number of bytes required from the current position /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if fewer than `needed` bytes remain. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if fewer than `needed` bytes remain. /// /// # Examples /// @@ -1015,7 +1015,7 @@ impl<'a> Parser<'a> { /// * `length` - The length to add to the current position /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if the calculation would overflow + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if the calculation would overflow /// or if the resulting position exceeds the data length. /// /// # Examples @@ -1052,7 +1052,7 @@ impl<'a> Parser<'a> { /// * `length` - The number of bytes to read /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading `length` bytes would exceed the data. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading `length` bytes would exceed the data. /// /// # Examples /// @@ -1080,8 +1080,8 @@ impl<'a> Parser<'a> { /// character strings in .NET metadata. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if reading would exceed the data length or - /// [`crate::Error::Malformed`] for invalid UTF-16 encoding or odd byte length. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if reading would exceed the data length or + /// [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] for invalid UTF-16 encoding or odd byte length. /// /// # Examples /// diff --git a/dotscope/src/formatting/helpers.rs b/dotscope/src/formatting/helpers.rs index 21b1f763..8cf15ffb 100644 --- a/dotscope/src/formatting/helpers.rs +++ b/dotscope/src/formatting/helpers.rs @@ -598,7 +598,7 @@ pub(super) fn format_method_call_sig( pub(super) fn format_typespec_from_blob(asm: &CilObject, token: &Token) -> Option { let tables = asm.tables()?; let table = tables.table::()?; - let row = table.get(token.row())?; + let row = table.get(token.row()).ok().flatten()?; let blob = asm.blob()?; let sig_data = blob.get(row.signature as usize).ok()?; let parsed = parse_type_spec_signature(sig_data).ok()?; diff --git a/dotscope/src/formatting/method_body.rs b/dotscope/src/formatting/method_body.rs index f7c2f8a8..391689c8 100644 --- a/dotscope/src/formatting/method_body.rs +++ b/dotscope/src/formatting/method_body.rs @@ -149,7 +149,7 @@ fn format_locals(w: &mut dyn Write, method: &Method, asm: &CilObject) -> io::Res let rid = body.local_var_sig_token & 0x00FF_FFFF; let tables = asm.tables()?; let table = tables.table::()?; - let row = table.get(rid)?; + let row = table.get(rid).ok().flatten()?; let blob = asm.blob()?; let sig_data = blob.get(row.signature as usize).ok()?; parse_local_var_signature(sig_data).ok() diff --git a/dotscope/src/formatting/mod.rs b/dotscope/src/formatting/mod.rs index f3dd5bed..a4b4d1ff 100644 --- a/dotscope/src/formatting/mod.rs +++ b/dotscope/src/formatting/mod.rs @@ -418,7 +418,7 @@ mod tests { let cil_type = asm .query_types() .defined() - .filter(|t| t.fullname() == type_name) + .filter(|t| &*t.fullname() == type_name) .find_all() .into_iter() .next() @@ -550,14 +550,14 @@ mod tests { // DerivedClass has MetadataTestAttribute with fixed args let output = format_type_by_name(&asm, "DerivedClass"); // Find the comment and .custom lines - if let Some(comment_pos) = output.find("// (int32(100)") { - if let Some(custom_pos) = output[comment_pos..].find(".custom ") { - // Comment should come before .custom in the same region - assert!( - custom_pos > 0, - "Comment should appear before .custom directive" - ); - } + if let Some((_, after_comment)) = output.split_once("// (int32(100)") { + // Searching the text that *follows* the comment states the property directly. The + // previous form sliced from the comment's own offset and asserted the directive's + // index was non-zero, which no match could ever violate. + assert!( + after_comment.contains(".custom "), + "Comment should appear before .custom directive" + ); } } diff --git a/dotscope/src/formatting/types.rs b/dotscope/src/formatting/types.rs index d8d5cd06..c3ba4403 100644 --- a/dotscope/src/formatting/types.rs +++ b/dotscope/src/formatting/types.rs @@ -171,7 +171,7 @@ pub(super) fn format_type_end(w: &mut dyn Write, cil_type: &CilType) -> io::Resu let display_name = if cil_type.enclosing_type().is_some() { cil_type.name.clone() } else { - cil_type.fullname() + cil_type.fullname().to_string() }; writeln!(w, "}} // end of class {display_name}")?; writeln!(w)?; diff --git a/dotscope/src/lib.rs b/dotscope/src/lib.rs index 67db3cc3..c840b16d 100644 --- a/dotscope/src/lib.rs +++ b/dotscope/src/lib.rs @@ -32,9 +32,19 @@ ) )] #![allow(dead_code)] -//#![deny(unsafe_code)] -// - 'userstring.rs' uses a transmute for converting a &[u8] to &[u16] -// - 'file/physical.rs' uses mmap to map a file into memory +// The crate's trusted computing base is exactly one `unsafe` block, and it carries a targeted +// `#[allow(unsafe_code)]` with its own SAFETY note: +// +// - `cilassembly/writer/output.rs` — `Mmap::map_mut` over the output file. +// +// The deny lint does not reach through dependencies, so note the other mapping here: `cowfile`'s +// `map_copy`, which the primary load path routes its mmap through. +#![deny(unsafe_code)] +// A broken intra-doc link renders as plain text on docs.rs, so a `# Errors` contract naming a +// variant that no longer exists still reads as authoritative. Denying the lint is what keeps +// the documented error taxonomy tied to the real one; CI sets `RUSTDOCFLAGS: -Dwarnings` so it +// is enforced on the doc build too. +#![deny(rustdoc::broken_intra_doc_links)] //! # dotscope //! @@ -91,7 +101,7 @@ //! //! ```toml //! [dependencies] -//! dotscope = "0.7.0" +//! dotscope = "0.9" //! ``` //! //! ### Using the Prelude diff --git a/dotscope/src/metadata/cilassemblyview.rs b/dotscope/src/metadata/cilassemblyview.rs index 1e7ab370..9366cf7e 100644 --- a/dotscope/src/metadata/cilassemblyview.rs +++ b/dotscope/src/metadata/cilassemblyview.rs @@ -121,7 +121,7 @@ use crate::{ identity::{AssemblyIdentity, AssemblyVersion, Identity, ProcessorArchitecture}, root::Root, streams::{Blob, Guid, StreamHeader, Strings, TablesHeader, UserStrings}, - tables::{AssemblyProcessorRaw, AssemblyRaw, AssemblyRefRaw, ModuleRaw}, + tables::{skip_unreadable, AssemblyProcessorRaw, AssemblyRaw, AssemblyRefRaw, ModuleRaw}, validation::ValidationEngine, }, Error, Result, ValidationConfig, @@ -198,7 +198,7 @@ impl<'a> CilAssemblyViewData<'a> { /// # Errors /// /// Returns [`crate::Error::NotSupported`] if the file is not a .NET assembly (missing CLR header). - /// Returns [`crate::Error::OutOfBounds`] if the file data is truncated or corrupted. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if the file data is truncated or corrupted. pub fn from_dotscope_file(file: Arc, data: &'a [u8]) -> Result { let (clr_rva, clr_size) = file.clr().ok_or(Error::NotSupported)?; @@ -431,7 +431,7 @@ impl CilAssemblyView { /// /// Returns [`crate::Error::Io`] if the file cannot be read. /// Returns [`crate::Error::NotSupported`] if the file is not a .NET assembly. - /// Returns [`crate::Error::OutOfBounds`] if the file data is corrupted. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if the file data is corrupted. /// /// # Examples /// @@ -475,7 +475,7 @@ impl CilAssemblyView { /// /// Returns [`crate::Error::Io`] if the file cannot be read. /// Returns [`crate::Error::NotSupported`] if the file is not a .NET assembly. - /// Returns [`crate::Error::OutOfBounds`] if the file data is corrupted. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if the file data is corrupted. /// Returns validation errors if validation checks fail. /// /// # Examples @@ -518,7 +518,7 @@ impl CilAssemblyView { /// # Errors /// /// Returns [`crate::Error::NotSupported`] if the data is not a .NET assembly. - /// Returns [`crate::Error::OutOfBounds`] if the data is corrupted or truncated. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if the data is corrupted or truncated. /// /// # Examples /// @@ -551,7 +551,7 @@ impl CilAssemblyView { /// # Errors /// /// Returns [`crate::Error::NotSupported`] if the data is not a .NET assembly. - /// Returns [`crate::Error::OutOfBounds`] if the data is corrupted or truncated. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if the data is corrupted or truncated. /// Returns validation errors if validation checks fail. /// /// # Examples @@ -1064,6 +1064,7 @@ impl CilAssemblyView { assembly_ref_table .into_iter() + .filter_map(skip_unreadable) .filter_map(|row| row.to_owned(strings, blobs).ok()) .map(|assembly_ref| AssemblyIdentity::from_assembly_ref(&assembly_ref)) .collect() @@ -1119,7 +1120,10 @@ impl CilAssemblyView { return Ok(None); }; - let Some(assembly_row) = assembly_table.iter().next() else { + // The Assembly table holds at most one row, at RID 1 (ECMA-335 II.22.2), so it is + // fetched by RID. Taking the first row that *parses* would silently promote a later + // row when RID 1 is malformed, and report another row's identity as the assembly's. + let Some(assembly_row) = assembly_table.get(1).ok().flatten() else { // Empty Assembly table - also a netmodule return Ok(None); }; @@ -1152,7 +1156,7 @@ impl CilAssemblyView { let processor_architecture = tables .table::() - .and_then(|proc_table| proc_table.iter().next()) + .and_then(|proc_table| proc_table.get(1).ok().flatten()) .and_then(|proc| ProcessorArchitecture::try_from(proc.processor).ok()); #[allow(clippy::cast_possible_truncation)] @@ -1184,7 +1188,8 @@ impl CilAssemblyView { let strings = self.strings()?; let module_table = tables.table::()?; - let module_row = module_table.iter().next()?; + // Module is RID 1 by definition (ECMA-335 II.22.30). + let module_row = module_table.get(1).ok().flatten()?; strings.get(module_row.name as usize).ok().map(String::from) } diff --git a/dotscope/src/metadata/cor20header.rs b/dotscope/src/metadata/cor20header.rs index 9171c7ab..65baf1be 100644 --- a/dotscope/src/metadata/cor20header.rs +++ b/dotscope/src/metadata/cor20header.rs @@ -257,8 +257,8 @@ impl Cor20Header { /// Returns a parsed and validated [`Cor20Header`] on success. /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if the data is too short, or - /// [`crate::Error::Malformed`] if any field validation fails per ECMA-335 requirements: + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if the data is too short, or + /// [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] if any field validation fails per ECMA-335 requirements: /// - Invalid header size (not 72 bytes) /// - Invalid runtime version (0 or > 10) /// - Zero metadata RVA or size diff --git a/dotscope/src/metadata/customattributes/encoder.rs b/dotscope/src/metadata/customattributes/encoder.rs index ddc07822..eebc09e8 100644 --- a/dotscope/src/metadata/customattributes/encoder.rs +++ b/dotscope/src/metadata/customattributes/encoder.rs @@ -129,7 +129,7 @@ use crate::{ /// /// # Errors /// -/// Returns [`crate::Error::Malformed`] in the following cases: +/// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] in the following cases: /// - Too many named arguments (exceeds u16 maximum of 65535) /// - Array length exceeds u32 maximum /// - String length exceeds compressed uint maximum (0x1FFFFFFF bytes) @@ -273,7 +273,7 @@ fn encode_named_arguments( /// /// # Errors /// -/// Returns [`crate::Error::Malformed`] in the following cases: +/// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] in the following cases: /// - `Char`: Character code point exceeds 0xFFFF (outside Basic Multilingual Plane) /// - `String`: String length exceeds compressed uint maximum (0x1FFFFFFF bytes) /// - `Array`: Array length exceeds u32 maximum (4,294,967,295 elements) diff --git a/dotscope/src/metadata/customattributes/parser.rs b/dotscope/src/metadata/customattributes/parser.rs index 76b720ad..00f6d008 100644 --- a/dotscope/src/metadata/customattributes/parser.rs +++ b/dotscope/src/metadata/customattributes/parser.rs @@ -197,8 +197,8 @@ const MAX_ATTRIBUTE_ARRAY_LENGTH: i32 = 65536; /// - `named_args` - Field and property assignments with names and values /// /// # Errors -/// Returns [`crate::Error::OutOfBounds`] if the index is invalid, or one of the following: -/// - [`crate::Error::Malformed`]: Invalid prolog (not 0x0001), insufficient data for declared arguments, or type/value mismatches in argument parsing +/// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if the index is invalid, or one of the following: +/// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid prolog (not 0x0001), insufficient data for declared arguments, or type/value mismatches in argument parsing /// - [`crate::Error::DepthLimitExceeded`]: Maximum nesting depth exceeded during parsing /// /// # Examples @@ -265,7 +265,7 @@ pub fn parse_custom_attribute_blob( /// /// # Errors /// Returns one of the following errors if the blob data doesn't conform to ECMA-335 format: -/// - [`crate::Error::Malformed`]: Invalid or missing prolog (must be 0x0001), insufficient data for the number of declared arguments, type mismatches between expected and actual argument types, invalid serialization type tags in named arguments, or truncated/corrupted blob data +/// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid or missing prolog (must be 0x0001), insufficient data for the number of declared arguments, type mismatches between expected and actual argument types, invalid serialization type tags in named arguments, or truncated/corrupted blob data /// - [`crate::Error::DepthLimitExceeded`]: Maximum nesting depth exceeded during complex type parsing /// /// # Examples @@ -478,7 +478,7 @@ impl<'a> CustomAttributeParser<'a> { /// A complete [`crate::metadata::customattributes::CustomAttributeValue`] with all parsed data. /// /// # Errors - /// Returns [`crate::Error::Malformed`] for various format violations: + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] for various format violations: /// - Invalid prolog (not 0x0001) /// - Insufficient data for declared arguments /// - Invalid serialization types in named arguments @@ -561,7 +561,7 @@ impl<'a> CustomAttributeParser<'a> { /// Vector of parsed arguments in constructor parameter order /// /// # Errors - /// Returns [`crate::Error::Malformed`] if: + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] if: /// - Constructor has parameters but no resolved types /// - Insufficient blob data for declared parameters /// - Parameter type parsing fails @@ -689,7 +689,7 @@ impl<'a> CustomAttributeParser<'a> { /// Parsed argument if successful, None if type is unsupported /// /// # Errors - /// Returns [`crate::Error::Malformed`] for invalid data or unsupported types + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] for invalid data or unsupported types fn parse_fixed_argument( &mut self, cil_type: &CilTypeRef, @@ -809,7 +809,7 @@ impl<'a> CustomAttributeParser<'a> { // BUT: Enum types can also appear as Class and should be handled as ValueType/Enum let type_name = type_ref.fullname(); - if type_name == "System.Type" { + if &*type_name == "System.Type" { // System.Type is stored as a string (type name) if self.parser.peek_byte()? == 0xFF { let _ = self.parser.read_le::()?; // consume null marker @@ -820,7 +820,7 @@ impl<'a> CustomAttributeParser<'a> { })?; Ok(Some(CustomAttributeArgument::Type(s))) } - } else if type_name == "System.String" { + } else if &*type_name == "System.String" { // System.String is stored as a string if self.parser.peek_byte()? == 0xFF { let _ = self.parser.read_le::()?; // consume null marker @@ -831,7 +831,7 @@ impl<'a> CustomAttributeParser<'a> { })?; Ok(Some(CustomAttributeArgument::String(s))) } - } else if type_name == "System.Object" { + } else if &*type_name == "System.Object" { // System.Object is stored as a tagged object - read type tag first let type_tag = self.parser.read_le::()?; let value = self.parse_argument_by_type_tag(type_tag)?; @@ -846,7 +846,8 @@ impl<'a> CustomAttributeParser<'a> { if EnumUtils::is_enum_type(&resolved_type, Some(registry)) { let underlying_type_size = EnumUtils::get_enum_underlying_type_size(&resolved_type); - return self.parse_enum(type_name, underlying_type_size); + return self + .parse_enum(type_name.to_string(), underlying_type_size); } } } @@ -865,7 +866,7 @@ impl<'a> CustomAttributeParser<'a> { EnumUtils::get_enum_underlying_type_size(&type_ref) }; - return self.parse_enum(type_name, underlying_type_size); + return self.parse_enum(type_name.to_string(), underlying_type_size); } // Stage 3: Fallback for unresolvable external types @@ -874,7 +875,7 @@ impl<'a> CustomAttributeParser<'a> { // determine inheritance. For custom attributes, external types that // aren't System.Type/String/Object are typically enums. Assume int32 // underlying type (the most common) to allow parsing to continue. - self.parse_enum(type_name, 4) + self.parse_enum(type_name.to_string(), 4) } } CilFlavor::ValueType => { @@ -901,7 +902,7 @@ impl<'a> CustomAttributeParser<'a> { if EnumUtils::is_enum_type(&resolved_type, Some(registry)) { let underlying_type_size = EnumUtils::get_enum_underlying_type_size(&resolved_type); - return self.parse_enum(type_name, underlying_type_size); + return self.parse_enum(type_name.to_string(), underlying_type_size); } } } @@ -920,7 +921,7 @@ impl<'a> CustomAttributeParser<'a> { EnumUtils::get_enum_underlying_type_size(&type_ref) }; - self.parse_enum(type_name, underlying_type_size) + self.parse_enum(type_name.to_string(), underlying_type_size) } else { // Stage 3: No resolution possible - missing dependencies Err(malformed_error!( @@ -1001,7 +1002,7 @@ impl<'a> CustomAttributeParser<'a> { /// Parsed named argument with name, type, and value, or None if no more data /// /// # Errors - /// Returns [`crate::Error::Malformed`] for invalid format or unsupported types + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] for invalid format or unsupported types fn parse_named_argument(&mut self) -> Result> { if !self.parser.has_more_data() { return Ok(None); @@ -1054,9 +1055,23 @@ impl<'a> CustomAttributeParser<'a> { } }; - // Read field/property name - let name_length = self.parser.read_compressed_uint()?; - let mut name = String::with_capacity(name_length as usize); + // Read field/property name. + // + // Each name byte consumes exactly one input byte, so a length beyond what remains in the + // blob can never be legitimate. Reject it before reserving, mirroring `parse_string`: + // `read_compressed_uint` reaches 0x1FFF_FFFF, so four attacker bytes would otherwise buy + // a ~512 MB reservation per row before the first read could fail. + let name_length = self.parser.read_compressed_uint()? as usize; + let available_data = self.parser.len().saturating_sub(self.parser.pos()); + if name_length > available_data { + return Err(malformed_error!( + "Named argument name length {} exceeds {} remaining blob bytes", + name_length, + available_data + )); + } + + let mut name = String::with_capacity(name_length); for _ in 0..name_length { name.push(char::from(self.parser.read_le::()?)); } @@ -1097,7 +1112,7 @@ impl<'a> CustomAttributeParser<'a> { /// /// # Errors /// - [`crate::Error::DepthLimitExceeded`]: Maximum nesting depth exceeded - /// - [`crate::Error::Malformed`]: Invalid type tags or malformed data format + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid type tags or malformed data format fn parse_argument_by_type_tag(&mut self, type_tag: u8) -> Result { /// Work item for iterative parsing stack enum WorkItem { @@ -1217,6 +1232,26 @@ impl<'a> CustomAttributeParser<'a> { "Invalid array length: {}", array_length )); + } else if array_length > MAX_ATTRIBUTE_ARRAY_LENGTH { + // Same cap the type-driven path applies. Without it this arm is + // an allocation bomb: the push loop below consumes no further + // input, so the blob's own length does not bound the work and + // ~12 bytes can demand gigabytes of `work_stack`. + return Err(malformed_error!( + "Custom attribute array too large: {} (max: {})", + array_length, + MAX_ATTRIBUTE_ARRAY_LENGTH + )); + } else if work_stack + .len() + .saturating_add(usize::try_from(array_length).unwrap_or(usize::MAX)) + > MAX_NESTING_DEPTH + { + // Bound the burst *before* it happens. The depth check at the top + // of this loop is only evaluated on the next pop, i.e. after the + // vector has already grown. Behaviour-preserving: any length past + // the depth limit would error on that next pop regardless. + return Err(DepthLimitExceeded(MAX_NESTING_DEPTH)); } else { // Schedule work to build array after parsing elements work_stack.push(WorkItem::BuildArray(array_length)); @@ -1367,7 +1402,7 @@ impl<'a> CustomAttributeParser<'a> { /// Parsed string (empty `String` for both null marker and zero length) /// /// # Errors - /// Returns [`crate::Error::Malformed`] if: + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] if: /// - No data available for reading /// - Declared length exceeds available data /// - Compressed length parsing fails @@ -1438,6 +1473,7 @@ mod tests { create_constructor_with_params, create_constructor_with_params_and_registry, create_empty_constructor, get_test_type_registry, }, + utils::LazyList, }; #[test] @@ -1812,7 +1848,7 @@ mod tests { modifiers: Arc::new(boxcar::Vec::new()), base: OnceLock::new(), is_by_ref: std::sync::atomic::AtomicBool::new(false), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); param.base.set(CilTypeRef::from(array_type)).ok(); method.params.push(param); diff --git a/dotscope/src/metadata/customattributes/types.rs b/dotscope/src/metadata/customattributes/types.rs index b1632a3f..22c73414 100644 --- a/dotscope/src/metadata/customattributes/types.rs +++ b/dotscope/src/metadata/customattributes/types.rs @@ -134,7 +134,7 @@ use std::{fmt, sync::Arc}; -use crate::metadata::typesystem::CilTypeReference; +use crate::{metadata::typesystem::CilTypeReference, utils::LazyList}; /// A reference-counted pointer to a [`CustomAttributeValue`] for efficient sharing. /// @@ -148,7 +148,7 @@ pub type CustomAttributeValueRc = Arc; /// Provides thread-safe storage for custom attribute collections on metadata objects. /// Uses [`boxcar::Vec`] for lock-free concurrent access and [`Arc`] for reference counting, /// enabling efficient metadata processing in multi-threaded scenarios. -pub type CustomAttributeValueList = Arc>; +pub type CustomAttributeValueList = LazyList; /// Represents a complete parsed custom attribute with fixed and named arguments. /// diff --git a/dotscope/src/metadata/dependencies/analyzer.rs b/dotscope/src/metadata/dependencies/analyzer.rs index c2ef1f51..a2366083 100644 --- a/dotscope/src/metadata/dependencies/analyzer.rs +++ b/dotscope/src/metadata/dependencies/analyzer.rs @@ -194,6 +194,13 @@ impl DependencyAnalyzer { let source_identity = Self::extract_current_assembly_identity(context)?; for row in table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; // Convert raw AssemblyRef to owned representation let assembly_ref = row.to_owned(strings, blobs)?; @@ -262,6 +269,13 @@ impl DependencyAnalyzer { let source_identity = Self::extract_current_assembly_identity(context)?; for row in table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; // Convert raw ModuleRef to owned representation let module_ref = row.to_owned(strings)?; @@ -376,6 +390,13 @@ impl DependencyAnalyzer { let source_identity = Self::extract_current_assembly_identity(context)?; for row in table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; // Convert raw File to owned representation let file_ref = row.to_owned(blobs, strings)?; diff --git a/dotscope/src/metadata/dependencies/graph.rs b/dotscope/src/metadata/dependencies/graph.rs index 490823f8..304aa1c5 100644 --- a/dotscope/src/metadata/dependencies/graph.rs +++ b/dotscope/src/metadata/dependencies/graph.rs @@ -539,7 +539,7 @@ impl AssemblyDependencyGraph { /// os_major_version: std::sync::atomic::AtomicU32::new(0), /// os_minor_version: std::sync::atomic::AtomicU32::new(0), /// processor: std::sync::atomic::AtomicU32::new(0), - /// custom_attributes: Arc::new(boxcar::Vec::new()), + /// custom_attributes: dotscope::metadata::customattributes::CustomAttributeValueList::new(), /// }); /// /// let dep = AssemblyDependency { diff --git a/dotscope/src/metadata/identity/cryptographic.rs b/dotscope/src/metadata/identity/cryptographic.rs index 99730b5e..cfccbd3c 100644 --- a/dotscope/src/metadata/identity/cryptographic.rs +++ b/dotscope/src/metadata/identity/cryptographic.rs @@ -199,7 +199,7 @@ impl Identity { /// - [`Identity::PubKey`] for other public key sizes /// /// # Errors - /// Returns [`crate::Error::OutOfBounds`] if: + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if: /// - Token creation requested but data has fewer than 8 bytes /// - Data cannot be read as little-endian `u64` /// diff --git a/dotscope/src/metadata/loader/context.rs b/dotscope/src/metadata/loader/context.rs index 3efa0306..c5d8f198 100644 --- a/dotscope/src/metadata/loader/context.rs +++ b/dotscope/src/metadata/loader/context.rs @@ -73,13 +73,17 @@ //! - [`crate::metadata::tables`] - All metadata table types and coded index resolution //! - [`crate::metadata::streams`] - Metadata stream access and heap operations -use std::sync::{Arc, OnceLock}; +use std::{ + fmt::Display, + result::Result as StdResult, + sync::{Arc, OnceLock}, +}; use crate::{ file::File, metadata::{ cor20header::Cor20Header, - diagnostics::{DiagnosticCategory, Diagnostics}, + diagnostics::{Diagnostic, DiagnosticCategory, DiagnosticSeverity, Diagnostics}, exports::Exports, imports::Imports, method::MethodMap, @@ -95,11 +99,12 @@ use crate::{ ImportScopeMap, InterfaceImplMap, LocalConstantMap, LocalScopeMap, LocalVariableMap, MemberRefMap, MethodDebugInformationMap, MethodImplMap, MethodPtrMap, MethodSemanticsMap, MethodSpecMap, ModuleRc, ModuleRefMap, NestedClassMap, ParamMap, - ParamPtrMap, PropertyMap, PropertyMapEntryMap, PropertyPtrMap, StandAloneSigMap, - StateMachineMethodMap, TableId, TypeSpecMap, + ParamPtrMap, PropertyMap, PropertyMapEntryMap, PropertyPtrMap, RowReadable, + StandAloneSigMap, StateMachineMethodMap, TableId, TypeSpecMap, }, typesystem::{CilTypeReference, TypeRegistry}, }, + Result, }; /// Centralized context for metadata table maps during assembly loading. @@ -569,12 +574,12 @@ impl LoaderContext<'_> { /// // Use owned value /// } /// ``` - pub fn handle_result String>( + pub fn handle_result String>( &self, - result: std::result::Result, + result: StdResult, category: DiagnosticCategory, context_msg: F, - ) -> crate::Result> { + ) -> Result> { match result { Ok(value) => Ok(Some(value)), Err(e) => { @@ -590,6 +595,69 @@ impl LoaderContext<'_> { } } + /// Handles a *row-parse* failure according to the lenient/strict loading mode. + /// + /// The counterpart to [`handle_result`](Self::handle_result) for the one failure that + /// happens before there is anything to build a message from. Every loader routes the + /// failures it can describe — `to_owned`, reference resolution — through `handle_result` + /// with a message keyed on `row.token`. A row that failed to *parse* has no token, so it + /// cannot use that path, and the natural `let row = row?;` propagates straight out of + /// `load()` and aborts the whole [`CilObject`](crate::CilObject) load. That defeats + /// [`lenient`](Self::lenient) for exactly the malformed input the flag exists to tolerate: + /// one truncated row at the end of one table would take the file down. + /// + /// Identity comes from the table id and RID instead, recorded structurally on the + /// diagnostic via [`Diagnostic::with_table_row`] so consumers can locate the row without + /// parsing the message. + /// + /// # Arguments + /// + /// * `row` - The row-parse result, as yielded by the table iterators + /// * `index` - The row's 0-based position in the iteration. The table iterators are + /// range-based and read each row at an offset derived from its own RID, so nothing is + /// dropped on a parse failure and `index + 1` is the true RID — unlike a filtered + /// sequence, where position stops tracking RID. + /// + /// The table is taken from [`RowReadable::TABLE_ID`], so it cannot drift from the row + /// type being loaded. + /// + /// # Returns + /// + /// * `Ok(Some(row))` - The row parsed + /// * `Ok(None)` - The row did not parse and we are in lenient mode; skip it + /// * `Err(e)` - The row did not parse and we are in strict mode + /// + /// # Errors + /// + /// Returns [`crate::Error::Malformed`] when a row fails to parse in strict mode. + pub fn handle_row(&self, row: Result, index: usize) -> Result> { + match row { + Ok(row) => Ok(Some(row)), + Err(e) => { + let table = T::TABLE_ID; + let rid = u32::try_from(index).unwrap_or(u32::MAX).saturating_add(1); + if self.lenient { + self.diagnostics.push( + Diagnostic::new( + DiagnosticSeverity::Warning, + DiagnosticCategory::Table, + format!("Failed to parse {table:?} row {rid}: {e}"), + ) + .with_table_row(table as u8, rid), + ); + Ok(None) + } else { + Err(malformed_error!( + "Failed to parse {:?} row {}: {}", + table, + rid, + e + )) + } + } + } + } + /// Handles an error according to the lenient/strict loading mode. /// /// In strict mode (default), errors are propagated immediately. @@ -605,12 +673,12 @@ impl LoaderContext<'_> { /// /// * `Ok(())` - Operation succeeded or failed in lenient mode /// * `Err(e)` - Operation failed in strict mode - pub fn handle_error String>( + pub fn handle_error String>( &self, - result: std::result::Result<(), E>, + result: StdResult<(), E>, category: DiagnosticCategory, context_msg: F, - ) -> crate::Result<()> { + ) -> Result<()> { match result { Ok(()) => Ok(()), Err(e) => { diff --git a/dotscope/src/metadata/loader/inheritance.rs b/dotscope/src/metadata/loader/inheritance.rs index 9b69b0ff..c0aea7e1 100644 --- a/dotscope/src/metadata/loader/inheritance.rs +++ b/dotscope/src/metadata/loader/inheritance.rs @@ -8,6 +8,8 @@ //! scattered across individual table loaders, providing a unified and more robust //! approach to handling circular dependencies. +use rayon::iter::ParallelIterator; + use crate::{ metadata::{ loader::{LoaderContext, MetadataLoader}, @@ -49,6 +51,7 @@ impl MetadataLoader for InheritanceResolver { typedef_table .par_iter() .try_for_each(|raw_typedef| -> Result<()> { + let raw_typedef = raw_typedef?; if raw_typedef.extends.row == 0 { return Ok(()); } @@ -65,6 +68,37 @@ impl MetadataLoader for InheritanceResolver { )) })?; + // A type that extends itself is malformed, and it is the + // cheapest inheritance cycle for a hostile assembly to + // express. Leave the base unset rather than recording the + // self-edge: every base walk is bounded by + // `CilType::base_chain`, but keeping the trivial cycle out + // of the graph entirely means consumers never observe it. + // + // This deliberately does not fail the load. Analysing + // malformed and obfuscated binaries is the point of this + // library, so the condition is reported by validation + // rather than by refusing the file. + // + // Eliding the edge has a consequence that is easy to miss: + // any validator that looks for inheritance cycles by + // *walking the graph* can no longer see this one, because + // the edge it would follow was never recorded. + // `OwnedCircularityValidator` therefore re-reads the raw + // `extends` column specifically for the self-edge; see + // `validate_self_referential_bases`. Longer cycles + // (`A -> B -> A`) are not elided and are still found by the + // walk. If this elision is ever removed, that check becomes + // redundant rather than wrong. + if base_type_ref.token == type_def.token { + log::warn!( + "InheritanceResolver: type {} extends itself; \ + ignoring the self-referential base", + type_def.token + ); + return Ok(()); + } + // Use the resolved base type directly by its token. // Do NOT lookup by fullname as that can return the wrong type // when multiple types share the same name (e.g., nested types diff --git a/dotscope/src/metadata/method/body.rs b/dotscope/src/metadata/method/body.rs index bfce74f3..e474d140 100644 --- a/dotscope/src/metadata/method/body.rs +++ b/dotscope/src/metadata/method/body.rs @@ -105,6 +105,8 @@ use std::io::Write; +use log::debug; + use crate::{ metadata::method::{ encode_exception_handlers, ExceptionHandler, ExceptionHandlerFlags, MethodBodyFlags, @@ -470,7 +472,38 @@ impl MethodBody { let first_duo = read_le::(data)?; - let size_header = (first_duo >> 12).wrapping_mul(4); + // ECMA-335 II.25.4.3: the top nibble is the header size in dwords, and the + // spec mandates 3 (12 bytes). + // + // Only a value *below* 3 is actually dangerous: `size_code` is read from + // offset 4 and `localVarSigTok` from offset 8, so a shorter header would make + // the body start inside its own header and shift every downstream offset. That + // is refused. + // + // A value above 3 is non-conformant but not ambiguous, and it is what real + // runtimes accept: CoreCLR and dnlib both locate the IL at + // `4 * (flags >> 12)` rather than assuming 12. Rejecting those outright would + // make dotscope refuse assemblies the CLR happily executes — the wrong trade + // for a malware-analysis tool, whose input is by definition not + // spec-conformant. The declared size is honoured and the deviation is reported. + let header_dwords = first_duo >> 12; + if header_dwords < 3 { + return Err(malformed_error!( + "fat method header declares {} dwords, ECMA-335 II.25.4.3 requires at least 3", + header_dwords + )); + } + if header_dwords > 3 { + debug!( + "fat method header declares {header_dwords} dwords; ECMA-335 II.25.4.3 \ + mandates 3, honouring the declared size as CoreCLR and dnlib do" + ); + } + + // Every downstream offset — the IL start, the `size_code` bounds check and the + // EH cursor — is derived from this, so it must be the declared size rather + // than a hard-coded 12. + let size_header: u16 = header_dwords.saturating_mul(4); let size_code = read_le::(data.get(4..).ok_or(out_of_bounds_error!())?)?; let total = (size_code as usize) .checked_add(size_header as usize) @@ -1301,4 +1334,70 @@ mod tests { assert_eq!(raw_body.exception_handlers.len(), 1); assert_eq!(raw_body.exception_handlers[0].try_offset, 0xFFFF); } + + /// Builds a fat method body whose header-size nibble is `header_dwords`. + /// + /// `size_code` bytes of IL follow the header. The IL is a run of `nop` so a mis-located + /// start is visible as a wrong `size_code`/`size_header` rather than as a decode error. + fn build_fat_body_with_header_dwords(header_dwords: u16, size_code: u32) -> Vec { + let mut data = Vec::new(); + // flags (fat, init-locals) in the low 12 bits, header size in the top nibble + let first_duo = (header_dwords << 12) | 0x0013; + data.extend_from_slice(&first_duo.to_le_bytes()); + data.extend_from_slice(&1u16.to_le_bytes()); // max_stack + data.extend_from_slice(&size_code.to_le_bytes()); + data.extend_from_slice(&0u32.to_le_bytes()); // local_var_sig_token + + // Pad out to the declared header length, then the IL. + let header_len = (header_dwords as usize).saturating_mul(4); + data.resize(header_len, 0); + data.resize(header_len + size_code as usize, 0x00); + data + } + + /// A header shorter than 3 dwords would overlap the code it precedes, so it is refused. + /// + /// `size_code` is read from offset 4 and `localVarSigTok` from offset 8; with a header + /// under 12 bytes the body starts inside its own header and every downstream offset is + /// wrong. This is the hazard the nibble check exists for. + #[test] + fn fat_header_below_three_dwords_is_refused() { + for dwords in 0..3u16 { + let data = build_fat_body_with_header_dwords(dwords, 4); + assert!( + MethodBody::from(&data).is_err(), + "a {dwords}-dword fat header must be refused" + ); + } + } + + /// A header larger than the mandated 3 dwords is honoured, not rejected and not normalised. + /// + /// ECMA-335 II.25.4.3 fixes the value at 3, but CoreCLR and dnlib both locate the IL at + /// `4 * (flags >> 12)`, so such a body is loadable in practice and a malware-analysis tool + /// has to read it the same way. Rejecting it refused files the CLR runs; normalising to 12 + /// shifted the IL start, the `size_code` bounds check and the EH cursor by `4 * (n - 3)` + /// and decoded garbage. + #[test] + fn fat_header_above_three_dwords_is_honoured() { + let data = build_fat_body_with_header_dwords(4, 8); + + let body = MethodBody::from(&data).expect("a 4-dword fat header must load"); + + assert!(body.is_fat); + assert_eq!(body.size_header, 16, "the declared header size is honoured"); + assert_eq!(body.size_code, 8); + assert_eq!(body.max_stack, 1); + } + + /// The mandated 3-dword header is unaffected by the above. + #[test] + fn fat_header_of_three_dwords_is_unchanged() { + let data = build_fat_body_with_header_dwords(3, 8); + + let body = MethodBody::from(&data).expect("a conformant fat header must load"); + + assert_eq!(body.size_header, 12); + assert_eq!(body.size_code, 8); + } } diff --git a/dotscope/src/metadata/method/mod.rs b/dotscope/src/metadata/method/mod.rs index 742e887f..6b821609 100644 --- a/dotscope/src/metadata/method/mod.rs +++ b/dotscope/src/metadata/method/mod.rs @@ -125,7 +125,7 @@ use crate::{ CilModifier, CilTypeRc, CilTypeRef, CilTypeReference, TypeRegistry, TypeResolver, }, }, - utils::VisitedMap, + utils::{LazyList, VisitedMap}, CilObject, Error::SsaError, Result, @@ -136,7 +136,7 @@ pub type MethodMap = SkipMap; /// A vector that holds several parsed `Method`s. pub type MethodList = Arc>; /// A vector that holds `MethodRef` instances (weak references) -pub type MethodRefList = Arc>; +pub type MethodRefList = LazyList; /// A reference-counted pointer to a `Method`. pub type MethodRc = Arc; @@ -271,7 +271,7 @@ impl MethodRef { /// /// # Errors /// - /// Returns [`crate::Error::Malformed`] if the underlying method has been dropped. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] if the underlying method has been dropped. /// /// # Examples /// @@ -1162,7 +1162,7 @@ impl Method { /// Returns the fully-qualified name of the declaring type, if available. #[must_use] pub fn declaring_type_fullname(&self) -> Option { - self.declaring_type_rc().map(|t| t.fullname()) + self.declaring_type_rc().map(|t| t.fullname().to_string()) } /// Returns the fully-qualified method name in `"DeclaringType::MethodName"` format. @@ -1330,7 +1330,7 @@ impl Method { return Ok(()); }; - let local_var_sig_data = match sigs_table.get(body.local_var_sig_token & 0x00FF_FFFF) { + let local_var_sig_data = match sigs_table.get(body.local_var_sig_token & 0x00FF_FFFF)? { Some(var_sig_row) => blobs.get(var_sig_row.signature as usize)?, None => { return Err(malformed_error!( diff --git a/dotscope/src/metadata/query.rs b/dotscope/src/metadata/query.rs index 09171dd6..acfcdc43 100644 --- a/dotscope/src/metadata/query.rs +++ b/dotscope/src/metadata/query.rs @@ -162,7 +162,7 @@ impl<'a> TypeQuery<'a> { /// Filters to types whose fullname matches exactly. #[must_use] pub fn fullname(mut self, fqn: &'a str) -> Self { - self.filters.push(Box::new(move |t| t.fullname() == fqn)); + self.filters.push(Box::new(move |t| &*t.fullname() == fqn)); self } diff --git a/dotscope/src/metadata/resolver.rs b/dotscope/src/metadata/resolver.rs index f40cf5a6..f94adcb7 100644 --- a/dotscope/src/metadata/resolver.rs +++ b/dotscope/src/metadata/resolver.rs @@ -462,16 +462,22 @@ impl<'a> TokenResolver<'a> { self.assembly.types().resolve(&type_token) } 0x06 => { - for type_info in self.assembly.types().all_types() { - for (_, method_ref) in type_info.methods.iter() { - if let Some(method) = method_ref.upgrade() { - if method.token == method_token { - return self.assembly.types().get(&type_info.token); - } - } + // O(1) via the method's own back-pointer where it is set. This lookup is on the + // emulator's hottest paths — every `call`/`callvirt`/`newobj` — so the former + // scan over every type and every method made dispatch O(total methods). + if let Ok(method) = self.assembly.method(&method_token) { + if let Some(declaring) = method.declaring_type_rc() { + return Some(declaring); } } - None + + // Fallback: the `OnceLock` back-pointer is legitimately unset for synthetic + // methods (Reflection.Emit), so the registry's memoised index answers those. + let type_token = self + .assembly + .types() + .declaring_type_token_of_method(method_token)?; + self.assembly.types().get(&type_token) } 0x2B => { let method_spec = self.assembly.method_spec(&method_token).ok()?; @@ -484,10 +490,9 @@ impl<'a> TokenResolver<'a> { /// Finds the declaring type of a field token. /// - /// Scans all types in the registry to find which type's field list contains - /// the given field token. Handles both FieldDef (0x04) and MemberRef (0x0A) - /// tokens — for MemberRef, the `declaredby` field is used for O(1) lookup; - /// for FieldDef, an O(n) scan over all types is performed. + /// Handles both FieldDef (0x04) and MemberRef (0x0A) tokens. MemberRef resolves through + /// the `declaredby` field; FieldDef resolves through the type registry's memoised + /// `field_to_type` index. /// /// # Arguments /// @@ -500,8 +505,9 @@ impl<'a> TokenResolver<'a> { /// /// # Performance /// - /// For MemberRef tokens, resolution is O(1) via the `declaredby` field. - /// For FieldDef tokens, this performs an O(n) scan over all types. + /// O(1) for MemberRef via `declaredby`. For FieldDef, the first lookup of a given token + /// populates the registry index and every later lookup is O(1); this is on the + /// `ldsfld`/`stsfld` path, so the repeat case is the one that matters. /// /// # Examples /// @@ -527,15 +533,13 @@ impl<'a> TokenResolver<'a> { self.resolve_declaring_type(&member.declaredby) } 0x04 => { - // FieldDef: scan all types - for type_info in self.assembly.types().all_types() { - for (_, field_rc) in type_info.fields.iter() { - if field_rc.token == field_token { - return self.assembly.types().get(&type_info.token); - } - } - } - None + // FieldDef: memoised index rather than a scan over every type's field list. + // Reached on every `ldsfld`/`stsfld`. + let type_token = self + .assembly + .types() + .declaring_type_token_of_field(field_token)?; + self.assembly.types().get(&type_token) } _ => None, } diff --git a/dotscope/src/metadata/root.rs b/dotscope/src/metadata/root.rs index 9568371f..28de531c 100644 --- a/dotscope/src/metadata/root.rs +++ b/dotscope/src/metadata/root.rs @@ -342,8 +342,8 @@ impl Root { /// /// # Errors /// - /// - [`crate::Error::OutOfBounds`]: If the data slice is too short for the required fields - /// - [`crate::Error::Malformed`]: If the magic signature is invalid, version string is malformed, + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]: If the data slice is too short for the required fields + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: If the magic signature is invalid, version string is malformed, /// or stream directory is inconsistent /// - [`crate::Error::TypeError`]: If integer overflow occurs during parsing /// diff --git a/dotscope/src/metadata/security/builders.rs b/dotscope/src/metadata/security/builders.rs index c26ca23d..e7eb048f 100644 --- a/dotscope/src/metadata/security/builders.rs +++ b/dotscope/src/metadata/security/builders.rs @@ -284,8 +284,8 @@ impl PermissionSetBuilder { /// # Errors /// /// Returns [`crate::Error`] in the following cases: - /// - [`crate::Error::Malformed`] - When permission data contains unsupported types - /// - [`crate::Error::Malformed`] - When the target format is [`crate::metadata::security::PermissionSetFormat::Unknown`] + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] - When permission data contains unsupported types + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] - When the target format is [`crate::metadata::security::PermissionSetFormat::Unknown`] /// /// # Examples /// diff --git a/dotscope/src/metadata/signatures/mod.rs b/dotscope/src/metadata/signatures/mod.rs index 30b9c2d5..eee7774d 100644 --- a/dotscope/src/metadata/signatures/mod.rs +++ b/dotscope/src/metadata/signatures/mod.rs @@ -832,7 +832,7 @@ mod tests { ]) .unwrap(); assert!(matches!(result.base, TypeSignature::SzArray(_))); - if let TypeSignature::SzArray(inner) = result.base { + if let TypeSignature::SzArray(inner) = &result.base { assert_eq!(*inner.base, TypeSignature::String); } } @@ -908,8 +908,8 @@ mod tests { ]) .unwrap(); assert!(matches!(result.base, TypeSignature::GenericInst(_, _))); - if let TypeSignature::GenericInst(class, args) = result.base { - assert!(matches!(*class, TypeSignature::Class(_))); + if let TypeSignature::GenericInst(class, args) = &result.base { + assert!(matches!(**class, TypeSignature::Class(_))); assert_eq!(args.len(), 1); assert_eq!(args[0], TypeSignature::I4); } diff --git a/dotscope/src/metadata/signatures/parser.rs b/dotscope/src/metadata/signatures/parser.rs index 7349d916..877e6adf 100644 --- a/dotscope/src/metadata/signatures/parser.rs +++ b/dotscope/src/metadata/signatures/parser.rs @@ -89,7 +89,8 @@ //! //! ## Nesting Depth Protection //! The parser includes protection against stack overflow from malformed signatures: -//! - Maximum nesting depth of 10,000 levels using iterative parsing +//! - Bounded nesting: 1000 levels per parse plus a 10,000-node budget across the whole +//! signature, using iterative parsing //! - Explicit stack-based processing prevents call stack exhaustion //! - Early termination on depth limit exceeded //! - Clear error reporting for nesting depth limits @@ -187,10 +188,35 @@ pub mod CALLING_CONVENTION { /// through circular type references or deeply nested generic types. The iterative parser /// uses an explicit stack which is tracked against this limit. /// -/// The limit of 10,000 levels accommodates even the most complex real-world .NET assemblies -/// with deep generic hierarchies while still preventing resource exhaustion. +/// 1000 levels accommodates even the most complex real-world .NET assemblies with deep generic +/// hierarchies while still preventing resource exhaustion. +/// +/// This bound is **per invocation** of `parse_type_inner`, which is not sufficient on its own — +/// see [`MAX_SIGNATURE_NODES`]. const MAX_NESTING_DEPTH: usize = 1000; +/// Maximum number of type nodes a single [`SignatureParser`] will construct, across every +/// nested invocation. +/// +/// [`MAX_NESTING_DEPTH`] bounds one invocation's heap work stack, and +/// [`MAX_NATIVE_RECURSION_DEPTH`] bounds native re-entry — but the two **multiply** rather than +/// compose. `work_stack` and `result_stack` are allocated fresh per `parse_type_inner` call +/// while `self.depth` advances by only 1 per native re-entry, so at native depth *d* an +/// invocation may still build `MAX_NESTING_DEPTH - d` levels. Summing over the permitted native +/// depths yields roughly 61 000 nested `Box` levels in the *returned* value from +/// a ~60 KB blob of `PTR` runs separated by `FNPTR` headers. +/// +/// Construction survives that; the consumers do not. Drop glue, `Display` and the reference +/// scanners all recurse over the finished structure, and signature blobs are parsed on rayon +/// workers with the default 2 MiB stack. The result is a SIGSEGV while merely opening an +/// assembly — not a catchable panic, since `deny(panic)` does not apply to drop glue. +/// +/// This counter is never reset for the lifetime of a parser, so it bounds the total structure +/// regardless of how native recursion and heap work stacks interleave. It is deliberately +/// larger than [`MAX_NESTING_DEPTH`] because it counts *breadth* as well as depth: a method +/// signature with many shallow parameters is legitimate and must not trip it. +const MAX_SIGNATURE_NODES: usize = 10_000; + /// Maximum depth of **native** (call-stack) recursion in this parser. /// /// Deliberately far tighter than [`MAX_NESTING_DEPTH`], because the two bound @@ -226,8 +252,16 @@ const MAX_GENERIC_ARGS: u32 = 256; /// Maximum number of local variables in a method. /// -/// While some generated code may have many locals, 65536 is a reasonable upper bound -/// that prevents allocation attacks while supporting legitimate complex methods. +/// This is an early reject on the *declared* count, before any allocation is made against it, +/// so a bogus length field costs nothing. It matches the ECMA-335 ceiling. +/// +/// It is **not** the effective limit on a local-variable signature. +/// [`MAX_SIGNATURE_NODES`] bounds the total type nodes a parser constructs and is an order of +/// magnitude smaller, so it fires first for any signature declaring more than ~10 000 locals — +/// each local contributes at least one node. That is deliberate: the node budget is the bound +/// that actually protects against the recursive-drop stack overflow, and it counts breadth as +/// well as depth precisely so a wide signature cannot evade it. A signature declaring 65 536 +/// locals is rejected, just by the other limit and with a different error. const MAX_LOCAL_VARIABLES: u32 = 65536; /// Binary signature parser for all .NET metadata signature types according to ECMA-335. @@ -375,6 +409,13 @@ pub struct SignatureParser<'a> { /// Both are attacker-reachable: signature blobs come straight from the /// metadata heap of the .NET file being analyzed. depth: usize, + + /// Total type nodes constructed by this parser, across every nested invocation. + /// + /// Never reset — see [`MAX_SIGNATURE_NODES`] for why a per-invocation bound cannot catch + /// the case this exists for. Parsers are single-use by contract, so this is a per-signature + /// total. + nodes: usize, } impl<'a> SignatureParser<'a> { @@ -411,7 +452,22 @@ impl<'a> SignatureParser<'a> { SignatureParser { parser: Parser::new(data), depth: 0, + nodes: 0, + } + } + + /// Charges one constructed type node against [`MAX_SIGNATURE_NODES`]. + /// + /// # Errors + /// + /// Returns [`crate::Error::DepthLimitExceeded`] once the parser has built more nodes than + /// the budget allows. + fn charge_node(&mut self) -> Result<()> { + self.nodes = self.nodes.saturating_add(1); + if self.nodes > MAX_SIGNATURE_NODES { + return Err(DepthLimitExceeded(MAX_SIGNATURE_NODES)); } + Ok(()) } /// Parse a single type signature from the current position in the signature blob. @@ -452,8 +508,8 @@ impl<'a> SignatureParser<'a> { /// /// # Errors /// - [`crate::error::Error::DepthLimitExceeded`]: Maximum nesting depth exceeded - /// - [`crate::Error::Malformed`]: Invalid element type or malformed signature data - /// - [`crate::error::Error::OutOfBounds`]: Truncated signature data + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid element type or malformed signature data + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]: Truncated signature data /// /// # Implementation Notes /// @@ -515,6 +571,11 @@ impl<'a> SignatureParser<'a> { return Err(DepthLimitExceeded(MAX_NESTING_DEPTH)); } + // Global budget across every nested invocation. The check above is re-armed on each + // native re-entry because the work stacks are per-invocation, so it alone permits a + // structure tens of thousands of levels deep; this one does not reset. + self.charge_node()?; + match work { WorkItem::Seek(pos) => { self.parser.seek(pos)?; @@ -583,7 +644,21 @@ impl<'a> SignatureParser<'a> { self.parse_type_simple()?; // Skip element type // Read array metadata + // + // `rank` is bounded here and not merely stored: it is the only + // ceiling on `num_lo_bounds` below, so an unbounded rank makes that + // check permissive rather than protective, and the lower-bound + // extension loop then grows `dimensions` to the declared count + // before any read can run out of input. let rank = self.parser.read_compressed_uint()?; + if rank > MAX_ARRAY_DIMENSIONS { + return Err(malformed_error!( + "Array signature has too many dimensions: rank {} (max: {})", + rank, + MAX_ARRAY_DIMENSIONS + )); + } + let num_sizes = self.parser.read_compressed_uint()?; if num_sizes > MAX_ARRAY_DIMENSIONS { return Err(malformed_error!( @@ -947,8 +1022,8 @@ impl<'a> SignatureParser<'a> { /// The vector is empty if no custom modifiers are present. /// /// # Errors - /// - [`crate::Error::Malformed`]: Invalid compressed token encoding - /// - [`crate::error::Error::OutOfBounds`]: Truncated modifier data + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid compressed token encoding + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]: Truncated modifier data /// /// # Performance Notes /// - Modifiers are relatively uncommon in most .NET code @@ -1030,9 +1105,9 @@ impl<'a> SignatureParser<'a> { /// - Complete type signature information /// /// # Errors - /// - [`crate::Error::Malformed`]: Invalid parameter encoding + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid parameter encoding /// - [`crate::Error::DepthLimitExceeded`]: Parameter type parsing exceeds nesting depth limit - /// - [`crate::error::Error::OutOfBounds`]: Truncated parameter data + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]: Truncated parameter data /// /// # Usage Notes /// @@ -1158,9 +1233,9 @@ impl<'a> SignatureParser<'a> { /// - Variable argument list (if applicable) /// /// # Errors - /// - [`crate::Error::Malformed`]: Invalid calling convention or parameter encoding + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid calling convention or parameter encoding /// - [`crate::Error::DepthLimitExceeded`]: Parameter type parsing exceeds nesting depth limit - /// - [`crate::error::Error::OutOfBounds`]: Truncated signature data + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]: Truncated signature data /// /// # Performance Notes /// - Parameter vectors are pre-allocated based on parameter count @@ -1382,9 +1457,9 @@ impl<'a> SignatureParser<'a> { /// - Type constraints and annotations /// /// # Errors - /// - [`crate::Error::Malformed`]: Invalid field signature header (not 0x06) + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid field signature header (not 0x06) /// - [`crate::Error::DepthLimitExceeded`]: Field type parsing exceeds nesting depth limit - /// - [`crate::error::Error::OutOfBounds`]: Truncated field signature data + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]: Truncated field signature data /// /// # Custom Modifier Applications /// @@ -1530,9 +1605,9 @@ impl<'a> SignatureParser<'a> { /// - Complete type and modifier information /// /// # Errors - /// - [`crate::Error::Malformed`]: Invalid property signature header (missing PROPERTY bit) + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid property signature header (missing PROPERTY bit) /// - [`crate::Error::DepthLimitExceeded`]: Property or parameter type parsing exceeds nesting depth limit - /// - [`crate::error::Error::OutOfBounds`]: Truncated property signature data + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]: Truncated property signature data /// /// # Indexer Design Patterns /// @@ -1723,9 +1798,9 @@ impl<'a> SignatureParser<'a> { /// - Custom modifier information /// /// # Errors - /// - [`crate::Error::Malformed`]: Invalid local variable signature header (not 0x07) + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid local variable signature header (not 0x07) /// - [`crate::Error::DepthLimitExceeded`]: Local variable type parsing exceeds nesting depth limit - /// - [`crate::error::Error::OutOfBounds`]: Truncated local variable signature data + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]: Truncated local variable signature data /// /// # Memory Management Implications /// @@ -1965,8 +2040,8 @@ impl<'a> SignatureParser<'a> { /// /// # Errors /// - [`crate::Error::DepthLimitExceeded`]: Type parsing exceeds maximum nesting depth - /// - [`crate::Error::Malformed`]: Invalid type encoding or format - /// - [`crate::error::Error::OutOfBounds`]: Truncated type specification data + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid type encoding or format + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]: Truncated type specification data /// /// # Performance Notes /// - Type specifications often involve complex recursive parsing @@ -1989,13 +2064,17 @@ impl<'a> SignatureParser<'a> { // Then parse the base type signature let base_type_sig = self.parse_type()?; - // If we got a ModifiedRequired/Optional from parse_type, we need to handle it specially - match base_type_sig { + // If we got a ModifiedRequired/Optional from parse_type, we need to handle it specially. + // + // Matched by reference and drained rather than moved out: `TypeSignature` implements + // `Drop` (for iterative teardown), which makes moving a field out of it illegal. + let mut base_type_sig = base_type_sig; + match &mut base_type_sig { TypeSignature::ModifiedRequired(mod_modifiers) | TypeSignature::ModifiedOptional(mod_modifiers) => { // Combine the modifiers from parse_custom_mods() and from the ModifiedRequired let mut all_modifiers = modifiers; - all_modifiers.extend(mod_modifiers); + all_modifiers.append(mod_modifiers); // Parse the base type that follows the modifiers let base_type = self.parse_type()?; @@ -2145,10 +2224,10 @@ impl<'a> SignatureParser<'a> { /// - Ready for runtime method instantiation /// /// # Errors - /// - [`crate::Error::Malformed`]: Invalid method specification header (not 0x0A) + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Invalid method specification header (not 0x0A) /// - [`crate::Error::DepthLimitExceeded`]: Type argument parsing exceeds nesting depth limit - /// - [`crate::error::Error::OutOfBounds`]: Truncated method specification data - /// - [`crate::error::Error::Malformed`]: Mismatched type argument count + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]: Truncated method specification data + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`]: Mismatched type argument count /// /// # Performance Notes /// - Type argument parsing cost is linear in the number of arguments @@ -2256,7 +2335,7 @@ mod tests { let result = parser.parse_type().unwrap(); assert!(matches!(result, TypeSignature::SzArray(_))); - if let TypeSignature::SzArray(inner) = result { + if let TypeSignature::SzArray(inner) = &result { assert_eq!(*inner.base, TypeSignature::I4); } @@ -2271,12 +2350,28 @@ mod tests { let result = parser.parse_type().unwrap(); assert!(matches!(result, TypeSignature::Array(_))); - if let TypeSignature::Array(array) = result { + if let TypeSignature::Array(array) = &result { assert_eq!(*array.base, TypeSignature::I4); assert_eq!(array.rank, 2); assert_eq!(array.dimensions.len(), 0) } + // An overlarge rank is rejected on the declared value, before anything is sized from + // it. `rank` is the only ceiling on `num_lo_bounds`, and the lower-bound extension + // loop grows `dimensions` to that count before any read can run out of input -- so an + // unbounded rank is an allocation bomb, not just an implausible array. + let mut parser = SignatureParser::new(&[ + 0x14, // ARRAY + 0x08, // I4 (element type) + 0xC0, 0x40, 0x00, 0x00, // rank 0x400000, a 4-byte compressed uint + 0x00, // num_sizes 0 + 0x00, // num_lo_bounds 0 + ]); + assert!( + parser.parse_type().is_err(), + "an array rank above MAX_ARRAY_DIMENSIONS must be rejected" + ); + // Multi-dimensional array int[2,3] with rank 2, with sizes let mut parser = SignatureParser::new(&[ 0x14, // ARRAY @@ -2290,7 +2385,7 @@ mod tests { let result = parser.parse_type().unwrap(); assert!(matches!(result, TypeSignature::Array(_))); - if let TypeSignature::Array(array) = result { + if let TypeSignature::Array(array) = &result { assert_eq!(*array.base, TypeSignature::I4); assert_eq!(array.rank, 2); assert_eq!(array.dimensions.len(), 2); @@ -2301,6 +2396,51 @@ mod tests { } } + #[test] + fn deeply_nested_signature_is_refused_not_built() { + // A long run of ELEMENT_TYPE_PTR (0x0F) terminated by I4. Each byte is one nesting + // level and consumes no other input, which is what makes this cheap to weaponise: a + // ~60 KB blob describes a ~61 000-level structure, deep enough that recursive drop + // glue overflows a 2 MiB rayon worker stack. + let mut blob = vec![0x0F_u8; MAX_SIGNATURE_NODES.saturating_mul(2)]; + blob.push(0x08); // I4 + + let mut parser = SignatureParser::new(&blob); + assert!( + parser.parse_type().is_err(), + "a signature past the node budget must be refused during parsing" + ); + } + + /// Signatures of ordinary depth must still parse — the budget must not be so tight that it + /// rejects real assemblies. + #[test] + fn moderately_nested_signature_still_parses() { + // int**...* nested 100 levels: far beyond anything a real compiler emits, well inside + // the budget. + let mut blob = vec![0x0F_u8; 100]; + blob.push(0x08); // I4 + + let mut parser = SignatureParser::new(&blob); + assert!(parser.parse_type().is_ok()); + } + + /// Dropping a deep signature must not recurse. + /// + /// Built by hand rather than parsed, because the parser now refuses to construct one this + /// deep; the `Drop` impl has to hold for any `TypeSignature` however it was produced. + #[test] + fn deep_signature_drops_without_recursing() { + let mut sig = TypeSignature::I4; + for _ in 0..200_000 { + sig = TypeSignature::Pinned(Box::new(sig)); + } + + // The assertion is simply that this returns: with drop glue recursing, 200 000 levels + // overflows the stack and takes the process down with SIGSEGV. + drop(sig); + } + #[test] fn test_parse_pointers_and_byrefs() { // Pointer to Int32 (int*) @@ -2308,7 +2448,7 @@ mod tests { let result = parser.parse_type().unwrap(); assert!(matches!(result, TypeSignature::Ptr(_))); - if let TypeSignature::Ptr(inner) = result { + if let TypeSignature::Ptr(inner) = &result { assert_eq!(*inner.base, TypeSignature::I4); } @@ -2317,8 +2457,8 @@ mod tests { let result = parser.parse_type().unwrap(); assert!(matches!(result, TypeSignature::ByRef(_))); - if let TypeSignature::ByRef(inner) = result { - assert_eq!(*inner, TypeSignature::I4); + if let TypeSignature::ByRef(inner) = &result { + assert_eq!(**inner, TypeSignature::I4); } } @@ -2336,8 +2476,8 @@ mod tests { let result = parser.parse_type().unwrap(); assert!(matches!(result, TypeSignature::GenericInst(_, _))); - if let TypeSignature::GenericInst(class, args) = result { - assert!(matches!(*class, TypeSignature::Class(_))); + if let TypeSignature::GenericInst(class, args) = &result { + assert!(matches!(**class, TypeSignature::Class(_))); assert_eq!(args.len(), 1); assert_eq!(args[0], TypeSignature::I4); } @@ -2355,8 +2495,8 @@ mod tests { let result = parser.parse_type().unwrap(); assert!(matches!(result, TypeSignature::GenericInst(_, _))); - if let TypeSignature::GenericInst(class, args) = result { - assert!(matches!(*class, TypeSignature::Class(_))); + if let TypeSignature::GenericInst(class, args) = &result { + assert!(matches!(**class, TypeSignature::Class(_))); assert_eq!(args.len(), 2); assert_eq!(args[0], TypeSignature::String); assert_eq!(args[1], TypeSignature::I4); diff --git a/dotscope/src/metadata/signatures/types.rs b/dotscope/src/metadata/signatures/types.rs index f4e689e8..0c9262d3 100644 --- a/dotscope/src/metadata/signatures/types.rs +++ b/dotscope/src/metadata/signatures/types.rs @@ -3350,6 +3350,63 @@ pub struct SignatureMethodSpec { pub generic_args: Vec, } +impl Drop for TypeSignature { + /// Tears the signature down iteratively rather than letting drop glue recurse. + /// + /// A signature blob is attacker-supplied and can describe a type nested thousands of levels + /// deep (long runs of `ELEMENT_TYPE_PTR`). Compiler-generated drop glue walks that chain + /// with one stack frame per level, and signature blobs are parsed on rayon workers with the + /// default 2 MiB stack, so a deep enough signature exhausts the stack while the value is + /// merely going out of scope. That is a SIGSEGV, not a panic — `deny(panic)` and + /// `Result`-returning parsers offer no protection against drop glue. + /// + /// This moves each node's children onto a heap worklist before the node is dropped, so + /// every individual drop is shallow and stack usage stays constant regardless of depth. + /// It is defence in depth alongside the parser's node budget: the budget stops such a + /// signature being built here, but this holds for any `TypeSignature` however constructed. + fn drop(&mut self) { + let mut worklist: Vec = Vec::new(); + take_children(self, &mut worklist); + + while let Some(mut node) = worklist.pop() { + take_children(&mut node, &mut worklist); + // `node` is dropped here with its children already moved out, so its own drop + // glue — and the recursive call to this impl — bottoms out immediately. + } + } +} + +/// Moves every nested [`TypeSignature`] out of `node` and onto `out`, leaving leaves behind. +/// +/// Used by [`TypeSignature`]'s [`Drop`] to flatten a chain that drop glue would otherwise walk +/// recursively. Replacing each child with [`TypeSignature::Void`] is what makes the subsequent +/// drop shallow. +fn take_children(node: &mut TypeSignature, out: &mut Vec) { + /// Takes a boxed child, leaving a leaf in its place. + fn take(slot: &mut TypeSignature) -> TypeSignature { + std::mem::replace(slot, TypeSignature::Void) + } + + match node { + TypeSignature::Ptr(ptr) => out.push(take(&mut ptr.base)), + TypeSignature::SzArray(arr) => out.push(take(&mut arr.base)), + TypeSignature::Array(arr) => out.push(take(&mut arr.base)), + TypeSignature::ByRef(inner) | TypeSignature::Pinned(inner) => out.push(take(inner)), + TypeSignature::GenericInst(base, args) => { + out.push(take(base)); + out.append(args); + } + TypeSignature::FnPtr(method) => { + out.push(take(&mut method.return_type.base)); + out.extend(method.params.drain(..).map(|p| p.base)); + out.extend(method.varargs.drain(..).map(|p| p.base)); + } + // Remaining variants are leaves, or hold only tokens and flags. `ModifiedRequired` and + // `ModifiedOptional` carry `CustomModifier`s, which are a `Token` plus a `bool`. + _ => {} + } +} + impl TypeSignature { /// Check if a constant primitive value is compatible with this type signature /// diff --git a/dotscope/src/metadata/streams/blob.rs b/dotscope/src/metadata/streams/blob.rs index d3051771..902cbf2b 100644 --- a/dotscope/src/metadata/streams/blob.rs +++ b/dotscope/src/metadata/streams/blob.rs @@ -400,7 +400,7 @@ impl<'a> Blob<'a> { /// /// // Access null blob (always empty) /// let null_blob = blob_heap.get(0)?; - /// assert_eq!(null_blob, &[]); + /// assert_eq!(null_blob, &[] as &[u8]); /// /// // Access first real blob /// let first_blob = blob_heap.get(1)?; diff --git a/dotscope/src/metadata/streams/streamheader.rs b/dotscope/src/metadata/streams/streamheader.rs index e682fb1a..317e73c3 100644 --- a/dotscope/src/metadata/streams/streamheader.rs +++ b/dotscope/src/metadata/streams/streamheader.rs @@ -387,7 +387,7 @@ impl StreamHeader { /// # Errors /// /// Returns [`crate::Error`] in the following cases: - /// - **[`crate::Error::OutOfBounds`]**: Data slice too short (< 9 bytes) + /// - **[`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]**: Data slice too short (< 9 bytes) /// - **Malformed data**: Stream size not 4-byte aligned (ECMA-335 violation) /// - **Range error**: Offset or size exceeds 0x7FFFFFFF (integer overflow protection) /// - **Invalid name**: Stream name doesn't match standard ECMA-335 identifiers diff --git a/dotscope/src/metadata/streams/strings.rs b/dotscope/src/metadata/streams/strings.rs index fa51b668..91ba2d71 100644 --- a/dotscope/src/metadata/streams/strings.rs +++ b/dotscope/src/metadata/streams/strings.rs @@ -574,7 +574,7 @@ impl<'a> Strings<'a> { /// # Errors /// /// Returns [`crate::Error`] in the following cases: - /// - **[`crate::Error::OutOfBounds`]**: Index exceeds heap data length + /// - **[`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]**: Index exceeds heap data length /// - **Malformed string**: No null terminator found within remaining heap data /// - **Invalid UTF-8**: String bytes do not form a valid UTF-8 sequence /// - **Encoding errors**: Non-ASCII bytes that violate UTF-8 encoding rules diff --git a/dotscope/src/metadata/streams/tablesheader.rs b/dotscope/src/metadata/streams/tablesheader.rs index 69467098..f428f9c8 100644 --- a/dotscope/src/metadata/streams/tablesheader.rs +++ b/dotscope/src/metadata/streams/tablesheader.rs @@ -105,7 +105,10 @@ //! println!("Found {} type definitions", typedef_table.row_count); //! //! // Examine first few types +//! // Iteration yields `Result`: a row that fails to parse is reported +//! // rather than silently ending the iteration. //! for (index, type_def) in typedef_table.iter().enumerate().take(5) { +//! let type_def = type_def?; //! println!("Type {}: flags={:#x}, name_idx={}, namespace_idx={}", //! index, type_def.flags, type_def.type_name, type_def.type_namespace); //! } @@ -117,6 +120,7 @@ //! //! // Find methods by characteristics //! let static_methods = method_table.iter() +//! .filter_map(Result::ok) // skip rows that fail to parse //! .filter(|method| method.flags & 0x0010 != 0) // MethodAttributes.Static //! .count(); //! println!("Static methods: {}", static_methods); @@ -138,12 +142,14 @@ //! tables.table::() //! ) { //! for (type_idx, type_def) in typedef_table.iter().enumerate().take(10) { +//! let type_def = type_def?; //! // Calculate field range for this type //! let field_start = type_def.field_list.saturating_sub(1) as usize; -//! -//! // Find field range end (next type's field_list or table end) +//! +//! // Find field range end (next type's field_list or table end). `get` returns +//! // `Result>`: outer for a malformed row, inner for a RID past the end. //! let field_end = if type_idx + 1 < typedef_table.row_count as usize { -//! typedef_table.get((type_idx + 1) as u32) +//! typedef_table.get((type_idx + 1) as u32)? //! .map(|next_type| next_type.field_list.saturating_sub(1) as usize) //! .unwrap_or(field_table.row_count as usize) //! } else { @@ -173,6 +179,7 @@ //! //! // Parallel analysis using rayon //! let attribute_stats = ca_table.par_iter() +//! .filter_map(Result::ok) // parallel iteration yields `Result` //! .map(|attr| { //! // Analyze attribute type and parent //! let parent_table = attr.parent.tag; @@ -207,7 +214,7 @@ //! // Process chunk without loading entire table into memory //! let mut external_refs = 0; //! for i in chunk_start..chunk_end { -//! if let Some(member_ref) = memberref_table.get(i) { +//! if let Ok(Some(member_ref)) = memberref_table.get(i) { //! // Analyze member reference //! if member_ref.class.tag == TableId::TypeRef { //! external_refs += 1; @@ -279,7 +286,8 @@ //! - **Lifetime enforcement**: Rust borrow checker prevents use-after-free //! - **Type safety**: Generic type parameters prevent incorrect table access //! - **Bounds verification**: All array and slice access bounds-checked -//! - **No unsafe aliasing**: Careful pointer management in type casting +//! - **No unsafe code**: Table access goes through the safe accessor layer; there is no pointer +//! casting here to manage //! //! # See Also //! - [`crate::metadata::tables`]: Individual metadata table definitions and structures @@ -385,7 +393,10 @@ use crate::{ /// println!("Assembly defines {} types", typedef_table.row_count); /// /// // Analyze type characteristics +/// // Iteration yields `Result`: a row that fails to parse is reported +/// // rather than silently ending the iteration. /// for (index, type_def) in typedef_table.iter().enumerate().take(10) { +/// let type_def = type_def?; /// let is_public = type_def.flags & 0x00000007 == 0x00000001; /// let is_sealed = type_def.flags & 0x00000100 != 0; /// let is_abstract = type_def.flags & 0x00000080 != 0; @@ -412,9 +423,11 @@ use crate::{ /// tables.table::() /// ) { /// for (type_idx, type_def) in typedef_table.iter().enumerate().take(5) { -/// // Calculate member ranges for this type -/// let next_type = typedef_table.get((type_idx + 1) as u32); -/// +/// let type_def = type_def?; +/// // Calculate member ranges for this type. `get` returns `Result>`: +/// // outer for a malformed row, inner for a RID past the end of the table. +/// let next_type = typedef_table.get((type_idx + 1) as u32)?; +/// /// let field_start = type_def.field_list.saturating_sub(1); /// let field_end = next_type.as_ref() /// .map(|t| t.field_list.saturating_sub(1)) @@ -492,7 +505,7 @@ use crate::{ /// let chunk_end = (chunk_start + CHUNK_SIZE).min(total_rows); /// /// for i in chunk_start..chunk_end { -/// if let Some(member_ref) = memberref_table.get(i) { +/// if let Ok(Some(member_ref)) = memberref_table.get(i) { /// // Analyze member reference type and parent /// let is_method = true; // Simplified: check signature /// let is_external = true; // Simplified: check class reference @@ -578,7 +591,7 @@ use crate::{ /// - **Lifetime enforcement**: Rust borrow checker prevents use-after-free vulnerabilities /// - **Type safety**: Generic parameters prevent incorrect table type access /// - **Bounds verification**: All array and slice access is bounds-checked -/// - **Controlled unsafe**: Minimal unsafe code with careful pointer management +/// - **No unsafe code**: The safe accessor layer replaced the pointer casting this once needed /// /// # ECMA-335 Compliance /// @@ -608,7 +621,7 @@ use crate::{ /// println!("TypeDef table has {} rows", typedef_table.row_count); /// /// // Access individual rows by index (0-based) -/// if let Some(first_type) = typedef_table.get(0) { +/// if let Ok(Some(first_type)) = typedef_table.get(1) { /// println!("First type: flags={}, name_idx={}, namespace_idx={}", /// first_type.flags, first_type.type_name, first_type.type_namespace); /// } @@ -626,6 +639,7 @@ use crate::{ /// // Iterate over all methods in the assembly /// if let Some(method_table) = tables_header.table::() { /// for (index, method) in method_table.iter().enumerate() { +/// let method = method?; /// println!("Method {}: RVA={:#x}, impl_flags={}, flags={}, name_idx={}", /// index, method.rva, method.impl_flags, method.flags, method.name); /// @@ -646,6 +660,7 @@ use crate::{ /// // Process field metadata in parallel /// if let Some(field_table) = tables_header.table::() { /// let field_count = field_table.par_iter() +/// .filter_map(Result::ok) // parallel iteration yields `Result` /// .filter(|field| field.flags & 0x0010 != 0) // FieldAttributes.Static /// .count(); /// @@ -666,6 +681,7 @@ use crate::{ /// tables_header.table::() /// ) { /// for (type_idx, type_def) in typedef_table.iter().enumerate().take(5) { +/// let type_def = type_def?; /// println!("Type {}: methods {}-{}", /// type_idx, type_def.method_list, /// type_def.method_list.saturating_add(10)); // Simplified example @@ -717,7 +733,7 @@ use crate::{ /// let chunk_end = (chunk_start + CHUNK_SIZE).min(total_rows); /// /// for i in chunk_start..chunk_end { -/// if let Some(attr) = ca_table.get(i) { +/// if let Ok(Some(attr)) = ca_table.get(i) { /// // Process individual custom attribute /// // attr.parent, attr.type_def, attr.value are available /// // without copying the entire table into memory @@ -849,7 +865,7 @@ impl<'a> TablesHeader<'a> { /// # Errors /// /// Returns [`crate::Error`] in the following cases: - /// - **[`crate::Error::OutOfBounds`]**: Data too short for complete header (< 24 bytes) + /// - **[`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`]**: Data too short for complete header (< 24 bytes) /// - **Malformed data**: No valid tables present (all bits in `valid` are 0) /// - **Version error**: Unsupported major/minor version combination /// - **Format error**: Invalid table data or corrupted stream structure @@ -1083,6 +1099,16 @@ impl<'a> TablesHeader<'a> { )?; } + // Checked again after the loop, not only at the top of each iteration: the check above + // catches an overshoot on the *next* variant, which leaves the last one — presently + // `CustomDebugInformation` — with nothing behind it to notice. + if current_offset > data.len() { + return Err(ParseFailure::OutOfBounds { + stage: ParseStage::TildeStream, + } + .into()); + } + Ok(tables_header) } @@ -1136,12 +1162,13 @@ impl<'a> TablesHeader<'a> { /// if let Some(typedef_table) = tables.table::() { /// // Efficient access to all type definitions /// for type_def in typedef_table.iter().take(5) { + /// let type_def = type_def?; /// println!("Type: flags={:#x}, name_idx={}, namespace_idx={}", /// type_def.flags, type_def.type_name, type_def.type_namespace); /// } - /// - /// // Random access to specific rows - /// if let Some(first_type) = typedef_table.get(0) { + /// + /// // Random access to specific rows, by RID (1-based) + /// if let Ok(Some(first_type)) = typedef_table.get(1) { /// println!("First type name index: {}", first_type.type_name); /// } /// } diff --git a/dotscope/src/metadata/streams/userstrings.rs b/dotscope/src/metadata/streams/userstrings.rs index 30ac050c..92823408 100644 --- a/dotscope/src/metadata/streams/userstrings.rs +++ b/dotscope/src/metadata/streams/userstrings.rs @@ -40,7 +40,7 @@ //! # Reference //! - [ECMA-335 II.24.2.4](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) -use widestring::U16Str; +use widestring::U16String; use crate::{ utils::{read_compressed_int, read_compressed_int_at}, @@ -111,7 +111,7 @@ impl<'a> UserStrings<'a> { /// * `Ok(UserStrings)` - Valid heap accessor /// /// # Errors - /// * [`crate::Error::OutOfBounds`] - If data is empty or doesn't start with null byte + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] - If data is empty or doesn't start with null byte /// /// # Examples /// @@ -145,10 +145,10 @@ impl<'a> UserStrings<'a> { /// * `index` - The byte offset within the heap (typically from metadata table references) /// /// # Returns - /// * `Ok(&U16Str)` - Reference to the UTF-16 string at the specified offset + /// * `Ok(U16String)` - The UTF-16 string at the specified offset /// /// # Errors - /// * [`crate::Error::OutOfBounds`] - If index is out of bounds + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] - If index is out of bounds /// * [`crate::Error`] - If string data is malformed or has invalid UTF-16 length /// /// # Examples @@ -164,15 +164,12 @@ impl<'a> UserStrings<'a> { /// ``` /// /// # Platform Notes - /// This method performs unaligned memory access which is well-supported on all modern - /// platforms (x86/x64, aarch64/ARMv8, most ARMv7+). On platforms without hardware - /// unaligned access support, there may be a performance penalty but no correctness issues. - /// - /// # Panics - /// This function will not panic under normal circumstances - the internal `unwrap()` is - /// used on a raw pointer conversion that is guaranteed to succeed when the input slice - /// is valid. - pub fn get(&self, index: usize) -> Result<&'a U16Str> { + /// ECMA-335 stores `#US` entries little-endian and guarantees no alignment for them, so the + /// code points are decoded explicitly from byte pairs. Returning an owned string rather than + /// a borrowed `&U16Str` is what makes that possible: a `&[u16]` borrowed from the heap would + /// have to be reinterpreted from a possibly-odd byte offset, which is undefined behaviour + /// regardless of what the hardware tolerates, and would read native-endian. + pub fn get(&self, index: usize) -> Result { let oob = || ParseFailure::HeapOutOfBounds { heap: HeapKind::UserStrings, index: u32::try_from(index).unwrap_or(u32::MAX), @@ -195,8 +192,7 @@ impl<'a> UserStrings<'a> { } if total_bytes == 1 { - static EMPTY_U16: [u16; 0] = []; - return Ok(U16Str::from_slice(&EMPTY_U16)); + return Ok(U16String::new()); } // Total bytes includes UTF-16 data + terminator byte (1 byte) @@ -224,35 +220,26 @@ impl<'a> UserStrings<'a> { .get(data_start..utf16_data_end) .ok_or_else(|| Error::from(oob()))?; - // Convert byte slice to u16 slice for UTF-16 string construction. + // Decode the code points from byte pairs rather than reinterpreting the bytes as + // `&[u16]`. // - // SAFETY: - // - `utf16_data.len()` is guaranteed to be even (checked above with `utf16_length % 2 != 0`) - // - The resulting slice length is exactly `utf16_data.len() / 2` u16 elements - // - The pointer comes from a valid `&[u8]` slice that outlives this function + // A cast would be wrong twice over. `data_start` is attacker-influenced — the compressed + // length prefix is 1, 2 or 4 bytes depending on its first byte, so an odd offset is + // reachable without touching the stream offset at all — and building a *reference* from a + // pointer that is not aligned to its referent is undefined behaviour in Rust however the + // hardware behaves. It also reads native-endian, while ECMA-335 §II.24.2.4 specifies + // little-endian, so the same input would decode differently on a big-endian target. // - // Alignment considerations: - // - x86/x64: Hardware supports unaligned access natively - // - aarch64 (ARMv8): Hardware supports unaligned access (may be slightly slower) - // - arm (ARMv7+): Most implementations support unaligned access via hardware or kernel trap - // - Older ARM/MIPS/PowerPC without unaligned support: May fault - // - // The ECMA-335 specification does not guarantee alignment of #US heap entries. - // However, unaligned access works on all modern platforms where .NET runs. - // We accept the potential performance penalty on ARM rather than failing entirely. - // - // Note: We suppress the clippy warning because unaligned u16 access is safe (not UB) - // on all platforms we target - it may be slow but won't cause memory corruption. - let str_slice = unsafe { - let ptr = utf16_data.as_ptr(); - - #[allow(clippy::cast_ptr_alignment)] - core::ptr::slice_from_raw_parts(ptr.cast::(), utf16_data.len() / 2) - .as_ref() - .ok_or_else(|| corrupt("null pointer in user string slice conversion".into()))? - }; - - Ok(U16Str::from_slice(str_slice)) + // `utf16_data.len()` is even (rejected above otherwise), so `chunks_exact(2)` consumes it + // fully with no remainder. + let code_units: Vec = utf16_data + .chunks_exact(2) + // `chunks_exact(2)` yields only two-byte chunks, so the conversion cannot fail; the + // fallback exists so this stays free of indexing that could panic. + .map(|pair| <[u8; 2]>::try_from(pair).map_or(0, u16::from_le_bytes)) + .collect(); + + Ok(U16String::from_vec(code_units)) } /// Returns an iterator over all user strings in the heap @@ -401,7 +388,7 @@ impl<'a> UserStrings<'a> { } impl<'a> IntoIterator for &'a UserStrings<'a> { - type Item = (usize, &'a U16Str); + type Item = (usize, U16String); type IntoIter = UserStringsIterator<'a>; /// Create an iterator over the user strings heap. @@ -429,7 +416,7 @@ impl<'a> IntoIterator for &'a UserStrings<'a> { /// Iterator over entries in the `#US` (`UserStrings`) heap /// /// Provides zero-copy access to UTF-16 user strings with their byte offsets. -/// Each iteration returns a `(usize, &U16Str)` containing the offset and string content. +/// Each iteration returns a `(usize, U16String)` containing the offset and string content. /// The iterator automatically handles length prefixes and string format validation. /// /// # Iteration Behavior @@ -458,7 +445,7 @@ impl<'a> UserStringsIterator<'a> { } impl<'a> Iterator for UserStringsIterator<'a> { - type Item = (usize, &'a U16Str); + type Item = (usize, U16String); /// Get the next user string from the heap /// @@ -544,6 +531,58 @@ mod tests { assert_eq!(us_str.get(1).unwrap(), u16str!("Hello, World!")); } + /// A user string whose payload starts at an *odd* byte offset. + /// + /// This is the case that makes a zero-copy `&[u16]` unsound: building one from + /// `utf16_data.as_ptr()` creates a reference from a pointer that is not aligned to its + /// referent, which is UB in Rust no matter how tolerant the hardware is. Nothing upstream + /// constrains the parity — a two-byte compressed length prefix at index 1 puts the payload + /// at index 3 without touching the stream offset at all — so the decoder copies. + /// + /// A 0x81-byte entry forces that two-byte prefix: values >= 0x80 cannot use the one-byte form. + #[test] + fn payload_at_odd_offset_decodes() { + const TOTAL_BYTES: usize = 0x81; // 128 bytes of UTF-16 + 1 terminator byte + const CODE_UNITS: usize = 64; + + let mut data = vec![0x00]; + // Compressed-int encoding of 0x81 in the two-byte form. + data.push(0x80); + data.push(0x81); + assert_eq!(data.len(), 3, "payload must begin at an odd offset"); + + // 64 x 'A' (U+0041), little-endian as ECMA-335 requires. + for _ in 0..CODE_UNITS { + data.push(0x41); + data.push(0x00); + } + data.push(0x00); // terminator byte + + assert_eq!(data.len(), 3 + TOTAL_BYTES); + + let heap = UserStrings::from(&data).unwrap(); + let s = heap.get(1).unwrap(); + + assert_eq!(s.len(), CODE_UNITS); + assert_eq!(s.to_string_lossy(), "A".repeat(CODE_UNITS)); + } + + /// Code units must be read little-endian regardless of host byte order (ECMA-335 II.24.2.4). + /// + /// U+3042 is asymmetric, so a native-endian read on a big-endian host would yield U+4230 + /// ('B0' as two bytes) instead. Reading it back correctly pins the byte order down. + #[test] + fn code_units_are_little_endian() { + // 0x05 = 4 payload bytes + terminator; U+3042 U+3044 encoded little-endian. + let data: [u8; 8] = [0x00, 0x05, 0x42, 0x30, 0x44, 0x30, 0x00, 0x00]; + + let heap = UserStrings::from(&data).unwrap(); + let s = heap.get(1).unwrap(); + + assert_eq!(s.as_slice(), &[0x3042, 0x3044]); + assert_eq!(s.to_string_lossy(), "\u{3042}\u{3044}"); + } + #[test] fn invalid() { let data_empty = []; diff --git a/dotscope/src/metadata/tables/assembly/loader.rs b/dotscope/src/metadata/tables/assembly/loader.rs index 5c409c59..c0ede501 100644 --- a/dotscope/src/metadata/tables/assembly/loader.rs +++ b/dotscope/src/metadata/tables/assembly/loader.rs @@ -55,7 +55,7 @@ impl MetadataLoader for AssemblyLoader { let Some(table) = header.table::() else { return Ok(()); }; - let Some(row) = table.get(1) else { + let Some(row) = table.get(1)? else { return Ok(()); }; diff --git a/dotscope/src/metadata/tables/assembly/raw.rs b/dotscope/src/metadata/tables/assembly/raw.rs index de569779..1b2e8a63 100644 --- a/dotscope/src/metadata/tables/assembly/raw.rs +++ b/dotscope/src/metadata/tables/assembly/raw.rs @@ -28,6 +28,7 @@ use crate::{ tables::{Assembly, AssemblyFlags, AssemblyRc, HashAlgorithmId, TableInfoRef, TableRow}, token::Token, }, + utils::LazyList, Result, }; @@ -160,7 +161,7 @@ impl AssemblyRaw { Some(strings.get(self.culture as usize)?.to_string()) }, security: OnceLock::new(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } diff --git a/dotscope/src/metadata/tables/assembly/reader.rs b/dotscope/src/metadata/tables/assembly/reader.rs index 12f588e9..f8c5a98c 100644 --- a/dotscope/src/metadata/tables/assembly/reader.rs +++ b/dotscope/src/metadata/tables/assembly/reader.rs @@ -133,12 +133,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -181,12 +182,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/assemblyos/loader.rs b/dotscope/src/metadata/tables/assemblyos/loader.rs index d9a4c0f4..6a258d8a 100644 --- a/dotscope/src/metadata/tables/assemblyos/loader.rs +++ b/dotscope/src/metadata/tables/assemblyos/loader.rs @@ -56,7 +56,7 @@ impl MetadataLoader for AssemblyOsLoader { let Some(table) = header.table::() else { return Ok(()); }; - let Some(row) = table.get(1) else { + let Some(row) = table.get(1)? else { return Ok(()); }; diff --git a/dotscope/src/metadata/tables/assemblyos/reader.rs b/dotscope/src/metadata/tables/assemblyos/reader.rs index 11f2d736..b5aedb6c 100644 --- a/dotscope/src/metadata/tables/assemblyos/reader.rs +++ b/dotscope/src/metadata/tables/assemblyos/reader.rs @@ -111,12 +111,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/assemblyprocessor/loader.rs b/dotscope/src/metadata/tables/assemblyprocessor/loader.rs index dcbf3cdc..2047a8e6 100644 --- a/dotscope/src/metadata/tables/assemblyprocessor/loader.rs +++ b/dotscope/src/metadata/tables/assemblyprocessor/loader.rs @@ -86,7 +86,7 @@ impl MetadataLoader for AssemblyProcessorLoader { let Some(table) = header.table::() else { return Ok(()); }; - let Some(row) = table.get(1) else { + let Some(row) = table.get(1)? else { return Ok(()); }; diff --git a/dotscope/src/metadata/tables/assemblyprocessor/reader.rs b/dotscope/src/metadata/tables/assemblyprocessor/reader.rs index 0caa314e..07e12069 100644 --- a/dotscope/src/metadata/tables/assemblyprocessor/reader.rs +++ b/dotscope/src/metadata/tables/assemblyprocessor/reader.rs @@ -103,12 +103,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/assemblyref/loader.rs b/dotscope/src/metadata/tables/assemblyref/loader.rs index f9ae97a2..1dbb03f0 100644 --- a/dotscope/src/metadata/tables/assemblyref/loader.rs +++ b/dotscope/src/metadata/tables/assemblyref/loader.rs @@ -58,6 +58,8 @@ //! //! - [ECMA-335 II.22.5](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `AssemblyRef` table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -118,7 +120,10 @@ impl MetadataLoader for AssemblyRefLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("assembly ref 0x{:08x}", row.token.value()); let Some(res) = context.handle_result( diff --git a/dotscope/src/metadata/tables/assemblyref/raw.rs b/dotscope/src/metadata/tables/assemblyref/raw.rs index 7254af53..6ae0f03d 100644 --- a/dotscope/src/metadata/tables/assemblyref/raw.rs +++ b/dotscope/src/metadata/tables/assemblyref/raw.rs @@ -53,6 +53,7 @@ use crate::{ }, token::Token, }, + utils::LazyList, Result, }; @@ -201,7 +202,7 @@ impl AssemblyRefRaw { os_major_version: AtomicU32::new(0), os_minor_version: AtomicU32::new(0), processor: AtomicU32::new(0), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } diff --git a/dotscope/src/metadata/tables/assemblyref/reader.rs b/dotscope/src/metadata/tables/assemblyref/reader.rs index 16be8b20..88d27e85 100644 --- a/dotscope/src/metadata/tables/assemblyref/reader.rs +++ b/dotscope/src/metadata/tables/assemblyref/reader.rs @@ -134,12 +134,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -182,12 +183,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/assemblyrefos/loader.rs b/dotscope/src/metadata/tables/assemblyrefos/loader.rs index e60f8ee5..38c477d9 100644 --- a/dotscope/src/metadata/tables/assemblyrefos/loader.rs +++ b/dotscope/src/metadata/tables/assemblyrefos/loader.rs @@ -41,6 +41,8 @@ //! //! - [ECMA-335 II.22.7](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `AssemblyRefOS` table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -99,7 +101,10 @@ impl MetadataLoader for AssemblyRefOsLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("assembly ref os 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/assemblyrefos/reader.rs b/dotscope/src/metadata/tables/assemblyrefos/reader.rs index b1e95424..94402be4 100644 --- a/dotscope/src/metadata/tables/assemblyrefos/reader.rs +++ b/dotscope/src/metadata/tables/assemblyrefos/reader.rs @@ -100,12 +100,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -138,12 +139,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/assemblyrefprocessor/loader.rs b/dotscope/src/metadata/tables/assemblyrefprocessor/loader.rs index 6b0c0034..609e6869 100644 --- a/dotscope/src/metadata/tables/assemblyrefprocessor/loader.rs +++ b/dotscope/src/metadata/tables/assemblyrefprocessor/loader.rs @@ -39,6 +39,8 @@ //! //! - [ECMA-335 II.22.8](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `AssemblyRefProcessor` table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -105,7 +107,10 @@ impl MetadataLoader for AssemblyRefProcessorLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("assembly ref processor 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/assemblyrefprocessor/reader.rs b/dotscope/src/metadata/tables/assemblyrefprocessor/reader.rs index f515691b..7d43f8bb 100644 --- a/dotscope/src/metadata/tables/assemblyrefprocessor/reader.rs +++ b/dotscope/src/metadata/tables/assemblyrefprocessor/reader.rs @@ -95,12 +95,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -132,12 +133,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/classlayout/loader.rs b/dotscope/src/metadata/tables/classlayout/loader.rs index 3c1df081..3c82e83e 100644 --- a/dotscope/src/metadata/tables/classlayout/loader.rs +++ b/dotscope/src/metadata/tables/classlayout/loader.rs @@ -48,6 +48,8 @@ //! //! - [ECMA-335 II.22.8](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `ClassLayout` table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -116,7 +118,10 @@ impl MetadataLoader for ClassLayoutLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("class layout 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/classlayout/reader.rs b/dotscope/src/metadata/tables/classlayout/reader.rs index e7f4f493..0a625f4a 100644 --- a/dotscope/src/metadata/tables/classlayout/reader.rs +++ b/dotscope/src/metadata/tables/classlayout/reader.rs @@ -61,12 +61,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -85,8 +86,7 @@ mod tests { true, true, )); - let table = - MetadataTable::::new(&data, u16::MAX as u32 + 3, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: ClassLayoutRaw| { assert_eq!(row.rid, 1); @@ -98,12 +98,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/constant/loader.rs b/dotscope/src/metadata/tables/constant/loader.rs index a6d23e0b..70933c8c 100644 --- a/dotscope/src/metadata/tables/constant/loader.rs +++ b/dotscope/src/metadata/tables/constant/loader.rs @@ -50,6 +50,8 @@ //! //! - [ECMA-335 II.22.9](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - Constant table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -119,7 +121,10 @@ impl MetadataLoader for ConstantLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("constant 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/constant/reader.rs b/dotscope/src/metadata/tables/constant/reader.rs index 18bb4cd7..c30a57c7 100644 --- a/dotscope/src/metadata/tables/constant/reader.rs +++ b/dotscope/src/metadata/tables/constant/reader.rs @@ -64,12 +64,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -89,7 +90,7 @@ mod tests { true, true, )); - let table = MetadataTable::::new(&data, u16::MAX as u32 + 3, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: ConstantRaw| { assert_eq!(row.rid, 1); @@ -104,12 +105,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/customattribute/loader.rs b/dotscope/src/metadata/tables/customattribute/loader.rs index 2fd82bd8..a870549d 100644 --- a/dotscope/src/metadata/tables/customattribute/loader.rs +++ b/dotscope/src/metadata/tables/customattribute/loader.rs @@ -55,6 +55,8 @@ //! //! - [ECMA-335 II.22.10](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `CustomAttribute` table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -125,7 +127,10 @@ impl MetadataLoader for CustomAttributeLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("custom attribute 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/customattribute/reader.rs b/dotscope/src/metadata/tables/customattribute/reader.rs index 2ac88fc6..ddd41590 100644 --- a/dotscope/src/metadata/tables/customattribute/reader.rs +++ b/dotscope/src/metadata/tables/customattribute/reader.rs @@ -68,12 +68,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -95,8 +96,7 @@ mod tests { true, true, )); - let table = - MetadataTable::::new(&data, u16::MAX as u32 + 3, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: CustomAttributeRaw| { assert_eq!(row.rid, 1); @@ -122,12 +122,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/customdebuginformation/loader.rs b/dotscope/src/metadata/tables/customdebuginformation/loader.rs index e99f4750..7652f0be 100644 --- a/dotscope/src/metadata/tables/customdebuginformation/loader.rs +++ b/dotscope/src/metadata/tables/customdebuginformation/loader.rs @@ -58,6 +58,8 @@ //! //! - [Portable PDB v1.1](https://github.com/dotnet/corefx/blob/master/src/System.Reflection.Metadata/specs/PortablePdb-Metadata.md#customdebuginformation-table-0x37) - `CustomDebugInformation` table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -128,7 +130,10 @@ impl MetadataLoader for CustomDebugInformationLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("custom debug info 0x{:08x}", row.token.value()); let Some(custom_debug_info) = context.handle_result( diff --git a/dotscope/src/metadata/tables/customdebuginformation/reader.rs b/dotscope/src/metadata/tables/customdebuginformation/reader.rs index bb1689dc..f1f5d939 100644 --- a/dotscope/src/metadata/tables/customdebuginformation/reader.rs +++ b/dotscope/src/metadata/tables/customdebuginformation/reader.rs @@ -89,12 +89,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -129,12 +130,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/declsecurity/loader.rs b/dotscope/src/metadata/tables/declsecurity/loader.rs index 72713362..704acc25 100644 --- a/dotscope/src/metadata/tables/declsecurity/loader.rs +++ b/dotscope/src/metadata/tables/declsecurity/loader.rs @@ -28,6 +28,8 @@ //! - [ECMA-335 II.22.11](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `DeclSecurity` table specification //! - [ECMA-335 II.23.1.16](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `SecurityAction` enumeration +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -89,7 +91,10 @@ impl MetadataLoader for DeclSecurityLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("decl security 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/declsecurity/raw.rs b/dotscope/src/metadata/tables/declsecurity/raw.rs index 9a3c07db..b091c46f 100644 --- a/dotscope/src/metadata/tables/declsecurity/raw.rs +++ b/dotscope/src/metadata/tables/declsecurity/raw.rs @@ -44,6 +44,7 @@ use crate::{ token::Token, typesystem::CilTypeReference, }, + utils::LazyList, Result, }; @@ -223,7 +224,7 @@ impl DeclSecurityRaw { action, parent, permission_set, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } } diff --git a/dotscope/src/metadata/tables/declsecurity/reader.rs b/dotscope/src/metadata/tables/declsecurity/reader.rs index e6d33634..629b2159 100644 --- a/dotscope/src/metadata/tables/declsecurity/reader.rs +++ b/dotscope/src/metadata/tables/declsecurity/reader.rs @@ -83,12 +83,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -111,8 +112,7 @@ mod tests { true, true, )); - let table = - MetadataTable::::new(&data, u16::MAX as u32 + 3, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: DeclSecurityRaw| { assert_eq!(row.rid, 1); @@ -127,12 +127,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/document/loader.rs b/dotscope/src/metadata/tables/document/loader.rs index 4e41d5cd..58a979af 100644 --- a/dotscope/src/metadata/tables/document/loader.rs +++ b/dotscope/src/metadata/tables/document/loader.rs @@ -63,7 +63,11 @@ impl MetadataLoader for DocumentLoader { table .par_iter() - .map(|row| { + .enumerate() + .map(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("document 0x{:08x}", row.token.value()); let Some(document) = context.handle_result( diff --git a/dotscope/src/metadata/tables/enclog/loader.rs b/dotscope/src/metadata/tables/enclog/loader.rs index 538906f8..ea555b56 100644 --- a/dotscope/src/metadata/tables/enclog/loader.rs +++ b/dotscope/src/metadata/tables/enclog/loader.rs @@ -19,6 +19,8 @@ //! # Reference //! - [ECMA-335 II.22.12](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `EncLog` table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -57,7 +59,10 @@ impl MetadataLoader for EncLogLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("enc log 0x{:08x}", row.token.value()); let Some(owned) = diff --git a/dotscope/src/metadata/tables/enclog/reader.rs b/dotscope/src/metadata/tables/enclog/reader.rs index 1eac56f1..f6a90859 100644 --- a/dotscope/src/metadata/tables/enclog/reader.rs +++ b/dotscope/src/metadata/tables/enclog/reader.rs @@ -66,12 +66,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/encmap/loader.rs b/dotscope/src/metadata/tables/encmap/loader.rs index 4340b6a0..4afa40ad 100644 --- a/dotscope/src/metadata/tables/encmap/loader.rs +++ b/dotscope/src/metadata/tables/encmap/loader.rs @@ -18,6 +18,8 @@ //! # Reference //! - [ECMA-335 II.22.13](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `EncMap` table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -55,7 +57,10 @@ impl MetadataLoader for EncMapLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("enc map 0x{:08x}", row.token.value()); let Some(owned) = diff --git a/dotscope/src/metadata/tables/encmap/reader.rs b/dotscope/src/metadata/tables/encmap/reader.rs index 8998d2ec..5acd6f2f 100644 --- a/dotscope/src/metadata/tables/encmap/reader.rs +++ b/dotscope/src/metadata/tables/encmap/reader.rs @@ -66,12 +66,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/event/loader.rs b/dotscope/src/metadata/tables/event/loader.rs index f4a43f7e..3102a7e8 100644 --- a/dotscope/src/metadata/tables/event/loader.rs +++ b/dotscope/src/metadata/tables/event/loader.rs @@ -27,6 +27,8 @@ //! # Reference //! - [ECMA-335 II.22.13](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - Event table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -74,7 +76,10 @@ impl MetadataLoader for EventLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("event 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/event/raw.rs b/dotscope/src/metadata/tables/event/raw.rs index 5ea7f50c..206d86a4 100644 --- a/dotscope/src/metadata/tables/event/raw.rs +++ b/dotscope/src/metadata/tables/event/raw.rs @@ -29,6 +29,7 @@ use crate::{ token::Token, typesystem::TypeRegistry, }, + utils::LazyList, Result, }; @@ -128,7 +129,7 @@ impl EventRaw { fn_on_other: OnceLock::new(), fn_on_raise: OnceLock::new(), fn_on_remove: OnceLock::new(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } diff --git a/dotscope/src/metadata/tables/event/reader.rs b/dotscope/src/metadata/tables/event/reader.rs index c40ca346..b6b8ec4e 100644 --- a/dotscope/src/metadata/tables/event/reader.rs +++ b/dotscope/src/metadata/tables/event/reader.rs @@ -68,12 +68,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -96,7 +97,7 @@ mod tests { true, true, )); - let table = MetadataTable::::new(&data, u16::MAX as u32 + 3, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: EventRaw| { assert_eq!(row.rid, 1); @@ -111,12 +112,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/eventmap/loader.rs b/dotscope/src/metadata/tables/eventmap/loader.rs index da7c3062..a5466d92 100644 --- a/dotscope/src/metadata/tables/eventmap/loader.rs +++ b/dotscope/src/metadata/tables/eventmap/loader.rs @@ -14,6 +14,8 @@ //! # Reference //! - [ECMA-335 II.22.12](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `EventMap` table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -61,7 +63,10 @@ impl MetadataLoader for EventMapLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("event map 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/eventmap/raw.rs b/dotscope/src/metadata/tables/eventmap/raw.rs index ecabd5c2..3007a374 100644 --- a/dotscope/src/metadata/tables/eventmap/raw.rs +++ b/dotscope/src/metadata/tables/eventmap/raw.rs @@ -141,7 +141,7 @@ impl EventMapRaw { let end = if next_row_id > map.row_count { events.len().saturating_add(1) } else { - match map.get(next_row_id) { + match map.get(next_row_id)? { Some(next_row) => next_row.event_list as usize, None => { return Err(malformed_error!( @@ -158,35 +158,23 @@ impl EventMapRaw { let event_list = Arc::new(boxcar::Vec::with_capacity(end.saturating_sub(start))); for counter in start..end { + let event_rid = u32::try_from(counter) + .map_err(|_| malformed_error!("Event row index out of range: {}", counter))?; let actual_event_token = if event_ptr.is_empty() { - let token_value = counter | 0x1400_0000; - Token::new(u32::try_from(token_value).map_err(|_| { - malformed_error!("Token value {} exceeds u32 range", token_value) - })?) + Token::from_parts(TableId::Event, event_rid) } else { - let event_ptr_token_value = u32::try_from(counter | 0x0D00_0000).map_err(|_| { - malformed_error!("EventPtr token value too large: {}", counter | 0x0D00_0000) - })?; - let event_ptr_token = Token::new(event_ptr_token_value); + // Built from the TableId enum rather than a hand-written prefix so the table + // id cannot drift from the value `EventPtrReader` keys rows under. + let event_ptr_token = Token::from_parts(TableId::EventPtr, event_rid); match event_ptr.get(&event_ptr_token) { Some(event_ptr_entry) => { - let actual_event_rid = event_ptr_entry.value().event; - let actual_event_token_value = u32::try_from( - actual_event_rid as usize | 0x1400_0000, - ) - .map_err(|_| { - malformed_error!( - "Event token value too large: {}", - actual_event_rid as usize | 0x1400_0000 - ) - })?; - Token::new(actual_event_token_value) + Token::from_parts(TableId::Event, event_ptr_entry.value().event) } None => { return Err(malformed_error!( "Failed to resolve EventPtr - {}", - counter | 0x0D00_0000 + event_ptr_token.value() )) } } diff --git a/dotscope/src/metadata/tables/eventmap/reader.rs b/dotscope/src/metadata/tables/eventmap/reader.rs index 8e6d2528..a56d697a 100644 --- a/dotscope/src/metadata/tables/eventmap/reader.rs +++ b/dotscope/src/metadata/tables/eventmap/reader.rs @@ -80,12 +80,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -106,7 +107,7 @@ mod tests { true, true, )); - let table = MetadataTable::::new(&data, u16::MAX as u32 + 3, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: EventMapRaw| { assert_eq!(row.rid, 1); @@ -117,12 +118,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/eventptr/loader.rs b/dotscope/src/metadata/tables/eventptr/loader.rs index e5a1fb12..5e9bec74 100644 --- a/dotscope/src/metadata/tables/eventptr/loader.rs +++ b/dotscope/src/metadata/tables/eventptr/loader.rs @@ -70,7 +70,10 @@ impl MetadataLoader for EventPtrLoader { return Ok(()); }; - for row in table { + for (index, row) in table.into_iter().enumerate() { + let Some(row) = context.handle_row(row, index)? else { + continue; + }; let token_msg = || format!("event ptr 0x{:08x}", row.token.value()); let Some(owned) = diff --git a/dotscope/src/metadata/tables/eventptr/reader.rs b/dotscope/src/metadata/tables/eventptr/reader.rs index 0bce0b0c..e59ebc79 100644 --- a/dotscope/src/metadata/tables/eventptr/reader.rs +++ b/dotscope/src/metadata/tables/eventptr/reader.rs @@ -72,12 +72,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -104,12 +105,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/exportedtype/loader.rs b/dotscope/src/metadata/tables/exportedtype/loader.rs index 3b37cf57..0a99449a 100644 --- a/dotscope/src/metadata/tables/exportedtype/loader.rs +++ b/dotscope/src/metadata/tables/exportedtype/loader.rs @@ -22,6 +22,8 @@ //! # Reference //! - [ECMA-335 II.22.14](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `ExportedType` table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -70,28 +72,37 @@ impl MetadataLoader for ExportedTypeLoader { }; // First pass: create exported type entries - table.par_iter().try_for_each(|row| { - let token_msg = || format!("exported type 0x{:08x}", row.token.value()); + table + .par_iter() + .enumerate() + .try_for_each(|(index, row)| -> Result<()> { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; + let token_msg = || format!("exported type 0x{:08x}", row.token.value()); - let Some(owned) = context.handle_result( - row.to_owned(|coded_index| context.get_ref(coded_index), strings, true), - DiagnosticCategory::Type, - token_msg, - )? - else { - return Ok(()); - }; + let Some(owned) = context.handle_result( + row.to_owned(|coded_index| context.get_ref(coded_index), strings, true), + DiagnosticCategory::Type, + token_msg, + )? + else { + return Ok(()); + }; - context.handle_result( - context.exported_type.insert(row.token, owned.clone()), - DiagnosticCategory::Type, - token_msg, - )?; - Ok(()) - })?; + context.handle_result( + context.exported_type.insert(row.token, owned.clone()), + DiagnosticCategory::Type, + token_msg, + )?; + Ok(()) + })?; // Second pass: resolve implementations - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("exported type impl 0x{:08x}", row.token.value()); if let Some(implementation) = diff --git a/dotscope/src/metadata/tables/exportedtype/raw.rs b/dotscope/src/metadata/tables/exportedtype/raw.rs index eda1e8c6..5778f3d8 100644 --- a/dotscope/src/metadata/tables/exportedtype/raw.rs +++ b/dotscope/src/metadata/tables/exportedtype/raw.rs @@ -53,6 +53,7 @@ use crate::{ token::Token, typesystem::CilTypeReference, }, + utils::LazyList, Result, }; @@ -199,7 +200,7 @@ impl ExportedTypeRaw { Some(string.get(self.namespace as usize)?.to_string()) }, implementation: implementation_lock, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } diff --git a/dotscope/src/metadata/tables/exportedtype/reader.rs b/dotscope/src/metadata/tables/exportedtype/reader.rs index 4ebd9aeb..83022867 100644 --- a/dotscope/src/metadata/tables/exportedtype/reader.rs +++ b/dotscope/src/metadata/tables/exportedtype/reader.rs @@ -92,12 +92,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -139,12 +140,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/field/loader.rs b/dotscope/src/metadata/tables/field/loader.rs index eb183399..f0558d00 100644 --- a/dotscope/src/metadata/tables/field/loader.rs +++ b/dotscope/src/metadata/tables/field/loader.rs @@ -22,6 +22,8 @@ //! # Reference //! - [ECMA-335 II.22.15](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - Field table specification +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -65,7 +67,10 @@ impl MetadataLoader for FieldLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("field 0x{:08x}", row.token.value()); let Some(res) = context.handle_result( diff --git a/dotscope/src/metadata/tables/field/raw.rs b/dotscope/src/metadata/tables/field/raw.rs index f907de9c..f071f2c5 100644 --- a/dotscope/src/metadata/tables/field/raw.rs +++ b/dotscope/src/metadata/tables/field/raw.rs @@ -22,6 +22,7 @@ use crate::{ tables::{Field, FieldAttributes, FieldRc, TableInfoRef, TableRow}, token::Token, }, + utils::LazyList, Result, }; @@ -123,7 +124,7 @@ impl FieldRaw { rva: OnceLock::new(), layout: OnceLock::new(), marshal: OnceLock::new(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), declaring_type: OnceLock::new(), })) } diff --git a/dotscope/src/metadata/tables/field/reader.rs b/dotscope/src/metadata/tables/field/reader.rs index a35fcf07..19daeb6a 100644 --- a/dotscope/src/metadata/tables/field/reader.rs +++ b/dotscope/src/metadata/tables/field/reader.rs @@ -55,12 +55,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -91,12 +92,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/fieldlayout/loader.rs b/dotscope/src/metadata/tables/fieldlayout/loader.rs index 6dd00e46..06c6f4f5 100644 --- a/dotscope/src/metadata/tables/fieldlayout/loader.rs +++ b/dotscope/src/metadata/tables/fieldlayout/loader.rs @@ -17,6 +17,8 @@ //! # ECMA-335 Reference //! See ECMA-335, Partition II, §22.16 for the `FieldLayout` table specification. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -69,7 +71,10 @@ impl MetadataLoader for FieldLayoutLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("field layout 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/fieldlayout/reader.rs b/dotscope/src/metadata/tables/fieldlayout/reader.rs index 8ed11272..4bed4bb2 100644 --- a/dotscope/src/metadata/tables/fieldlayout/reader.rs +++ b/dotscope/src/metadata/tables/fieldlayout/reader.rs @@ -72,12 +72,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -95,8 +96,7 @@ mod tests { true, true, )); - let table = - MetadataTable::::new(&data, u16::MAX as u32 + 3, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: FieldLayoutRaw| { assert_eq!(row.rid, 1); @@ -107,12 +107,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/fieldmarshal/loader.rs b/dotscope/src/metadata/tables/fieldmarshal/loader.rs index 96d05328..d68766c6 100644 --- a/dotscope/src/metadata/tables/fieldmarshal/loader.rs +++ b/dotscope/src/metadata/tables/fieldmarshal/loader.rs @@ -19,6 +19,8 @@ //! # ECMA-335 Reference //! See ECMA-335, Partition II, §22.17 for the `FieldMarshal` table specification. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -78,7 +80,10 @@ impl MetadataLoader for FieldMarshalLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("field marshal 0x{:08x}", row.token.value()); let Some(res) = context.handle_result( diff --git a/dotscope/src/metadata/tables/fieldmarshal/reader.rs b/dotscope/src/metadata/tables/fieldmarshal/reader.rs index d027b99d..5802f16a 100644 --- a/dotscope/src/metadata/tables/fieldmarshal/reader.rs +++ b/dotscope/src/metadata/tables/fieldmarshal/reader.rs @@ -74,12 +74,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -100,8 +101,7 @@ mod tests { true, true, )); - let table = - MetadataTable::::new(&data, u16::MAX as u32 + 3, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: FieldMarshalRaw| { assert_eq!(row.rid, 1); @@ -115,12 +115,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/fieldptr/loader.rs b/dotscope/src/metadata/tables/fieldptr/loader.rs index a90bd7e1..4bf6d561 100644 --- a/dotscope/src/metadata/tables/fieldptr/loader.rs +++ b/dotscope/src/metadata/tables/fieldptr/loader.rs @@ -66,7 +66,10 @@ impl MetadataLoader for FieldPtrLoader { return Ok(()); }; - for row in table { + for (index, row) in table.into_iter().enumerate() { + let Some(row) = context.handle_row(row, index)? else { + continue; + }; let token_msg = || format!("field ptr 0x{:08x}", row.token.value()); let Some(owned) = diff --git a/dotscope/src/metadata/tables/fieldptr/reader.rs b/dotscope/src/metadata/tables/fieldptr/reader.rs index 064747b3..ab0ac26c 100644 --- a/dotscope/src/metadata/tables/fieldptr/reader.rs +++ b/dotscope/src/metadata/tables/fieldptr/reader.rs @@ -49,12 +49,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -81,12 +82,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/fieldrva/loader.rs b/dotscope/src/metadata/tables/fieldrva/loader.rs index 6db751a2..6b6f343c 100644 --- a/dotscope/src/metadata/tables/fieldrva/loader.rs +++ b/dotscope/src/metadata/tables/fieldrva/loader.rs @@ -25,6 +25,8 @@ //! # ECMA-335 Reference //! See ECMA-335, Partition II, §22.19 for the `FieldRva` table specification. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -77,7 +79,10 @@ impl MetadataLoader for FieldRvaLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("field rva 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/fieldrva/reader.rs b/dotscope/src/metadata/tables/fieldrva/reader.rs index 974a11a2..96309ceb 100644 --- a/dotscope/src/metadata/tables/fieldrva/reader.rs +++ b/dotscope/src/metadata/tables/fieldrva/reader.rs @@ -52,12 +52,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -89,12 +90,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/file/loader.rs b/dotscope/src/metadata/tables/file/loader.rs index b45af5f9..9d0b706b 100644 --- a/dotscope/src/metadata/tables/file/loader.rs +++ b/dotscope/src/metadata/tables/file/loader.rs @@ -30,6 +30,8 @@ //! # ECMA-335 Reference //! See ECMA-335, Partition II, §22.19 for the File table specification. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -85,7 +87,10 @@ impl MetadataLoader for FileLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("file 0x{:08x}", row.token.value()); let Some(res) = context.handle_result( diff --git a/dotscope/src/metadata/tables/file/raw.rs b/dotscope/src/metadata/tables/file/raw.rs index dda5cf0a..97d3b730 100644 --- a/dotscope/src/metadata/tables/file/raw.rs +++ b/dotscope/src/metadata/tables/file/raw.rs @@ -32,6 +32,7 @@ use crate::{ tables::{AssemblyRefHash, File, FileAttributes, FileRc, TableInfoRef, TableRow}, token::Token, }, + utils::LazyList, Result, }; @@ -144,7 +145,7 @@ impl FileRaw { flags: FileAttributes::new(self.flags), name: strings.get(self.name as usize)?.to_string(), hash_value: AssemblyRefHash::new(blob.get(self.hash_value as usize)?)?, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } diff --git a/dotscope/src/metadata/tables/file/reader.rs b/dotscope/src/metadata/tables/file/reader.rs index d6886da7..f1b07c9c 100644 --- a/dotscope/src/metadata/tables/file/reader.rs +++ b/dotscope/src/metadata/tables/file/reader.rs @@ -55,12 +55,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -86,12 +87,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/genericparam/loader.rs b/dotscope/src/metadata/tables/genericparam/loader.rs index c27a67d5..ffd1776c 100644 --- a/dotscope/src/metadata/tables/genericparam/loader.rs +++ b/dotscope/src/metadata/tables/genericparam/loader.rs @@ -29,6 +29,8 @@ //! # ECMA-335 Reference //! See ECMA-335, Partition II, §22.20 for the `GenericParam` table specification. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -83,22 +85,28 @@ impl MetadataLoader for GenericParamLoader { return Ok(()); }; - generics.par_iter().try_for_each(|row| { - let token_msg = || format!("generic param 0x{:08x}", row.token.value()); + generics + .par_iter() + .enumerate() + .try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; + let token_msg = || format!("generic param 0x{:08x}", row.token.value()); - let Some(owned) = context.handle_result( - row.to_owned(|coded_index| context.get_ref(coded_index), strings), - DiagnosticCategory::Type, - token_msg, - )? - else { - return Ok(()); - }; + let Some(owned) = context.handle_result( + row.to_owned(|coded_index| context.get_ref(coded_index), strings), + DiagnosticCategory::Type, + token_msg, + )? + else { + return Ok(()); + }; - context.handle_error(owned.apply(), DiagnosticCategory::Type, token_msg)?; - context.generic_param.insert(row.token, owned.clone()); - Ok(()) - }) + context.handle_error(owned.apply(), DiagnosticCategory::Type, token_msg)?; + context.generic_param.insert(row.token, owned.clone()); + Ok(()) + }) } /// Returns the table identifier for the `GenericParam` table. diff --git a/dotscope/src/metadata/tables/genericparam/mod.rs b/dotscope/src/metadata/tables/genericparam/mod.rs index 96a19fe3..729ceded 100644 --- a/dotscope/src/metadata/tables/genericparam/mod.rs +++ b/dotscope/src/metadata/tables/genericparam/mod.rs @@ -68,7 +68,7 @@ use std::sync::Arc; use crossbeam_skiplist::SkipMap; -use crate::metadata::token::Token; +use crate::{metadata::token::Token, utils::LazyList}; mod builder; mod loader; @@ -94,7 +94,7 @@ pub type GenericParamMap = SkipMap; /// This collection provides ordered access to generic parameter entries, useful for /// sequential processing and bulk operations during generic type analysis and /// parameter enumeration. -pub type GenericParamList = Arc>; +pub type GenericParamList = LazyList; /// Reference-counted generic parameter entry. /// diff --git a/dotscope/src/metadata/tables/genericparam/owned.rs b/dotscope/src/metadata/tables/genericparam/owned.rs index 2d92b743..161210ad 100644 --- a/dotscope/src/metadata/tables/genericparam/owned.rs +++ b/dotscope/src/metadata/tables/genericparam/owned.rs @@ -269,7 +269,7 @@ impl fmt::Display for GenericParam { // Type/interface constraints for (_, constraint_ref) in self.constraints.iter() { if let Some(constraint_type) = constraint_ref.upgrade() { - constraints.push(constraint_type.fullname()); + constraints.push(constraint_type.fullname().to_string()); } } diff --git a/dotscope/src/metadata/tables/genericparam/raw.rs b/dotscope/src/metadata/tables/genericparam/raw.rs index 0bb0eac8..943e8d41 100644 --- a/dotscope/src/metadata/tables/genericparam/raw.rs +++ b/dotscope/src/metadata/tables/genericparam/raw.rs @@ -34,6 +34,7 @@ use crate::{ token::Token, typesystem::CilTypeReference, }, + utils::LazyList, Result, }; @@ -169,7 +170,7 @@ impl GenericParamRaw { owner, constraints: Arc::new(boxcar::Vec::new()), name: strings.get(self.name as usize)?.to_string(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } } diff --git a/dotscope/src/metadata/tables/genericparam/reader.rs b/dotscope/src/metadata/tables/genericparam/reader.rs index 59c355c1..084a561c 100644 --- a/dotscope/src/metadata/tables/genericparam/reader.rs +++ b/dotscope/src/metadata/tables/genericparam/reader.rs @@ -65,12 +65,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -110,12 +111,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/genericparamconstraint/loader.rs b/dotscope/src/metadata/tables/genericparamconstraint/loader.rs index a1ad41b2..593bc26d 100644 --- a/dotscope/src/metadata/tables/genericparamconstraint/loader.rs +++ b/dotscope/src/metadata/tables/genericparamconstraint/loader.rs @@ -31,6 +31,8 @@ //! # ECMA-335 Reference //! See ECMA-335, Partition II, §22.21 for the `GenericParamConstraint` table specification. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -94,7 +96,10 @@ impl MetadataLoader for GenericParamConstraintLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("generic param constraint 0x{:08x}", row.token.value()); let Some(res) = context.handle_result( diff --git a/dotscope/src/metadata/tables/genericparamconstraint/raw.rs b/dotscope/src/metadata/tables/genericparamconstraint/raw.rs index 244fe67c..43f62019 100644 --- a/dotscope/src/metadata/tables/genericparamconstraint/raw.rs +++ b/dotscope/src/metadata/tables/genericparamconstraint/raw.rs @@ -30,6 +30,7 @@ use crate::{ token::Token, typesystem::TypeRegistry, }, + utils::LazyList, Result, }; @@ -212,7 +213,7 @@ impl GenericParamConstraintRaw { )) } }, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } } diff --git a/dotscope/src/metadata/tables/genericparamconstraint/reader.rs b/dotscope/src/metadata/tables/genericparamconstraint/reader.rs index c1eb35a1..a95628e0 100644 --- a/dotscope/src/metadata/tables/genericparamconstraint/reader.rs +++ b/dotscope/src/metadata/tables/genericparamconstraint/reader.rs @@ -64,12 +64,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -107,12 +108,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/implmap/loader.rs b/dotscope/src/metadata/tables/implmap/loader.rs index dbd2957a..c8b7638f 100644 --- a/dotscope/src/metadata/tables/implmap/loader.rs +++ b/dotscope/src/metadata/tables/implmap/loader.rs @@ -28,6 +28,8 @@ //! //! # ECMA-335 Reference //! See ECMA-335, Partition II, §22.22 for the `ImplMap` table specification. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -65,7 +67,10 @@ impl MetadataLoader for ImplMapLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("impl map 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/implmap/reader.rs b/dotscope/src/metadata/tables/implmap/reader.rs index 0a091459..c8cfde04 100644 --- a/dotscope/src/metadata/tables/implmap/reader.rs +++ b/dotscope/src/metadata/tables/implmap/reader.rs @@ -86,12 +86,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -132,12 +133,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/importscope/loader.rs b/dotscope/src/metadata/tables/importscope/loader.rs index b29b7ff6..c178c17e 100644 --- a/dotscope/src/metadata/tables/importscope/loader.rs +++ b/dotscope/src/metadata/tables/importscope/loader.rs @@ -4,6 +4,8 @@ //! `ImportScope` table data during metadata loading. The loader handles parallel //! processing and integration with the broader loader context. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -44,7 +46,10 @@ impl MetadataLoader for ImportScopeLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("import scope 0x{:08x}", row.token.value()); let Some(import_scope) = diff --git a/dotscope/src/metadata/tables/importscope/reader.rs b/dotscope/src/metadata/tables/importscope/reader.rs index 28fae389..33983e8e 100644 --- a/dotscope/src/metadata/tables/importscope/reader.rs +++ b/dotscope/src/metadata/tables/importscope/reader.rs @@ -52,12 +52,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -86,12 +87,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/interfaceimpl/loader.rs b/dotscope/src/metadata/tables/interfaceimpl/loader.rs index 9ec715bd..225f3036 100644 --- a/dotscope/src/metadata/tables/interfaceimpl/loader.rs +++ b/dotscope/src/metadata/tables/interfaceimpl/loader.rs @@ -27,6 +27,8 @@ //! # ECMA-335 Reference //! See ECMA-335, Partition II, §22.23 for the `InterfaceImpl` table specification. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -65,7 +67,10 @@ impl MetadataLoader for InterfaceImplLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("interface impl 0x{:08x}", row.token.value()); let Some(res) = context.handle_result( diff --git a/dotscope/src/metadata/tables/interfaceimpl/raw.rs b/dotscope/src/metadata/tables/interfaceimpl/raw.rs index 5722641d..f6d3fbe1 100644 --- a/dotscope/src/metadata/tables/interfaceimpl/raw.rs +++ b/dotscope/src/metadata/tables/interfaceimpl/raw.rs @@ -23,6 +23,7 @@ use crate::{ token::Token, typesystem::{CilTypeRef, TypeRegistry}, }, + utils::LazyList, Result, }; @@ -111,7 +112,7 @@ impl InterfaceImplRaw { Some(class) => { class.interfaces.push(InterfaceEntry { interface: CilTypeRef::new(&interface), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); Ok(()) } @@ -163,7 +164,7 @@ impl InterfaceImplRaw { )) } }, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } } diff --git a/dotscope/src/metadata/tables/interfaceimpl/reader.rs b/dotscope/src/metadata/tables/interfaceimpl/reader.rs index 66cc21d5..26a71595 100644 --- a/dotscope/src/metadata/tables/interfaceimpl/reader.rs +++ b/dotscope/src/metadata/tables/interfaceimpl/reader.rs @@ -72,12 +72,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -95,8 +96,7 @@ mod tests { true, true, )); - let table = - MetadataTable::::new(&data, u16::MAX as u32 + 2, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: InterfaceImplRaw| { assert_eq!(row.rid, 1); @@ -109,7 +109,7 @@ mod tests { }; { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/localconstant/loader.rs b/dotscope/src/metadata/tables/localconstant/loader.rs index f7e8867d..c59b7e3d 100644 --- a/dotscope/src/metadata/tables/localconstant/loader.rs +++ b/dotscope/src/metadata/tables/localconstant/loader.rs @@ -4,6 +4,8 @@ //! `LocalConstant` table data during metadata loading. The loader handles parallel //! processing and integration with the broader loader context. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -47,7 +49,10 @@ impl MetadataLoader for LocalConstantLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("local constant 0x{:08x}", row.token.value()); let Some(local_constant) = context.handle_result( diff --git a/dotscope/src/metadata/tables/localconstant/reader.rs b/dotscope/src/metadata/tables/localconstant/reader.rs index d23bbca1..8b12b8cc 100644 --- a/dotscope/src/metadata/tables/localconstant/reader.rs +++ b/dotscope/src/metadata/tables/localconstant/reader.rs @@ -52,12 +52,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -86,12 +87,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/localscope/loader.rs b/dotscope/src/metadata/tables/localscope/loader.rs index c2ce8185..2b22a7fe 100644 --- a/dotscope/src/metadata/tables/localscope/loader.rs +++ b/dotscope/src/metadata/tables/localscope/loader.rs @@ -4,6 +4,8 @@ //! ``LocalScope`` table data during metadata loading. The loader handles parallel //! processing and integration with the broader loader context. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -47,7 +49,10 @@ impl MetadataLoader for LocalScopeLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("local scope 0x{:08x}", row.token.value()); let Some(local_scope) = context.handle_result( diff --git a/dotscope/src/metadata/tables/localscope/raw.rs b/dotscope/src/metadata/tables/localscope/raw.rs index 7bc9eec9..39439f71 100644 --- a/dotscope/src/metadata/tables/localscope/raw.rs +++ b/dotscope/src/metadata/tables/localscope/raw.rs @@ -158,7 +158,7 @@ impl LocalScopeRaw { } else { let start = self.variable_list; - let end = if let Some(next_scope) = scope_table.get(next_rid) { + let end = if let Some(next_scope) = scope_table.get(next_rid)? { if next_scope.variable_list != 0 { next_scope.variable_list } else { @@ -195,7 +195,7 @@ impl LocalScopeRaw { } else { let start = self.constant_list; - let end = if let Some(next_scope) = scope_table.get(next_rid) { + let end = if let Some(next_scope) = scope_table.get(next_rid)? { if next_scope.constant_list != 0 { next_scope.constant_list } else { diff --git a/dotscope/src/metadata/tables/localscope/reader.rs b/dotscope/src/metadata/tables/localscope/reader.rs index 54c02f53..439baafd 100644 --- a/dotscope/src/metadata/tables/localscope/reader.rs +++ b/dotscope/src/metadata/tables/localscope/reader.rs @@ -70,12 +70,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -118,12 +119,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/localvariable/loader.rs b/dotscope/src/metadata/tables/localvariable/loader.rs index b2436e29..4bf38271 100644 --- a/dotscope/src/metadata/tables/localvariable/loader.rs +++ b/dotscope/src/metadata/tables/localvariable/loader.rs @@ -4,6 +4,8 @@ //! `LocalVariable` table data during metadata loading. The loader handles parallel //! processing and integration with the broader loader context. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -45,7 +47,10 @@ impl MetadataLoader for LocalVariableLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("local variable 0x{:08x}", row.token.value()); let Some(local_variable) = context.handle_result( diff --git a/dotscope/src/metadata/tables/localvariable/reader.rs b/dotscope/src/metadata/tables/localvariable/reader.rs index 41a40d5f..a0f79ef2 100644 --- a/dotscope/src/metadata/tables/localvariable/reader.rs +++ b/dotscope/src/metadata/tables/localvariable/reader.rs @@ -58,12 +58,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -94,12 +95,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/manifestresource/loader.rs b/dotscope/src/metadata/tables/manifestresource/loader.rs index c596cc60..f237807c 100644 --- a/dotscope/src/metadata/tables/manifestresource/loader.rs +++ b/dotscope/src/metadata/tables/manifestresource/loader.rs @@ -27,6 +27,8 @@ //! # ECMA-335 Reference //! See ECMA-335, Partition II, §22.24 for the `ManifestResource` table specification. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -65,7 +67,10 @@ impl MetadataLoader for ManifestResourceLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("manifest resource 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/manifestresource/raw.rs b/dotscope/src/metadata/tables/manifestresource/raw.rs index 986df167..f528e484 100644 --- a/dotscope/src/metadata/tables/manifestresource/raw.rs +++ b/dotscope/src/metadata/tables/manifestresource/raw.rs @@ -160,7 +160,7 @@ impl ManifestResourceRaw { self.rid ) })?; - data_size = if let Some(next_res) = table.get(next_rid) { + data_size = if let Some(next_res) = table.get(next_rid)? { (next_res.offset_field as usize) .checked_sub(self.offset_field as usize) .ok_or_else(|| { diff --git a/dotscope/src/metadata/tables/manifestresource/reader.rs b/dotscope/src/metadata/tables/manifestresource/reader.rs index 31d2c919..0d453eee 100644 --- a/dotscope/src/metadata/tables/manifestresource/reader.rs +++ b/dotscope/src/metadata/tables/manifestresource/reader.rs @@ -68,12 +68,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -114,12 +115,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/memberref/loader.rs b/dotscope/src/metadata/tables/memberref/loader.rs index 2320d0e4..b3910c05 100644 --- a/dotscope/src/metadata/tables/memberref/loader.rs +++ b/dotscope/src/metadata/tables/memberref/loader.rs @@ -32,6 +32,8 @@ //! //! [`MemberRefLoader`]: crate::metadata::tables::MemberRefLoader +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -72,7 +74,10 @@ impl MetadataLoader for MemberRefLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("member ref 0x{:08x}", row.token.value()); let Some(res) = context.handle_result( diff --git a/dotscope/src/metadata/tables/memberref/raw.rs b/dotscope/src/metadata/tables/memberref/raw.rs index 7fa04f16..4563c95f 100644 --- a/dotscope/src/metadata/tables/memberref/raw.rs +++ b/dotscope/src/metadata/tables/memberref/raw.rs @@ -27,6 +27,7 @@ use crate::{ token::Token, typesystem::{CilTypeReference, TypeRegistry}, }, + utils::LazyList, Result, }; @@ -138,7 +139,7 @@ impl MemberRefRaw { modifiers: Arc::new(boxcar::Vec::new()), base: OnceLock::new(), is_by_ref: AtomicBool::new(method_sig.return_type.by_ref), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); params.push(return_param); @@ -157,7 +158,7 @@ impl MemberRefRaw { modifiers: Arc::new(boxcar::Vec::new()), base: OnceLock::new(), is_by_ref: AtomicBool::new(param_sig.by_ref), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); params.push(param); } @@ -320,7 +321,7 @@ impl MemberRefRaw { name: strings.get(self.name as usize)?.to_string(), signature, params, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); Ok(member_ref) diff --git a/dotscope/src/metadata/tables/memberref/reader.rs b/dotscope/src/metadata/tables/memberref/reader.rs index 6ff6350e..258c8cd2 100644 --- a/dotscope/src/metadata/tables/memberref/reader.rs +++ b/dotscope/src/metadata/tables/memberref/reader.rs @@ -58,12 +58,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -82,7 +83,7 @@ mod tests { true, true, )); - let table = MetadataTable::::new(&data, u16::MAX as u32 + 3, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: MemberRefRaw| { assert_eq!(row.rid, 1); @@ -97,12 +98,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/methoddebuginformation/loader.rs b/dotscope/src/metadata/tables/methoddebuginformation/loader.rs index f30d6bb4..3c53640e 100644 --- a/dotscope/src/metadata/tables/methoddebuginformation/loader.rs +++ b/dotscope/src/metadata/tables/methoddebuginformation/loader.rs @@ -20,6 +20,8 @@ //! # Reference //! * [Portable PDB Format - MethodDebugInformation Table](https://github.com/dotnet/core/blob/main/Documentation/diagnostics/portable_pdb.md#methoddebuginformation-table-0x31) +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -60,7 +62,10 @@ impl MetadataLoader for MethodDebugInformationLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("method debug info 0x{:08x}", row.token.value()); let Some(method_debug_info) = diff --git a/dotscope/src/metadata/tables/methoddebuginformation/reader.rs b/dotscope/src/metadata/tables/methoddebuginformation/reader.rs index 19bde398..b12b31a3 100644 --- a/dotscope/src/metadata/tables/methoddebuginformation/reader.rs +++ b/dotscope/src/metadata/tables/methoddebuginformation/reader.rs @@ -52,12 +52,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -89,12 +90,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/methoddef/loader.rs b/dotscope/src/metadata/tables/methoddef/loader.rs index 1e593fb0..68c9e795 100644 --- a/dotscope/src/metadata/tables/methoddef/loader.rs +++ b/dotscope/src/metadata/tables/methoddef/loader.rs @@ -73,6 +73,8 @@ //! - Table ID: 0x06 //! - Purpose: Define method implementations within types +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -150,7 +152,10 @@ impl MetadataLoader for MethodDefLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("method 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/methoddef/raw.rs b/dotscope/src/metadata/tables/methoddef/raw.rs index 0e293916..c1a118c3 100644 --- a/dotscope/src/metadata/tables/methoddef/raw.rs +++ b/dotscope/src/metadata/tables/methoddef/raw.rs @@ -72,6 +72,7 @@ use crate::{ tables::{MetadataTable, ParamMap, ParamPtrMap, TableId, TableInfoRef, TableRow}, token::Token, }, + utils::LazyList, Result, }; @@ -254,7 +255,7 @@ impl MethodDefRaw { let end = if next_row_id > table.row_count { params_map.len().saturating_add(1) } else { - match table.get(next_row_id) { + match table.get(next_row_id)? { Some(next_row) => next_row.param_list as usize, None => { return Err(malformed_error!( @@ -270,38 +271,27 @@ impl MethodDefRaw { } else { let type_params = Arc::new(boxcar::Vec::with_capacity(end.saturating_sub(start))); for counter in start..end { + let param_rid = u32::try_from(counter).map_err(|_| { + malformed_error!("Param row index out of range: {}", counter) + })?; let actual_param_token = if param_ptr_map.is_empty() { - let token_value = u32::try_from(counter | 0x0800_0000).map_err(|_| { - malformed_error!("Token value too large: {}", counter | 0x0800_0000) - })?; - Token::new(token_value) + Token::from_parts(TableId::Param, param_rid) } else { - let param_ptr_token_value = - u32::try_from(counter | 0x0A00_0000).map_err(|_| { - malformed_error!( - "ParamPtr token value too large: {}", - counter | 0x0A00_0000 - ) - })?; - let param_ptr_token = Token::new(param_ptr_token_value); + // Built from the TableId enum rather than a hand-written prefix so the + // table id cannot drift from the value `ParamPtrReader` keys rows under. + // It had: this was `0x0A00_0000`, which is MemberRef, while ParamPtr is + // `0x07` — so the lookup below never hit and every assembly carrying a + // ParamPtr table failed to load. + let param_ptr_token = Token::from_parts(TableId::ParamPtr, param_rid); match param_ptr_map.get(¶m_ptr_token) { Some(param_ptr) => { - let actual_param_rid = param_ptr.value().param; - let actual_param_token_value = - u32::try_from(actual_param_rid as usize | 0x0800_0000) - .map_err(|_| { - malformed_error!( - "Param token value too large: {}", - actual_param_rid as usize | 0x0800_0000 - ) - })?; - Token::new(actual_param_token_value) + Token::from_parts(TableId::Param, param_ptr.value().param) } None => { return Err(malformed_error!( "Failed to resolve ParamPtr - {}", - counter | 0x0A00_0000 + param_ptr_token.value() )) } } @@ -335,17 +325,17 @@ impl MethodDefRaw { flags_pinvoke: AtomicU32::new(0), params: type_params, varargs: Arc::new(boxcar::Vec::new()), - generic_params: Arc::new(boxcar::Vec::new()), - generic_args: Arc::new(boxcar::Vec::new()), + generic_params: LazyList::new(), + generic_args: LazyList::new(), signature, rva: if self.rva == 0 { None } else { Some(self.rva) }, body: OnceLock::new(), local_vars: Arc::new(boxcar::Vec::new()), overrides: Arc::new(boxcar::Vec::new()), - interface_impls: Arc::new(boxcar::Vec::new()), + interface_impls: LazyList::new(), security: OnceLock::new(), blocks: OnceLock::new(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), declaring_type: OnceLock::new(), })) } diff --git a/dotscope/src/metadata/tables/methoddef/reader.rs b/dotscope/src/metadata/tables/methoddef/reader.rs index dc194def..2fa3013c 100644 --- a/dotscope/src/metadata/tables/methoddef/reader.rs +++ b/dotscope/src/metadata/tables/methoddef/reader.rs @@ -64,12 +64,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -91,7 +92,7 @@ mod tests { true, true, )); - let table = MetadataTable::::new(&data, u16::MAX as u32 + 2, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: MethodDefRaw| { assert_eq!(row.rid, 1); @@ -105,7 +106,7 @@ mod tests { }; { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/methodimpl/loader.rs b/dotscope/src/metadata/tables/methodimpl/loader.rs index 93f745f9..f333efeb 100644 --- a/dotscope/src/metadata/tables/methodimpl/loader.rs +++ b/dotscope/src/metadata/tables/methodimpl/loader.rs @@ -73,6 +73,8 @@ //! - Partition II, §22.27 for the `MethodImpl` table specification //! - Table ID: 0x19 //! - Purpose: Define method implementation mappings for interface and virtual method resolution +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -149,7 +151,10 @@ impl MetadataLoader for MethodImplLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("method impl 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/methodimpl/reader.rs b/dotscope/src/metadata/tables/methodimpl/reader.rs index 9d328428..923e54b7 100644 --- a/dotscope/src/metadata/tables/methodimpl/reader.rs +++ b/dotscope/src/metadata/tables/methodimpl/reader.rs @@ -71,12 +71,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -118,12 +119,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/methodptr/loader.rs b/dotscope/src/metadata/tables/methodptr/loader.rs index 690eba6e..9dec6f46 100644 --- a/dotscope/src/metadata/tables/methodptr/loader.rs +++ b/dotscope/src/metadata/tables/methodptr/loader.rs @@ -74,7 +74,10 @@ impl MetadataLoader for MethodPtrLoader { return Ok(()); }; - for row in table { + for (index, row) in table.into_iter().enumerate() { + let Some(row) = context.handle_row(row, index)? else { + continue; + }; let token_msg = || format!("method ptr 0x{:08x}", row.token.value()); let Some(owned) = diff --git a/dotscope/src/metadata/tables/methodptr/reader.rs b/dotscope/src/metadata/tables/methodptr/reader.rs index 118bdf42..01d69ce8 100644 --- a/dotscope/src/metadata/tables/methodptr/reader.rs +++ b/dotscope/src/metadata/tables/methodptr/reader.rs @@ -49,12 +49,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -81,12 +82,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/methodsemantics/loader.rs b/dotscope/src/metadata/tables/methodsemantics/loader.rs index 97a3e90d..abeda6c6 100644 --- a/dotscope/src/metadata/tables/methodsemantics/loader.rs +++ b/dotscope/src/metadata/tables/methodsemantics/loader.rs @@ -25,6 +25,8 @@ //! - Semantic relationships conflict (e.g., duplicate setters) //! - Required dependency tables are missing or malformed +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -62,7 +64,10 @@ impl MetadataLoader for MethodSemanticsLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("method semantics 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/methodsemantics/reader.rs b/dotscope/src/metadata/tables/methodsemantics/reader.rs index 4631ea5a..201b4c4c 100644 --- a/dotscope/src/metadata/tables/methodsemantics/reader.rs +++ b/dotscope/src/metadata/tables/methodsemantics/reader.rs @@ -86,12 +86,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -130,12 +131,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/methodspec/loader.rs b/dotscope/src/metadata/tables/methodspec/loader.rs index b7ca5caa..6e6b772d 100644 --- a/dotscope/src/metadata/tables/methodspec/loader.rs +++ b/dotscope/src/metadata/tables/methodspec/loader.rs @@ -19,6 +19,8 @@ //! - [`crate::metadata::tables::MethodDefRaw`] - For method definition resolution //! - [`crate::metadata::tables::MemberRef`] - For member reference resolution +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -57,7 +59,10 @@ impl MetadataLoader for MethodSpecLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("method spec 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/methodspec/mod.rs b/dotscope/src/metadata/tables/methodspec/mod.rs index 6d348bc7..04a68977 100644 --- a/dotscope/src/metadata/tables/methodspec/mod.rs +++ b/dotscope/src/metadata/tables/methodspec/mod.rs @@ -60,7 +60,10 @@ use std::sync::Arc; use crossbeam_skiplist::SkipMap; -use crate::metadata::{token::Token, typesystem::CilTypeReference}; +use crate::{ + metadata::{token::Token, typesystem::CilTypeReference}, + utils::LazyList, +}; mod builder; mod loader; @@ -84,7 +87,7 @@ pub type MethodSpecMap = SkipMap; /// /// Uses a lock-free vector implementation for efficient concurrent access to /// the collection of all method specification entries in the metadata. -pub type MethodSpecList = Arc>; +pub type MethodSpecList = LazyList; /// Reference-counted pointer to a [`MethodSpec`] entry. /// diff --git a/dotscope/src/metadata/tables/methodspec/raw.rs b/dotscope/src/metadata/tables/methodspec/raw.rs index 680ae24c..4b5d4feb 100644 --- a/dotscope/src/metadata/tables/methodspec/raw.rs +++ b/dotscope/src/metadata/tables/methodspec/raw.rs @@ -13,6 +13,7 @@ use crate::{ token::Token, typesystem::{CilTypeReference, TypeRegistry, TypeResolver}, }, + utils::LazyList, Result, }; @@ -141,7 +142,7 @@ impl MethodSpecRaw { offset: self.offset, method: method.clone(), instantiation, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), generic_args, }); diff --git a/dotscope/src/metadata/tables/methodspec/reader.rs b/dotscope/src/metadata/tables/methodspec/reader.rs index f07444d1..0f59b544 100644 --- a/dotscope/src/metadata/tables/methodspec/reader.rs +++ b/dotscope/src/metadata/tables/methodspec/reader.rs @@ -79,12 +79,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -120,12 +121,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/module/loader.rs b/dotscope/src/metadata/tables/module/loader.rs index 099a34cc..eb01e88a 100644 --- a/dotscope/src/metadata/tables/module/loader.rs +++ b/dotscope/src/metadata/tables/module/loader.rs @@ -69,7 +69,7 @@ impl MetadataLoader for ModuleLoader { || "module".to_string(), ); }; - let Some(row) = table.get(1) else { + let Some(row) = table.get(1)? else { return context.handle_error( Err(malformed_error!( "Module table is present but contains no rows" diff --git a/dotscope/src/metadata/tables/module/raw.rs b/dotscope/src/metadata/tables/module/raw.rs index 1c70e30f..7db493fd 100644 --- a/dotscope/src/metadata/tables/module/raw.rs +++ b/dotscope/src/metadata/tables/module/raw.rs @@ -11,6 +11,7 @@ use crate::{ tables::{Module, ModuleRc, TableInfoRef, TableRow}, token::Token, }, + utils::LazyList, Result, }; @@ -138,7 +139,7 @@ impl ModuleRaw { Some(guids.get(self.encbaseid as usize)?) }, imports: Vec::new(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } diff --git a/dotscope/src/metadata/tables/module/reader.rs b/dotscope/src/metadata/tables/module/reader.rs index 805054bc..e5d616b3 100644 --- a/dotscope/src/metadata/tables/module/reader.rs +++ b/dotscope/src/metadata/tables/module/reader.rs @@ -138,12 +138,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -178,12 +179,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/moduleref/loader.rs b/dotscope/src/metadata/tables/moduleref/loader.rs index 5911e0a4..ff539b2a 100644 --- a/dotscope/src/metadata/tables/moduleref/loader.rs +++ b/dotscope/src/metadata/tables/moduleref/loader.rs @@ -24,6 +24,8 @@ //! - `ModuleRef` table contains invalid data //! - Token conflicts occur during storage +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -61,7 +63,10 @@ impl MetadataLoader for ModuleRefLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("module ref 0x{:08x}", row.token.value()); let Some(res) = context.handle_result( diff --git a/dotscope/src/metadata/tables/moduleref/raw.rs b/dotscope/src/metadata/tables/moduleref/raw.rs index 323a4983..5c3a059b 100644 --- a/dotscope/src/metadata/tables/moduleref/raw.rs +++ b/dotscope/src/metadata/tables/moduleref/raw.rs @@ -10,6 +10,7 @@ use crate::{ tables::{ModuleRef, ModuleRefRc, TableInfoRef, TableRow}, token::Token, }, + utils::LazyList, Result, }; @@ -102,7 +103,7 @@ impl ModuleRefRaw { token: self.token, offset: self.offset, name: strings.get(self.name as usize)?.to_string(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } diff --git a/dotscope/src/metadata/tables/moduleref/reader.rs b/dotscope/src/metadata/tables/moduleref/reader.rs index dcf552de..738cbbc6 100644 --- a/dotscope/src/metadata/tables/moduleref/reader.rs +++ b/dotscope/src/metadata/tables/moduleref/reader.rs @@ -118,12 +118,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -150,12 +151,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/nestedclass/loader.rs b/dotscope/src/metadata/tables/nestedclass/loader.rs index ad91ed86..6546c528 100644 --- a/dotscope/src/metadata/tables/nestedclass/loader.rs +++ b/dotscope/src/metadata/tables/nestedclass/loader.rs @@ -27,6 +27,8 @@ //! - Circular nesting relationships are detected //! - Token conflicts occur during storage //! +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -66,22 +68,28 @@ impl MetadataLoader for NestedClassLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { - let token_msg = || format!("nested class 0x{:08x}", row.token.value()); + table + .par_iter() + .enumerate() + .try_for_each(|(index, row)| -> Result<()> { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; + let token_msg = || format!("nested class 0x{:08x}", row.token.value()); - let Some(owned) = context.handle_result( - row.to_owned(context.types), - DiagnosticCategory::Type, - token_msg, - )? - else { - return Ok(()); - }; + let Some(owned) = context.handle_result( + row.to_owned(context.types), + DiagnosticCategory::Type, + token_msg, + )? + else { + return Ok(()); + }; - context.handle_error(owned.apply(), DiagnosticCategory::Type, token_msg)?; - context.nested_class.insert(row.token, owned); - Ok(()) - })?; + context.handle_error(owned.apply(), DiagnosticCategory::Type, token_msg)?; + context.nested_class.insert(row.token, owned); + Ok(()) + })?; // Rebuild the fullname index now that enclosing type relationships are established. // `set_enclosing_type()` invalidates the cached fullname on each nested type, so diff --git a/dotscope/src/metadata/tables/nestedclass/raw.rs b/dotscope/src/metadata/tables/nestedclass/raw.rs index b225c0af..a544f7ea 100644 --- a/dotscope/src/metadata/tables/nestedclass/raw.rs +++ b/dotscope/src/metadata/tables/nestedclass/raw.rs @@ -111,6 +111,7 @@ impl NestedClassRaw { let mut mapping: BTreeMap> = BTreeMap::new(); for row in classes { + let row = row?; mapping .entry(row.enclosing_class | 0x0200_0000) .or_default() diff --git a/dotscope/src/metadata/tables/nestedclass/reader.rs b/dotscope/src/metadata/tables/nestedclass/reader.rs index 625d8a59..b2db9119 100644 --- a/dotscope/src/metadata/tables/nestedclass/reader.rs +++ b/dotscope/src/metadata/tables/nestedclass/reader.rs @@ -122,12 +122,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -159,12 +160,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/param/loader.rs b/dotscope/src/metadata/tables/param/loader.rs index 13f8abb6..7fc3bc41 100644 --- a/dotscope/src/metadata/tables/param/loader.rs +++ b/dotscope/src/metadata/tables/param/loader.rs @@ -25,6 +25,8 @@ //! - Param table contains invalid or corrupted data //! - Token conflicts occur during storage +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -75,7 +77,10 @@ impl MetadataLoader for ParamLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("param 0x{:08x}", row.token.value()); let Some(res) = context.handle_result( diff --git a/dotscope/src/metadata/tables/param/raw.rs b/dotscope/src/metadata/tables/param/raw.rs index a12b025b..6dd2b8fa 100644 --- a/dotscope/src/metadata/tables/param/raw.rs +++ b/dotscope/src/metadata/tables/param/raw.rs @@ -11,6 +11,7 @@ use crate::{ tables::{Param, ParamAttributes, ParamRc, TableInfoRef, TableRow}, token::Token, }, + utils::LazyList, Result, }; @@ -157,7 +158,7 @@ impl ParamRaw { modifiers: Arc::new(boxcar::Vec::new()), base: OnceLock::new(), is_by_ref: AtomicBool::new(false), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } } diff --git a/dotscope/src/metadata/tables/param/reader.rs b/dotscope/src/metadata/tables/param/reader.rs index 8fb072e2..e7c0114f 100644 --- a/dotscope/src/metadata/tables/param/reader.rs +++ b/dotscope/src/metadata/tables/param/reader.rs @@ -137,12 +137,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -173,12 +174,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/paramptr/loader.rs b/dotscope/src/metadata/tables/paramptr/loader.rs index cc97957f..3c1b4b8c 100644 --- a/dotscope/src/metadata/tables/paramptr/loader.rs +++ b/dotscope/src/metadata/tables/paramptr/loader.rs @@ -57,7 +57,10 @@ impl MetadataLoader for ParamPtrLoader { return Ok(()); }; - for row in table { + for (index, row) in table.into_iter().enumerate() { + let Some(row) = context.handle_row(row, index)? else { + continue; + }; let token_msg = || format!("param ptr 0x{:08x}", row.token.value()); let Some(owned) = diff --git a/dotscope/src/metadata/tables/paramptr/reader.rs b/dotscope/src/metadata/tables/paramptr/reader.rs index 795cf227..e1c6dc44 100644 --- a/dotscope/src/metadata/tables/paramptr/reader.rs +++ b/dotscope/src/metadata/tables/paramptr/reader.rs @@ -79,8 +79,8 @@ impl RowReadable for ParamPtrRaw { /// /// ## Errors /// - /// * [`crate::error::Error::OutOfBounds`] - Insufficient data for complete entry - /// * [`crate::error::Error::Malformed`] - Malformed table entry structure + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] - Insufficient data for complete entry + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] - Malformed table entry structure fn row_read(data: &[u8], offset: &mut usize, rid: u32, sizes: &TableInfoRef) -> Result { Ok(ParamPtrRaw { rid, @@ -120,12 +120,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -152,12 +153,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/property/loader.rs b/dotscope/src/metadata/tables/property/loader.rs index 910e1b33..dd77814a 100644 --- a/dotscope/src/metadata/tables/property/loader.rs +++ b/dotscope/src/metadata/tables/property/loader.rs @@ -24,6 +24,8 @@ //! - [`crate::metadata::tables::PropertyRaw`] - Raw table entry structure //! - [`crate::metadata::tables::Property`] - Owned table entry type +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -68,7 +70,10 @@ impl MetadataLoader for PropertyLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("property 0x{:08x}", row.token.value()); let Some(res) = context.handle_result( diff --git a/dotscope/src/metadata/tables/property/raw.rs b/dotscope/src/metadata/tables/property/raw.rs index 8ec31291..9fdb35a9 100644 --- a/dotscope/src/metadata/tables/property/raw.rs +++ b/dotscope/src/metadata/tables/property/raw.rs @@ -12,6 +12,7 @@ use crate::{ tables::{Property, PropertyAttributes, PropertyRc, TableInfoRef, TableRow}, token::Token, }, + utils::LazyList, Result, }; @@ -112,8 +113,8 @@ impl PropertyRaw { /// /// ## Errors /// - /// * [`crate::error::Error::OutOfBounds`] - Invalid string or blob heap index - /// * [`crate::error::Error::Malformed`] - Malformed property signature + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] - Invalid string or blob heap index + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] - Malformed property signature pub fn to_owned(&self, strings: &Strings, blob: &Blob) -> Result { Ok(Arc::new(Property { token: self.token, @@ -124,7 +125,7 @@ impl PropertyRaw { fn_setter: OnceLock::new(), fn_getter: OnceLock::new(), fn_other: OnceLock::new(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), })) } diff --git a/dotscope/src/metadata/tables/property/reader.rs b/dotscope/src/metadata/tables/property/reader.rs index 6a299100..5c10fd69 100644 --- a/dotscope/src/metadata/tables/property/reader.rs +++ b/dotscope/src/metadata/tables/property/reader.rs @@ -82,8 +82,8 @@ impl RowReadable for PropertyRaw { /// /// ## Errors /// - /// * [`crate::error::Error::OutOfBounds`] - Insufficient data for complete entry - /// * [`crate::error::Error::Malformed`] - Malformed table entry structure + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] - Insufficient data for complete entry + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] - Malformed table entry structure fn row_read(data: &[u8], offset: &mut usize, rid: u32, sizes: &TableInfoRef) -> Result { Ok(PropertyRaw { rid, @@ -129,12 +129,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -165,12 +166,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/propertymap/loader.rs b/dotscope/src/metadata/tables/propertymap/loader.rs index 3639f4b5..3cb84b80 100644 --- a/dotscope/src/metadata/tables/propertymap/loader.rs +++ b/dotscope/src/metadata/tables/propertymap/loader.rs @@ -27,6 +27,8 @@ //! - [`crate::metadata::tables::PropertyMapRaw`] - Raw table entry structure //! - [`crate::metadata::tables::PropertyMap`] - Owned table entry type +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -69,7 +71,10 @@ impl MetadataLoader for PropertyMapLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("property map 0x{:08x}", row.token.value()); let Some(owned) = context.handle_result( diff --git a/dotscope/src/metadata/tables/propertymap/raw.rs b/dotscope/src/metadata/tables/propertymap/raw.rs index 11c935f8..d994d1f6 100644 --- a/dotscope/src/metadata/tables/propertymap/raw.rs +++ b/dotscope/src/metadata/tables/propertymap/raw.rs @@ -110,7 +110,7 @@ impl PropertyMapRaw { let end = if next_row_id > map.row_count { properties.len().saturating_add(1) } else { - match map.get(next_row_id) { + match map.get(next_row_id)? { Some(next_row) => next_row.property_list as usize, None => { return Err(malformed_error!( @@ -127,40 +127,23 @@ impl PropertyMapRaw { let property_list = Arc::new(boxcar::Vec::with_capacity(end.saturating_sub(start))); for counter in start..end { + let property_rid = u32::try_from(counter) + .map_err(|_| malformed_error!("Property row index out of range: {}", counter))?; let actual_property_token = if property_ptr.is_empty() { - let token_value = counter | 0x1700_0000; - Token::new( - u32::try_from(token_value) - .map_err(|_| malformed_error!("Property counter overflow"))?, - ) + Token::from_parts(TableId::Property, property_rid) } else { - let property_ptr_token_value = - u32::try_from(counter | 0x0E00_0000).map_err(|_| { - malformed_error!( - "PropertyPtr token value too large: {}", - counter | 0x0E00_0000 - ) - })?; - let property_ptr_token = Token::new(property_ptr_token_value); + // Built from the TableId enum rather than a hand-written prefix so the table + // id cannot drift from the value `PropertyPtrReader` keys rows under. + let property_ptr_token = Token::from_parts(TableId::PropertyPtr, property_rid); match property_ptr.get(&property_ptr_token) { Some(property_ptr_entry) => { - let actual_property_rid = property_ptr_entry.value().property; - let actual_property_token_value = u32::try_from( - actual_property_rid as usize | 0x1700_0000, - ) - .map_err(|_| { - malformed_error!( - "Property token value too large: {}", - actual_property_rid as usize | 0x1700_0000 - ) - })?; - Token::new(actual_property_token_value) + Token::from_parts(TableId::Property, property_ptr_entry.value().property) } None => { return Err(malformed_error!( "Failed to resolve PropertyPtr - {}", - counter | 0x0E00_0000 + property_ptr_token.value() )) } } diff --git a/dotscope/src/metadata/tables/propertymap/reader.rs b/dotscope/src/metadata/tables/propertymap/reader.rs index c739f25b..4ceba39d 100644 --- a/dotscope/src/metadata/tables/propertymap/reader.rs +++ b/dotscope/src/metadata/tables/propertymap/reader.rs @@ -130,12 +130,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -156,8 +157,7 @@ mod tests { true, true, )); - let table = - MetadataTable::::new(&data, u16::MAX as u32 + 3, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: PropertyMapRaw| { assert_eq!(row.rid, 1); @@ -168,12 +168,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/propertyptr/loader.rs b/dotscope/src/metadata/tables/propertyptr/loader.rs index 1c9c6dd4..c3c49836 100644 --- a/dotscope/src/metadata/tables/propertyptr/loader.rs +++ b/dotscope/src/metadata/tables/propertyptr/loader.rs @@ -68,7 +68,10 @@ impl MetadataLoader for PropertyPtrLoader { return Ok(()); }; - for row in table { + for (index, row) in table.into_iter().enumerate() { + let Some(row) = context.handle_row(row, index)? else { + continue; + }; let token_msg = || format!("property ptr 0x{:08x}", row.token.value()); let Some(owned) = diff --git a/dotscope/src/metadata/tables/propertyptr/reader.rs b/dotscope/src/metadata/tables/propertyptr/reader.rs index 781a30d0..8df500fe 100644 --- a/dotscope/src/metadata/tables/propertyptr/reader.rs +++ b/dotscope/src/metadata/tables/propertyptr/reader.rs @@ -71,8 +71,8 @@ impl RowReadable for PropertyPtrRaw { /// /// ## Errors /// - /// * [`crate::error::Error::OutOfBounds`] - Insufficient data for complete entry - /// * [`crate::error::Error::Malformed`] - Malformed table entry structure + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] - Insufficient data for complete entry + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] - Malformed table entry structure fn row_read(data: &[u8], offset: &mut usize, rid: u32, sizes: &TableInfoRef) -> Result { Ok(PropertyPtrRaw { rid, @@ -112,12 +112,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -144,12 +145,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/standalonesig/reader.rs b/dotscope/src/metadata/tables/standalonesig/reader.rs index 8ff997b1..2ccfe71e 100644 --- a/dotscope/src/metadata/tables/standalonesig/reader.rs +++ b/dotscope/src/metadata/tables/standalonesig/reader.rs @@ -72,8 +72,8 @@ impl RowReadable for StandAloneSigRaw { /// /// ## Errors /// - /// * [`crate::error::Error::OutOfBounds`] - Insufficient data for complete entry - /// * [`crate::error::Error::Malformed`] - Malformed table entry structure + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] - Insufficient data for complete entry + /// * [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] - Malformed table entry structure fn row_read(data: &[u8], offset: &mut usize, rid: u32, sizes: &TableInfoRef) -> Result { let offset_org = *offset; @@ -112,12 +112,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -129,8 +130,7 @@ mod tests { ]; let sizes = Arc::new(TableInfo::new_test(&[], true, true, true)); - let table = - MetadataTable::::new(&data, u16::MAX as u32 + 3, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: StandAloneSigRaw| { assert_eq!(row.rid, 1); @@ -140,12 +140,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/statemachinemethod/loader.rs b/dotscope/src/metadata/tables/statemachinemethod/loader.rs index 7d9fbaae..0f1925c2 100644 --- a/dotscope/src/metadata/tables/statemachinemethod/loader.rs +++ b/dotscope/src/metadata/tables/statemachinemethod/loader.rs @@ -5,6 +5,8 @@ //! The loader follows the established `MetadataLoader` pattern for consistent parallel //! processing and efficient memory utilization. +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -53,7 +55,10 @@ impl MetadataLoader for StateMachineMethodLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("state machine method 0x{:08x}", row.token.value()); let Some(state_machine_method) = context.handle_result( diff --git a/dotscope/src/metadata/tables/statemachinemethod/reader.rs b/dotscope/src/metadata/tables/statemachinemethod/reader.rs index 7e87c304..6b9140a6 100644 --- a/dotscope/src/metadata/tables/statemachinemethod/reader.rs +++ b/dotscope/src/metadata/tables/statemachinemethod/reader.rs @@ -92,12 +92,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -129,12 +130,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/typedef/loader.rs b/dotscope/src/metadata/tables/typedef/loader.rs index e2ec5724..c79e9136 100644 --- a/dotscope/src/metadata/tables/typedef/loader.rs +++ b/dotscope/src/metadata/tables/typedef/loader.rs @@ -50,6 +50,8 @@ //! - [`crate::metadata::tables::TypeDefRaw`] - Raw table entry structure //! - [`crate::metadata::typesystem::CilType`] - Type system integration +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -130,7 +132,10 @@ impl MetadataLoader for TypeDefLoader { .get() .map(|assembly| CilTypeReference::Assembly(assembly.clone())); - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("type 0x{:08x}", row.token.value()); let Some(type_def) = context.handle_result( diff --git a/dotscope/src/metadata/tables/typedef/raw.rs b/dotscope/src/metadata/tables/typedef/raw.rs index a6e2e957..0ffde0b6 100644 --- a/dotscope/src/metadata/tables/typedef/raw.rs +++ b/dotscope/src/metadata/tables/typedef/raw.rs @@ -28,6 +28,7 @@ use crate::{ token::Token, typesystem::{CilType, CilTypeRc, CilTypeRef, CilTypeReference}, }, + utils::LazyList, Result, }; @@ -164,7 +165,7 @@ impl TypeDefRaw { .ok_or_else(|| malformed_error!("Method count overflow: {}", methods.len()))?; (fields_end, methods_end) } else { - match defs.get(next_rid) { + match defs.get(next_rid)? { Some(next_row) => (next_row.field_list as usize, next_row.method_list as usize), None => { return Err(malformed_error!( @@ -197,38 +198,21 @@ impl TypeDefRaw { })?; let type_fields = Arc::new(boxcar::Vec::with_capacity(capacity)); for counter in start_fields..end_fields { + let field_rid = u32::try_from(counter) + .map_err(|_| malformed_error!("Field row index out of range: {}", counter))?; let actual_field_token = if field_ptr.is_empty() { - Token::new(u32::try_from(counter | 0x0400_0000).map_err(|_| { - malformed_error!("Field token overflow: {}", counter | 0x0400_0000) - })?) + Token::from_parts(TableId::Field, field_rid) } else { - let field_ptr_token_value = - u32::try_from(counter | 0x0300_0000).map_err(|_| { - malformed_error!( - "FieldPtr token value too large: {}", - counter | 0x0300_0000 - ) - })?; - let field_ptr_token = Token::new(field_ptr_token_value); + let field_ptr_token = Token::from_parts(TableId::FieldPtr, field_rid); match field_ptr.get(&field_ptr_token) { Some(field_ptr_entry) => { - let actual_field_rid = field_ptr_entry.value().field; - let actual_field_token_value = u32::try_from( - actual_field_rid as usize | 0x0400_0000, - ) - .map_err(|_| { - malformed_error!( - "Field token value too large: {}", - actual_field_rid as usize | 0x0400_0000 - ) - })?; - Token::new(actual_field_token_value) + Token::from_parts(TableId::Field, field_ptr_entry.value().field) } None => { return Err(malformed_error!( "Failed to resolve FieldPtr - {}", - counter | 0x0300_0000 + field_ptr_token.value() )) } } @@ -259,7 +243,7 @@ impl TypeDefRaw { || start_methods > methods.len() || end_methods < start_methods { - Arc::new(boxcar::Vec::new()) + LazyList::new() } else { let capacity = end_methods.checked_sub(start_methods).ok_or_else(|| { malformed_error!( @@ -268,40 +252,27 @@ impl TypeDefRaw { start_methods ) })?; - let type_methods = Arc::new(boxcar::Vec::with_capacity(capacity)); + // `LazyList` has no with_capacity: the point is to not allocate until first push. + let _ = capacity; + let type_methods = LazyList::new(); for counter in start_methods..end_methods { + let method_rid = u32::try_from(counter) + .map_err(|_| malformed_error!("Method row index out of range: {}", counter))?; let actual_method_token = if method_ptr.is_empty() { - Token::new(u32::try_from(counter | 0x0600_0000).map_err(|_| { - malformed_error!("Method token overflow: {}", counter | 0x0600_0000) - })?) + Token::from_parts(TableId::MethodDef, method_rid) } else { - let method_ptr_token_value = - u32::try_from(counter | 0x0900_0000).map_err(|_| { - malformed_error!( - "MethodPtr token value too large: {}", - counter | 0x0900_0000 - ) - })?; - let method_ptr_token = Token::new(method_ptr_token_value); + // Built from the TableId enum rather than a hand-written prefix so the + // table id cannot drift from the value `MethodPtrReader` keys rows under. + let method_ptr_token = Token::from_parts(TableId::MethodPtr, method_rid); match method_ptr.get(&method_ptr_token) { Some(method_ptr_entry) => { - let actual_method_rid = method_ptr_entry.value().method; - let actual_method_token_value = u32::try_from( - actual_method_rid as usize | 0x0600_0000, - ) - .map_err(|_| { - malformed_error!( - "Method token value too large: {}", - actual_method_rid as usize | 0x0600_0000 - ) - })?; - Token::new(actual_method_token_value) + Token::from_parts(TableId::MethodDef, method_ptr_entry.value().method) } None => { return Err(malformed_error!( "Failed to resolve MethodPtr - {}", - counter | 0x0900_0000 + method_ptr_token.value() )) } } diff --git a/dotscope/src/metadata/tables/typedef/reader.rs b/dotscope/src/metadata/tables/typedef/reader.rs index 0c338090..ea134191 100644 --- a/dotscope/src/metadata/tables/typedef/reader.rs +++ b/dotscope/src/metadata/tables/typedef/reader.rs @@ -135,12 +135,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -168,7 +169,7 @@ mod tests { true, true, )); - let table = MetadataTable::::new(&data, u16::MAX as u32 + 2, sizes).unwrap(); + let table = MetadataTable::::new(&data, 1, sizes).unwrap(); let eval = |row: TypeDefRaw| { assert_eq!(row.rid, 1); @@ -185,7 +186,7 @@ mod tests { }; { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/typeref/builder.rs b/dotscope/src/metadata/tables/typeref/builder.rs index 0ae63456..5dcd0d5d 100644 --- a/dotscope/src/metadata/tables/typeref/builder.rs +++ b/dotscope/src/metadata/tables/typeref/builder.rs @@ -153,8 +153,8 @@ impl TypeRefBuilder { // Create the TypeRefRaw entry let typeref_raw = TypeRefRaw { rid, - token: Token::new(rid | 0x0100_0000), // TypeRef table token prefix - offset: 0, // Will be set during binary generation + token: Token::from_parts(TableId::TypeRef, rid), + offset: 0, // Will be set during binary generation resolution_scope, type_name: name_index, type_namespace: namespace_index, diff --git a/dotscope/src/metadata/tables/typeref/loader.rs b/dotscope/src/metadata/tables/typeref/loader.rs index 3486cd09..4e0a2feb 100644 --- a/dotscope/src/metadata/tables/typeref/loader.rs +++ b/dotscope/src/metadata/tables/typeref/loader.rs @@ -57,7 +57,10 @@ impl MetadataLoader for TypeRefLoader { // ECMA-335 guarantees parent entries appear before children in the table. // Using parallel iteration could process children before parents, // causing resolution scope lookups to fail. - for row in table { + for (index, row) in table.into_iter().enumerate() { + let Some(row) = context.handle_row(row, index)? else { + continue; + }; let token_msg = || format!("type ref 0x{:08x}", row.token.value()); let Some(new_entry) = context.handle_result( diff --git a/dotscope/src/metadata/tables/typeref/raw.rs b/dotscope/src/metadata/tables/typeref/raw.rs index 88370c9c..cc439e3b 100644 --- a/dotscope/src/metadata/tables/typeref/raw.rs +++ b/dotscope/src/metadata/tables/typeref/raw.rs @@ -29,6 +29,7 @@ use crate::{ token::Token, typesystem::{CilType, CilTypeRc, CilTypeReference}, }, + utils::LazyList, Result, }; @@ -155,7 +156,7 @@ impl TypeRefRaw { None, TypeAttributes::ZERO, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), None, ))) } diff --git a/dotscope/src/metadata/tables/typeref/reader.rs b/dotscope/src/metadata/tables/typeref/reader.rs index 11d77b31..1ad6b174 100644 --- a/dotscope/src/metadata/tables/typeref/reader.rs +++ b/dotscope/src/metadata/tables/typeref/reader.rs @@ -122,12 +122,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -168,12 +169,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/tables/types/common/codedindex.rs b/dotscope/src/metadata/tables/types/common/codedindex.rs index 133372a4..f4758a15 100644 --- a/dotscope/src/metadata/tables/types/common/codedindex.rs +++ b/dotscope/src/metadata/tables/types/common/codedindex.rs @@ -364,61 +364,7 @@ impl CodedIndex { tag, row, ci_type, - token: match tag { - TableId::Module => Token::new(row), - TableId::TypeRef => Token::new(row | 0x0100_0000), - TableId::TypeDef => Token::new(row | 0x0200_0000), - TableId::FieldPtr => Token::new(row | 0x0300_0000), - TableId::Field => Token::new(row | 0x0400_0000), - TableId::MethodPtr => Token::new(row | 0x0500_0000), - TableId::MethodDef => Token::new(row | 0x0600_0000), - TableId::ParamPtr => Token::new(row | 0x0700_0000), - TableId::Param => Token::new(row | 0x0800_0000), - TableId::InterfaceImpl => Token::new(row | 0x0900_0000), - TableId::MemberRef => Token::new(row | 0x0A00_0000), - TableId::Constant => Token::new(row | 0x0B00_0000), - TableId::CustomAttribute => Token::new(row | 0x0C00_0000), - TableId::FieldMarshal => Token::new(row | 0x0D00_0000), - TableId::DeclSecurity => Token::new(row | 0x0E00_0000), - TableId::ClassLayout => Token::new(row | 0x0F00_0000), - TableId::FieldLayout => Token::new(row | 0x1000_0000), - TableId::StandAloneSig => Token::new(row | 0x1100_0000), - TableId::EventMap => Token::new(row | 0x1200_0000), - TableId::EventPtr => Token::new(row | 0x1300_0000), - TableId::Event => Token::new(row | 0x1400_0000), - TableId::PropertyMap => Token::new(row | 0x1500_0000), - TableId::PropertyPtr => Token::new(row | 0x1600_0000), - TableId::Property => Token::new(row | 0x1700_0000), - TableId::MethodSemantics => Token::new(row | 0x1800_0000), - TableId::MethodImpl => Token::new(row | 0x1900_0000), - TableId::ModuleRef => Token::new(row | 0x1A00_0000), - TableId::TypeSpec => Token::new(row | 0x1B00_0000), - TableId::ImplMap => Token::new(row | 0x1C00_0000), - TableId::FieldRVA => Token::new(row | 0x1D00_0000), - TableId::EncLog => Token::new(row | 0x1E00_0000), - TableId::EncMap => Token::new(row | 0x1F00_0000), - TableId::Assembly => Token::new(row | 0x2000_0000), - TableId::AssemblyProcessor => Token::new(row | 0x2100_0000), - TableId::AssemblyOS => Token::new(row | 0x2200_0000), - TableId::AssemblyRef => Token::new(row | 0x2300_0000), - TableId::AssemblyRefProcessor => Token::new(row | 0x2400_0000), - TableId::AssemblyRefOS => Token::new(row | 0x2500_0000), - TableId::File => Token::new(row | 0x2600_0000), - TableId::ExportedType => Token::new(row | 0x2700_0000), - TableId::ManifestResource => Token::new(row | 0x2800_0000), - TableId::NestedClass => Token::new(row | 0x2900_0000), - TableId::GenericParam => Token::new(row | 0x2A00_0000), - TableId::MethodSpec => Token::new(row | 0x2B00_0000), - TableId::GenericParamConstraint => Token::new(row | 0x2C00_0000), - TableId::Document => Token::new(row | 0x3000_0000), - TableId::MethodDebugInformation => Token::new(row | 0x3100_0000), - TableId::LocalScope => Token::new(row | 0x3200_0000), - TableId::LocalVariable => Token::new(row | 0x3300_0000), - TableId::LocalConstant => Token::new(row | 0x3400_0000), - TableId::ImportScope => Token::new(row | 0x3500_0000), - TableId::StateMachineMethod => Token::new(row | 0x3600_0000), - TableId::CustomDebugInformation => Token::new(row | 0x3700_0000), - }, + token: Token::from_parts(tag, row), } } diff --git a/dotscope/src/metadata/tables/types/common/info.rs b/dotscope/src/metadata/tables/types/common/info.rs index 6d3852a6..d491ddff 100644 --- a/dotscope/src/metadata/tables/types/common/info.rs +++ b/dotscope/src/metadata/tables/types/common/info.rs @@ -172,6 +172,15 @@ pub struct TableInfo { /// Determined by bit 2 of the heap size flags in the metadata tables header. /// When `true`, all blob heap references use 4 bytes; when `false`, 2 bytes. is_large_index_blob: bool, + + /// Bit mask of the tables this `TableInfo` actually describes. + /// + /// [`Self::row_count`] cannot answer this: it returns 0 both for a table that is present + /// and empty and for one this `TableInfo` says nothing about. Those are different facts, + /// and code that bounds a row against `row_count` needs the second one — a `TableInfo` + /// reconstructed for a writer round-trip or a unit test describes only the tables it was + /// given, so treating "not described" as "has no rows" rejects valid data. + declared: u64, } /// Shared reference to a [`TableInfo`] structure for efficient multi-threaded access. @@ -182,6 +191,19 @@ pub struct TableInfo { pub type TableInfoRef = Arc; impl TableInfo { + /// Bit mask of every table id that has a [`TableId`] variant. + /// + /// Derived from the enum itself rather than written out, so it cannot drift as tables are + /// added. The `#~` valid mask is checked against this before any row count is read; see + /// [`Self::new`]. + fn known_table_mask() -> u64 { + TableId::iter().fold(0u64, |mask, table_id| { + // Every discriminant is below 64 (the mask is a `u64` by specification), so the + // shift is always defined. + mask | (1u64 << (table_id as usize)) + }) + } + /// Constructs a new `TableInfo` from metadata tables header data. /// /// Parses the metadata tables header to extract table row counts and heap size flags, @@ -214,12 +236,31 @@ impl TableInfo { /// /// ## Errors /// - /// - [`crate::Error::OutOfBounds`] - Insufficient data to read required header fields + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] - Insufficient data to read required header fields /// /// ## Reference /// /// * [ECMA-335 Partition II, Section 24.2.6](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - #~ Stream pub fn new(data: &[u8], valid_bitvec: u64) -> Result { + // Reject bits that name no table before reading a single row count. + // + // `TableId` covers 0x00-0x2C and 0x30-0x37; the gaps at 0x2D-0x2F and 0x38-0x3F have no + // variant. The loop below consumes a `u32` only for bits that map to a known `TableId`, + // but `TablesHeader::from` computes where table data starts as + // `24 + valid_bitvec.count_ones() * 4` — which counts them. A single gap bit therefore + // makes the two disagree about how many row-count words the header holds, shifting + // every subsequent row count onto the wrong table and mis-locating the table data. + // Rejecting the mask keeps them in agreement by construction, rather than requiring + // both to implement the same tolerance. + let unknown = valid_bitvec & !Self::known_table_mask(); + if unknown != 0 { + return Err(malformed_error!( + "#~ valid mask names {} table(s) that do not exist (unknown bits 0x{:016X})", + unknown.count_ones(), + unknown + )); + } + let table_info_len = (TableId::CustomDebugInformation as usize) .checked_add(1) .ok_or_else(|| malformed_error!("Table info size overflow"))?; @@ -251,6 +292,9 @@ impl TableInfo { let heap_size_flags = read_le::(data.get(6..).ok_or(out_of_bounds_error!())?)?; let mut table_info = TableInfo { rows: table_info, + // Every bit survived the check above, so the mask is exactly the tables the + // header declares. + declared: valid_bitvec, coded_indexes: vec![0; CodedIndexType::COUNT], is_large_index_str: heap_size_flags & 1 == 1, is_large_index_guid: heap_size_flags & 2 == 2, @@ -287,6 +331,7 @@ impl TableInfo { ) -> Self { let mut table_info = TableInfo { rows: vec![TableRowInfo::default(); TableId::CustomDebugInformation as usize + 1], + declared: 0, coded_indexes: vec![0; CodedIndexType::COUNT], is_large_index_str: large_str, is_large_index_guid: large_guid, @@ -295,6 +340,7 @@ impl TableInfo { for valid_table in valid_tables { table_info.rows[valid_table.0 as usize] = TableRowInfo::new(valid_table.1); + table_info.declared |= 1u64 << (valid_table.0 as usize); } table_info.calculate_coded_index_bits(); @@ -333,7 +379,7 @@ impl TableInfo { /// /// ## Errors /// - /// - [`crate::Error::OutOfBounds`] - Tag value exceeds the number of tables in the coded index union + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] - Tag value exceeds the number of tables in the coded index union /// /// ## Reference /// @@ -396,7 +442,7 @@ impl TableInfo { /// /// ## Errors /// - /// - [`crate::Error::OutOfBounds`] - Table ID is not valid for the specified coded index type + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] - Table ID is not valid for the specified coded index type /// /// ## Reference /// @@ -420,8 +466,13 @@ impl TableInfo { _ => return Err(out_of_bounds_error!()), }; // CustomAttributeType uses 3 bits for the tag (5 entries -> ceil(log2(5)) = 3) - let encoded = (row << 3) | tag; - return Ok(encoded); + let encoded = row + .checked_shl(3) + .filter(|shifted| (shifted >> 3) == row) + .ok_or_else(|| { + malformed_error!("Row {} does not fit a CustomAttributeType coded index", row) + })?; + return Ok(encoded | tag); } let tables = coded_index_type.tables(); @@ -441,11 +492,26 @@ impl TableInfo { let tag_bits = (tables.len() as f32).log2().ceil() as u8; // Encode: (row << tag_bits) | tag + // + // A row too large for the bits left after the tag would shift its high bits off the + // top and encode a *different* row, so it is rejected rather than silently truncated. + // Round-tripping the shift is the check. + let encoded = row + .checked_shl(u32::from(tag_bits)) + .filter(|shifted| (shifted >> tag_bits) == row) + .ok_or_else(|| { + malformed_error!( + "Row {} does not fit a {}-bit-tagged coded index", + row, + tag_bits + ) + })?; + // Tag cast is safe as table count is limited by metadata format #[allow(clippy::cast_possible_truncation)] - let encoded = (row << tag_bits) | (tag as u32); + let tag = tag as u32; - Ok(encoded) + Ok(encoded | tag) } /// Checks whether a specific table requires large (4-byte) indices due to size. @@ -745,12 +811,26 @@ impl TableInfo { /// /// ## Returns /// - /// The number of rows in the specified table (0 if table is not present). + /// The number of rows in the specified table, or 0. + /// + /// **A 0 here is ambiguous**: it means either "this table is present and empty" or "this + /// `TableInfo` does not describe this table". Do not bound a row index against this value + /// without first asking [`Self::is_declared`] — a `TableInfo` built for a writer + /// round-trip or a test describes only some tables, and rejecting rows that reference the + /// others would reject valid metadata. #[must_use] pub fn row_count(&self, table_id: TableId) -> u32 { self.rows.get(table_id as usize).map_or(0, |info| info.rows) } + /// Whether this `TableInfo` describes the given table at all. + /// + /// See [`Self::row_count`] for why the two questions are separate. + #[must_use] + pub fn is_declared(&self, table_id: TableId) -> bool { + self.declared & (1u64 << (table_id as usize)) != 0 + } + /// Creates a new `TableInfo` with modified row counts for specified tables. /// /// This method is used when writing assemblies with modified table sizes (e.g., after @@ -792,6 +872,8 @@ impl TableInfo { let mut new_info = TableInfo { rows, + // A row count changing does not change which tables are described. + declared: self.declared, coded_indexes: vec![0; CodedIndexType::COUNT], is_large_index_str: self.is_large_index_str, is_large_index_guid: self.is_large_index_guid, @@ -801,4 +883,94 @@ impl TableInfo { new_info.calculate_coded_index_bits(); new_info } + + /// Returns a copy with heap index widths derived from the given **output** heap sizes. + /// + /// ECMA-335 II.24.2.6 selects a 4-byte heap index once a heap exceeds 0xFFFF bytes. The + /// widths on `self` come from the *input* file's `HeapSizes` byte and are copied verbatim by + /// [`with_modified_row_counts`](Self::with_modified_row_counts), so a writer that appends + /// past the threshold — new type and method names, decrypted literals, rebuilt signatures — + /// keeps emitting 2-byte indices while resolving offsets that no longer fit them. + /// + /// Call this with the byte counts actually written for `#Strings`, `#GUID` and `#Blob`, + /// before the tables stream is laid out. + /// + /// # Arguments + /// + /// * `strings_bytes` - Size of the emitted `#Strings` heap. + /// * `guid_bytes` - Size of the emitted `#GUID` heap. + /// * `blob_bytes` - Size of the emitted `#Blob` heap. + #[must_use] + pub fn with_modified_heap_sizes( + &self, + strings_bytes: usize, + guid_bytes: usize, + blob_bytes: usize, + ) -> TableInfo { + /// Offsets are indices into the heap, so a heap of exactly 0x10000 bytes still needs a + /// 4-byte index to address its last entry. + const LARGE_HEAP_THRESHOLD: usize = 0xFFFF; + + let mut new_info = TableInfo { + rows: self.rows.clone(), + declared: self.declared, + coded_indexes: vec![0; CodedIndexType::COUNT], + // Widths only ever widen here: an input that already declared 4-byte indices keeps + // them, because existing rows were read at that width. + is_large_index_str: self.is_large_index_str || strings_bytes > LARGE_HEAP_THRESHOLD, + is_large_index_guid: self.is_large_index_guid || guid_bytes > LARGE_HEAP_THRESHOLD, + is_large_index_blob: self.is_large_index_blob || blob_bytes > LARGE_HEAP_THRESHOLD, + }; + + new_info.calculate_coded_index_bits(); + new_info + } +} + +#[cfg(test)] +mod heap_size_tests { + use super::*; + + /// A heap that grew past 0xFFFF must promote its index to 4 bytes. + /// + /// Without this the emitted `HeapSizes` byte keeps the *input* file's widths, every + /// heap-reference field is written 2 bytes wide, and any offset at or above 0x10000 is + /// masked — producing a file that parses cleanly while its names and signatures point at + /// arbitrary earlier heap positions. + #[test] + fn heap_index_widens_past_the_threshold() { + let base = TableInfo::new_test(&[], false, false, false); + + let widened = base.with_modified_heap_sizes(0x1_0000, 0, 0); + assert!(widened.is_large_str(), "#Strings must widen"); + assert!(!widened.is_large_guid()); + assert!(!widened.is_large_blob()); + + let widened = base.with_modified_heap_sizes(0, 0x20_0000, 0x1_0000); + assert!(!widened.is_large_str()); + assert!(widened.is_large_guid(), "#GUID must widen"); + assert!(widened.is_large_blob(), "#Blob must widen"); + } + + /// Exactly at the boundary the 2-byte index still addresses every offset. + #[test] + fn heap_index_stays_small_at_the_boundary() { + let base = TableInfo::new_test(&[], false, false, false); + let same = base.with_modified_heap_sizes(0xFFFF, 0xFFFF, 0xFFFF); + + assert!(!same.is_large_str()); + assert!(!same.is_large_guid()); + assert!(!same.is_large_blob()); + } + + /// Widths never narrow: existing rows in the input were read at the declared width. + #[test] + fn heap_index_never_narrows() { + let base = TableInfo::new_test(&[], true, true, true); + let narrowed = base.with_modified_heap_sizes(16, 16, 16); + + assert!(narrowed.is_large_str()); + assert!(narrowed.is_large_guid()); + assert!(narrowed.is_large_blob()); + } } diff --git a/dotscope/src/metadata/tables/types/read/access.rs b/dotscope/src/metadata/tables/types/read/access.rs index f418c2a9..8db53e89 100644 --- a/dotscope/src/metadata/tables/types/read/access.rs +++ b/dotscope/src/metadata/tables/types/read/access.rs @@ -21,8 +21,10 @@ use crate::metadata::tables::{MetadataTable, RowReadable}; /// # fn example(tables: &TablesHeader) -> dotscope::Result<()> { /// // Type-safe access - no table ID needed /// if let Some(typedef_table) = tables.table::() { -/// // Work with the table safely +/// // Iteration yields `Result`: a row that fails to parse is reported +/// // rather than silently ending the iteration. /// for type_def in typedef_table.iter().take(5) { +/// let type_def = type_def?; /// println!("Type: {}", type_def.type_name); /// } /// } diff --git a/dotscope/src/metadata/tables/types/read/iter.rs b/dotscope/src/metadata/tables/types/read/iter.rs index 6abd3894..cedf9e02 100644 --- a/dotscope/src/metadata/tables/types/read/iter.rs +++ b/dotscope/src/metadata/tables/types/read/iter.rs @@ -32,15 +32,60 @@ //! - [`crate::metadata::tables::types::read::traits`] - Core parsing traits //! - [`crate::metadata::tables::types::read::access`] - Low-level access utilities -use std::sync::{Arc, Mutex}; - use rayon::iter::{plumbing, IndexedParallelIterator, ParallelIterator}; use crate::{ metadata::tables::{MetadataTable, RowReadable}, - Error, Result, + Result, }; +/// Discards a row that could not be parsed, after logging it. +/// +/// For consumers that genuinely cannot propagate — an adapter chain inside a function +/// returning `Option`, or a scan that has no error channel. Written as a named function +/// rather than `.flatten()` or `.filter_map(Result::ok)` so that dropping a row is greppable +/// and leaves a trace: a silently short table is how a malformed row turns into missing +/// analysis output rather than an error. +/// +/// Prefer `?` wherever the caller can carry an error. +/// +/// # Position is not RID +/// +/// Dropping a row shifts everything after it, so **the position of a row in a filtered +/// sequence is not its RID**. Never derive a RID from `enumerate()`, an index into a +/// collected `Vec`, `.first()`, or `.next()`: +/// +/// ```rust,ignore +/// // Wrong: one unreadable row ahead of the match names a different row entirely. +/// for (index, row) in table.iter().filter_map(skip_unreadable).enumerate() { +/// let rid = index as u32 + 1; +/// } +/// +/// // Right: the row carries its own RID. +/// for row in table.iter().filter_map(skip_unreadable) { +/// let rid = row.rid; +/// } +/// +/// // Right, when a specific RID is wanted: `get` derives each row's offset from its RID, +/// // so an unreadable row cannot displace any other. +/// let module = table.get(1).ok().flatten(); +/// ``` +/// +/// Every raw row struct carries a `rid` field, and +/// [`MetadataTable::get`](super::MetadataTable::get) fetches by RID directly. This mattered +/// most where a row was read positionally and written back under the original RID, which +/// copied one row's contents onto a different row. +#[must_use] +pub fn skip_unreadable(row: Result) -> Option { + match row { + Ok(row) => Some(row), + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + None + } + } +} + /// Sequential iterator for metadata table rows. /// /// This iterator provides lazy, on-demand access to table rows in sequential order. @@ -51,40 +96,52 @@ use crate::{ /// /// - **Lazy evaluation**: Rows are parsed only when accessed /// - **Memory efficient**: Constant memory usage regardless of table size -/// - **Error resilient**: Parsing errors result in `None` rather than panics +/// - **Honest about failure**: a row that does not parse is yielded as `Err`, not skipped /// - **Cache friendly**: Sequential access pattern optimizes memory locality +/// +/// ## Why the item type is a `Result` +/// +/// Yielding `None` on a parse error would end the iteration, silently dropping every +/// remaining row — a row-hiding primitive on a hostile file, and one that reaches further +/// than it looks: the writer rebuilds tables by iterating them, so a table truncated this way +/// is re-emitted without the rows that were skipped. Every row is therefore reported, and the +/// iterator always yields exactly `row_count` items. pub struct TableIterator<'a, T> { /// Reference to the table being iterated pub table: &'a MetadataTable<'a, T>, /// Current row number (0-based for internal tracking) pub current_row: u32, - /// Current byte offset in the table data - pub current_offset: usize, } impl Iterator for TableIterator<'_, T> { - type Item = T; + type Item = Result; fn next(&mut self) -> Option { if self.current_row >= self.table.row_count { return None; } - match T::row_read( - self.table.data, - &mut self.current_offset, - self.current_row.saturating_add(1), - &self.table.sizes, - ) { - Ok(row) => { - self.current_row = self.current_row.saturating_add(1); - Some(row) - } - Err(_) => None, - } + let rid = self.current_row.saturating_add(1); + self.current_row = rid; + + // `get` derives each row's offset from its index, so a failure here costs this row + // and no other — which is what makes continuing after one sound. + self.table.get(rid).transpose() + } + + fn size_hint(&self) -> (usize, Option) { + let remaining = self + .table + .row_count + .saturating_sub(self.current_row) + .try_into() + .unwrap_or(usize::MAX); + (remaining, Some(remaining)) } } +impl ExactSizeIterator for TableIterator<'_, T> {} + /// Parallel iterator for metadata table rows. /// /// This iterator enables concurrent processing of table rows across multiple threads @@ -116,64 +173,8 @@ pub struct TableParIterator<'a, T> { pub range: std::ops::Range, } -// Extension methods for more efficient parallel operations -impl<'a, T: RowReadable + Send + Sync + 'a> TableParIterator<'a, T> { - /// Processes the iterator in parallel with early error detection and termination. - /// - /// This method provides a parallel equivalent to the standard iterator's `try_for_each`, - /// executing the provided operation on each row concurrently while monitoring for - /// errors. If any operation fails, processing stops and the first error encountered - /// is returned. - /// - /// ## Arguments - /// - /// * `op` - A closure that takes each row and returns a [`Result`]. Must be `Send + Sync` - /// to enable safe parallel execution. - /// - /// ## Returns - /// - /// Returns `Ok(())` if all operations complete successfully, or the first error - /// encountered during parallel processing. - /// - /// # Errors - /// - /// Returns an error if any operation applied to an item returns an error. The first error encountered is returned. - pub fn try_for_each(self, op: F) -> Result<()> - where - F: Fn(T) -> Result<()> + Send + Sync, - { - let error = Arc::new(Mutex::new(None)); - - self.for_each(|item| { - if let Ok(guard) = error.lock() { - if guard.is_some() { - return; - } - } - - if let Err(e) = op(item) { - if let Ok(mut guard) = error.lock() { - if guard.is_none() { - *guard = Some(e); - } - } - } - }); - - let mutex = Arc::into_inner(error) - .ok_or_else(|| Error::LockError("Arc still has references".into()))?; - match mutex - .into_inner() - .map_err(|e| Error::LockError(format!("iterator error lock: {e}")))? - { - Some(e) => Err(e), - None => Ok(()), - } - } -} - impl ParallelIterator for TableParIterator<'_, T> { - type Item = T; + type Item = Result; fn drive_unindexed(self, consumer: C) -> C::Result where @@ -233,7 +234,7 @@ struct TableProducer<'a, T> { } impl<'a, T: RowReadable + Send + Sync> rayon::iter::plumbing::Producer for TableProducer<'a, T> { - type Item = T; + type Item = Result; type IntoIter = TableProducerIterator<'a, T>; fn into_iter(self) -> Self::IntoIter { @@ -285,7 +286,7 @@ struct TableProducerIterator<'a, T> { } impl Iterator for TableProducerIterator<'_, T> { - type Item = T; + type Item = Result; fn next(&mut self) -> Option { if self.range.start >= self.range.end { @@ -297,7 +298,11 @@ impl Iterator for TableProducerIterator<'_, T> { // Get the row directly from the table // +1 because row indices start at 1 - self.table.get(row_index.saturating_add(1)) + // + // A failing row is yielded as `Err`, so this chunk still produces exactly + // `range.len()` items — which is what `ExactSizeIterator` below promises rayon, and + // what `IndexedParallelIterator::len` reports. + self.table.get(row_index.saturating_add(1)).transpose() } fn size_hint(&self) -> (usize, Option) { @@ -319,6 +324,6 @@ impl DoubleEndedIterator for TableProducerIterator // Get the row directly from the table // +1 because row indices start at 1 - self.table.get(self.range.end.saturating_add(1)) + self.table.get(self.range.end.saturating_add(1)).transpose() } } diff --git a/dotscope/src/metadata/tables/types/read/mod.rs b/dotscope/src/metadata/tables/types/read/mod.rs index e9f14fe3..6fea7b95 100644 --- a/dotscope/src/metadata/tables/types/read/mod.rs +++ b/dotscope/src/metadata/tables/types/read/mod.rs @@ -29,6 +29,6 @@ mod traits; pub(crate) use access::TableAccess; pub use data::TableData; -pub use iter::{TableIterator, TableParIterator}; +pub use iter::{skip_unreadable, TableIterator, TableParIterator}; pub use table::MetadataTable; pub use traits::RowReadable; diff --git a/dotscope/src/metadata/tables/types/read/table.rs b/dotscope/src/metadata/tables/types/read/table.rs index 7b4d146c..3065c9f6 100644 --- a/dotscope/src/metadata/tables/types/read/table.rs +++ b/dotscope/src/metadata/tables/types/read/table.rs @@ -66,12 +66,12 @@ use crate::{ /// let table: MetadataTable = MetadataTable::new(data, 100, table_info)?; /// /// // Access specific rows -/// if let Some(first_row) = table.get(1) { +/// if let Some(first_row) = table.get(1)? { /// println!("First row ID: {}", first_row.id); /// } /// /// // Sequential iteration -/// for (index, row) in table.iter().enumerate() { +/// for (index, row) in table.iter().filter_map(skip_unreadable).enumerate() { /// println!("Row {}: ID = {}", index + 1, row.id); /// } /// # Ok(()) @@ -97,7 +97,8 @@ use crate::{ /// /// // Parallel processing with automatic error handling /// table.par_iter().try_for_each(|row| { -/// // Process each row in parallel +/// // A row that failed to parse is reported, not skipped. +/// let row = row?; /// println!("Processing row: {}", row.id); /// Ok(()) /// })?; @@ -142,10 +143,37 @@ impl<'a, T: RowReadable> MetadataTable<'a, T> { /// - The table configuration is invalid or inconsistent /// - Row size calculation fails due to invalid size parameters pub fn new(data: &'a [u8], row_count: u32, sizes: TableInfoRef) -> Result { + let row_size = T::row_size(&sizes); + + // `data` is the remainder of the whole `#~` stream, not this table's own extent, so a + // row count larger than the table really has would read rows belonging to the tables + // that follow it. Widened to `u64` because `row_count * row_size` overflows `u32` at + // counts the raw validator still permits. + let declared = u64::from(row_count) + .checked_mul(u64::from(row_size)) + .ok_or_else(|| malformed_error!("Table extent overflows"))?; + let available = u64::try_from(data.len()) + .map_err(|_| malformed_error!("Table buffer length exceeds the address width"))?; + if declared > available { + return Err(malformed_error!( + "Table declares {} rows of {} bytes ({}) but only {} bytes remain", + row_count, + row_size, + declared, + available + )); + } + + // Truncate to the declared extent so no read can cross into the next table, whatever + // the iterators do. + let data = data + .get(..usize::try_from(declared).unwrap_or(data.len())) + .unwrap_or(data); + Ok(MetadataTable { data, row_count, - row_size: T::row_size(&sizes), + row_size, sizes, _phantom: Arc::new(PhantomData), }) @@ -176,14 +204,23 @@ impl<'a, T: RowReadable> MetadataTable<'a, T> { /// /// ## Returns /// - /// Returns `Some(T)` if the row exists and can be parsed successfully, - /// or `None` if the index is out of bounds or parsing fails. - #[must_use] - pub fn get(&self, index: u32) -> Option { + /// `Ok(Some(row))` when the row exists and parses, `Ok(None)` when the index is outside + /// the table, and `Err` when the row exists but does not parse. Those are three different + /// facts and callers act on them differently — a row that will not parse is a defect in + /// the input, not an absent row. + /// + /// # Errors + /// + /// Returns an error if the row's bytes cannot be decoded. + pub fn get(&self, index: u32) -> Result> { if index == 0 || self.row_count < index { - return None; + return Ok(None); } + // The offset is derived from the index rather than carried across calls, so a row + // that fails to parse cannot desynchronise any other row. That independence is what + // lets the iterators report a failure and keep going, instead of having to stop + // because they no longer know where the next row begins. T::row_read( self.data, &mut (index as usize) @@ -192,7 +229,7 @@ impl<'a, T: RowReadable> MetadataTable<'a, T> { index, &self.sizes, ) - .ok() + .map(Some) } /// Creates a sequential iterator over all rows in the table. @@ -210,7 +247,6 @@ impl<'a, T: RowReadable> MetadataTable<'a, T> { TableIterator { table: self, current_row: 0, - current_offset: 0, } } @@ -234,7 +270,7 @@ impl<'a, T: RowReadable> MetadataTable<'a, T> { } impl<'a, T: RowReadable> IntoIterator for &'a MetadataTable<'a, T> { - type Item = T; + type Item = Result; type IntoIter = TableIterator<'a, T>; fn into_iter(self) -> Self::IntoIter { diff --git a/dotscope/src/metadata/tables/typespec/loader.rs b/dotscope/src/metadata/tables/typespec/loader.rs index 1ffcecd8..6ac8ef16 100644 --- a/dotscope/src/metadata/tables/typespec/loader.rs +++ b/dotscope/src/metadata/tables/typespec/loader.rs @@ -24,6 +24,8 @@ //! //! * [ECMA-335 Partition II, Section 22.39](https://ecma-international.org/wp-content/uploads/ECMA-335_6th_edition_june_2012.pdf) - `TypeSpec` Table +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + use crate::{ metadata::{ diagnostics::DiagnosticCategory, @@ -58,7 +60,7 @@ impl MetadataLoader for TypeSpecLoader { /// /// ## Errors /// - /// - [`crate::Error::Malformed`] - Malformed type signature in blob heap + /// - [`crate::Error::Parse`] carrying [`crate::ParseFailure::Other`] - Malformed type signature in blob heap /// - [`crate::Error::TypeNotFound`] - Referenced type cannot be resolved /// - [`crate::Error::TypeError`] - Type specification violates semantic rules fn load(&self, context: &LoaderContext) -> Result<()> { @@ -69,7 +71,10 @@ impl MetadataLoader for TypeSpecLoader { return Ok(()); }; - table.par_iter().try_for_each(|row| { + table.par_iter().enumerate().try_for_each(|(index, row)| { + let Some(row) = context.handle_row(row, index)? else { + return Ok(()); + }; let token_msg = || format!("type spec 0x{:08x}", row.token.value()); let Some(owned) = diff --git a/dotscope/src/metadata/tables/typespec/reader.rs b/dotscope/src/metadata/tables/typespec/reader.rs index 0c6650eb..fa376f0c 100644 --- a/dotscope/src/metadata/tables/typespec/reader.rs +++ b/dotscope/src/metadata/tables/typespec/reader.rs @@ -106,12 +106,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } @@ -138,12 +139,13 @@ mod tests { { for row in table.iter() { + let row = row.expect("row parses"); eval(row); } } { - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); eval(row); } } diff --git a/dotscope/src/metadata/token.rs b/dotscope/src/metadata/token.rs index 9e1c6f3f..f4bcb2cf 100644 --- a/dotscope/src/metadata/token.rs +++ b/dotscope/src/metadata/token.rs @@ -598,4 +598,65 @@ mod tests { assert!(!null_token.is_table(TableId::MethodDef)); // Null token assert!(null_token.is_table(TableId::Module)); // Null token matches table 0 } + + // The `*Ptr` indirection tables are keyed by tokens built in their readers from literal + // prefixes, while the consumers that look rows up build the same tokens independently. + // These assert both halves agree, so a lookup key can never silently address the wrong + // table — a mismatch makes every lookup miss, which reads as "the assembly has no methods" + // rather than as an error. + + #[test] + fn indirection_table_tokens_match_reader_prefixes() { + // Left side: what `Token::from_parts` produces for each indirection table. + // Right side: the literal prefix the corresponding reader keys its rows under. + assert_eq!( + Token::from_parts(TableId::FieldPtr, 1).value(), + 0x0300_0001, + "FieldPtr" + ); + assert_eq!( + Token::from_parts(TableId::MethodPtr, 1).value(), + 0x0500_0001, + "MethodPtr" + ); + assert_eq!( + Token::from_parts(TableId::EventPtr, 1).value(), + 0x1300_0001, + "EventPtr" + ); + assert_eq!( + Token::from_parts(TableId::PropertyPtr, 1).value(), + 0x1600_0001, + "PropertyPtr" + ); + } + + #[test] + fn indirection_targets_use_their_own_table_ids() { + // The tables an indirection row resolves *to* are distinct from the indirection + // tables themselves; transposing the two is the mistake this guards against. + assert_eq!(Token::from_parts(TableId::Field, 1).value(), 0x0400_0001); + assert_eq!( + Token::from_parts(TableId::MethodDef, 1).value(), + 0x0600_0001 + ); + assert_eq!(Token::from_parts(TableId::Event, 1).value(), 0x1400_0001); + assert_eq!(Token::from_parts(TableId::Property, 1).value(), 0x1700_0001); + } + + #[test] + fn indirection_and_target_tables_are_distinct() { + for (ptr, target) in [ + (TableId::FieldPtr, TableId::Field), + (TableId::MethodPtr, TableId::MethodDef), + (TableId::EventPtr, TableId::Event), + (TableId::PropertyPtr, TableId::Property), + ] { + assert_ne!( + Token::from_parts(ptr, 1).value(), + Token::from_parts(target, 1).value(), + "{ptr:?} and {target:?} must not share a token space" + ); + } + } } diff --git a/dotscope/src/metadata/typesystem/base.rs b/dotscope/src/metadata/typesystem/base.rs index 1ec1779b..26d6ea23 100644 --- a/dotscope/src/metadata/typesystem/base.rs +++ b/dotscope/src/metadata/typesystem/base.rs @@ -352,7 +352,7 @@ impl CilTypeRef { /// See the [Convenience Accessors](#convenience-accessors) section above for performance notes. #[must_use] pub fn fullname(&self) -> Option { - self.upgrade().map(|t| t.fullname()) + self.upgrade().map(|t| t.fullname().to_string()) } /// Gets a clone of the nested types collection, or [`None`] if dropped. @@ -1100,8 +1100,16 @@ pub enum CilFlavor { Pinned, /// Function pointer type with a specific method signature FnPtr { - /// The method signature this function pointer must match - signature: SignatureMethod, + /// The method signature this function pointer must match. + /// + /// Boxed deliberately, matching [`TypeSignature::FnPtr`]. An inline `SignatureMethod` + /// is ~144 bytes and, as the widest variant, sets the size of `CilFlavor` and + /// everything built on it — `SymbolicValue`, and through it every `EmValue` moved on + /// the interpreter's evaluation stack, local slot, argument slot and array element. + /// Function pointers are rare in real assemblies, so the indirection is off the hot + /// path while the size saving applies to every value. See the `EmValue` size + /// assertion in `emulation::value::emvalue`. + signature: Box, }, /// Generic parameter from a type or method definition GenericParameter { @@ -1692,7 +1700,7 @@ impl From<&TypeSignature> for CilFlavor { // Function pointer - carry the full signature TypeSignature::FnPtr(method_sig) => CilFlavor::FnPtr { - signature: (**method_sig).clone(), + signature: method_sig.clone(), }, // Single-dimensional zero-based array @@ -1770,5 +1778,5 @@ impl From<&TypeSignature> for CilFlavor { /// /// Per ECMA-335, `native int` and `native uint` (`System.IntPtr` / /// `System.UIntPtr`) are pointer-sized: 4 bytes on PE32, 8 bytes on PE32+. -/// Use [`PointerSize::from_pe`] to derive from the PE header. +/// Use [`PointerSize::from_is_64bit`] to derive from the PE header. pub use analyssa::PointerSize; diff --git a/dotscope/src/metadata/typesystem/builder.rs b/dotscope/src/metadata/typesystem/builder.rs index 87158b27..1e58b353 100644 --- a/dotscope/src/metadata/typesystem/builder.rs +++ b/dotscope/src/metadata/typesystem/builder.rs @@ -116,6 +116,7 @@ use crate::{ CompleteTypeSpec, TypeRegistry, TypeSource, }, }, + utils::{truncate_chars, LazyList}, Error::TypeError, Result, }; @@ -796,7 +797,9 @@ impl TypeBuilder { let fn_ptr_type = self.registry.get_or_create_type(&CompleteTypeSpec { token_init: self.token_init.take(), - flavor: CilFlavor::FnPtr { signature }, + flavor: CilFlavor::FnPtr { + signature: Box::new(signature), + }, namespace: String::new(), name, source: self.source.clone(), @@ -919,7 +922,7 @@ impl TypeBuilder { instantiation: SignatureMethodSpec { generic_args: vec![], }, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), generic_args: { let type_ref_list = Arc::new(boxcar::Vec::with_capacity(1)); type_ref_list.push(arg.clone().into()); @@ -991,11 +994,7 @@ impl TypeBuilder { let return_info = format!("{:?}", signature.return_type.base).replace(' ', ""); // Truncate return_info to avoid extremely long names - let return_short = if return_info.len() > 16 { - &return_info[..16] - } else { - &return_info - }; + let return_short = truncate_chars(&return_info, 16); format!("FnPtr_{calling_convention}_{param_count}_{return_short}") } @@ -1217,7 +1216,7 @@ mod tests { constraints: Arc::new(boxcar::Vec::new()), rid: 0, offset: 0, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); list_type.generic_params.push(generic_param); @@ -1243,7 +1242,13 @@ mod tests { assert_eq!(list_int_instance.generic_args.count(), 1); assert_eq!( - list_int_instance.generic_args[0].generic_args[0] + list_int_instance + .generic_args + .get(0) + .unwrap() + .generic_args + .get(0) + .unwrap() .name() .unwrap(), "Int32" @@ -1525,7 +1530,7 @@ mod tests { constraints: Arc::new(boxcar::Vec::new()), rid: 0, offset: 0, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); let value_param = Arc::new(GenericParam { @@ -1537,7 +1542,7 @@ mod tests { constraints: Arc::new(boxcar::Vec::new()), rid: 1, offset: 1, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); dict_type.generic_params.push(key_param); @@ -1565,13 +1570,25 @@ mod tests { assert_eq!(dict_instance.generic_args.count(), 2); assert_eq!( - dict_instance.generic_args[0].generic_args[0] + dict_instance + .generic_args + .get(0) + .unwrap() + .generic_args + .get(0) + .unwrap() .name() .unwrap(), "String" ); assert_eq!( - dict_instance.generic_args[1].generic_args[0] + dict_instance + .generic_args + .get(1) + .unwrap() + .generic_args + .get(0) + .unwrap() .name() .unwrap(), "Int32" @@ -1580,13 +1597,25 @@ mod tests { // With the simplified approach, we only store the resolved types // The order corresponds to the generic parameter order (0=TKey, 1=TValue) assert_eq!( - dict_instance.generic_args[0].generic_args[0] + dict_instance + .generic_args + .get(0) + .unwrap() + .generic_args + .get(0) + .unwrap() .name() .unwrap(), "String" ); // TKey -> String assert_eq!( - dict_instance.generic_args[1].generic_args[0] + dict_instance + .generic_args + .get(1) + .unwrap() + .generic_args + .get(0) + .unwrap() .name() .unwrap(), "Int32" diff --git a/dotscope/src/metadata/typesystem/mod.rs b/dotscope/src/metadata/typesystem/mod.rs index da7b628b..b13474ee 100644 --- a/dotscope/src/metadata/typesystem/mod.rs +++ b/dotscope/src/metadata/typesystem/mod.rs @@ -75,6 +75,7 @@ use crate::{ }, token::Token, }, + utils::LazyList, Error, Result, }; @@ -90,6 +91,47 @@ pub type CilTypeList = Arc>; /// while maintaining thread safety for concurrent access scenarios. pub type CilTypeRc = Arc; +/// Maximum number of ancestors [`CilType::base_chain`] will yield before stopping. +/// +/// Real inheritance chains are shallow — the deepest in the .NET base class library is well +/// under 20. This bound exists solely to terminate walks over malformed metadata, and is set +/// high enough that no legitimate assembly can reach it. It matches the limit already applied +/// to the virtual-dispatch walks in the emulation engine. +pub const MAX_BASE_WALK_DEPTH: usize = 256; + +/// Iterator over a type's inheritance chain, hardened against attacker-supplied cycles. +/// +/// Created by [`CilType::base_chain`]; see that method for why walking `base()` by hand is +/// unsafe on untrusted input. +pub struct BaseChain { + /// Next ancestor to yield, or `None` once the chain is exhausted. + current: Option, + /// Tokens already yielded, seeded with the starting type so self-extends terminates. + visited: HashSet, + /// Remaining depth budget. + remaining: usize, +} + +impl Iterator for BaseChain { + type Item = CilTypeRc; + + fn next(&mut self) -> Option { + if self.remaining == 0 { + return None; + } + let current = self.current.take()?; + // A token we have already yielded means the chain closed into a cycle. + if !self.visited.insert(current.token) { + return None; + } + self.remaining = self.remaining.saturating_sub(1); + self.current = current.base(); + Some(current) + } +} + +impl std::iter::FusedIterator for BaseChain {} + /// Represents a unified type definition combining information from `TypeDef`, `TypeRef`, and `TypeSpec` tables. /// /// `CilType` provides a complete representation of a .NET type, merging metadata from multiple @@ -160,7 +202,7 @@ pub struct CilType { /// Enclosing type for nested types - used for reverse lookup to build hierarchical names pub enclosing_type: OnceLock, /// Cached full name to avoid expensive recomputation. - fullname: RwLock>, + fullname: RwLock>>, // vtable // security // default_constructor: Option @@ -197,6 +239,7 @@ impl CilType { /// /// ```rust,no_run /// use dotscope::metadata::{ + /// method::MethodRefList, /// tables::TypeAttributes, /// typesystem::{CilType, CilFlavor}, /// token::Token, @@ -211,7 +254,7 @@ impl CilType { /// None, // No base type specified yet /// TypeAttributes::new(0x00100001), // TypeAttributes flags /// Arc::new(boxcar::Vec::new()), // Empty fields list - /// Arc::new(boxcar::Vec::new()), // Empty methods list + /// MethodRefList::new(), // Empty methods list, allocated on first push /// Some(CilFlavor::Class), // Explicit class flavor /// ); /// ``` @@ -257,9 +300,9 @@ impl CilType { interfaces: Arc::new(boxcar::Vec::new()), overwrites: Arc::new(boxcar::Vec::new()), nested_types: Arc::new(boxcar::Vec::new()), - generic_params: Arc::new(boxcar::Vec::new()), - generic_args: Arc::new(boxcar::Vec::new()), - custom_attributes: Arc::new(boxcar::Vec::new()), + generic_params: LazyList::new(), + generic_args: LazyList::new(), + custom_attributes: LazyList::new(), packing_size: OnceLock::new(), class_size: OnceLock::new(), spec: OnceLock::new(), @@ -393,6 +436,49 @@ impl CilType { } } + /// Iterates the type's inheritance chain, from the immediate base type towards the root. + /// + /// The starting type itself is **not** yielded — the first item is `self.base()`. + /// + /// # Why this exists + /// + /// [`CilType::base`] is populated directly from the TypeDef `extends` coded index, which is + /// attacker-controlled. A malformed assembly can describe a type that extends itself, or an + /// `A -> B -> A` cycle, and such an assembly still loads under + /// [`ValidationConfig::analysis`](crate::metadata::validation::ValidationConfig::analysis) + /// because that preset reports circular inheritance as a diagnostic rather than refusing the + /// file. A hand-written `while let Some(t) = t.base()` loop over such a graph never + /// terminates, and the recursive equivalent exhausts the native stack — an abort that no + /// caller can catch and that no emulation budget can interrupt. + /// + /// This iterator is the only sanctioned way to walk a base chain. It is bounded twice over: + /// by a visited set seeded with the starting type's token, so a cycle stops the walk the + /// first time it closes, and by [`MAX_BASE_WALK_DEPTH`], so even an acyclic-but-absurd chain + /// terminates. Both bounds end the iteration quietly rather than erroring, matching the + /// permissive posture the rest of the analysis path takes towards malformed input. + /// + /// # Examples + /// + /// ```rust,no_run + /// # use dotscope::metadata::typesystem::CilType; + /// # fn example(cil_type: &CilType) { + /// // Safe against attacker-supplied inheritance cycles. + /// let derives_from_exception = cil_type + /// .base_chain() + /// .any(|ancestor| &*ancestor.fullname() == "System.Exception"); + /// # } + /// ``` + #[must_use] + pub fn base_chain(&self) -> BaseChain { + let mut visited = HashSet::new(); + visited.insert(self.token); + BaseChain { + current: self.base(), + visited, + remaining: MAX_BASE_WALK_DEPTH, + } + } + /// Set the enclosing type for nested types. /// /// This method allows setting the enclosing type for nested types, establishing the @@ -656,14 +742,14 @@ impl CilType { let base_fullname = base_type.fullname(); // Direct well-known base types - if base_fullname == wellknown::names::VALUE_TYPE - || base_fullname == wellknown::names::ENUM + if &*base_fullname == wellknown::names::VALUE_TYPE + || &*base_fullname == wellknown::names::ENUM { return Some(CilFlavor::ValueType); } - if base_fullname == wellknown::names::DELEGATE - || base_fullname == wellknown::names::MULTICAST_DELEGATE + if &*base_fullname == wellknown::names::DELEGATE + || &*base_fullname == wellknown::names::MULTICAST_DELEGATE { return Some(CilFlavor::Class); // Delegates are reference types but special classes } @@ -731,19 +817,19 @@ impl CilType { let ancestor_name = ancestor.fullname(); // Check for well-known ancestor types - if ancestor_name == wellknown::names::VALUE_TYPE - || ancestor_name == wellknown::names::ENUM + if &*ancestor_name == wellknown::names::VALUE_TYPE + || &*ancestor_name == wellknown::names::ENUM { return Some(CilFlavor::ValueType); } - if ancestor_name == wellknown::names::DELEGATE - || ancestor_name == wellknown::names::MULTICAST_DELEGATE + if &*ancestor_name == wellknown::names::DELEGATE + || &*ancestor_name == wellknown::names::MULTICAST_DELEGATE { return Some(CilFlavor::Class); } - if ancestor_name == wellknown::names::OBJECT { + if &*ancestor_name == wellknown::names::OBJECT { // Reached the root - this is a reference type class return Some(CilFlavor::Class); } @@ -939,7 +1025,7 @@ impl CilType { #[must_use] pub fn is_enum(&self) -> bool { self.base() - .is_some_and(|b| b.fullname() == wellknown::names::ENUM) + .is_some_and(|b| &*b.fullname() == wellknown::names::ENUM) } /// Returns true if this type is a delegate (inherits from `System.Delegate` or `System.MulticastDelegate`). @@ -947,7 +1033,7 @@ impl CilType { pub fn is_delegate(&self) -> bool { self.base().is_some_and(|b| { let name = b.fullname(); - name == wellknown::names::MULTICAST_DELEGATE || name == wellknown::names::DELEGATE + &*name == wellknown::names::MULTICAST_DELEGATE || &*name == wellknown::names::DELEGATE }) } @@ -1086,16 +1172,16 @@ impl CilType { /// /// # Caching /// The result is cached after first computation for performance. - pub fn fullname(&self) -> String { + pub fn fullname(&self) -> Arc { if let Ok(guard) = self.fullname.read() { if let Some(cached) = guard.as_ref() { - return cached.clone(); + return Arc::clone(cached); } } - let fullname = self.compute_fullname(); + let fullname: Arc = Arc::from(self.compute_fullname()); if let Ok(mut guard) = self.fullname.write() { - *guard = Some(fullname.clone()); + *guard = Some(Arc::clone(&fullname)); } fullname } @@ -1266,37 +1352,35 @@ impl CilType { } /// Check if this type is a subtype of (inherits from) the target type + /// + /// Walks via [`CilType::base_chain`], so a cyclic `extends` graph terminates the search + /// instead of spinning forever. fn is_subtype_of(&self, target: &CilType) -> bool { - let mut current = self.base(); - while let Some(base_type) = current { - if base_type.token == target.token + self.base_chain().any(|base_type| { + base_type.token == target.token || (base_type.namespace == target.namespace && base_type.name == target.name) - { - return true; - } - current = base_type.base(); - } - false + }) } /// Check if this type implements the specified interface + /// + /// Walks via [`CilType::base_chain`] rather than recursing, so a cyclic `extends` graph + /// terminates the search instead of exhausting the native stack. fn implements_interface(&self, interface: &CilType) -> bool { - for (_, entry) in self.interfaces.iter() { - if let Some(impl_type) = entry.interface.upgrade() { - if impl_type.token == interface.token - || (impl_type.namespace == interface.namespace - && impl_type.name == interface.name) - { - return true; - } - } - } - - if let Some(base_type) = self.base() { - return base_type.implements_interface(interface); + fn declares(candidate: &CilType, interface: &CilType) -> bool { + candidate.interfaces.iter().any(|(_, entry)| { + entry.interface.upgrade().is_some_and(|impl_type| { + impl_type.token == interface.token + || (impl_type.namespace == interface.namespace + && impl_type.name == interface.name) + }) + }) } - false + declares(self, interface) + || self + .base_chain() + .any(|base_type| declares(&base_type, interface)) } /// Check if a constant value is compatible with this type @@ -1573,3 +1657,116 @@ impl fmt::Display for CilType { write!(f, "{}", self.fullname()) } } + +/// Tests that inheritance-chain walks terminate on attacker-controlled `extends` cycles. +/// +/// Each case constructs an inheritance graph that no valid assembly contains but that a +/// malformed one can describe, then walks it. A regression shows up as a hung or killed test +/// process rather than an assertion failure, since the failure mode being guarded against is +/// an infinite loop or native stack exhaustion. +#[cfg(test)] +mod tests { + use super::*; + use crate::test::builders::CilTypeBuilder; + + /// Builds a type with a distinct token so the visited set can tell instances apart. + fn typ(rid: u32, name: &str) -> CilTypeRc { + CilTypeBuilder::simple_class("Test", name) + .with_token(Token::new(0x0200_0000 | rid)) + .build() + } + + #[test] + fn base_chain_terminates_on_self_extends() { + let a = typ(1, "SelfRef"); + a.set_base(&CilTypeRef::new(&a)).expect("set_base"); + + // The starting token seeds the visited set, so the self-edge yields nothing at all. + assert_eq!(a.base_chain().count(), 0); + } + + #[test] + fn base_chain_terminates_on_two_type_cycle() { + let a = typ(1, "A"); + let b = typ(2, "B"); + a.set_base(&CilTypeRef::new(&b)).expect("set_base a -> b"); + b.set_base(&CilTypeRef::new(&a)).expect("set_base b -> a"); + + // Walking from A yields B, then stops: A's token is already in the visited set. + let chain: Vec = a.base_chain().map(|t| t.token).collect(); + assert_eq!(chain, vec![b.token]); + } + + #[test] + fn base_chain_terminates_on_three_type_cycle() { + let a = typ(1, "A"); + let b = typ(2, "B"); + let c = typ(3, "C"); + a.set_base(&CilTypeRef::new(&b)).expect("set_base a -> b"); + b.set_base(&CilTypeRef::new(&c)).expect("set_base b -> c"); + c.set_base(&CilTypeRef::new(&a)).expect("set_base c -> a"); + + let chain: Vec = a.base_chain().map(|t| t.token).collect(); + assert_eq!(chain, vec![b.token, c.token]); + } + + #[test] + fn base_chain_yields_acyclic_ancestors_in_order() { + let object = typ(1, "Object"); + let mid = typ(2, "Mid"); + let leaf = typ(3, "Leaf"); + mid.set_base(&CilTypeRef::new(&object)) + .expect("set_base mid -> object"); + leaf.set_base(&CilTypeRef::new(&mid)) + .expect("set_base leaf -> mid"); + + let chain: Vec = leaf.base_chain().map(|t| t.token).collect(); + assert_eq!(chain, vec![mid.token, object.token]); + } + + #[test] + fn base_chain_is_bounded_even_without_a_cycle() { + // A chain longer than the depth budget, with every token distinct so only the depth + // bound can stop it. + let types: Vec = (0..MAX_BASE_WALK_DEPTH + 50) + .map(|i| { + #[allow(clippy::cast_possible_truncation)] + typ(i as u32 + 1, "Deep") + }) + .collect(); + for pair in types.windows(2) { + if let [derived, base] = pair { + derived.set_base(&CilTypeRef::new(base)).expect("set_base"); + } + } + + let first = types.first().expect("non-empty"); + assert_eq!(first.base_chain().count(), MAX_BASE_WALK_DEPTH); + } + + #[test] + fn is_subtype_of_terminates_on_cycle() { + let a = typ(1, "A"); + let b = typ(2, "B"); + let unrelated = typ(3, "Unrelated"); + a.set_base(&CilTypeRef::new(&b)).expect("set_base a -> b"); + b.set_base(&CilTypeRef::new(&a)).expect("set_base b -> a"); + + // Searching for a type that is nowhere in the cycle is the non-terminating case. + assert!(!a.is_subtype_of(&unrelated)); + assert!(a.is_subtype_of(&b)); + } + + #[test] + fn implements_interface_terminates_on_cycle() { + let a = typ(1, "A"); + let b = typ(2, "B"); + let iface = CilTypeBuilder::interface("Test", "IUnimplemented") + .with_token(Token::new(0x0200_0003)) + .build(); + a.set_base(&CilTypeRef::new(&b)).expect("set_base a -> b"); + b.set_base(&CilTypeRef::new(&a)).expect("set_base b -> a"); + + assert!(!a.implements_interface(&iface)); + } +} diff --git a/dotscope/src/metadata/typesystem/primitives.rs b/dotscope/src/metadata/typesystem/primitives.rs index 8b55de7a..befb0ac8 100644 --- a/dotscope/src/metadata/typesystem/primitives.rs +++ b/dotscope/src/metadata/typesystem/primitives.rs @@ -1330,7 +1330,7 @@ impl CilPrimitive { /// /// # Errors /// Returns [`crate::Error::TypeNotPrimitive`] if the primitive type is invalid. - /// Returns [`crate::Error::OutOfBounds`] or other errors if the blob data is insufficient or invalid. + /// Returns [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] or other errors if the blob data is insufficient or invalid. pub fn from_blob(p_type: u8, blob: &[u8]) -> Result { Ok(CilPrimitive { kind: CilPrimitiveKind::from_byte(p_type)?, diff --git a/dotscope/src/metadata/typesystem/registry.rs b/dotscope/src/metadata/typesystem/registry.rs index 79ea8745..744a2fc2 100644 --- a/dotscope/src/metadata/typesystem/registry.rs +++ b/dotscope/src/metadata/typesystem/registry.rs @@ -65,6 +65,7 @@ //! //! ```rust,no_run //! use dotscope::metadata::typesystem::{TypeRegistry, CilType, TypeSource}; +//! use dotscope::metadata::method::MethodRefList; //! use dotscope::metadata::tables::TypeAttributes; //! use dotscope::metadata::identity::AssemblyIdentity; //! use dotscope::metadata::token::Token; @@ -83,7 +84,7 @@ //! None, // No base type yet //! TypeAttributes::new(0x00100001), // Public class //! Arc::new(boxcar::Vec::new()), // Empty fields -//! Arc::new(boxcar::Vec::new()), // Empty methods +//! MethodRefList::new(), // Empty methods, allocated on first push //! None, // Flavor will be computed //! ); //! @@ -106,9 +107,12 @@ //! - Generic type instantiations //! - Cross-assembly type resolution -use std::sync::{ - atomic::{AtomicU32, Ordering}, - Arc, +use std::{ + collections::BTreeSet, + sync::{ + atomic::{AtomicU32, Ordering}, + Arc, + }, }; use crossbeam_skiplist::SkipMap; @@ -127,6 +131,7 @@ use crate::{ CilTypeReference, PointerSize, TypeSignatureHash, }, }, + utils::LazyList, Error::TypeNotFound, Result, }; @@ -377,6 +382,12 @@ struct SourceRegistry { assembly_refs: DashMap, /// File references indexed by their metadata tokens files: DashMap, + /// Memoised `TypeSource` per defining-assembly token. + /// + /// Deriving an `AssemblyIdentity` from an `Assembly` row rebuilds its name, culture and + /// public-key blob. Every locally-defined type resolves to the same one, so it is + /// derived once per assembly rather than once per type. + assembly_sources: DashMap, } impl SourceRegistry { @@ -393,6 +404,7 @@ impl SourceRegistry { module_refs: DashMap::new(), assembly_refs: DashMap::new(), files: DashMap::new(), + assembly_sources: DashMap::new(), } } @@ -429,7 +441,17 @@ impl SourceRegistry { TypeSource::AssemblyRef(assembly_ref.token) } CilTypeReference::Assembly(assembly) => { - TypeSource::Assembly(AssemblyIdentity::from_assembly(assembly)) + // Reached once per type that names the defining assembly as its source. + // `from_assembly` is worse than a clone — it re-clones the name and culture + // and rebuilds the public-key blob — and the result is identical for every + // one of them, so it is derived once and reused. + if let Some(cached) = self.assembly_sources.get(&assembly.token) { + return cached.value().clone(); + } + + let source = TypeSource::Assembly(AssemblyIdentity::from_assembly(assembly)); + self.assembly_sources.insert(assembly.token, source.clone()); + source } CilTypeReference::File(file) => { self.files.insert(file.token, file.clone()); @@ -559,16 +581,41 @@ pub struct TypeRegistry { /// Identity of the assembly this registry represents current_assembly: AssemblyIdentity, /// Secondary index: types grouped by their origin source - types_by_source: DashMap>, + /// + /// The four indices below hold `BTreeSet`s rather than `Vec`s because the TypeRef + /// redirect sweep removes a token from each of them per resolved TypeRef, and a linear + /// `retain` over lists that keep growing makes that sweep quadratic in TypeRef count. + /// Ordered rather than hashed so that iteration — which decides *which* type a + /// duplicated name resolves to — stays deterministic across runs. + types_by_source: DashMap>, /// Secondary index: types indexed by full name (namespace.name) - types_by_fullname: DashMap>, + types_by_fullname: DashMap, BTreeSet>, /// Secondary index: types indexed by simple name (may have duplicates) - types_by_name: DashMap>, + types_by_name: DashMap>, /// Secondary index: types grouped by namespace - types_by_namespace: DashMap>, + types_by_namespace: DashMap>, + /// Secondary index: nested types keyed by the last component of their fullname. + /// + /// A TypeRef may name a nested type by its inner name alone (`Enumerator`2`) while the + /// TypeDef carries the full path (`NS.Partitioner/Enumerator`2`). Resolving that without + /// an index means an `ends_with` scan of every registered fullname on *every* exact-name + /// miss, and the custom-attribute parser drives two such misses per fixed argument. + types_by_nested_suffix: DashMap, BTreeSet>, /// Registered external TypeRegistries for cross-assembly type resolution /// Maps AssemblyIdentity to external TypeRegistry for cross-assembly lookups external_registries: DashMap>, + /// Secondary index: MethodDef token to the token of its declaring type. + /// + /// Without this, `TypeResolver::declaring_type` answers the MethodDef case by scanning + /// every type and, inside that, every method — upgrading a `Weak` per candidate. That + /// lookup sits on the emulator's hottest paths (`call`/`callvirt`/`newobj`), so an O(1) + /// answer here is what keeps dispatch from degrading to O(total methods in assembly). + method_to_type: DashMap, + /// Secondary index: FieldDef token to the token of its declaring type. + /// + /// Same rationale as [`method_to_type`](Self::method_to_type), for the `ldsfld`/`stsfld` + /// paths. + field_to_type: DashMap, } impl TypeRegistry { @@ -625,7 +672,10 @@ impl TypeRegistry { types_by_fullname: DashMap::new(), types_by_name: DashMap::new(), types_by_namespace: DashMap::new(), + types_by_nested_suffix: DashMap::new(), external_registries: DashMap::new(), + method_to_type: DashMap::new(), + field_to_type: DashMap::new(), }; registry.initialize_primitives()?; @@ -694,7 +744,7 @@ impl TypeRegistry { None, TypeAttributes::ZERO, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), Some(flavor), )); @@ -780,24 +830,41 @@ impl TypeRegistry { self.types_by_source .entry(source) .or_default() - .push(type_rc.token); + .insert(type_rc.token); if !type_rc.namespace.is_empty() { self.types_by_namespace .entry(type_rc.namespace.clone()) .or_default() - .push(type_rc.token); + .insert(type_rc.token); } self.types_by_name .entry(type_rc.name.clone()) .or_default() - .push(type_rc.token); + .insert(type_rc.token); + let fullname = type_rc.fullname(); + self.index_nested_suffix(&fullname, type_rc.token); self.types_by_fullname - .entry(type_rc.fullname()) + .entry(fullname) .or_default() - .push(type_rc.token); + .insert(type_rc.token); + } + + /// Indexes a nested type under the last component of its fullname. + /// + /// Non-nested names are skipped: they are already answered exactly by + /// `types_by_fullname`, and indexing them would make the suffix map a second copy of it. + fn index_nested_suffix(&self, fullname: &str, token: Token) { + if let Some(last) = fullname.rsplit('/').next() { + if last.len() != fullname.len() { + self.types_by_nested_suffix + .entry(Arc::from(last)) + .or_default() + .insert(token); + } + } } /// Insert a `CilType` into the registry @@ -828,7 +895,7 @@ impl TypeRegistry { None, TypeAttributes::ZERO, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), None, )); @@ -854,7 +921,7 @@ impl TypeRegistry { None, TypeAttributes::ZERO, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), Some(flavor), )); @@ -1014,7 +1081,7 @@ impl TypeRegistry { } } - if let Some(tokens) = self.types_by_fullname.get(&fullname) { + if let Some(tokens) = self.types_by_fullname.get(fullname.as_str()) { if let Some(&token) = tokens.first() { return self.types.get(&token).map(|res| res.value().clone()); } @@ -1157,37 +1224,33 @@ impl TypeRegistry { } } } else { - // Fallback: try suffix matching for nested types - // This handles cases where TypeRef has incomplete name (e.g., "DynamicPartitionEnumerator_Abstract`2") - // but TypeDef has complete name (e.g., "Partitioner/DynamicPartitionEnumerator_Abstract`2") - let mut candidates = Vec::new(); - for key_entry in &self.types_by_fullname { - let key = key_entry.key(); - let tokens = key_entry.value(); - - // Check if this key ends with our target fullname (handling nested types) - if key.ends_with(fullname) && key != fullname { - // Additional check: ensure it's a proper nested type match (contains '/') - if key.contains('/') { - // Check tokens for TypeDef or TypeSpec - for token in tokens { - if let Some(entry) = self.types.get(token) { - let type_rc = entry.value().clone(); - if type_rc.token.is_table(TableId::TypeDef) - || type_rc.token.is_table(TableId::TypeSpec) - { - candidates.push(type_rc); - break; // Take first TypeDef/TypeSpec found - } - } - } + // Fallback: suffix matching for nested types. A TypeRef may carry the inner name + // alone (`DynamicPartitionEnumerator_Abstract`2`) where the TypeDef carries the + // whole path (`Partitioner/DynamicPartitionEnumerator_Abstract`2`). + // + // `types_by_nested_suffix` narrows this to the types whose fullname ends in the + // same *component*; the `ends_with` below then confirms the full requested + // suffix, so a multi-component request still matches exactly as before. Only + // matches that cut a component in half — `Bar` against `A/FooBar` — are no + // longer accepted, and those were never nested-type resolutions. + let requested_suffix = fullname.rsplit('/').next().unwrap_or(fullname); + if let Some(tokens) = self.types_by_nested_suffix.get(requested_suffix) { + for token in tokens.value() { + let Some(entry) = self.types.get(token) else { + continue; + }; + let type_rc = entry.value().clone(); + if !type_rc.token.is_table(TableId::TypeDef) + && !type_rc.token.is_table(TableId::TypeSpec) + { + continue; } - } - } - // Return first candidate (could be enhanced with disambiguation logic) - if let Some(candidate) = candidates.first() { - return Some(candidate.clone()); + let candidate = type_rc.fullname(); + if candidate.ends_with(fullname) && &*candidate != fullname { + return Some(type_rc); + } + } } } @@ -1467,7 +1530,7 @@ impl TypeRegistry { None, flags, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), Some(spec.flavor.clone()), )); @@ -1555,7 +1618,7 @@ impl TypeRegistry { instantiation: SignatureMethodSpec { generic_args: vec![], }, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), generic_args: { let type_ref_list = Arc::new(boxcar::Vec::with_capacity(1)); type_ref_list.push(arg_type.clone().into()); @@ -1585,6 +1648,9 @@ impl TypeRegistry { } /// Get all types in the registry + /// + /// Prefer [`iter`](Self::iter) where a borrow suffices: this clones an `Arc` per type, so + /// each call is a heap allocation plus two atomic RMWs per type in the assembly. pub fn all_types(&self) -> Vec { self.types .iter() @@ -1592,6 +1658,59 @@ impl TypeRegistry { .collect() } + /// Returns the token of the type declaring `method_token`, or `None` if no type claims it. + /// + /// The first lookup for a given token scans the registry; the answer is then memoised, so + /// repeated dispatch to the same method — the emulator's actual access pattern, since a hot + /// loop calls the same handful of methods — is O(1) thereafter. + /// + /// Memoisation is on demand rather than built once at load because a type's method list is + /// populated *after* the type itself is registered; an eagerly built index could be + /// permanently cached while incomplete. + #[must_use] + pub fn declaring_type_token_of_method(&self, method_token: Token) -> Option { + if let Some(cached) = self.method_to_type.get(&method_token) { + return Some(*cached); + } + + for entry in self.types.iter() { + let type_token = *entry.key(); + for (_, method_ref) in entry.value().methods.iter() { + if let Some(method) = method_ref.upgrade() { + if method.token == method_token { + self.method_to_type.insert(method_token, type_token); + return Some(type_token); + } + } + } + } + + None + } + + /// Returns the token of the type declaring `field_token`, or `None` if no type claims it. + /// + /// Memoised on the same terms as + /// [`declaring_type_token_of_method`](Self::declaring_type_token_of_method). + #[must_use] + pub fn declaring_type_token_of_field(&self, field_token: Token) -> Option { + if let Some(cached) = self.field_to_type.get(&field_token) { + return Some(*cached); + } + + for entry in self.types.iter() { + let type_token = *entry.key(); + for (_, field) in entry.value().fields.iter() { + if field.token == field_token { + self.field_to_type.insert(field_token, type_token); + return Some(type_token); + } + } + } + + None + } + /// Get types from a specific source /// /// ## Arguments @@ -1755,28 +1874,32 @@ impl TypeRegistry { if let Some(external) = original_typeref.external() { let source = self.register_source(external); if let Some(mut list) = self.types_by_source.get_mut(&source) { - list.retain(|&token| token != typeref_token); + list.remove(&typeref_token); } } else { let current_source = self.current_assembly_source(); if let Some(mut list) = self.types_by_source.get_mut(¤t_source) { - list.retain(|&token| token != typeref_token); + list.remove(&typeref_token); } } if !original_typeref.namespace.is_empty() { if let Some(mut list) = self.types_by_namespace.get_mut(&original_typeref.namespace) { - list.retain(|&token| token != typeref_token); + list.remove(&typeref_token); } } if let Some(mut list) = self.types_by_name.get_mut(&original_typeref.name) { - list.retain(|&token| token != typeref_token); + list.remove(&typeref_token); } let old_fullname = original_typeref.fullname(); - if let Some(mut list) = self.types_by_fullname.get_mut(&old_fullname) { - list.retain(|&token| token != typeref_token); + let old_suffix = old_fullname.rsplit('/').next().unwrap_or(&old_fullname); + if let Some(mut list) = self.types_by_nested_suffix.get_mut(old_suffix) { + list.remove(&typeref_token); + } + if let Some(mut list) = self.types_by_fullname.get_mut(&*old_fullname) { + list.remove(&typeref_token); } // Add TypeRef token to the TypeDef's indexes (new metadata) @@ -1785,25 +1908,27 @@ impl TypeRegistry { self.types_by_source .entry(source) .or_default() - .push(typeref_token); + .insert(typeref_token); } if !resolved_typedef.namespace.is_empty() { self.types_by_namespace .entry(resolved_typedef.namespace.clone()) .or_default() - .push(typeref_token); + .insert(typeref_token); } self.types_by_name .entry(resolved_typedef.name.clone()) .or_default() - .push(typeref_token); + .insert(typeref_token); + let new_fullname = resolved_typedef.fullname(); + self.index_nested_suffix(&new_fullname, typeref_token); self.types_by_fullname - .entry(resolved_typedef.fullname()) + .entry(new_fullname) .or_default() - .push(typeref_token); + .insert(typeref_token); self.types.insert(typeref_token, resolved_typedef.clone()); true @@ -1823,15 +1948,17 @@ impl TypeRegistry { /// - Before InheritanceResolver runs (which needs to look up nested types) pub fn build_fullnames(&self) { self.types_by_fullname.clear(); + self.types_by_nested_suffix.clear(); for entry in &self.types { let type_rc = entry.value(); let current_fullname = type_rc.fullname(); + self.index_nested_suffix(¤t_fullname, type_rc.token); self.types_by_fullname .entry(current_fullname) .or_default() - .push(type_rc.token); + .insert(type_rc.token); } } } @@ -1854,6 +1981,79 @@ mod tests { AssemblyFlags, AssemblyRef, AssemblyRefHash, File, FileAttributes, Module, ModuleRef, }; + /// Builds a type whose *name* already contains the nested path, which is the shape + /// `compute_fullname` produces for a type with an enclosing type. + fn nested_type(token: Token, namespace: &str, name: &str) -> CilTypeRc { + Arc::new(CilType::new( + token, + namespace.to_string(), + name.to_string(), + None, + None, + TypeAttributes::ZERO, + Arc::new(boxcar::Vec::new()), + LazyList::new(), + None, + )) + } + + /// A TypeRef may name a nested type by its inner name alone while the TypeDef carries + /// the whole path. Resolving that is what the suffix index replaced a full scan of every + /// registered fullname with; the resolution itself must be unchanged. + #[test] + fn nested_types_resolve_by_their_inner_name() { + let identity = AssemblyIdentity::parse("TestAssembly, Version=1.0.0.0").unwrap(); + let registry = TypeRegistry::new(identity).unwrap(); + + let token = Token::new(0x0200_1234); + registry.insert(&nested_type( + token, + "NS", + "Partitioner/DynamicPartitionEnumerator`2", + )); + + // The inner name alone resolves to the nested type... + let found = registry + .get_by_fullname("DynamicPartitionEnumerator`2", false) + .expect("nested type must resolve by its inner name"); + assert_eq!(found.token, token); + + // ...as does a multi-component suffix. + let found = registry + .get_by_fullname("Partitioner/DynamicPartitionEnumerator`2", false) + .expect("nested type must resolve by a multi-component suffix"); + assert_eq!(found.token, token); + + // The exact fullname still goes through the exact index, not the fallback. + assert!(registry + .get_by_fullname("NS.Partitioner/DynamicPartitionEnumerator`2", false) + .is_some()); + + // An unrelated name must not resolve. + assert!(registry.get_by_fullname("SomethingElse", false).is_none()); + } + + /// The fullname index is set-valued so the TypeRef redirect sweep can remove in `O(log n)` + /// instead of a linear `retain` over lists that keep growing. Registering the + /// same token twice must therefore not duplicate it. + #[test] + fn fullname_index_holds_each_token_once() { + let identity = AssemblyIdentity::parse("TestAssembly, Version=1.0.0.0").unwrap(); + let registry = TypeRegistry::new(identity).unwrap(); + + let token = Token::new(0x0200_4321); + let type_rc = nested_type(token, "NS", "Outer/Inner"); + registry.insert(&type_rc); + registry.build_fullnames(); + registry.build_fullnames(); + + let tokens = registry + .types_by_fullname + .get("NS.Outer/Inner") + .expect("fullname must be indexed"); + assert_eq!(tokens.value().len(), 1); + } + #[test] fn test_registry_primitives() { let test_identity = AssemblyIdentity::parse("TestAssembly, Version=1.0.0.0").unwrap(); @@ -2046,7 +2246,7 @@ mod tests { None, TypeAttributes::ZERO, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), Some(CilFlavor::Class), )); @@ -2077,7 +2277,7 @@ mod tests { generation: 0, encbaseid: None, imports: Vec::new(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); let module_ref = Arc::new(ModuleRef { @@ -2085,7 +2285,7 @@ mod tests { name: "ReferenceModule".to_string(), rid: 0, offset: 0, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); let assembly_ref = Arc::new(AssemblyRef { @@ -2105,7 +2305,7 @@ mod tests { os_major_version: AtomicU32::new(0), os_minor_version: AtomicU32::new(0), processor: AtomicU32::new(0), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); let file = Arc::new(File { @@ -2115,7 +2315,7 @@ mod tests { rid: 0, offset: 0, hash_value: AssemblyRefHash::new(&[0xCC, 0xCC]).unwrap(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); let module_source = registry.register_source(&CilTypeReference::Module(module.clone())); diff --git a/dotscope/src/metadata/typesystem/resolver.rs b/dotscope/src/metadata/typesystem/resolver.rs index bf689a6d..53b9289d 100644 --- a/dotscope/src/metadata/typesystem/resolver.rs +++ b/dotscope/src/metadata/typesystem/resolver.rs @@ -151,6 +151,7 @@ use crate::{ CompleteTypeSpec, TypeRegistry, TypeSource, }, }, + utils::{truncate_chars, LazyList}, Error::{RecursionLimit, TypeError, TypeMissingParent, TypeNotFound}, Result, }; @@ -522,11 +523,7 @@ impl TypeResolver { let return_info = format!("{:?}", signature.return_type.base).replace(' ', ""); // Truncate return_info to avoid extremely long names - let return_short = if return_info.len() > 16 { - &return_info[..16] - } else { - &return_info - }; + let return_short = truncate_chars(&return_info, 16); format!("FnPtr_{calling_convention}_{param_count}_{return_short}") } @@ -764,7 +761,7 @@ impl TypeResolver { let fnptr_type = self.registry.get_or_create_type(&CompleteTypeSpec { token_init: self.token_init.take(), flavor: CilFlavor::FnPtr { - signature: *fn_ptr.clone(), + signature: fn_ptr.clone(), }, namespace: String::new(), name, @@ -854,7 +851,7 @@ impl TypeResolver { instantiation: SignatureMethodSpec { generic_args: vec![], }, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), generic_args: { let type_ref_list = Arc::new(boxcar::Vec::with_capacity(1)); type_ref_list.push(arg_type.into()); @@ -1219,7 +1216,7 @@ mod tests { constraints: Arc::new(boxcar::Vec::new()), rid: 1, offset: 1, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); list_type.generic_params.push(type_param); @@ -1238,7 +1235,15 @@ mod tests { assert_eq!(list_int.generic_args.count(), 1); assert_eq!( - list_int.generic_args[0].generic_args[0].name().unwrap(), + list_int + .generic_args + .get(0) + .unwrap() + .generic_args + .get(0) + .unwrap() + .name() + .unwrap(), "Int32" ); } diff --git a/dotscope/src/metadata/validation/config.rs b/dotscope/src/metadata/validation/config.rs index 4b2531e2..69776c2a 100644 --- a/dotscope/src/metadata/validation/config.rs +++ b/dotscope/src/metadata/validation/config.rs @@ -229,8 +229,14 @@ impl ValidationConfig { /// Creates a disabled validation configuration. /// /// **Warning**: This disables ALL validation checks, including basic structural - /// validation. Use only when you absolutely trust the assembly format. Malformed - /// assemblies may cause panics or undefined behavior. + /// validation. Use only when you already trust the assembly format. + /// + /// It does not disable bounds checking: parsers are individually bounds-checked and the + /// crate denies `panic`, `unwrap_used`, `expect_used`, `indexing_slicing` and + /// `arithmetic_side_effects`, so a malformed assembly still yields an `Err` rather than a + /// panic. What is lost is *semantic* rejection — a file with contradictory metadata is + /// accepted and analysed as if it were coherent, so anything derived from it (type + /// hierarchies, signatures, method bodies) may be wrong without any error being reported. /// /// # Returns /// @@ -244,9 +250,9 @@ impl ValidationConfig { /// /// # Risks /// - /// - No protection against malformed metadata - /// - Potential for crashes on invalid data + /// - No protection against semantically malformed metadata /// - Silent acceptance of ECMA-335 violations + /// - Analysis results derived from incoherent metadata, with no diagnostic /// /// # Field values /// diff --git a/dotscope/src/metadata/validation/result.rs b/dotscope/src/metadata/validation/result.rs index 4b32a3d4..502a3ee0 100644 --- a/dotscope/src/metadata/validation/result.rs +++ b/dotscope/src/metadata/validation/result.rs @@ -370,12 +370,23 @@ impl ValidationResult { let errors = self.errors().into_iter().cloned().collect::>(); let error_count = errors.len(); - let validator_names: Vec<_> = failures.iter().map(|f| f.validator_name()).collect(); + // Name *and* message. The errors are carried in `errors` either way, but the + // summary is what reaches a caller that only prints the error, so a name alone + // makes a failure report which validator objected without saying to what. + // The error's own Display already names the validator, so it is used alone rather + // than prefixed with the name again. + let details: Vec = failures + .iter() + .map(|f| match f.error() { + Some(err) => err.to_string(), + None => f.validator_name().to_string(), + }) + .collect(); let summary = format!( "{} of {} validators failed: {}", error_count, self.validator_count, - validator_names.join(", ") + details.join("; ") ); Err(Error::ValidationStage2Failed { diff --git a/dotscope/src/metadata/validation/scanner.rs b/dotscope/src/metadata/validation/scanner.rs index 1871c269..93629011 100644 --- a/dotscope/src/metadata/validation/scanner.rs +++ b/dotscope/src/metadata/validation/scanner.rs @@ -309,6 +309,13 @@ impl ReferenceScanner { dispatch_table_type!(table_id, |RawType| { if let Some(table) = tables.table::() { for row in table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let token = Token::new(table_token_base | row.rid); self.valid_tokens.insert(token); } @@ -333,6 +340,13 @@ impl ReferenceScanner { if let Some(table) = tables.table::() { let token_base = u32::from(table_id.token_type()) << 24; for row in table { + let row = match row { + Ok(row) => row, + Err(e) => { + log::warn!("skipping unreadable metadata row: {e}"); + continue; + } + }; let from_token = Token::new(token_base | row.rid); self.extract_row_references(table_id, from_token, &row); } diff --git a/dotscope/src/metadata/validation/validators/owned/constraints/types.rs b/dotscope/src/metadata/validation/validators/owned/constraints/types.rs index 40021dd5..c6f06b53 100644 --- a/dotscope/src/metadata/validation/validators/owned/constraints/types.rs +++ b/dotscope/src/metadata/validation/validators/owned/constraints/types.rs @@ -261,8 +261,8 @@ impl OwnedTypeConstraintValidator { // Allow System value types and enums let type_name = constraint_type.fullname(); if type_name.starts_with("System.") - || type_name == "System.ValueType" - || type_name == "System.Enum" + || &*type_name == "System.ValueType" + || &*type_name == "System.Enum" { Ok(()) } else { diff --git a/dotscope/src/metadata/validation/validators/owned/relationships/circularity.rs b/dotscope/src/metadata/validation/validators/owned/relationships/circularity.rs index 32193103..3eb8d8fa 100644 --- a/dotscope/src/metadata/validation/validators/owned/relationships/circularity.rs +++ b/dotscope/src/metadata/validation/validators/owned/relationships/circularity.rs @@ -74,6 +74,7 @@ use rustc_hash::{FxHashMap, FxHashSet}; use crate::{ metadata::{ + tables::{skip_unreadable, TypeDefRaw}, typesystem::CilType, validation::{ context::{OwnedValidationContext, ValidationContext}, @@ -119,6 +120,64 @@ impl OwnedCircularityValidator { Self } + /// Detects a TypeDef whose `extends` names the row itself. + /// + /// This reads the raw `extends` column rather than walking `base()`, and it has to: + /// [`InheritanceResolver`] deliberately **elides** a self-referential base instead of + /// recording it, so that no consumer can ever observe the trivial inheritance cycle. An + /// edge that was never recorded is structurally invisible to + /// [`validate_inheritance_cycles`](Self::validate_inheritance_cycles), which walks the + /// resolved graph — so without this check the most easily forged inheritance cycle in the + /// format would load silently and be reported by nothing. + /// + /// Longer cycles (`A -> B -> A`) are *not* elided and are caught by the graph walk. Only + /// the self-edge needs reading back from the metadata. + /// + /// The two halves are a matched pair: if the elision in [`InheritanceResolver`] is ever + /// removed, this becomes redundant rather than wrong, and the graph walk will report the + /// same condition. + /// + /// # Arguments + /// + /// * `context` - Owned validation context, used here for its raw metadata tables + /// + /// # Returns + /// + /// * `Ok(())` - No TypeDef extends itself + /// * `Err(`[`crate::Error::ValidationOwnedFailed`]`)` - A self-referential base was found + /// + /// [`InheritanceResolver`]: crate::metadata::loader + fn validate_self_referential_bases(&self, context: &OwnedValidationContext) -> Result<()> { + let Some(tables) = context.object().tables() else { + return Ok(()); + }; + let Some(typedefs) = tables.table::() else { + return Ok(()); + }; + + for row in typedefs.iter().filter_map(skip_unreadable) { + // `extends` is a TypeDefOrRef coded index; row 0 means "no base type". + if row.extends.token.row() == 0 { + continue; + } + + // A self-extends is only expressible as a TypeDef pointing at its own row. The + // table byte is part of the token, so a TypeRef or TypeSpec can never compare + // equal here even when it resolves to this same type. + if row.extends.token == row.token { + return Err(Error::ValidationOwnedFailed { + validator: self.name().to_string(), + message: format!( + "Circular inheritance relationship detected: type (token 0x{:08X}) extends itself", + row.token.value() + ), + }); + } + } + + Ok(()) + } + /// Validates inheritance cycles across type relationships. /// /// Detects circular inheritance patterns where types form cycles through their @@ -432,6 +491,8 @@ impl OwnedCircularityValidator { impl OwnedValidator for OwnedCircularityValidator { fn validate_owned(&self, context: &OwnedValidationContext) -> Result<()> { + // First, because it is the one inheritance cycle the graph walk below cannot see. + self.validate_self_referential_bases(context)?; self.validate_inheritance_cycles(context)?; self.validate_interface_implementation_cycles(context)?; self.validate_cross_reference_cycles(context)?; @@ -465,9 +526,13 @@ mod tests { use crate::{ metadata::validation::ValidationConfig, test::{ - factories::validation::circularity::owned_circularity_validator_file_factory, - owned_validator_test, + factories::validation::circularity::{ + create_assembly_with_self_referential_type, + owned_circularity_validator_file_factory, + }, + owned_validator_test, TestAssemblySource, }, + CilObject, }; #[test] @@ -487,4 +552,63 @@ mod tests { |context| validator.validate_owned(context), ) } + + /// The self-edge really is absent from the resolved graph. + /// + /// This is the half of the contract that `test_owned_circularity_validator` cannot see. + /// That test only asserts the validator *rejects* the assembly — which it would also do if + /// the edge were present and the ordinary graph walk caught it. Pinning `base() == None` + /// here is what makes the raw-column check in + /// [`validate_self_referential_bases`](OwnedCircularityValidator::validate_self_referential_bases) + /// demonstrably necessary rather than merely redundant, and it pins + /// `InheritanceResolver`'s two other claims at the same time: that a self-extends does not + /// fail the load, and that the cycle never reaches a consumer. + /// + /// Without this, deleting the elision and deleting the raw check would cancel out and no + /// test would notice. + #[test] + #[cfg(not(feature = "skip-expensive-tests"))] + fn self_referential_base_is_elided_from_the_resolved_graph() -> Result<()> { + let assembly = create_assembly_with_self_referential_type()?; + let TestAssemblySource::Memory(data) = &assembly.source else { + return Err(Error::Other( + "factory is expected to produce an in-memory assembly".to_string(), + )); + }; + + // Claim 1: a self-referential `extends` does not fail the load. + let object = CilObject::from_mem_with_validation(data.clone(), ValidationConfig::minimal()) + .map_err(|e| { + Error::Other(format!("a self-extends assembly must still load, got: {e}")) + })?; + + // Claim 2: the type exists, and its base was left unset rather than pointing at itself. + let self_ref = object + .types() + .iter() + .find(|entry| entry.value().name == "SelfReferentialType") + .map(|entry| entry.value().clone()) + .ok_or_else(|| Error::Other("SelfReferentialType was not registered".to_string()))?; + + assert!( + self_ref.base().is_none(), + "the self-referential base must be elided, not recorded; found a base of {:?}", + self_ref.base().map(|b| b.token) + ); + + // Claim 3: and because it is elided, walking the graph finds nothing — which is + // exactly why the raw-column check has to exist. + let validator = OwnedCircularityValidator::new(); + let mut visited = FxHashSet::default(); + let mut visiting = FxHashSet::default(); + validator + .check_inheritance_cycle_relationships(&self_ref, &mut visited, &mut visiting) + .map_err(|e| { + Error::Other(format!( + "the graph walk unexpectedly saw the elided self-edge: {e}" + )) + })?; + + Ok(()) + } } diff --git a/dotscope/src/metadata/validation/validators/owned/relationships/dependency.rs b/dotscope/src/metadata/validation/validators/owned/relationships/dependency.rs index cfa52c39..824a024a 100644 --- a/dotscope/src/metadata/validation/validators/owned/relationships/dependency.rs +++ b/dotscope/src/metadata/validation/validators/owned/relationships/dependency.rs @@ -321,7 +321,16 @@ impl OwnedDependencyValidator { for type_entry in context.target_assembly_types() { // Validate inheritance ordering if let Some(base_type) = type_entry.base() { - // Check for self-referential inheritance + // Backstop for a self-referential base reached through some path other than + // the TypeDef `extends` column — a synthetic type built by + // `typesystem::resolver`, say. + // + // It does *not* catch the metadata condition. `InheritanceResolver` elides a + // self-referential `extends` rather than recording it, so for a type loaded + // from metadata `base()` is `None` here and this comparison cannot fire. + // `OwnedCircularityValidator::validate_self_referential_bases` reads the raw + // column and is what actually reports that case. + // // IMPORTANT: Only flag self-reference if both types are from the same assembly. // Tokens are only unique within an assembly, so we need to check assembly context. // A type is local (from target assembly) if get_external() returns None. diff --git a/dotscope/src/metadata/validation/validators/owned/types/definition.rs b/dotscope/src/metadata/validation/validators/owned/types/definition.rs index 07a9845f..a37fc21e 100644 --- a/dotscope/src/metadata/validation/validators/owned/types/definition.rs +++ b/dotscope/src/metadata/validation/validators/owned/types/definition.rs @@ -177,22 +177,17 @@ impl OwnedTypeDefinitionValidator { // Validate special naming patterns (but allow legitimate compiler-generated types) if type_entry.name.starts_with('<') && !type_entry.name.ends_with('>') { - // Allow compiler-generated patterns: - // - '<>c' (closures) - // - 'd__N' (async state machines) - // - '<>c__DisplayClassN' (closure display classes) - // - 'b__N' (lambda expressions) - // - 'e__FixedBuffer' (fixed buffer struct) - // - '<g__LocalFunction|N_N>d' (async local function state machines) - // - '' (global/module type) - let is_compiler_generated = type_entry.name.starts_with("<>") - || type_entry.name.contains(">d__") - || type_entry.name.contains(">b__") - || type_entry.name.contains(">c__") - || type_entry.name.contains(">e__FixedBuffer") - || type_entry.name.contains(">g__") // Local function patterns - || type_entry.name.ends_with(">d") // Async state machines ending with >d - || type_entry.name == ""; // Global module type + // A leading '<' opens a compiler- or tool-generated name, and the property + // that makes such a name well formed is that the bracket is *closed*. What + // follows the '>' is chosen by whatever emitted it, so enumerating suffixes + // cannot be exhaustive: Roslyn appends 'd__N'/'b__N'/'c__DisplayClassN', and + // both Roslyn and .NET Reactor append '{GUID}' to '' + // and ''. The previous allowlist rejected that last shape, which is + // present in shipped assemblies and in every .NET Reactor sample -- so real + // input failed validation, before any rewriting was involved. + // + // An unterminated '<' is still rejected, which is what this check is for. + let is_compiler_generated = type_entry.name.contains('>'); if !is_compiler_generated { return Err(Error::ValidationOwnedFailed { @@ -342,9 +337,9 @@ impl OwnedTypeDefinitionValidator { // Value types should typically inherit from System.ValueType or System.Enum if let Some(base_type) = type_entry.base() { let base_fullname = base_type.fullname(); - if base_fullname != "System.ValueType" - && base_fullname != "System.Enum" - && base_fullname != "System.Object" + if &*base_fullname != "System.ValueType" + && &*base_fullname != "System.Enum" + && &*base_fullname != "System.Object" { // Object is allowed for primitives // Allow some flexibility for special cases diff --git a/dotscope/src/metadata/validation/validators/owned/types/dependency.rs b/dotscope/src/metadata/validation/validators/owned/types/dependency.rs index cee366b1..e72a736f 100644 --- a/dotscope/src/metadata/validation/validators/owned/types/dependency.rs +++ b/dotscope/src/metadata/validation/validators/owned/types/dependency.rs @@ -421,7 +421,7 @@ impl OwnedTypeDependencyValidator { .find(|t| Arc::as_ptr(t) as usize == type_key) .map_or_else( || format!("Unknown type at address 0x{type_key:X}"), - |t| t.fullname(), + |t| t.fullname().to_string(), ); return Err(Error::ValidationOwnedFailed { diff --git a/dotscope/src/metadata/validation/validators/owned/types/inheritance.rs b/dotscope/src/metadata/validation/validators/owned/types/inheritance.rs index 5dafa075..5260ee8e 100644 --- a/dotscope/src/metadata/validation/validators/owned/types/inheritance.rs +++ b/dotscope/src/metadata/validation/validators/owned/types/inheritance.rs @@ -269,8 +269,8 @@ impl OwnedInheritanceValidator { let is_array_relationship = type_entry.is_array_of(&base_fullname); let is_system_type = base_type.namespace.starts_with("System"); - let is_value_type_inheritance = base_type.fullname() == "System.ValueType" - || base_type.fullname() == "System.Enum"; + let is_value_type_inheritance = &*base_type.fullname() == "System.ValueType" + || &*base_type.fullname() == "System.Enum"; if !is_system_type && !is_value_type_inheritance @@ -473,7 +473,7 @@ impl OwnedInheritanceValidator { (CilFlavor::Array { .. }, CilFlavor::Class | CilFlavor::ValueType | CilFlavor::Interface) | // Arrays can inherit from their element types (CilFlavor::GenericInstance, _) => Ok(()), // Generic instances can inherit from any type (CilFlavor::ValueType, CilFlavor::Object) => { - if base_type.fullname() == "System.Object" { + if &*base_type.fullname() == "System.Object" { Ok(()) } else { Err(Error::ValidationOwnedFailed { diff --git a/dotscope/src/metadata/validation/validators/raw/constraints/generic.rs b/dotscope/src/metadata/validation/validators/raw/constraints/generic.rs index fbb7dbd6..04de4aa1 100644 --- a/dotscope/src/metadata/validation/validators/raw/constraints/generic.rs +++ b/dotscope/src/metadata/validation/validators/raw/constraints/generic.rs @@ -150,6 +150,7 @@ impl RawGenericConstraintValidator { if let Some(generic_param_table) = tables.table::() { for generic_param in generic_param_table { + let generic_param = generic_param?; if generic_param.flags > 0xFFFF { return Err(malformed_error!( "GenericParam RID {} has invalid flags value {} exceeding maximum", @@ -205,6 +206,7 @@ impl RawGenericConstraintValidator { let generic_param_table = tables.table::(); for constraint in constraint_table { + let constraint = constraint?; if constraint.owner == 0 { return Err(malformed_error!( "GenericParamConstraint RID {} has null owner reference", @@ -263,9 +265,14 @@ impl RawGenericConstraintValidator { tables.table::(), ) { for constraint in constraint_table { - let param_found = generic_param_table - .iter() - .any(|param| param.rid == constraint.owner); + let constraint = constraint?; + let mut param_found = false; + for param in generic_param_table.iter() { + if param?.rid == constraint.owner { + param_found = true; + break; + } + } if !param_found { return Err(malformed_error!( @@ -307,6 +314,7 @@ impl RawGenericConstraintValidator { if let Some(constraint_table) = tables.table::() { for constraint in constraint_table { + let constraint = constraint?; let constraint_tables = constraint.constraint.ci_type.tables(); let constraint_table_type = match constraint_tables { [single] if constraint_tables.len() == 1 => *single, @@ -405,6 +413,7 @@ impl RawGenericConstraintValidator { if let Some(generic_param_table) = tables.table::() { for generic_param in generic_param_table { + let generic_param = generic_param?; let flags = GenericParamAttributes::new(generic_param.flags); if flags.contains(GenericParamAttributes::COVARIANT) diff --git a/dotscope/src/metadata/validation/validators/raw/constraints/layout.rs b/dotscope/src/metadata/validation/validators/raw/constraints/layout.rs index 98ab6fab..04d72d2f 100644 --- a/dotscope/src/metadata/validation/validators/raw/constraints/layout.rs +++ b/dotscope/src/metadata/validation/validators/raw/constraints/layout.rs @@ -76,7 +76,7 @@ use rustc_hash::FxHashMap; use crate::{ metadata::{ cilassemblyview::CilAssemblyView, - tables::{ClassLayoutRaw, FieldLayoutRaw, FieldRaw, TypeDefRaw}, + tables::{skip_unreadable, ClassLayoutRaw, FieldLayoutRaw, FieldRaw, TypeDefRaw}, validation::{ context::{RawValidationContext, ValidationContext}, shared::err_no_metadata_tables, @@ -152,6 +152,7 @@ impl RawLayoutConstraintValidator { let mut field_offsets: FxHashMap> = FxHashMap::default(); for field_layout in field_layout_table { + let field_layout = field_layout?; if field_layout.field == 0 { return Err(malformed_error!( "FieldLayout RID {} has null field reference", @@ -231,6 +232,7 @@ impl RawLayoutConstraintValidator { let typedef_table = tables.table::(); for class_layout in class_layout_table { + let class_layout = class_layout?; let packing_size = class_layout.packing_size; if packing_size != 0 && !packing_size.is_power_of_two() { return Err(malformed_error!( @@ -311,10 +313,12 @@ impl RawLayoutConstraintValidator { ) { let mut class_layouts: FxHashMap = FxHashMap::default(); for class_layout in class_layout_table { + let class_layout = class_layout?; class_layouts.insert(class_layout.parent, class_layout.rid); } for field_layout in field_layout_table { + let field_layout = field_layout?; if field_layout.field_offset == 0x7FFF_FFFF { return Err(malformed_error!( "FieldLayout RID {} has field offset at maximum boundary - potential overflow", @@ -327,19 +331,23 @@ impl RawLayoutConstraintValidator { continue; } - let typedef_rows: Vec<_> = typedef_table.iter().collect(); + // A TypeDef owns fields `[self.field_list, next_row.field_list)`, where + // "next row" means the next *RID* — so both ends are fetched by RID rather + // than read off a filtered `Vec`. `skip_unreadable` drops unparseable rows, + // and in a collected vector that silently makes one type's field range + // swallow the following type's fields, mis-attributing the parent. let mut parent_typedef_rid = None; - for (index, typedef_entry) in typedef_rows.iter().enumerate() { + for rid in 1..=typedef_table.row_count { + let Ok(Some(typedef_entry)) = typedef_table.get(rid) else { + continue; + }; let start_field = typedef_entry.field_list; - let next_index = index.saturating_add(1); - let end_field = if next_index < typedef_rows.len() { - typedef_rows - .get(next_index) - .ok_or(out_of_bounds_error!())? - .field_list - } else { - u32::MAX + let end_field = match typedef_table.get(rid.saturating_add(1)) { + Ok(Some(next)) => next.field_list, + // Last row, or the next row is unreadable: leave the range open + // rather than guessing an end that could exclude a real field. + _ => u32::MAX, }; if field_layout.field >= start_field && field_layout.field < end_field { @@ -352,10 +360,12 @@ impl RawLayoutConstraintValidator { if let Some(parent_rid) = parent_typedef_rid { if let Some(&class_layout_rid) = class_layouts.get(&parent_rid) { // Find the actual class layout to validate field offset against class size - if let Some(parent_class_layout) = class_layout_table + let parent_class_layout = class_layout_table .iter() - .find(|cl| cl.rid == class_layout_rid) - { + .collect::>>()? + .into_iter() + .find(|cl| cl.rid == class_layout_rid); + if let Some(parent_class_layout) = parent_class_layout { // Validate field offset is reasonable (but allow flexibility for legitimate .NET patterns) // Note: In legitimate .NET assemblies, field offsets can exceed declared class size // due to explicit layout, union types, interop scenarios, inheritance, etc. @@ -376,9 +386,14 @@ impl RawLayoutConstraintValidator { } for class_layout in class_layout_table { - let typedef_found = typedef_table - .iter() - .any(|typedef| typedef.rid == class_layout.parent); + let class_layout = class_layout?; + let mut typedef_found = false; + for typedef in typedef_table.iter() { + if typedef?.rid == class_layout.parent { + typedef_found = true; + break; + } + } if !typedef_found { return Err(malformed_error!( @@ -421,6 +436,7 @@ impl RawLayoutConstraintValidator { (tables.table::(), tables.table::()) { for field_layout in field_layout_table { + let field_layout = field_layout?; let field_offset = field_layout.field_offset; if (field_offset % 4 == 1 || field_offset % 4 == 3) && field_offset > 65536 { @@ -481,10 +497,13 @@ impl RawLayoutConstraintValidator { tables.table::(), ) { for class_layout in class_layout_table { - if let Some(typedef_entry) = typedef_table + let class_layout = class_layout?; + let typedef_entry = typedef_table .iter() - .find(|td| td.rid == class_layout.parent) - { + .collect::>>()? + .into_iter() + .find(|td| td.rid == class_layout.parent); + if let Some(typedef_entry) = typedef_entry { const SEALED_FLAG: u32 = 0x0100; const SERIALIZABLE_FLAG: u32 = 0x2000; @@ -535,7 +554,10 @@ impl RawLayoutConstraintValidator { let tables = assembly_view.tables().ok_or_else(err_no_metadata_tables)?; if let Some(field_layout_table) = tables.table::() { - let field_layouts: Vec<_> = field_layout_table.iter().collect(); + let field_layouts: Vec<_> = field_layout_table + .iter() + .filter_map(skip_unreadable) + .collect(); let mut type_field_layouts: FxHashMap> = FxHashMap::default(); for field_layout in field_layouts { diff --git a/dotscope/src/metadata/validation/validators/raw/structure/signature.rs b/dotscope/src/metadata/validation/validators/raw/structure/signature.rs index bf252bae..c83cba66 100644 --- a/dotscope/src/metadata/validation/validators/raw/structure/signature.rs +++ b/dotscope/src/metadata/validation/validators/raw/structure/signature.rs @@ -444,6 +444,7 @@ impl RawValidator for RawSignatureValidator { if let Some(table) = tables.table::() { for method in table { + let method = method?; if let Some(blob_heap) = assembly_view.blobs() { if let Ok(blob_data) = blob_heap.get(method.signature as usize) { if let Some(&calling_convention) = blob_data.first() { @@ -471,6 +472,7 @@ impl RawValidator for RawSignatureValidator { if let Some(table) = tables.table::() { for field in table { + let field = field?; Self::validate_signature_blob_integrity( assembly_view, field.signature, @@ -481,6 +483,7 @@ impl RawValidator for RawSignatureValidator { if let Some(table) = tables.table::() { for property in table { + let property = property?; Self::validate_signature_blob_integrity( assembly_view, property.signature, @@ -491,6 +494,7 @@ impl RawValidator for RawSignatureValidator { if let Some(table) = tables.table::() { for standalone_sig in table { + let standalone_sig = standalone_sig?; if let Some(blob_heap) = assembly_view.blobs() { if let Ok(blob_data) = blob_heap.get(standalone_sig.signature as usize) { if let Some(&calling_convention) = blob_data.first() { @@ -513,6 +517,7 @@ impl RawValidator for RawSignatureValidator { if let Some(table) = tables.table::() { for type_spec in table { + let type_spec = type_spec?; Self::validate_signature_blob_integrity( assembly_view, type_spec.signature, @@ -523,6 +528,7 @@ impl RawValidator for RawSignatureValidator { if let Some(table) = tables.table::() { for member_ref in table { + let member_ref = member_ref?; if let Some(blob_heap) = assembly_view.blobs() { if let Ok(blob_data) = blob_heap.get(member_ref.signature as usize) { if let Some(&calling_convention) = blob_data.first() { diff --git a/dotscope/src/metadata/validation/validators/raw/structure/table.rs b/dotscope/src/metadata/validation/validators/raw/structure/table.rs index a81e0a26..5ba4b1bf 100644 --- a/dotscope/src/metadata/validation/validators/raw/structure/table.rs +++ b/dotscope/src/metadata/validation/validators/raw/structure/table.rs @@ -242,6 +242,7 @@ impl RawTableValidator { (tables.table::(), tables.table::()) { for typedef_row in typedef_table { + let typedef_row = typedef_row?; if typedef_row.field_list != 0 && typedef_row.field_list > field_table.row_count.saturating_add(1) { @@ -259,6 +260,7 @@ impl RawTableValidator { (tables.table::(), tables.table::()) { for typedef_row in typedef_table { + let typedef_row = typedef_row?; if typedef_row.method_list != 0 && typedef_row.method_list > method_table.row_count.saturating_add(1) { diff --git a/dotscope/src/metadata/validation/validators/raw/structure/token.rs b/dotscope/src/metadata/validation/validators/raw/structure/token.rs index 7da3a98d..0f2e05e3 100644 --- a/dotscope/src/metadata/validation/validators/raw/structure/token.rs +++ b/dotscope/src/metadata/validation/validators/raw/structure/token.rs @@ -142,6 +142,7 @@ impl RawTokenValidator { if let Some(tables) = assembly_view.tables() { if let Some(table) = tables.table::() { for typedef in table { + let typedef = typedef?; if typedef.extends.row != 0 { referenced_tokens.push(typedef.extends.token); } @@ -150,6 +151,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for interface_impl in table { + let interface_impl = interface_impl?; token_validator.validate_table_row(TableId::TypeDef, interface_impl.class)?; referenced_tokens.push(interface_impl.interface.token); } @@ -157,12 +159,14 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for memberref in table { + let memberref = memberref?; referenced_tokens.push(memberref.class.token); } } if let Some(table) = tables.table::() { for attr in table { + let attr = attr?; referenced_tokens.push(attr.parent.token); referenced_tokens.push(attr.constructor.token); } @@ -170,6 +174,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for nested in table { + let nested = nested?; token_validator.validate_table_row(TableId::TypeDef, nested.nested_class)?; token_validator.validate_table_row(TableId::TypeDef, nested.enclosing_class)?; } @@ -177,18 +182,21 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for genparam in table { + let genparam = genparam?; referenced_tokens.push(genparam.owner.token); } } if let Some(table) = tables.table::() { for methodspec in table { + let methodspec = methodspec?; referenced_tokens.push(methodspec.method.token); } } if let Some(table) = tables.table::() { for constraint in table { + let constraint = constraint?; token_validator.validate_table_row(TableId::GenericParam, constraint.owner)?; referenced_tokens.push(constraint.constraint.token); } @@ -196,6 +204,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for method_impl in table { + let method_impl = method_impl?; token_validator.validate_table_row(TableId::TypeDef, method_impl.class)?; referenced_tokens.push(method_impl.method_body.token); referenced_tokens.push(method_impl.method_declaration.token); @@ -204,18 +213,21 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for constant in table { + let constant = constant?; referenced_tokens.push(constant.parent.token); } } if let Some(table) = tables.table::() { for marshal in table { + let marshal = marshal?; referenced_tokens.push(marshal.parent.token); } } if let Some(table) = tables.table::() { for security in table { + let security = security?; referenced_tokens.push(security.parent.token); } } @@ -359,6 +371,7 @@ impl RawTokenValidator { if let Some(tables) = assembly_view.tables() { if let Some(table) = tables.table::() { for typedef in table { + let typedef = typedef?; Self::validate_coded_index_field( &typedef.extends, &token_validator, @@ -369,6 +382,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for interface_impl in table { + let interface_impl = interface_impl?; Self::validate_coded_index_field( &interface_impl.interface, &token_validator, @@ -379,6 +393,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for memberref in table { + let memberref = memberref?; Self::validate_coded_index_field( &memberref.class, &token_validator, @@ -389,6 +404,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for attr in table { + let attr = attr?; Self::validate_coded_index_field( &attr.parent, &token_validator, @@ -404,6 +420,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for genparam in table { + let genparam = genparam?; Self::validate_coded_index_field( &genparam.owner, &token_validator, @@ -414,6 +431,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for methodspec in table { + let methodspec = methodspec?; Self::validate_coded_index_field( &methodspec.method, &token_validator, @@ -424,6 +442,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for constraint in table { + let constraint = constraint?; Self::validate_coded_index_field( &constraint.constraint, &token_validator, @@ -434,6 +453,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for constant in table { + let constant = constant?; Self::validate_coded_index_field( &constant.parent, &token_validator, @@ -444,6 +464,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for marshal in table { + let marshal = marshal?; Self::validate_coded_index_field( &marshal.parent, &token_validator, @@ -454,6 +475,7 @@ impl RawTokenValidator { if let Some(table) = tables.table::() { for security in table { + let security = security?; Self::validate_coded_index_field( &security.parent, &token_validator, @@ -472,6 +494,7 @@ impl RawTokenValidator { reference_validator: &ReferenceValidator, ) -> Result<()> { for interface_impl in table { + let interface_impl = interface_impl?; token_validator.validate_table_row(TableId::TypeDef, interface_impl.class)?; let interface_token = interface_impl.interface.token; @@ -490,6 +513,7 @@ impl RawTokenValidator { reference_validator: &ReferenceValidator, ) -> Result<()> { for memberref in table { + let memberref = memberref?; let class_token = memberref.class.token; let allowed_tables = memberref.class.ci_type.tables(); token_validator.validate_typed_token(class_token, allowed_tables)?; @@ -506,6 +530,7 @@ impl RawTokenValidator { reference_validator: &ReferenceValidator, ) -> Result<()> { for attr in table { + let attr = attr?; let parent_token = attr.parent.token; token_validator.validate_token_bounds(parent_token)?; reference_validator.validate_token_integrity(parent_token)?; @@ -525,6 +550,7 @@ impl RawTokenValidator { reference_validator: &ReferenceValidator, ) -> Result<()> { for nested in table { + let nested = nested?; token_validator.validate_table_row(TableId::TypeDef, nested.nested_class)?; token_validator.validate_table_row(TableId::TypeDef, nested.enclosing_class)?; @@ -547,6 +573,7 @@ impl RawTokenValidator { reference_validator: &ReferenceValidator, ) -> Result<()> { for genparam in table { + let genparam = genparam?; let owner_token = genparam.owner.token; let allowed_tables = genparam.owner.ci_type.tables(); token_validator.validate_typed_token(owner_token, allowed_tables)?; @@ -562,6 +589,7 @@ impl RawTokenValidator { reference_validator: &ReferenceValidator, ) -> Result<()> { for methodspec in table { + let methodspec = methodspec?; let method_token = methodspec.method.token; let allowed_tables = methodspec.method.ci_type.tables(); token_validator.validate_typed_token(method_token, allowed_tables)?; @@ -577,6 +605,7 @@ impl RawTokenValidator { reference_validator: &ReferenceValidator, ) -> Result<()> { for constraint in table { + let constraint = constraint?; token_validator.validate_table_row(TableId::GenericParam, constraint.owner)?; let constraint_token = constraint.constraint.token; @@ -594,6 +623,7 @@ impl RawTokenValidator { reference_validator: &ReferenceValidator, ) -> Result<()> { for method_impl in table { + let method_impl = method_impl?; token_validator.validate_table_row(TableId::TypeDef, method_impl.class)?; let body_token = method_impl.method_body.token; @@ -616,6 +646,7 @@ impl RawTokenValidator { reference_validator: &ReferenceValidator, ) -> Result<()> { for constant in table { + let constant = constant?; let parent_token = constant.parent.token; token_validator.validate_token_bounds(parent_token)?; reference_validator.validate_token_integrity(parent_token)?; @@ -630,6 +661,7 @@ impl RawTokenValidator { reference_validator: &ReferenceValidator, ) -> Result<()> { for marshal in table { + let marshal = marshal?; let parent_token = marshal.parent.token; let allowed_tables = marshal.parent.ci_type.tables(); token_validator.validate_typed_token(parent_token, allowed_tables)?; @@ -645,6 +677,7 @@ impl RawTokenValidator { reference_validator: &ReferenceValidator, ) -> Result<()> { for security in table { + let security = security?; let parent_token = security.parent.token; let allowed_tables = security.parent.ci_type.tables(); token_validator.validate_typed_token(parent_token, allowed_tables)?; diff --git a/dotscope/src/prelude.rs b/dotscope/src/prelude.rs index 566928d0..fc923491 100644 --- a/dotscope/src/prelude.rs +++ b/dotscope/src/prelude.rs @@ -135,7 +135,10 @@ //! if let Some(tables) = assembly.tables() { //! if let Some(typedef_table) = tables.table::() { //! let row_index = typedef_token.row(); -//! if let Some(typedef) = typedef_table.get(row_index) { +//! // `get` reports a malformed row rather than hiding it as "absent", so the +//! // result is `Result>`: outer for parse failure, inner for a RID +//! // past the end of the table. +//! if let Ok(Some(typedef)) = typedef_table.get(row_index) { //! println!("Type name index: {}", typedef.type_name); //! } //! } diff --git a/dotscope/src/project/mod.rs b/dotscope/src/project/mod.rs index 65d76895..249409c5 100644 --- a/dotscope/src/project/mod.rs +++ b/dotscope/src/project/mod.rs @@ -209,7 +209,7 @@ impl CilProject { // Check if this assembly defines the type locally for entry in assembly.types().iter() { let type_instance = entry.value(); - if type_instance.fullname() == full_name { + if &*type_instance.fullname() == full_name { // Check if this is a TypeDef (0x02) or TypeRef (0x01) if type_instance.token.is_table(TableId::TypeDef) { // TypeDef - actual definition, return immediately @@ -471,7 +471,7 @@ impl CilProject { for (identity, assembly) in self.iter() { for entry in assembly.types().iter() { let type_instance = entry.value(); - if type_instance.fullname() == type_name + if &*type_instance.fullname() == type_name && type_instance.token.is_table(TableId::TypeDef) { results.push((identity.clone(), type_instance.clone())); diff --git a/dotscope/src/test/builders/assembly.rs b/dotscope/src/test/builders/assembly.rs index 9b2f27cd..df1ecb9f 100644 --- a/dotscope/src/test/builders/assembly.rs +++ b/dotscope/src/test/builders/assembly.rs @@ -5,9 +5,12 @@ use std::sync::{atomic::AtomicU32, Arc}; -use crate::metadata::{ - tables::{AssemblyFlags, AssemblyRef, AssemblyRefHash, AssemblyRefRc}, - token::Token, +use crate::{ + metadata::{ + tables::{AssemblyFlags, AssemblyRef, AssemblyRefHash, AssemblyRefRc}, + token::Token, + }, + utils::LazyList, }; /// Builder for creating mock AssemblyRef instances with realistic metadata @@ -97,7 +100,7 @@ impl AssemblyRefBuilder { os_major_version: AtomicU32::new(0), os_minor_version: AtomicU32::new(0), processor: AtomicU32::new(0), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }) } } diff --git a/dotscope/src/test/builders/fields.rs b/dotscope/src/test/builders/fields.rs index f2c4363c..13a71ec0 100644 --- a/dotscope/src/test/builders/fields.rs +++ b/dotscope/src/test/builders/fields.rs @@ -5,12 +5,15 @@ use std::sync::{Arc, OnceLock}; -use crate::metadata::{ - marshalling::{MarshallingInfo, NativeType}, - signatures::{SignatureField, TypeSignature}, - tables::{Field, FieldAttributes, FieldRc, TypeAttributes}, - token::Token, - typesystem::{CilFlavor, CilPrimitive, CilTypeRc}, +use crate::{ + metadata::{ + marshalling::{MarshallingInfo, NativeType}, + signatures::{SignatureField, TypeSignature}, + tables::{Field, FieldAttributes, FieldRc, TypeAttributes}, + token::Token, + typesystem::{CilFlavor, CilPrimitive, CilTypeRc}, + }, + utils::LazyList, }; /// Field layout types for explicit field positioning @@ -216,7 +219,7 @@ impl FieldBuilder { None, TypeAttributes::ZERO, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), Some(CilFlavor::I4), )); Self::new(name, i4_type) @@ -232,7 +235,7 @@ impl FieldBuilder { None, TypeAttributes::ZERO, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), Some(CilFlavor::String), )); Self::new(name, string_type) @@ -248,7 +251,7 @@ impl FieldBuilder { None, TypeAttributes::ZERO, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), Some(CilFlavor::Boolean), )); Self::new(name, bool_type) @@ -264,7 +267,7 @@ impl FieldBuilder { None, TypeAttributes::ZERO, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), Some(CilFlavor::R4), )); Self::new(name, r4_type) @@ -280,7 +283,7 @@ impl FieldBuilder { None, TypeAttributes::ZERO, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), Some(CilFlavor::Object), )); Self::new(name, object_type) @@ -338,7 +341,7 @@ impl FieldBuilder { rva: OnceLock::new(), layout: OnceLock::new(), marshal: OnceLock::new(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), declaring_type: OnceLock::new(), }; @@ -420,7 +423,7 @@ impl Default for FieldBuilder { None, TypeAttributes::ZERO, Arc::new(boxcar::Vec::new()), - Arc::new(boxcar::Vec::new()), + LazyList::new(), Some(CilFlavor::I4), )); diff --git a/dotscope/src/test/builders/files.rs b/dotscope/src/test/builders/files.rs index a64a8be9..ca55ef3e 100644 --- a/dotscope/src/test/builders/files.rs +++ b/dotscope/src/test/builders/files.rs @@ -5,17 +5,20 @@ use std::sync::Arc; -use crate::metadata::{ - customattributes::CustomAttributeValue, - tables::{AssemblyRefHash, File, FileAttributes, FileRc, ModuleRef, ModuleRefRc}, - token::Token, +use crate::{ + metadata::{ + customattributes::CustomAttributeValue, + tables::{AssemblyRefHash, File, FileAttributes, FileRc, ModuleRef, ModuleRefRc}, + token::Token, + }, + utils::LazyList, }; /// Builder for creating mock ModuleRef instances with various configurations pub struct ModuleRefBuilder { rid: u32, name: String, - custom_attributes: Option>>>, + custom_attributes: Option>>, } impl ModuleRefBuilder { @@ -43,9 +46,7 @@ impl ModuleRefBuilder { offset: self.rid as usize, token: Token::new(0x1A000000 + self.rid), name: self.name, - custom_attributes: self - .custom_attributes - .unwrap_or_else(|| Arc::new(boxcar::Vec::>::new())), + custom_attributes: self.custom_attributes.unwrap_or_default(), }) } } @@ -99,7 +100,7 @@ impl FileBuilder { hash_value: self .hash_value .unwrap_or_else(|| AssemblyRefHash::new(&[1, 2, 3, 4]).unwrap()), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }) } } diff --git a/dotscope/src/test/builders/methods.rs b/dotscope/src/test/builders/methods.rs index ac82d773..df3b1e66 100644 --- a/dotscope/src/test/builders/methods.rs +++ b/dotscope/src/test/builders/methods.rs @@ -16,6 +16,7 @@ use crate::{ typesystem::{CilFlavor, TypeRegistry}, }, test::builders::params::ParamBuilder, + utils::LazyList, }; /// Builder for creating mock Method instances with complex configurations @@ -245,8 +246,8 @@ impl MethodBuilder { name: self.name, params, varargs: Arc::new(boxcar::Vec::new()), - generic_params: Arc::new(boxcar::Vec::new()), - generic_args: Arc::new(boxcar::Vec::new()), + generic_params: LazyList::new(), + generic_args: LazyList::new(), signature: self.signature.unwrap_or_else(|| SignatureMethod { has_this: false, explicit_this: false, @@ -270,10 +271,10 @@ impl MethodBuilder { body: OnceLock::new(), local_vars: Arc::new(boxcar::Vec::new()), overrides: Arc::new(boxcar::Vec::new()), - interface_impls: Arc::new(boxcar::Vec::new()), + interface_impls: LazyList::new(), security: OnceLock::new(), blocks: OnceLock::new(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), declaring_type: OnceLock::new(), }) } diff --git a/dotscope/src/test/builders/params.rs b/dotscope/src/test/builders/params.rs index 48ced618..c088dbce 100644 --- a/dotscope/src/test/builders/params.rs +++ b/dotscope/src/test/builders/params.rs @@ -16,6 +16,7 @@ use crate::{ }, prelude::{CilPrimitiveData, CilPrimitiveKind, ParamAttributes}, test::{builders::FieldBuilder, factories::metadata::customattributes::get_test_type_registry}, + utils::LazyList, }; /// Parameter attribute flags for various parameter characteristics @@ -219,7 +220,7 @@ impl ParamBuilder { modifiers: Arc::new(boxcar::Vec::new()), base: OnceLock::new(), is_by_ref: AtomicBool::new(false), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }); // Set default value if provided diff --git a/dotscope/src/test/builders/properties.rs b/dotscope/src/test/builders/properties.rs index 96a4fa05..f20e5989 100644 --- a/dotscope/src/test/builders/properties.rs +++ b/dotscope/src/test/builders/properties.rs @@ -13,6 +13,7 @@ use crate::{ typesystem::CilPrimitive, }, prelude::SignatureParameter, + utils::LazyList, }; /// Property constant value types for default values @@ -138,7 +139,7 @@ impl PropertyBuilder { fn_setter: OnceLock::new(), fn_getter: OnceLock::new(), fn_other: OnceLock::new(), - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }; let property_rc = Arc::new(property); diff --git a/dotscope/src/test/builders/types.rs b/dotscope/src/test/builders/types.rs index a1b37514..9b6912a5 100644 --- a/dotscope/src/test/builders/types.rs +++ b/dotscope/src/test/builders/types.rs @@ -12,6 +12,7 @@ use crate::{ typesystem::{CilFlavor, CilType, CilTypeRc, CilTypeReference}, }, test::FileBuilder, + utils::LazyList, }; /// Builder for creating mock CilType instances with various characteristics @@ -102,7 +103,7 @@ impl CilTypeBuilder { None, // base type self.flags, Arc::new(boxcar::Vec::new()), // fields - Arc::new(boxcar::Vec::new()), // methods + LazyList::new(), // methods self.flavor, )) } @@ -135,6 +136,6 @@ pub fn create_exportedtype(dummy_type: CilTypeRc) -> ExportedTypeRc { name: "ExportedType".to_string(), namespace: Some("Test.Namespace".to_string()), implementation: implementation_lock, - custom_attributes: Arc::new(boxcar::Vec::new()), + custom_attributes: LazyList::new(), }) } diff --git a/dotscope/src/test/crafted2.rs b/dotscope/src/test/crafted2.rs index 03ff6c70..30f820a5 100644 --- a/dotscope/src/test/crafted2.rs +++ b/dotscope/src/test/crafted2.rs @@ -624,7 +624,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 1); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.generation, 0); assert_eq!(row.name, 0x9CF); @@ -645,7 +645,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 65); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.token.value(), 0x01000001); assert_eq!( @@ -664,7 +664,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 36); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 0); assert_eq!(row.type_name, 0x1FD); @@ -685,7 +685,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 48); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 0x26); assert_eq!(row.name, 0xABB); @@ -700,7 +700,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 97); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.rva, 0x2050); assert_eq!(row.impl_flags, 0); @@ -718,7 +718,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 72); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 0); assert_eq!(row.sequence, 1); @@ -733,7 +733,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 5); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.class, 7); assert_eq!( @@ -750,7 +750,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 67); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!( row.class, @@ -768,7 +768,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 7); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.base, 0xA); assert_eq!( @@ -787,7 +787,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 88); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!( row.parent, @@ -808,7 +808,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 1); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!( row.parent, @@ -825,7 +825,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 2); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.action, 8); assert_eq!( @@ -843,7 +843,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 3); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.packing_size, 0); assert_eq!(row.class_size, 0x10); @@ -858,7 +858,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 3); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.field_offset, 0); assert_eq!(row.field, 0xC); @@ -872,7 +872,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 11); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.signature, 0x53); } @@ -885,7 +885,7 @@ fn verify_tableheader(asm: &CilObject) { Some(module) => { assert_eq!(module.row_count, 2); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.parent, 0x7); assert_eq!(row.event_list, 0x1); @@ -899,7 +899,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 3); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 0); assert_eq!(row.name, 0x102); @@ -917,7 +917,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 8); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.parent, 0x4); assert_eq!(row.property_list, 1); @@ -931,7 +931,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 13); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 0); assert_eq!(row.name, 0x9B4); @@ -946,7 +946,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 31); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.semantics, 8); assert_eq!(row.method, 0xE); @@ -964,7 +964,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 4); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.class, 0xE); assert_eq!( @@ -985,7 +985,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 2); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.name, 0xA33); } @@ -998,7 +998,7 @@ fn verify_tableheader(asm: &CilObject) { Some(module) => { assert_eq!(module.row_count, 16); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.signature, 0x40); } @@ -1011,7 +1011,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 2); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.mapping_flags, 0x104); assert_eq!( @@ -1030,7 +1030,7 @@ fn verify_tableheader(asm: &CilObject) { Some(module) => { assert_eq!(module.row_count, 1); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rva, 0x5410); assert_eq!(row.field, 0x1E); } @@ -1043,7 +1043,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 1); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.hash_alg_id, 0x8004); assert_eq!(row.major_version, 1); @@ -1063,7 +1063,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 2); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.major_version, 4); assert_eq!(row.minor_version, 0); @@ -1083,7 +1083,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 10); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.nested_class, 0x1B); assert_eq!(row.enclosing_class, 0xE); @@ -1097,7 +1097,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 19); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.number, 0); assert_eq!(row.flags, 4); @@ -1116,7 +1116,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 4); - let row = table.get(1).unwrap(); + let row = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!( row.method_body, @@ -1136,7 +1136,7 @@ fn verify_tableheader(asm: &CilObject) { Some(table) => { assert_eq!(table.row_count, 16); - let row: GenericParamConstraintRaw = table.get(1).unwrap(); + let row: GenericParamConstraintRaw = table.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.owner, 0x3); assert_eq!( @@ -1182,6 +1182,7 @@ fn verify_assembly_custom_attributes(asm: &CilObject) { let mut assembly_attr_count = 0; for attr_row in custom_attr_table.iter() { + let attr_row = attr_row.expect("row parses"); // Check if this attribute is on the assembly (target token 0x20000001) if attr_row.parent.token.value() == 0x20000001 { assembly_attr_count += 1; @@ -1203,6 +1204,7 @@ fn verify_module_custom_attributes(asm: &CilObject) { let mut module_attr_count = 0; for attr_row in custom_attr_table.iter() { + let attr_row = attr_row.expect("row parses"); // Check if this attribute is on the module (target token 0x00000001) if attr_row.parent.token.value() == 0x00000001 { module_attr_count += 1; @@ -1385,7 +1387,7 @@ fn verify_specialized_attribute_tables(asm: &CilObject) { // Test marshalling descriptor parsing - this verifies our marshalling implementation // against real .NET assembly data - let row = field_marshal_table.get(1).unwrap(); + let row = field_marshal_table.get(1).unwrap().unwrap(); let blob_heap = asm.blob().expect("Expected blob heap to be present"); let descriptor_blob = blob_heap.get(row.native_type as usize).unwrap(); @@ -2254,6 +2256,7 @@ fn test_nested_type_relationships(asm: &CilObject) { ); for nested_row in nested_table.iter().take(5) { + let nested_row = nested_row.expect("row parses"); println!( " NestedClass: nested=0x{:X}, enclosing=0x{:X}", nested_row.nested_class, nested_row.enclosing_class @@ -2318,7 +2321,11 @@ fn test_enum_and_constant_validation(asm: &CilObject) { ); // Look for constants associated with our enum - let constant_rows: Vec<_> = constant_table.iter().take(5).collect(); + let constant_rows: Vec<_> = constant_table + .iter() + .take(5) + .map(|row| row.expect("row parses")) + .collect(); for constant_row in constant_rows { println!( " Constant: type=0x{:X}, parent=0x{:08X}, value=0x{:X}", @@ -2555,6 +2562,7 @@ fn test_pinvoke_and_security_validation(asm: &CilObject) { ); for implmap_row in implmap_table.iter() { + let implmap_row = implmap_row.expect("row parses"); println!( " ImplMap: flags=0x{:X}, member=0x{:08X}, name=0x{:X}, scope=0x{:X}", implmap_row.mapping_flags, @@ -2579,6 +2587,7 @@ fn test_pinvoke_and_security_validation(asm: &CilObject) { ); for security_row in declsecurity_table.iter() { + let security_row = security_row.expect("row parses"); println!( " Security: action={}, parent=0x{:08X}, permission_set=0x{:X}", security_row.action, @@ -2992,7 +3001,7 @@ fn test_assembly_metadata_validation(asm: &CilObject) { println!("Assembly table has {assembly_count} entries"); assert_eq!(assembly_count, 1, "Should have exactly 1 assembly entry"); - if let Some(assembly_row) = assembly_table.get(1) { + if let Some(assembly_row) = assembly_table.get(1).ok().flatten() { println!("Assembly metadata:"); println!(" Major version: {}", assembly_row.major_version); println!(" Minor version: {}", assembly_row.minor_version); @@ -3022,7 +3031,7 @@ fn test_assembly_metadata_validation(asm: &CilObject) { println!("Module table has {module_count} entries"); assert!(module_count >= 1, "Should have at least 1 module"); - if let Some(module_row) = module_table.get(1) { + if let Some(module_row) = module_table.get(1).ok().flatten() { println!(" Module generation: {}", module_row.generation); println!(" ✓ Module metadata validated"); } @@ -3098,6 +3107,7 @@ fn test_xml_permission_set_parsing(asm: &CilObject) { // Iterate through DeclSecurity entries for security_row in decl_security_table.iter() { + let security_row = security_row.expect("row parses"); // Get the permission set blob from the blob stream if let Some(blob_heap) = asm.blob() { if let Ok(blob_data) = blob_heap.get(security_row.permission_set as usize) { @@ -3201,6 +3211,7 @@ fn test_xml_permission_set_parsing(asm: &CilObject) { // Let's test that we can at least parse the binary permission sets for security_row in decl_security_table.iter() { + let security_row = security_row.expect("row parses"); if let Some(blob_heap) = asm.blob() { if let Ok(blob_data) = blob_heap.get(security_row.permission_set as usize) { let permission_set = PermissionSet::new(blob_data).unwrap(); diff --git a/dotscope/src/test/helpers/dependencies.rs b/dotscope/src/test/helpers/dependencies.rs index 7ee70e63..389d6dca 100644 --- a/dotscope/src/test/helpers/dependencies.rs +++ b/dotscope/src/test/helpers/dependencies.rs @@ -6,16 +6,17 @@ use std::sync::{atomic::AtomicU32, Arc}; -use boxcar::Vec as BoxcarVec; - -use crate::metadata::{ - dependencies::{ - AssemblyDependency, DependencyResolutionState, DependencySource, DependencyType, - VersionRequirement, +use crate::{ + metadata::{ + dependencies::{ + AssemblyDependency, DependencyResolutionState, DependencySource, DependencyType, + VersionRequirement, + }, + identity::{AssemblyIdentity, AssemblyVersion}, + tables::{AssemblyFlags, AssemblyRef, AssemblyRefHash, File, FileAttributes, ModuleRef}, + token::Token, }, - identity::{AssemblyIdentity, AssemblyVersion}, - tables::{AssemblyFlags, AssemblyRef, AssemblyRefHash, File, FileAttributes, ModuleRef}, - token::Token, + utils::LazyList, }; /// Create a test assembly identity with basic version information. @@ -66,7 +67,7 @@ pub fn create_test_assembly_ref_with_culture( offset: 0, rid: 1, token: Token::new(0x23000001), // AssemblyRef table token (0x23 = table, 1 = row) - custom_attributes: Arc::new(BoxcarVec::new()), + custom_attributes: LazyList::new(), hash: None, os_platform_id: AtomicU32::new(0), os_major_version: AtomicU32::new(0), @@ -82,7 +83,7 @@ pub fn create_test_module_ref(name: &str) -> Arc { offset: 0, rid: 1, token: Token::new(0x1A000001), // ModuleRef table token (0x1A = table, 1 = row) - custom_attributes: Arc::new(BoxcarVec::new()), + custom_attributes: LazyList::new(), }) } @@ -95,7 +96,7 @@ pub fn create_test_file(name: &str) -> Arc { offset: 0, rid: 1, token: Token::new(0x26000001), // File table token (0x26 = table, 1 = row) - custom_attributes: Arc::new(BoxcarVec::new()), + custom_attributes: LazyList::new(), }) } diff --git a/dotscope/src/test/mono/disassembly.rs b/dotscope/src/test/mono/disassembly.rs index 1de46d74..3b6cb499 100644 --- a/dotscope/src/test/mono/disassembly.rs +++ b/dotscope/src/test/mono/disassembly.rs @@ -96,15 +96,9 @@ impl DisassemblyResult { // Extract method name from lines like: // .method public static void Main() cil managed // Find the last word before the parenthesis - if let Some(paren_pos) = line.find('(') { - let before_paren = &line[..paren_pos]; - before_paren - .split_whitespace() - .last() - .map(|s| s.to_string()) - } else { - None - } + line.split_once('(') + .and_then(|(before_paren, _)| before_paren.split_whitespace().last()) + .map(ToString::to_string) }) .collect() } diff --git a/dotscope/src/test/windowsbase.rs b/dotscope/src/test/windowsbase.rs index d424e892..8c087e92 100644 --- a/dotscope/src/test/windowsbase.rs +++ b/dotscope/src/test/windowsbase.rs @@ -133,7 +133,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 1); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.generation, 0); assert_eq!(row.name, 0x1E026); @@ -150,7 +150,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 472); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.token.value(), 0x01000001); assert_eq!( @@ -160,7 +160,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.type_name, 0x18C2C); assert_eq!(row.type_namespace, 0x277D8); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.token.value(), 0x01000005); assert_eq!( @@ -170,7 +170,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.type_name, 0x27A21); assert_eq!(row.type_namespace, 0); - let row = module.get(320).unwrap(); + let row = module.get(320).unwrap().unwrap(); assert_eq!(row.rid, 320); assert_eq!(row.token.value(), 0x01000140); assert_eq!( @@ -189,7 +189,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 820); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 0); assert_eq!(row.type_name, 0x1495); @@ -201,7 +201,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.field_list, 1); assert_eq!(row.method_list, 1); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.flags, 0x100180); assert_eq!(row.type_name, 0x26FB4); @@ -213,7 +213,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.field_list, 0x2C); assert_eq!(row.method_list, 1); - let row = module.get(320).unwrap(); + let row = module.get(320).unwrap().unwrap(); assert_eq!(row.rid, 320); assert_eq!(row.flags, 0x100000); assert_eq!(row.type_name, 0x1238D); @@ -234,19 +234,19 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 6241); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 0x8056); assert_eq!(row.name, 0x2747E); assert_eq!(row.signature, 0x7F1); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.flags, 0x8056); assert_eq!(row.name, 0x5FD4); assert_eq!(row.signature, 0x7F1); - let row = module.get(320).unwrap(); + let row = module.get(320).unwrap().unwrap(); assert_eq!(row.rid, 320); assert_eq!(row.flags, 0x8056); assert_eq!(row.name, 0x5B7A); @@ -261,7 +261,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 6496); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.rva, 0xEAA84); assert_eq!(row.impl_flags, 0); @@ -270,7 +270,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.signature, 0xB041); assert_eq!(row.param_list, 1); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.rva, 0xEAB5C); assert_eq!(row.impl_flags, 0); @@ -279,7 +279,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.signature, 0xB041); assert_eq!(row.param_list, 9); - let row = module.get(320).unwrap(); + let row = module.get(320).unwrap().unwrap(); assert_eq!(row.rid, 320); assert_eq!(row.rva, 0xEB604); assert_eq!(row.impl_flags, 0); @@ -297,19 +297,19 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 7877); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 0); assert_eq!(row.sequence, 1); assert_eq!(row.name, 0x14593); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.flags, 0); assert_eq!(row.sequence, 1); assert_eq!(row.name, 0x16E17); - let row = module.get(320).unwrap(); + let row = module.get(320).unwrap().unwrap(); assert_eq!(row.rid, 320); assert_eq!(row.flags, 0); assert_eq!(row.sequence, 3); @@ -324,7 +324,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 122); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.class, 0xB); assert_eq!( @@ -332,7 +332,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { CodedIndex::new(TableId::TypeRef, 64, CodedIndexType::TypeDefOrRef) ); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.class, 0x10); assert_eq!( @@ -340,7 +340,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { CodedIndex::new(TableId::TypeSpec, 3, CodedIndexType::TypeDefOrRef) ); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!(row.class, 0x308); assert_eq!( @@ -357,7 +357,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 1762); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!( row.class, @@ -366,7 +366,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.name, 0x26F0B); assert_eq!(row.signature, 1); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!( row.class, @@ -375,7 +375,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.name, 0x26F0B); assert_eq!(row.signature, 0x10); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!( row.class, @@ -393,7 +393,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 4213); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.base, 0xE); assert_eq!( @@ -402,7 +402,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { ); assert_eq!(row.value, 0x3BB9); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.base, 0xE); assert_eq!( @@ -411,7 +411,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { ); assert_eq!(row.value, 0x3C1D); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!(row.base, 8); assert_eq!( @@ -430,7 +430,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 914); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!( row.parent, @@ -442,7 +442,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { ); assert_eq!(row.value, 0x4015); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!( row.parent, @@ -454,7 +454,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { ); assert_eq!(row.value, 0xFC8F); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!( row.parent, @@ -479,7 +479,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 620); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!( row.parent, @@ -487,7 +487,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { ); assert_eq!(row.native_type, 0xA56F); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!( row.parent, @@ -495,7 +495,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { ); assert_eq!(row.native_type, 0xA58F); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!( row.parent, @@ -512,7 +512,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 1); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.action, 8); assert_eq!( @@ -530,19 +530,19 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 13); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.packing_size, 0); assert_eq!(row.class_size, 0x10); assert_eq!(row.parent, 0x28); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.packing_size, 1); assert_eq!(row.class_size, 0); assert_eq!(row.parent, 0x24C); - let row = module.get(10).unwrap(); + let row = module.get(10).unwrap().unwrap(); assert_eq!(row.rid, 10); assert_eq!(row.packing_size, 1); assert_eq!(row.class_size, 0x34); @@ -557,17 +557,17 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 83); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.field_offset, 0); assert_eq!(row.field, 0x808); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.field_offset, 0); assert_eq!(row.field, 0x9A4); - let row = module.get(50).unwrap(); + let row = module.get(50).unwrap().unwrap(); assert_eq!(row.rid, 50); assert_eq!(row.field_offset, 0); assert_eq!(row.field, 0xC74); @@ -581,15 +581,15 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 668); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.signature, 0x220); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.signature, 0x280); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!(row.signature, 0xB9D); } @@ -602,17 +602,17 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 18); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.parent, 0xC); assert_eq!(row.event_list, 0x1); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.parent, 0x120); assert_eq!(row.event_list, 5); - let row = module.get(10).unwrap(); + let row = module.get(10).unwrap().unwrap(); assert_eq!(row.rid, 10); assert_eq!(row.parent, 0x158); assert_eq!(row.event_list, 0xD); @@ -626,7 +626,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 47); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 0); assert_eq!(row.name, 0xEF15); @@ -635,7 +635,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { CodedIndex::new(TableId::TypeRef, 69, CodedIndexType::TypeDefOrRef) ); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.flags, 0); assert_eq!(row.name, 0x1BBAA); @@ -644,7 +644,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { CodedIndex::new(TableId::TypeDef, 290, CodedIndexType::TypeDefOrRef) ); - let row = module.get(25).unwrap(); + let row = module.get(25).unwrap().unwrap(); assert_eq!(row.rid, 25); assert_eq!(row.flags, 0); assert_eq!(row.name, 0x13403); @@ -662,17 +662,17 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 234); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.parent, 0xC); assert_eq!(row.property_list, 1); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.parent, 0x15); assert_eq!(row.property_list, 0xB); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!(row.parent, 0x103); assert_eq!(row.property_list, 0x36D); @@ -686,19 +686,19 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 1511); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 0); assert_eq!(row.name, 0x1458B); assert_eq!(row.signature, 0xF6F5); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.flags, 0); assert_eq!(row.name, 0x4494); assert_eq!(row.signature, 0xF6FD); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!(row.flags, 0); assert_eq!(row.name, 0x1EEE3); @@ -713,7 +713,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 1848); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.semantics, 8); assert_eq!(row.method, 0x19C); @@ -722,7 +722,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { CodedIndex::new(TableId::Event, 1, CodedIndexType::HasSemantics) ); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.semantics, 0x10); assert_eq!(row.method, 0x336); @@ -731,7 +731,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { CodedIndex::new(TableId::Event, 2, CodedIndexType::HasSemantics) ); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!(row.semantics, 8); assert_eq!(row.method, 0x10FF); @@ -749,7 +749,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 174); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.class, 0xC); assert_eq!( @@ -761,7 +761,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { CodedIndex::new(TableId::MemberRef, 25, CodedIndexType::MethodDefOrRef) ); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.class, 0x5C); assert_eq!( @@ -773,7 +773,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { CodedIndex::new(TableId::MemberRef, 40, CodedIndexType::MethodDefOrRef) ); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!(row.class, 0x2E9); assert_eq!( @@ -794,15 +794,15 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 29); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.name, 0x1E036); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.name, 0x1DF93); - let row = module.get(25).unwrap(); + let row = module.get(25).unwrap().unwrap(); assert_eq!(row.rid, 25); assert_eq!(row.name, 0x1E09E); } @@ -815,15 +815,15 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 234); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.signature, 0x49); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.signature, 0x67); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!(row.signature, 0x1418); } @@ -836,7 +836,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 422); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.mapping_flags, 0x147); assert_eq!( @@ -846,7 +846,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.import_name, 0x2A5E1); assert_eq!(row.import_scope, 0x2); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.mapping_flags, 0x1120); assert_eq!( @@ -856,7 +856,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.import_name, 0x1BAE); assert_eq!(row.import_scope, 0x3); - let row = module.get(100).unwrap(); + let row = module.get(100).unwrap().unwrap(); assert_eq!(row.rid, 100); assert_eq!(row.mapping_flags, 0x1166); assert_eq!( @@ -875,12 +875,12 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 5); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.rva, 0x3CE40); assert_eq!(row.field, 0x119E); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.rva, 0x3CED0); assert_eq!(row.field, 0x11A2); @@ -894,7 +894,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 1); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.hash_alg_id, 0x8004); assert_eq!(row.major_version, 8); @@ -914,7 +914,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 32); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.major_version, 8); assert_eq!(row.minor_version, 0); @@ -925,7 +925,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.name, 0x15D67); assert_eq!(row.hash_value, 0); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.major_version, 8); assert_eq!(row.minor_version, 0); @@ -936,7 +936,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { assert_eq!(row.name, 0x289E2); assert_eq!(row.hash_value, 0); - let row = module.get(25).unwrap(); + let row = module.get(25).unwrap().unwrap(); assert_eq!(row.rid, 25); assert_eq!(row.major_version, 8); assert_eq!(row.minor_version, 0); @@ -956,7 +956,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 63); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 0x200000); assert_eq!(row.type_def_id, 0); @@ -967,7 +967,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { CodedIndex::new(TableId::AssemblyRef, 11, CodedIndexType::Implementation) ); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.flags, 0x200000); assert_eq!(row.type_def_id, 0); @@ -978,7 +978,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { CodedIndex::new(TableId::AssemblyRef, 11, CodedIndexType::Implementation) ); - let row = module.get(50).unwrap(); + let row = module.get(50).unwrap().unwrap(); assert_eq!(row.rid, 50); assert_eq!(row.flags, 0x200000); assert_eq!(row.type_def_id, 0); @@ -998,7 +998,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 1); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.flags, 1); assert_eq!(row.name, 0x279FC); @@ -1016,17 +1016,17 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 379); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.nested_class, 0x1BA); assert_eq!(row.enclosing_class, 2); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.nested_class, 0x1BE); assert_eq!(row.enclosing_class, 6); - let row = module.get(50).unwrap(); + let row = module.get(50).unwrap().unwrap(); assert_eq!(row.rid, 50); assert_eq!(row.nested_class, 0x1EB); assert_eq!(row.enclosing_class, 6); @@ -1040,7 +1040,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 60); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.number, 0); assert_eq!(row.flags, 0); @@ -1050,7 +1050,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { ); assert_eq!(row.name, 0xB6F1); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!(row.number, 0); assert_eq!(row.flags, 0); @@ -1060,7 +1060,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { ); assert_eq!(row.name, 0xB6F1); - let row = module.get(50).unwrap(); + let row = module.get(50).unwrap().unwrap(); assert_eq!(row.rid, 50); assert_eq!(row.number, 0); assert_eq!(row.flags, 0); @@ -1079,7 +1079,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 37); - let row = module.get(1).unwrap(); + let row = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!( row.method, @@ -1087,7 +1087,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { ); assert_eq!(row.instantiation, 0x343); - let row = module.get(5).unwrap(); + let row = module.get(5).unwrap().unwrap(); assert_eq!(row.rid, 5); assert_eq!( row.method, @@ -1095,7 +1095,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { ); assert_eq!(row.instantiation, 0x50C); - let row = module.get(25).unwrap(); + let row = module.get(25).unwrap().unwrap(); assert_eq!(row.rid, 25); assert_eq!( row.method, @@ -1112,7 +1112,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { Some(module) => { assert_eq!(module.row_count, 3); - let row: GenericParamConstraintRaw = module.get(1).unwrap(); + let row: GenericParamConstraintRaw = module.get(1).unwrap().unwrap(); assert_eq!(row.rid, 1); assert_eq!(row.owner, 0x11); assert_eq!( @@ -1120,7 +1120,7 @@ pub fn verify_tableheader(tables_header: &TablesHeader) { CodedIndex::new(TableId::TypeRef, 73, CodedIndexType::TypeDefOrRef) ); - let row = module.get(3).unwrap(); + let row = module.get(3).unwrap().unwrap(); assert_eq!(row.rid, 3); assert_eq!(row.owner, 0x32); assert_eq!( diff --git a/dotscope/src/utils/crypto.rs b/dotscope/src/utils/crypto.rs index b9b417ce..328ed9e5 100644 --- a/dotscope/src/utils/crypto.rs +++ b/dotscope/src/utils/crypto.rs @@ -46,6 +46,8 @@ use pbkdf2::pbkdf2_hmac; use sha1::{Digest as Sha1Digest, Sha1}; use sha2::{Digest as Sha2Digest, Sha256, Sha384, Sha512}; +use crate::{Error::TypeError, Result}; + // Type aliases for AES CBC modes type Aes128CbcEnc = Encryptor; type Aes128CbcDec = Decryptor; @@ -224,6 +226,24 @@ pub fn compute_hmac_sha512(key: &[u8], data: &[u8]) -> Vec { mac.finalize().into_bytes().to_vec() } +/// Upper bound on the PBKDF2 iteration count accepted by [`derive_pbkdf2_key`]. +/// +/// The iteration count reaching this function is a constant lifted out of an attacker-supplied +/// method body, so it is bounded only by `i32::MAX` (~2.1×10⁹) at the source. PBKDF2 runs its +/// full iteration loop once per output block, so iterations and key length multiply, and this +/// is native computation with no instruction budget or cancellation attached to it. +/// +/// Real-world obfuscators use 100–100 000. Ten million leaves three orders of magnitude of +/// headroom over observed usage while keeping the worst case to seconds rather than days. +pub(crate) const MAX_PBKDF2_ITERATIONS: u32 = 10_000_000; + +/// Upper bound on the derived key length accepted by [`derive_pbkdf2_key`], in bytes. +/// +/// Symmetric keys and IVs in the formats this crate handles are at most 32 and 16 bytes. +/// 1 KiB accommodates any plausible combined key+IV request and rejects the ~2 GB one that an +/// unvalidated `int32` allows. +pub(crate) const MAX_DERIVED_KEY_LEN: usize = 1024; + /// Derives a key using PBKDF2 (RFC 2898 Section 5.2). /// /// PBKDF2 is the algorithm used by .NET's `Rfc2898DeriveBytes` class. @@ -233,14 +253,25 @@ pub fn compute_hmac_sha512(key: &[u8], data: &[u8]) -> Vec { /// /// * `password` - The password bytes /// * `salt` - The salt bytes -/// * `iterations` - Number of iterations (minimum 1000 recommended) -/// * `key_len` - Desired length of the derived key +/// * `iterations` - Number of iterations (minimum 1000 recommended), capped at +/// [`MAX_PBKDF2_ITERATIONS`] +/// * `key_len` - Desired length of the derived key, capped at [`MAX_DERIVED_KEY_LEN`] /// * `hash_algorithm` - Hash algorithm: "SHA1", "SHA256", "SHA384", "SHA512" /// /// # Returns /// /// The derived key bytes of the specified length. /// +/// # Errors +/// +/// Returns [`crate::Error::NotSupported`] if `hash_algorithm` is not one of the recognised +/// names, or if it is "SHA1" in a build without the `legacy-crypto` feature. Returns +/// [`crate::Error::TypeError`] if `iterations` or `key_len` exceeds its ceiling. +/// +/// Unknown algorithms are an error rather than a silent fallback on purpose: substituting a +/// different hash yields plausible-looking bytes of the right length and the wrong value, which +/// an analyst then consumes as ground truth. +/// /// # Example /// /// ```rust,ignore @@ -248,7 +279,7 @@ pub fn compute_hmac_sha512(key: &[u8], data: &[u8]) -> Vec { /// /// let password = b"mypassword"; /// let salt = b"somesalt"; -/// let key = derive_pbkdf2_key(password, salt, 1000, 32, "SHA256"); +/// let key = derive_pbkdf2_key(password, salt, 1000, 32, "SHA256")?; /// assert_eq!(key.len(), 32); /// ``` pub fn derive_pbkdf2_key( @@ -257,7 +288,19 @@ pub fn derive_pbkdf2_key( iterations: u32, key_len: usize, hash_algorithm: &str, -) -> Vec { +) -> Result> { + if iterations > MAX_PBKDF2_ITERATIONS { + return Err(TypeError(format!( + "PBKDF2 iteration count {iterations} exceeds maximum {MAX_PBKDF2_ITERATIONS}" + ))); + } + + if key_len > MAX_DERIVED_KEY_LEN { + return Err(TypeError(format!( + "PBKDF2 derived key length {key_len} exceeds maximum {MAX_DERIVED_KEY_LEN}" + ))); + } + let mut key = vec![0u8; key_len]; match hash_algorithm.to_uppercase().as_str() { @@ -271,18 +314,28 @@ pub fn derive_pbkdf2_key( pbkdf2_hmac::(password, salt, iterations, &mut key); } #[cfg(feature = "legacy-crypto")] - _ => { - // Default to SHA1 (most common in .NET Framework) + "SHA1" => { pbkdf2_hmac::(password, salt, iterations, &mut key); } #[cfg(not(feature = "legacy-crypto"))] - _ => { - // Default to SHA256 when legacy-crypto is disabled - pbkdf2_hmac::(password, salt, iterations, &mut key); + "SHA1" => { + // Deriving with SHA-256 instead would return the right number of bytes and the + // wrong ones. `Rfc2898DeriveBytes` defaults to SHA-1, so this is the common path + // in such a build, not a corner case — name the feature that fixes it. + return Err(TypeError( + "PBKDF2-HMAC-SHA1 requires the `legacy-crypto` feature; \ + rebuild with it enabled to derive this key" + .to_string(), + )); + } + other => { + return Err(TypeError(format!( + "Unsupported PBKDF2 hash algorithm: {other}" + ))); } } - key + Ok(key) } /// Derives a key using PBKDF1 (RFC 2898 Section 5.1). @@ -785,14 +838,20 @@ impl Default for CryptoParameters { /// # Returns /// /// A tuple `(key, iv)` of the derived key material. +/// +/// # Errors +/// +/// Propagates any error from [`derive_pbkdf2_key`] — an iteration count or output length past +/// its ceiling, or an unsupported hash algorithm. Both are reachable from a crafted assembly, +/// since `params` is extracted from an attacker-supplied decryptor body. pub fn derive_key_iv( password: &[u8], salt: &[u8], params: &CryptoParameters, -) -> (Vec, Vec) { +) -> Result<(Vec, Vec)> { // Sizes come from configuration extracted via SSA from a decryptor body. - // Saturating addition avoids overflow on absurd values (in which case - // we degrade to empty key/iv rather than panicking). + // Saturating addition avoids overflow on absurd values; the resulting length is then + // bounds-checked by `derive_pbkdf2_key` rather than being allocated on trust. let output_len = params.key_size.saturating_add(params.iv_size); let derived = derive_pbkdf2_key( password, @@ -800,30 +859,30 @@ pub fn derive_key_iv( params.iterations, output_len, params.hash_algorithm, - ); + )?; let key = derived.get(..params.key_size).unwrap_or(&[]).to_vec(); let iv = derived .get(params.key_size..output_len) .unwrap_or(&[]) .to_vec(); - (key, iv) + Ok((key, iv)) } #[cfg(test)] mod tests { - use crate::utils::crypto::derive_pbkdf2_key; + use super::*; #[test] fn test_pbkdf2_sha256_basic() { let password = b"password"; let salt = b"salt"; - let key = derive_pbkdf2_key(password, salt, 1, 32, "SHA256"); + let key = derive_pbkdf2_key(password, salt, 1, 32, "SHA256").unwrap(); // Verify length assert_eq!(key.len(), 32); // The result should be deterministic - let key2 = derive_pbkdf2_key(password, salt, 1, 32, "SHA256"); + let key2 = derive_pbkdf2_key(password, salt, 1, 32, "SHA256").unwrap(); assert_eq!(key, key2); } @@ -831,7 +890,7 @@ mod tests { fn test_pbkdf2_sha384() { let password = b"test"; let salt = b"salt123"; - let key = derive_pbkdf2_key(password, salt, 1000, 48, "SHA384"); + let key = derive_pbkdf2_key(password, salt, 1000, 48, "SHA384").unwrap(); assert_eq!(key.len(), 48); } @@ -839,7 +898,7 @@ mod tests { fn test_pbkdf2_sha512() { let password = b"test"; let salt = b"salt123"; - let key = derive_pbkdf2_key(password, salt, 1000, 64, "SHA512"); + let key = derive_pbkdf2_key(password, salt, 1000, 64, "SHA512").unwrap(); assert_eq!(key.len(), 64); } @@ -848,8 +907,8 @@ mod tests { let password = b"password"; let salt = b"salt"; - let key_1 = derive_pbkdf2_key(password, salt, 1, 20, "SHA256"); - let key_1000 = derive_pbkdf2_key(password, salt, 1000, 20, "SHA256"); + let key_1 = derive_pbkdf2_key(password, salt, 1, 20, "SHA256").unwrap(); + let key_1000 = derive_pbkdf2_key(password, salt, 1000, 20, "SHA256").unwrap(); // Different iteration counts should produce different keys assert_ne!(key_1, key_1000); @@ -861,8 +920,8 @@ mod tests { let salt1 = b"salt1"; let salt2 = b"salt2"; - let key1 = derive_pbkdf2_key(password, salt1, 1000, 20, "SHA256"); - let key2 = derive_pbkdf2_key(password, salt2, 1000, 20, "SHA256"); + let key1 = derive_pbkdf2_key(password, salt1, 1000, 20, "SHA256").unwrap(); + let key2 = derive_pbkdf2_key(password, salt2, 1000, 20, "SHA256").unwrap(); assert_ne!(key1, key2); } @@ -872,9 +931,9 @@ mod tests { let password = b"test"; let salt = b"salt"; - let key_upper = derive_pbkdf2_key(password, salt, 1, 32, "SHA256"); - let key_lower = derive_pbkdf2_key(password, salt, 1, 32, "sha256"); - let key_mixed = derive_pbkdf2_key(password, salt, 1, 32, "Sha256"); + let key_upper = derive_pbkdf2_key(password, salt, 1, 32, "SHA256").unwrap(); + let key_lower = derive_pbkdf2_key(password, salt, 1, 32, "sha256").unwrap(); + let key_mixed = derive_pbkdf2_key(password, salt, 1, 32, "Sha256").unwrap(); assert_eq!(key_upper, key_lower); assert_eq!(key_upper, key_mixed); @@ -884,7 +943,7 @@ mod tests { fn test_pbkdf2_empty_password() { let password = b""; let salt = b"salt"; - let key = derive_pbkdf2_key(password, salt, 1000, 16, "SHA256"); + let key = derive_pbkdf2_key(password, salt, 1000, 16, "SHA256").unwrap(); assert_eq!(key.len(), 16); // Empty password should still produce non-zero key @@ -895,7 +954,7 @@ mod tests { fn test_pbkdf2_empty_salt() { let password = b"password"; let salt = b""; - let key = derive_pbkdf2_key(password, salt, 1000, 16, "SHA256"); + let key = derive_pbkdf2_key(password, salt, 1000, 16, "SHA256").unwrap(); assert_eq!(key.len(), 16); } @@ -906,14 +965,53 @@ mod tests { let salt = b"salt"; for len in [1, 16, 20, 32, 48, 64, 100] { - let key = derive_pbkdf2_key(password, salt, 1, len, "SHA256"); + let key = derive_pbkdf2_key(password, salt, 1, len, "SHA256").unwrap(); assert_eq!(key.len(), len); } } + /// The iteration ceiling is what stops a crafted decryptor body from wedging the process: + /// the count is an `i32` constant lifted straight out of attacker-supplied CIL. + #[test] + fn test_pbkdf2_rejects_excessive_iterations() { + let password = b"password"; + let salt = b"salt"; + + assert!(derive_pbkdf2_key(password, salt, MAX_PBKDF2_ITERATIONS, 1, "SHA256").is_ok()); + assert!(derive_pbkdf2_key( + password, + salt, + MAX_PBKDF2_ITERATIONS.saturating_add(1), + 32, + "SHA256" + ) + .is_err()); + // The value an unvalidated i32 allows. + assert!(derive_pbkdf2_key(password, salt, 2_147_483_647, 32, "SHA256").is_err()); + } + + /// The length ceiling bounds the output buffer and, with it, the number of PBKDF2 blocks — + /// the loop runs once per block, so length and iterations multiply. + #[test] + fn test_pbkdf2_rejects_excessive_key_length() { + let password = b"password"; + let salt = b"salt"; + + assert!(derive_pbkdf2_key(password, salt, 1, MAX_DERIVED_KEY_LEN, "SHA256").is_ok()); + assert!(derive_pbkdf2_key( + password, + salt, + 1, + MAX_DERIVED_KEY_LEN.saturating_add(1), + "SHA256" + ) + .is_err()); + // Roughly what a maximal i32 key size would request. + assert!(derive_pbkdf2_key(password, salt, 1, 2_000_000_000, "SHA256").is_err()); + } + // NOTE: PBKDF1 tests are in the legacy_tests module below // because they require the legacy-crypto feature (PBKDF1 uses SHA1) - use super::apply_crypto_transform; #[test] fn test_aes_128_encrypt_decrypt() { @@ -1077,19 +1175,19 @@ mod tests { mod legacy_tests { use sha1::Digest; - use crate::utils::crypto::{derive_pbkdf1_key, derive_pbkdf2_key}; + use super::*; #[test] fn test_pbkdf2_sha1_basic() { let password = b"password"; let salt = b"salt"; - let key = derive_pbkdf2_key(password, salt, 1, 20, "SHA1"); + let key = derive_pbkdf2_key(password, salt, 1, 20, "SHA1").unwrap(); // Verify length assert_eq!(key.len(), 20); // The result should be deterministic - let key2 = derive_pbkdf2_key(password, salt, 1, 20, "SHA1"); + let key2 = derive_pbkdf2_key(password, salt, 1, 20, "SHA1").unwrap(); assert_eq!(key, key2); } @@ -1097,22 +1195,25 @@ mod legacy_tests { fn test_pbkdf2_sha256_vs_sha1() { let password = b"password"; let salt = b"salt"; - let key_sha256 = derive_pbkdf2_key(password, salt, 1, 32, "SHA256"); - let key_sha1 = derive_pbkdf2_key(password, salt, 1, 32, "SHA1"); + let key_sha256 = derive_pbkdf2_key(password, salt, 1, 32, "SHA256").unwrap(); + let key_sha1 = derive_pbkdf2_key(password, salt, 1, 32, "SHA1").unwrap(); // Different algorithms produce different keys assert_ne!(key_sha256, key_sha1); } + /// An unrecognised algorithm name must be an error, not a silent substitution. + /// + /// A fallback to SHA-1 would return bytes of the correct length and the wrong value, which + /// an analyst then consumes as ground truth with no signal that the crypto was substituted. #[test] - fn test_pbkdf2_unknown_algorithm_defaults_to_sha1() { + fn test_pbkdf2_unknown_algorithm_is_rejected() { let password = b"test"; let salt = b"salt"; - let key_unknown = derive_pbkdf2_key(password, salt, 1000, 20, "UNKNOWN"); - let key_sha1 = derive_pbkdf2_key(password, salt, 1000, 20, "SHA1"); - - assert_eq!(key_unknown, key_sha1); + assert!(derive_pbkdf2_key(password, salt, 1000, 20, "UNKNOWN").is_err()); + assert!(derive_pbkdf2_key(password, salt, 1000, 20, "").is_err()); + assert!(derive_pbkdf2_key(password, salt, 1000, 20, "MD5").is_err()); } #[test] @@ -1244,7 +1345,7 @@ mod legacy_tests { let salt = b"salt"; let key1 = derive_pbkdf1_key(password, salt, 100, 20); - let key2 = derive_pbkdf2_key(password, salt, 100, 20, "SHA1"); + let key2 = derive_pbkdf2_key(password, salt, 100, 20, "SHA1").unwrap(); // PBKDF1 and PBKDF2 are different algorithms assert_ne!(key1, key2); @@ -1252,8 +1353,6 @@ mod legacy_tests { // DES encryption tests - use super::apply_crypto_transform; - #[test] fn test_des_encrypt_decrypt() { let key = [0u8; 8]; // 56-bit key (8 bytes with parity) diff --git a/dotscope/src/utils/decompress.rs b/dotscope/src/utils/decompress.rs index a967c44c..e8c523bd 100644 --- a/dotscope/src/utils/decompress.rs +++ b/dotscope/src/utils/decompress.rs @@ -40,6 +40,14 @@ pub enum DecompressError { DeflateError(String), /// Input buffer too small. BufferTooSmall, + /// Decompressed output exceeded the size limit. + /// + /// Compression ratios above 1000:1 are trivial to construct, so a few kilobytes of + /// attacker-supplied input can otherwise expand until the host runs out of memory. + OutputTooLarge { + /// The limit that was exceeded, in bytes. + limit: usize, + }, } impl std::fmt::Display for DecompressError { @@ -49,13 +57,16 @@ impl std::fmt::Display for DecompressError { Self::LzmaError(msg) => write!(f, "LZMA decompression error: {msg}"), Self::DeflateError(msg) => write!(f, "Deflate decompression error: {msg}"), Self::BufferTooSmall => write!(f, "Input buffer too small"), + Self::OutputTooLarge { limit } => { + write!(f, "Decompressed output exceeds the {limit} byte limit") + } } } } impl std::error::Error for DecompressError {} -/// Upper bound on a declared decompressed size, in bytes. +/// Upper bound on decompressed output, in bytes. /// /// This is a guard against corrupt or hostile length fields driving a large /// allocation, **not** a property of the format — LZMA itself allows any @@ -64,7 +75,7 @@ impl std::error::Error for DecompressError {} /// kilobytes for typical programs and single-digit megabytes for pathological /// ones, so 512 MB leaves several orders of magnitude of headroom while still /// rejecting a field that is plainly nonsense. -const MAX_DECOMPRESSED_SIZE: u64 = 512 * 1024 * 1024; +pub const MAX_DECOMPRESSED_BYTES: usize = 512 * 1024 * 1024; /// Header layouts used by ConfuserEx and its forks. /// @@ -75,6 +86,40 @@ const LZMA_HEADER_LAYOUTS: [usize; 2] = [ 9, // 5 props + 4-byte size — stock ConfuserEx's `Lzma.Decompress` ]; +/// A `Write` sink that refuses to grow past a byte ceiling. +/// +/// The LZMA decoder writes into a sink rather than being read from, so it cannot be bounded +/// with [`std::io::Read::take`] the way the Deflate and GZip paths are. This is the +/// equivalent: expansion stops *during* the decode instead of being measured afterwards, so a +/// decompression bomb never commits more than `limit` bytes of host memory. +/// +/// Bounding here rather than through LZMA's declared-size field is deliberate. lzma-rs treats +/// any size other than the all-ones "unknown" marker as an *exact* expected output length and +/// fails the stream when the decoded length differs, so writing a ceiling into that field +/// rejects every unknown-size stream whose real output is shorter — which is the normal case +/// for the ConfuserEx payloads this module decodes. +struct LimitedWriter { + buffer: Vec, + limit: usize, +} + +impl std::io::Write for LimitedWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + if self.buffer.len().saturating_add(buf.len()) > self.limit { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "decompressed output exceeds the configured limit", + )); + } + self.buffer.extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + /// Validates the 5 LZMA property bytes shared by every supported layout. /// /// Checks the properties byte encodes a legal `lc`/`lp`/`pb` triple and that @@ -101,7 +146,7 @@ fn valid_lzma_props(data: &[u8]) -> bool { /// Reads the declared uncompressed size for a given header length. /// /// Returns `None` when the buffer is too short, or the size is zero or beyond -/// [`MAX_DECOMPRESSED_SIZE`]. An all-ones field means "unknown" and maps to +/// [`MAX_DECOMPRESSED_BYTES`]. An all-ones field means "unknown" and maps to /// `Some(None)` — plausible, but with no length to verify against. fn declared_size(data: &[u8], header_len: usize) -> Option> { let size = match header_len { @@ -112,7 +157,9 @@ fn declared_size(data: &[u8], header_len: usize) -> Option> { if size == u64::MAX || size == u64::from(u32::MAX) { return Some(None); } - if size == 0 || size > MAX_DECOMPRESSED_SIZE { + // Widen the budget rather than narrowing `size`: on a 32-bit host a `u64` header field can + // exceed `usize::MAX`, and narrowing would wrap a nonsense value into an acceptable one. + if size == 0 || size > u64::try_from(MAX_DECOMPRESSED_BYTES).unwrap_or(u64::MAX) { return None; } Some(Some(size)) @@ -187,6 +234,23 @@ pub fn is_confuserex_lzma(data: &[u8]) -> bool { /// wrong shifts the payload by four bytes and corrupts the range coder, so a /// successful decode is itself the discriminator. pub fn decompress_confuserex_lzma(data: &[u8]) -> DecompressResult> { + decompress_confuserex_lzma_limited(data, MAX_DECOMPRESSED_BYTES) +} + +/// Decompresses a ConfuserEx LZMA payload, refusing output larger than `limit` bytes. +/// +/// # Arguments +/// +/// * `data` - The LZMA payload, including its header. +/// * `limit` - Maximum accepted output size in bytes. +/// +/// # Errors +/// +/// Returns [`DecompressError::OutputTooLarge`] if the declared or decoded size exceeds +/// `limit`, [`DecompressError::InvalidLzmaHeader`] if no header layout decodes cleanly, +/// [`DecompressError::LzmaError`] if the stream is malformed, or +/// [`DecompressError::BufferTooSmall`] if `data` is too short to contain a header. +pub fn decompress_confuserex_lzma_limited(data: &[u8], limit: usize) -> DecompressResult> { if data.len() < 9 { return Err(DecompressError::BufferTooSmall); } @@ -208,16 +272,41 @@ pub fn decompress_confuserex_lzma(data: &[u8]) -> DecompressResult> { continue; } + // A declared size above the limit is refused before decoding, so a header claiming a + // multi-gigabyte payload costs nothing. + if size.is_some_and(|declared| declared > limit as u64) { + last_err = Some(DecompressError::OutputTooLarge { limit }); + continue; + } + // lzma-rs expects the alone format: 5 props + 8-byte size + payload. + // + // The size field is passed through verbatim, including the all-ones "unknown" marker. + // lzma-rs treats *any* other value as an exact expected length and fails the stream + // when the decoded output differs (`decode::lzma`'s `len != output.len()` check), so + // substituting the limit for an undeclared size — which looks like a safe upper bound — + // actually rejects every unknown-size stream whose real output is shorter. That is the + // normal case for the ConfuserEx payloads this function exists to decode. + // + // The ceiling is enforced by `LimitedWriter` instead, which stops the expansion as it + // happens rather than after the fact. let mut lzma_stream = Vec::with_capacity(compressed.len().saturating_add(13)); lzma_stream.extend_from_slice(props); lzma_stream.extend_from_slice(&size.unwrap_or(u64::MAX).to_le_bytes()); lzma_stream.extend_from_slice(compressed); let mut cursor = Cursor::new(&lzma_stream); - let mut decompressed = Vec::new(); - match lzma_rs::lzma_decompress(&mut cursor, &mut decompressed) { + let mut sink = LimitedWriter { + buffer: Vec::new(), + limit, + }; + match lzma_rs::lzma_decompress(&mut cursor, &mut sink) { Ok(()) => { + let decompressed = sink.buffer; + if decompressed.len() > limit { + last_err = Some(DecompressError::OutputTooLarge { limit }); + continue; + } // A wrong layout can still decode into garbage of the wrong // length; hold the result to its declared size when known. if size.is_none_or(|expected| decompressed.len() as u64 == expected) { @@ -242,13 +331,34 @@ pub fn decompress_confuserex_lzma(data: &[u8]) -> DecompressResult> { /// /// The decompressed data, or an error if decompression fails. pub fn decompress_deflate(data: &[u8]) -> DecompressResult> { - let mut decoder = DeflateDecoder::new(data); + decompress_deflate_limited(data, MAX_DECOMPRESSED_BYTES) +} + +/// Decompresses Deflate data, refusing output larger than `limit` bytes. +/// +/// # Arguments +/// +/// * `data` - The Deflate compressed data. +/// * `limit` - Maximum accepted output size in bytes. +/// +/// # Errors +/// +/// Returns [`DecompressError::OutputTooLarge`] if the stream expands past `limit`, or +/// [`DecompressError::DeflateError`] if the stream is malformed. +pub fn decompress_deflate_limited(data: &[u8], limit: usize) -> DecompressResult> { + // `Read::take` bounds the decoder itself, so the expansion stops at the limit instead of + // being detected after the memory has already been committed. + let mut decoder = DeflateDecoder::new(data).take(limit.saturating_add(1) as u64); let mut decompressed = Vec::new(); decoder .read_to_end(&mut decompressed) .map_err(|e| DecompressError::DeflateError(e.to_string()))?; + if decompressed.len() > limit { + return Err(DecompressError::OutputTooLarge { limit }); + } + Ok(decompressed) } @@ -262,13 +372,32 @@ pub fn decompress_deflate(data: &[u8]) -> DecompressResult> { /// /// The decompressed data, or an error if decompression fails. pub fn decompress_gzip(data: &[u8]) -> DecompressResult> { - let mut decoder = GzDecoder::new(data); + decompress_gzip_limited(data, MAX_DECOMPRESSED_BYTES) +} + +/// Decompresses GZip data, refusing output larger than `limit` bytes. +/// +/// # Arguments +/// +/// * `data` - The GZip compressed data. +/// * `limit` - Maximum accepted output size in bytes. +/// +/// # Errors +/// +/// Returns [`DecompressError::OutputTooLarge`] if the stream expands past `limit`, or +/// [`DecompressError::DeflateError`] if the stream is malformed. +pub fn decompress_gzip_limited(data: &[u8], limit: usize) -> DecompressResult> { + let mut decoder = GzDecoder::new(data).take(limit.saturating_add(1) as u64); let mut decompressed = Vec::new(); decoder .read_to_end(&mut decompressed) .map_err(|e| DecompressError::DeflateError(e.to_string()))?; + if decompressed.len() > limit { + return Err(DecompressError::OutputTooLarge { limit }); + } + Ok(decompressed) } @@ -402,4 +531,122 @@ mod tests { let decompressed = decompress_gzip(&compressed).unwrap(); assert_eq!(&decompressed, original); } + + /// Compresses `len` zero bytes, which deflate reduces to a tiny stream. This is the + /// decompression-bomb shape: a few hundred input bytes expanding to megabytes. + fn deflate_zeros(len: usize) -> Vec { + let mut encoder = DeflateEncoder::new(Vec::new(), Compression::best()); + encoder.write_all(&vec![0u8; len]).unwrap(); + encoder.finish().unwrap() + } + + fn gzip_zeros(len: usize) -> Vec { + let mut encoder = GzEncoder::new(Vec::new(), Compression::best()); + encoder.write_all(&vec![0u8; len]).unwrap(); + encoder.finish().unwrap() + } + + #[test] + fn deflate_refuses_output_above_the_limit() { + let bomb = deflate_zeros(1024 * 1024); + assert!(bomb.len() < 4096, "input should be small: {}", bomb.len()); + assert!(matches!( + decompress_deflate_limited(&bomb, 64 * 1024), + Err(DecompressError::OutputTooLarge { .. }) + )); + } + + #[test] + fn deflate_accepts_output_at_the_limit() { + let payload = deflate_zeros(64 * 1024); + let out = decompress_deflate_limited(&payload, 64 * 1024).unwrap(); + assert_eq!(out.len(), 64 * 1024); + } + + #[test] + fn gzip_refuses_output_above_the_limit() { + let bomb = gzip_zeros(1024 * 1024); + assert!(matches!( + decompress_gzip_limited(&bomb, 64 * 1024), + Err(DecompressError::OutputTooLarge { .. }) + )); + } + + #[test] + fn gzip_accepts_output_at_the_limit() { + let payload = gzip_zeros(64 * 1024); + let out = decompress_gzip_limited(&payload, 64 * 1024).unwrap(); + assert_eq!(out.len(), 64 * 1024); + } + + #[test] + fn default_entry_points_carry_a_limit() { + // The unlimited-looking wrappers must still be bounded. + assert_eq!(MAX_DECOMPRESSED_BYTES, 512 * 1024 * 1024); + let payload = deflate_zeros(1024); + assert_eq!(decompress_deflate(&payload).unwrap().len(), 1024); + } + + #[test] + fn lzma_declared_size_above_limit_is_refused_before_decoding() { + // Valid property bytes, then an 8-byte declared size of 256 MB, against a 1 KB limit. + let mut data = vec![0x5D, 0x00, 0x00, 0x10, 0x00]; + data.extend_from_slice(&(256u64 * 1024 * 1024).to_le_bytes()); + data.extend_from_slice(&[0u8; 32]); + + assert!(matches!( + decompress_confuserex_lzma_limited(&data, 1024), + Err(DecompressError::OutputTooLarge { .. } | DecompressError::InvalidLzmaHeader) + )); + } + + /// Builds a real LZMA stream whose header carries the all-ones "unknown size" marker, + /// which is what `lzma_compress`'s default options emit. + fn unknown_size_stream(payload: &[u8]) -> Vec { + let mut compressed = Vec::new(); + lzma_rs::lzma_compress(&mut std::io::Cursor::new(payload), &mut compressed) + .expect("compressing a fixed payload cannot fail"); + + // Sanity: the header really does declare "unknown". + assert_eq!( + compressed.get(5..13), + Some(u64::MAX.to_le_bytes().as_slice()), + "expected the unknown-size marker in the 8-byte size field" + ); + compressed + } + + /// A stream that declares "unknown size" must still decode. + /// + /// The size field is passed to lzma-rs verbatim, and lzma-rs enforces any value other than + /// the all-ones marker as an *exact* output length. Substituting the byte limit there — an + /// apparently safe upper bound — therefore made every unknown-size stream whose real output + /// is shorter fail with a length mismatch, which silently disabled ConfuserEx constant + /// decryption. The previous test could not catch it: it fed a garbage payload and asserted + /// only `if let Ok(...)`, so the failing path satisfied it. + #[test] + fn lzma_unknown_declared_size_still_decodes() { + let payload = b"the quick brown fox jumps over the lazy dog".repeat(8); + let compressed = unknown_size_stream(&payload); + + let out = decompress_confuserex_lzma_limited(&compressed, 64 * 1024) + .expect("an unknown-size stream must decode"); + assert_eq!(out, payload); + } + + /// ...and is still bounded while doing so. + /// + /// With no declared size there is no length to pre-check, so the ceiling has to be enforced + /// during the decode. `LimitedWriter` is what does that; a limit below the true output must + /// refuse rather than return a truncated buffer. + #[test] + fn lzma_unknown_declared_size_is_still_bounded() { + let payload = vec![0x41u8; 8192]; + let compressed = unknown_size_stream(&payload); + + assert!( + decompress_confuserex_lzma_limited(&compressed, 1024).is_err(), + "an unknown-size stream expanding past the limit must be refused" + ); + } } diff --git a/dotscope/src/utils/io.rs b/dotscope/src/utils/io.rs index af5d7399..73accdcb 100644 --- a/dotscope/src/utils/io.rs +++ b/dotscope/src/utils/io.rs @@ -144,7 +144,7 @@ //! //! # Error Handling //! -//! All reading and writing functions return [`crate::Result`] and will return [`crate::Error::OutOfBounds`] +//! All reading and writing functions return [`crate::Result`] and will return [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] //! if there are insufficient bytes in the buffer to complete the operation. This ensures //! memory safety and prevents buffer overruns during parsing and generation. //! @@ -485,7 +485,7 @@ impl CilIO for isize { /// /// # Returns /// -/// Returns the decoded value or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns the decoded value or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// @@ -518,7 +518,7 @@ pub fn read_le(data: &[u8]) -> Result { /// /// # Returns /// -/// Returns the decoded value or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns the decoded value or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// @@ -569,7 +569,7 @@ pub fn read_le_at(data: &[u8], offset: &mut usize) -> Result { /// /// # Returns /// -/// Returns the decoded value as u32, or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns the decoded value as u32, or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// @@ -617,7 +617,7 @@ pub fn read_le_at_dyn(data: &[u8], offset: &mut usize, is_large: bool) -> Result /// /// # Returns /// -/// Returns the decoded value or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns the decoded value or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// @@ -651,7 +651,7 @@ pub fn read_be(data: &[u8]) -> Result { /// /// # Returns /// -/// Returns the decoded value or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns the decoded value or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// @@ -703,7 +703,7 @@ pub fn read_be_at(data: &[u8], offset: &mut usize) -> Result { /// /// # Returns /// -/// Returns the decoded value as u32, or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns the decoded value as u32, or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// @@ -751,7 +751,7 @@ pub fn read_be_at_dyn(data: &[u8], offset: &mut usize, is_large: bool) -> Result /// /// # Returns /// -/// Returns `Ok(())` on success or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns `Ok(())` on success or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// @@ -786,7 +786,7 @@ pub fn write_le(data: &mut [u8], value: T) -> Result<()> { /// /// # Returns /// -/// Returns `Ok(())` on success or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns `Ok(())` on success or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// @@ -837,7 +837,7 @@ pub fn write_le_at(data: &mut [u8], offset: &mut usize, value: T) -> R /// /// # Returns /// -/// Returns `Ok(())` on success or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns `Ok(())` on success or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// @@ -891,7 +891,7 @@ pub fn write_le_at_dyn( /// /// # Returns /// -/// Returns `Ok(())` on success or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns `Ok(())` on success or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// @@ -927,7 +927,7 @@ pub fn write_be(data: &mut [u8], value: T) -> Result<()> { /// /// # Returns /// -/// Returns `Ok(())` on success or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns `Ok(())` on success or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// @@ -979,7 +979,7 @@ pub fn write_be_at(data: &mut [u8], offset: &mut usize, value: T) -> R /// /// # Returns /// -/// Returns `Ok(())` on success or [`crate::Error::OutOfBounds`] if there are insufficient bytes. +/// Returns `Ok(())` on success or [`crate::Error::Parse`] carrying [`crate::ParseFailure::OutOfBounds`] if there are insufficient bytes. /// /// # Examples /// diff --git a/dotscope/src/utils/lazylist.rs b/dotscope/src/utils/lazylist.rs new file mode 100644 index 00000000..e277d9ba --- /dev/null +++ b/dotscope/src/utils/lazylist.rs @@ -0,0 +1,239 @@ +//! Lazily-allocated append-only list. +//! +//! [`LazyList`] wraps `Arc>` behind a [`OnceLock`], allocating the backing vector +//! on first push rather than at construction. +//! +//! # Why this exists +//! +//! `boxcar::Vec` is not a thin handle. It is `{ inflight: AtomicUsize, buckets: [AtomicPtr; 58], +//! count: AtomicUsize }` — 480 bytes inline, and `Buckets::new` materialises all 58 slots as +//! null rather than allocating them lazily. With the `Arc` header that is ~496 bytes for an +//! empty list. +//! +//! Metadata rows carry many such lists that are almost always empty: a method rarely has +//! varargs, generic arguments, overrides, local variables or interface implementations. Eagerly +//! constructing them made a 14-byte MethodDef row cost roughly 4 KB of zeroed bucket arrays — +//! a 300–500× amplification from table bytes to resident memory, driven entirely by attacker- +//! controlled row counts. +//! +//! # Concurrency +//! +//! Unchanged from the eager form. `OnceLock` makes first-push initialisation race-free, and +//! `boxcar::Vec` itself provides the append-only concurrent semantics; [`push`](LazyList::push) +//! still takes `&self`. + +use std::{ + fmt, + sync::{Arc, OnceLock}, +}; + +/// An append-only concurrent list whose backing storage is allocated on first use. +/// +/// See the [module docs](self) for why the deferral matters. +pub struct LazyList { + inner: OnceLock>>, +} + +impl LazyList { + /// Creates an empty list that has not allocated its backing storage. + #[must_use] + pub fn new() -> Self { + Self { + inner: OnceLock::new(), + } + } + + /// Returns the backing vector, allocating it if this is the first use. + fn materialize(&self) -> &Arc> { + self.inner.get_or_init(|| Arc::new(boxcar::Vec::new())) + } + + /// Appends a value, returning its index. + /// + /// Allocates the backing vector on the first call. + pub fn push(&self, value: T) -> usize { + self.materialize().push(value) + } + + /// Returns the number of elements. + /// + /// Does not allocate: an uninitialised list is empty by definition. + #[must_use] + pub fn count(&self) -> usize { + self.inner.get().map_or(0, |list| list.count()) + } + + /// Returns the number of elements. Alias of [`count`](Self::count), matching `boxcar::Vec`, + /// which exposes both. + #[must_use] + pub fn len(&self) -> usize { + self.count() + } + + /// Returns whether the list has no elements. + #[must_use] + pub fn is_empty(&self) -> bool { + self.count() == 0 + } + + /// Returns the element at `index`, or `None` if out of bounds. + #[must_use] + pub fn get(&self, index: usize) -> Option<&T> { + self.inner.get()?.get(index) + } + + /// Returns an iterator over `(index, &value)` pairs. + /// + /// Does not allocate when the list was never pushed to. + #[must_use] + pub fn iter(&self) -> LazyListIter<'_, T> { + LazyListIter { + inner: self.inner.get().map(|list| list.iter()), + } + } +} + +impl Default for LazyList { + fn default() -> Self { + Self::new() + } +} + +impl Clone for LazyList { + /// Clones the handle, sharing storage with the original. + /// + /// Deliberately materialises the backing vector: the field this type replaces was an `Arc`, + /// so clones shared their contents and a push through one was visible through the other. + /// Cloning an uninitialised list lazily would hand back two independent lists and silently + /// break that. Clones are rare compared with the empty lists this type exists to avoid, so + /// paying an allocation here is the right side of the trade. + fn clone(&self) -> Self { + let shared = Arc::clone(self.materialize()); + Self { + inner: OnceLock::from(shared), + } + } +} + +impl fmt::Debug for LazyList { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list().entries(self.iter().map(|(_, v)| v)).finish() + } +} + +impl<'a, T> IntoIterator for &'a LazyList { + type Item = (usize, &'a T); + type IntoIter = LazyListIter<'a, T>; + + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +impl FromIterator for LazyList { + fn from_iter>(iter: I) -> Self { + let list = Self::new(); + for value in iter { + list.push(value); + } + list + } +} + +/// Iterator over a [`LazyList`], yielding `(index, &value)`. +/// +/// Yields nothing when the list never allocated. +pub struct LazyListIter<'a, T> { + inner: Option>, +} + +impl<'a, T> Iterator for LazyListIter<'a, T> { + type Item = (usize, &'a T); + + fn next(&mut self) -> Option { + self.inner.as_mut()?.next() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn empty_list_does_not_allocate() { + let list: LazyList = LazyList::new(); + + assert!(list.is_empty()); + assert_eq!(list.count(), 0); + assert_eq!(list.iter().count(), 0); + assert!(list.get(0).is_none()); + assert!( + list.inner.get().is_none(), + "querying an empty list must not materialise the backing vector" + ); + } + + #[test] + fn push_materializes_and_stores() { + let list = LazyList::new(); + + assert_eq!(list.push(10), 0); + assert_eq!(list.push(20), 1); + + assert_eq!(list.count(), 2); + assert!(!list.is_empty()); + assert_eq!(list.get(0), Some(&10)); + assert_eq!(list.get(1), Some(&20)); + assert_eq!(list.get(2), None); + + let collected: Vec<(usize, u32)> = list.iter().map(|(i, v)| (i, *v)).collect(); + assert_eq!(collected, vec![(0, 10), (1, 20)]); + } + + /// Clones must share storage, matching the `Arc>` this replaces. + #[test] + fn clones_share_storage() { + let original = LazyList::new(); + original.push(1); + + let cloned = original.clone(); + cloned.push(2); + + assert_eq!( + original.count(), + 2, + "a push through the clone must be visible" + ); + assert_eq!(cloned.count(), 2); + } + + /// Sharing must hold even when the list had not been pushed to before cloning — the case a + /// naive lazy clone would get wrong. + #[test] + fn clones_of_empty_lists_still_share() { + let original: LazyList = LazyList::new(); + let cloned = original.clone(); + + cloned.push(7); + + assert_eq!(original.count(), 1); + assert_eq!(original.get(0), Some(&7)); + } + + #[test] + fn collects_from_iterator() { + let list: LazyList = (0..4).collect(); + assert_eq!(list.count(), 4); + assert_eq!(list.get(3), Some(&3)); + } + + #[test] + fn iterates_by_reference() { + let list: LazyList = (0..3).collect(); + let mut total = 0; + for (_, v) in &list { + total += *v; + } + assert_eq!(total, 3); + } +} diff --git a/dotscope/src/utils/mod.rs b/dotscope/src/utils/mod.rs index 802ed31f..a1271136 100644 --- a/dotscope/src/utils/mod.rs +++ b/dotscope/src/utils/mod.rs @@ -51,9 +51,11 @@ mod enums; mod hash; mod heap_calc; mod io; +mod lazylist; mod lebytes; mod math; mod synchronization; +mod text; mod visitedmap; pub use alignment::align_to; @@ -67,7 +69,10 @@ pub(crate) use crypto::derive_pbkdf1_key; #[cfg(feature = "deobfuscation")] pub(crate) use crypto::CryptoParameters; #[cfg(feature = "emulation")] -pub(crate) use crypto::{apply_crypto_transform, derive_pbkdf2_key, verify_rsa_pkcs1v15}; +pub(crate) use crypto::{ + apply_crypto_transform, derive_pbkdf2_key, verify_rsa_pkcs1v15, MAX_DERIVED_KEY_LEN, + MAX_PBKDF2_ITERATIONS, +}; #[cfg(feature = "emulation")] pub(crate) use crypto::{compute_hmac_sha256, compute_hmac_sha512}; #[cfg(feature = "legacy-crypto")] @@ -75,7 +80,9 @@ pub(crate) use crypto::{compute_md5, compute_sha1}; pub(crate) use crypto::{compute_sha256, compute_sha384, compute_sha512}; #[allow(unused_imports)] pub(crate) use decompress::{ - decompress_confuserex_lzma, decompress_deflate, decompress_gzip, is_confuserex_lzma, + decompress_confuserex_lzma, decompress_confuserex_lzma_limited, decompress_deflate, + decompress_deflate_limited, decompress_gzip, decompress_gzip_limited, is_confuserex_lzma, + MAX_DECOMPRESSED_BYTES, }; pub use dot::escape_dot; pub use enums::EnumUtils; @@ -90,10 +97,14 @@ pub use io::{ write_le_at_dyn, write_prefixed_string_utf16, write_prefixed_string_utf8, write_string_at, write_string_utf8, CilIO, }; +pub use lazylist::LazyList; +#[allow(unused_imports)] +pub use lazylist::LazyListIter; #[cfg(feature = "emulation")] pub use lebytes::LeBytes; #[cfg(feature = "emulation")] pub use math::to_i32_saturating; pub use math::to_u32; pub use synchronization::FailFastBarrier; +pub(crate) use text::truncate_chars; pub use visitedmap::VisitedMap; diff --git a/dotscope/src/utils/text.rs b/dotscope/src/utils/text.rs new file mode 100644 index 00000000..b5d6f23d --- /dev/null +++ b/dotscope/src/utils/text.rs @@ -0,0 +1,111 @@ +//! Character-boundary-safe text helpers. +//! +//! Rust's `str` indexing is by *byte* offset, so the natural-looking `&s[..n]` panics whenever +//! byte `n` lands inside a multi-byte UTF-8 character. Nearly every string this crate handles — +//! user-string literals, type and member names, emulated heap contents — comes from an +//! attacker-controlled assembly, so that panic is reachable input-driven behaviour rather than a +//! programming edge case. +//! +//! This is a blind spot in the crate's lint set: `clippy::indexing_slicing` only fires for types +//! that deref to a slice or array, and neither `str` nor `String` does. `clippy::string_slice` is +//! denied crate-wide to close it; use the helpers here instead of range-indexing text. + +/// Truncates `s` to at most `max_chars` characters without splitting a UTF-8 character. +/// +/// Returns `s` unchanged when it is already `max_chars` characters or shorter. +/// +/// # Arguments +/// +/// * `s` - The string to truncate. +/// * `max_chars` - Maximum number of characters to keep. +/// +/// # Returns +/// +/// A prefix of `s` containing at most `max_chars` characters. +/// +/// # Note +/// +/// The bound is in characters, not bytes, so the returned slice may be up to four times +/// `max_chars` bytes long. Callers truncating for display or for identifier length want the +/// character count; a caller that must respect a hard byte budget needs a different helper. +/// +/// # Examples +/// +/// ```rust,ignore +/// use dotscope::utils::truncate_chars; +/// +/// assert_eq!(truncate_chars("hello", 3), "hel"); +/// assert_eq!(truncate_chars("hello", 10), "hello"); +/// // A byte-offset slice would panic here; this splits between characters. +/// assert_eq!(truncate_chars("日本語", 2), "日本"); +/// ``` +pub(crate) fn truncate_chars(s: &str, max_chars: usize) -> &str { + match s.char_indices().nth(max_chars) { + Some((byte_idx, _)) => s.split_at(byte_idx).0, + None => s, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn shorter_than_limit_is_unchanged() { + assert_eq!(truncate_chars("abc", 10), "abc"); + assert_eq!(truncate_chars("", 5), ""); + } + + #[test] + fn exact_length_is_unchanged() { + assert_eq!(truncate_chars("abc", 3), "abc"); + } + + #[test] + fn ascii_truncates_at_limit() { + assert_eq!(truncate_chars("abcdef", 3), "abc"); + assert_eq!(truncate_chars("abc", 0), ""); + } + + /// The case that panics with `&s[..n]`: the limit falls inside a multi-byte character. + #[test] + fn multibyte_never_splits_a_character() { + // Each of these is 3 bytes, so byte offset 2 is mid-character. + assert_eq!(truncate_chars("日本語", 2), "日本"); + assert_eq!(truncate_chars("日本語", 1), "日"); + assert_eq!(truncate_chars("日本語", 3), "日本語"); + assert_eq!(truncate_chars("日本語", 99), "日本語"); + } + + /// Four-byte characters (astral plane) are the widest case. + #[test] + fn handles_four_byte_characters() { + let s = "𝄞𝄞𝄞"; + assert_eq!(s.len(), 12); + assert_eq!(truncate_chars(s, 1), "𝄞"); + assert_eq!(truncate_chars(s, 2).chars().count(), 2); + } + + /// A mixed string where the boundary lands differently than the byte count suggests. + #[test] + fn mixed_width_input() { + let s = "ab日cd"; + assert_eq!(s.len(), 7); + assert_eq!(truncate_chars(s, 3), "ab日"); + assert_eq!(truncate_chars(s, 4), "ab日c"); + } + + /// Every prefix length must be valid UTF-8 and round-trip through `chars()`. + #[test] + fn all_prefixes_are_valid_for_adversarial_input() { + let s = "a日b𝄞c\u{0301}d"; + for n in 0..=s.chars().count().saturating_add(3) { + let got = truncate_chars(s, n); + assert!( + s.starts_with(got), + "prefix {n} is not a prefix of the input" + ); + assert!(got.chars().count() <= n); + } + } +} diff --git a/dotscope/tests/bitmono.rs b/dotscope/tests/bitmono.rs index c17f6aaa..71433b54 100644 --- a/dotscope/tests/bitmono.rs +++ b/dotscope/tests/bitmono.rs @@ -37,7 +37,10 @@ use common::{ }; use dotscope::{ deobfuscation::{DeobfuscationEngine, DeobfuscationResult, EngineConfig}, - metadata::validation::ValidationConfig, + metadata::{ + tables::{AssemblyRefRaw, TypeRefRaw}, + validation::ValidationConfig, + }, CilObject, }; @@ -1258,8 +1261,6 @@ fn test_dotnethook_offset_diagnostic() { #[test] fn test_no_obfuscator_metadata_survives() { - use dotscope::metadata::tables::{AssemblyRefRaw, TypeRefRaw}; - let path = format!("{}/bitmono_maximum_il.exe", SAMPLES_DIR); if !std::path::Path::new(&path).exists() { eprintln!("Skipping: not found"); @@ -1279,6 +1280,7 @@ fn test_no_obfuscator_metadata_survives() { if let Some(aref_table) = tables.table::() { for aref in aref_table { + let aref = aref.expect("row parses"); let name = strings.get(aref.name as usize).unwrap_or("???"); // System.Private.CoreLib may survive legitimately if // /__StaticArrayInitTypeSize=N types @@ -1292,6 +1294,7 @@ fn test_no_obfuscator_metadata_survives() { if let Some(typeref_table) = tables.table::() { for tr in typeref_table { + let tr = tr.expect("row parses"); let name = strings.get(tr.type_name as usize).unwrap_or("???"); // System.ValueType may survive legitimately: __StaticArrayInitTypeSize=N // nested value types extend it for RuntimeHelpers.InitializeArray support. diff --git a/dotscope/tests/common/verification.rs b/dotscope/tests/common/verification.rs index 48958e15..f226e244 100644 --- a/dotscope/tests/common/verification.rs +++ b/dotscope/tests/common/verification.rs @@ -1183,7 +1183,7 @@ impl AssemblyStats { continue; } - type_names.insert(fullname); + type_names.insert(fullname.to_string()); // Count methods and fields on this type method_count += cil_type.methods().count(); diff --git a/dotscope/tests/fuzzer.rs b/dotscope/tests/fuzzer.rs index 7616f497..c1cbffbc 100644 --- a/dotscope/tests/fuzzer.rs +++ b/dotscope/tests/fuzzer.rs @@ -1,38 +1,96 @@ -//! Fuzzer corpus regression tests — load every file under `fuzz/corpus/` and -//! `fuzz/artifacts/` through `CilObject::from_path` and assert we don't panic. -//! Files are intentionally malformed; errors are expected, crashes are not. +#![allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::arithmetic_side_effects, + clippy::indexing_slicing, + missing_docs +)] -use std::{fs, path::PathBuf}; +//! Regression tests over adversarial inputs. +//! +//! Every file loaded here is malformed on purpose. A parse error is the expected outcome and +//! is not a failure; the property under test is that `CilObject::from_path` returns rather +//! than panicking, aborting, hanging, or exhausting memory. +//! +//! # Two corpora, two guarantees +//! +//! - [`fuzzer_crashes`] runs over `tests/samples/fuzz-regressions/`, which is **committed**. +//! Each file there once crashed the parser. This test is mandatory: if the directory is +//! missing or empty, that is itself a failure, because a regression suite that silently +//! finds nothing to check is worse than no suite at all. +//! - [`fuzzer_corpus`] runs over the local fuzzing corpus under `fuzz/corpus/`, which is too +//! large to commit and is therefore absent on a clean checkout. It skips loudly when +//! absent, and is exercised by anyone who has run the fuzzer locally. + +use std::{ + fs, + path::{Path, PathBuf}, +}; use dotscope::metadata::cilobject::CilObject; +/// Loads every input in the committed crash corpus. +/// +/// Fails if the directory is missing or empty — see the module docs for why. #[test] -/// Open all files from the fuzzer corpus, and load all of them without crashing. Can produce errors, as these are invalid files, -/// just don't crash -fn fuzzer_corpus() { - test_load_path(PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("fuzz/corpus/cilobject/")); +fn fuzzer_crashes() { + let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/samples/fuzz-regressions"); + let loaded = load_every_file(&path).unwrap_or_else(|e| { + panic!( + "crash-regression corpus at {} is unusable: {e}.\n\ + These inputs are committed to the repository; if they are missing, the checkout \ + is incomplete or the corpus was deleted.", + path.display() + ) + }); + + assert!( + loaded > 0, + "crash-regression corpus at {} contains no files — this test would otherwise pass \ + without checking anything", + path.display() + ); } +/// Loads every input in the local fuzzing corpus, when one is present. +/// +/// The corpus is hundreds of megabytes and is not committed, so this skips on a clean +/// checkout. It prints when it skips so a vacuous pass is visible in the test output rather +/// than indistinguishable from a real one. #[test] -/// Open all files from the fuzzer corpus that previously caused a crash, and load all of them without crashing. Can produce errors, -/// as these are invalid files, just don't crash. -fn fuzzer_crashes() { - test_load_path(PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("fuzz/artifacts/cilobject/")); +fn fuzzer_corpus() { + let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("fuzz/corpus/cilobject"); + match load_every_file(&path) { + Ok(loaded) => println!( + "fuzzer_corpus: loaded {loaded} inputs from {}", + path.display() + ), + Err(e) => println!( + "fuzzer_corpus: SKIPPED — no local corpus at {} ({e}). \ + Run `make fuzz` to generate one.", + path.display() + ), + } } -// #[test] -// /// Debug one specific test case -// fn debug() { -// test_load_path(PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("fuzz/artifacts/cilobject/leak-b96355f49d7918e856039d2c998de850e47ffdf3")); -// } - -fn test_load_path(path: PathBuf) { - if let Ok(artifacts_path) = fs::read_dir(path) { - for entry in artifacts_path.flatten() { - let path = entry.path(); - if path.is_file() { - let _ = CilObject::from_path(&path); - } +/// Loads every regular file in `dir` through the parser, returning how many were loaded. +/// +/// Parse errors are ignored by design; the assertion is that the call returns at all. +/// +/// # Errors +/// +/// Returns the underlying I/O error if `dir` cannot be read. +fn load_every_file(dir: &Path) -> std::io::Result { + let mut loaded = 0usize; + for entry in fs::read_dir(dir)? { + let path = entry?.path(); + if path.is_file() { + // Malformed input: an `Err` is the expected result and is discarded. What this + // exercises is that the call returns instead of taking the process down. + let _ = CilObject::from_path(&path); + loaded += 1; } } + Ok(loaded) } diff --git a/dotscope/tests/netreactor.rs b/dotscope/tests/netreactor.rs index e830c15d..12bd2350 100644 --- a/dotscope/tests/netreactor.rs +++ b/dotscope/tests/netreactor.rs @@ -418,6 +418,21 @@ fn test_all_netreactor_samples() { continue; } + // A native-stub build is an x86 executable that carries the .NET image + // as payload rather than being one, so it never reaches the point where + // deobfuscation could be judged — it is rejected at load. That is the + // documented limit of what this crate reads, not a deobfuscation + // result, and `test_netreactor_samples_load` already asserts the load + // error is the expected one. Measuring it here would only restate that. + if sample.expected_protections.has_native_exe { + eprintln!( + " [SKIP] {} — native exe stub, not a supported input format", + sample.filename + ); + skipped += 1; + continue; + } + results.push(run_deobfuscation_test( sample, SAMPLES_DIR, @@ -459,14 +474,35 @@ fn test_all_netreactor_samples() { continue; } - // Over-cleanup guard. Virtualized samples are not devirtualized by any - // technique, so the VM interpreter and its handler types remain live - // code — the stubs left in the virtualized methods still call into - // them. Cleanup reachability analysis must follow candidate-to-candidate - // call edges transitively to see that; a single-step version reads the - // handler cluster as isolated infrastructure and strips the assembly - // down to its application methods (854 -> 45 on reactor_virtualization). - if result.success && result.sample.expected_protections.has_virtualization { + // The contract cleanup must satisfy is structural validity, not a survival ratio. + // An assembly that fails validation has had something deleted out from under a + // reference that still names it, which is over-cleanup by definition. A ratio only + // correlates with that loosely — samples have been observed at 82% survival and + // still invalid, and it says nothing about *which* methods went. + assert!( + result.success, + "{}: deobfuscation did not complete: {}", + result.sample.filename, + result.error.as_deref().unwrap_or("unknown error") + ); + assert!( + result.assembly_valid, + "{}: deobfuscated output failed validation: {}", + result.sample.filename, + result.error.as_deref().unwrap_or("no error recorded") + ); + assert!( + result.roundtrip_ok, + "{}: deobfuscated output does not round-trip: {}", + result.sample.filename, + result.error.as_deref().unwrap_or("no error recorded") + ); + + // Coarse backstop, and only that. Validity catches deletions that leave a dangling + // reference; it cannot catch deleting code nothing statically references — a VM + // handler reached only through computed dispatch, say. Those samples are kept for + // later analysis, so losing their real code silently is the failure this guards. + if result.sample.expected_protections.has_virtualization { assert!( result.methods_after * 2 > result.methods_before, "{}: cleanup removed {} of {} methods — the VM runtime is still \ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-00556c4399794f6db44ffce728effc1ec3cd2fa1 b/dotscope/tests/samples/fuzz-regressions/crash-00556c4399794f6db44ffce728effc1ec3cd2fa1 new file mode 100644 index 00000000..943d17d9 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-00556c4399794f6db44ffce728effc1ec3cd2fa1 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-05e5dbbed85fdda41372c73d6e7eb80ce9a91546 b/dotscope/tests/samples/fuzz-regressions/crash-05e5dbbed85fdda41372c73d6e7eb80ce9a91546 new file mode 100644 index 00000000..f0afd654 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-05e5dbbed85fdda41372c73d6e7eb80ce9a91546 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-0ba1fe0ad577d8a505b2e34a29ecbec5d695cc81 b/dotscope/tests/samples/fuzz-regressions/crash-0ba1fe0ad577d8a505b2e34a29ecbec5d695cc81 new file mode 100644 index 00000000..b6a6aeae Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-0ba1fe0ad577d8a505b2e34a29ecbec5d695cc81 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-116ecc73c01734816489057f0c6081c787bd5b1d b/dotscope/tests/samples/fuzz-regressions/crash-116ecc73c01734816489057f0c6081c787bd5b1d new file mode 100644 index 00000000..4b5113e1 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-116ecc73c01734816489057f0c6081c787bd5b1d differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-13388014188607320f61641fd611342ff65a3e8c b/dotscope/tests/samples/fuzz-regressions/crash-13388014188607320f61641fd611342ff65a3e8c new file mode 100644 index 00000000..5e08ce00 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-13388014188607320f61641fd611342ff65a3e8c differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-13dfad7060839681c6f7c6bf1fe6acac3ebaa886 b/dotscope/tests/samples/fuzz-regressions/crash-13dfad7060839681c6f7c6bf1fe6acac3ebaa886 new file mode 100644 index 00000000..d679d38b Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-13dfad7060839681c6f7c6bf1fe6acac3ebaa886 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-1cbe9ef6db3eb6a372c269cb4c8e751565164a0e b/dotscope/tests/samples/fuzz-regressions/crash-1cbe9ef6db3eb6a372c269cb4c8e751565164a0e new file mode 100644 index 00000000..632553dc Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-1cbe9ef6db3eb6a372c269cb4c8e751565164a0e differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-1f924b9d1057f0e1957622ed59233e673aa72dff b/dotscope/tests/samples/fuzz-regressions/crash-1f924b9d1057f0e1957622ed59233e673aa72dff new file mode 100644 index 00000000..4cc6f175 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-1f924b9d1057f0e1957622ed59233e673aa72dff differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-22020b57bdfbba9c10cdfcde93865317b807b3b3 b/dotscope/tests/samples/fuzz-regressions/crash-22020b57bdfbba9c10cdfcde93865317b807b3b3 new file mode 100644 index 00000000..782a46b3 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-22020b57bdfbba9c10cdfcde93865317b807b3b3 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-272d0d545dbf696680104d5c85937babf4faf76e b/dotscope/tests/samples/fuzz-regressions/crash-272d0d545dbf696680104d5c85937babf4faf76e new file mode 100644 index 00000000..083c3f28 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-272d0d545dbf696680104d5c85937babf4faf76e differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-291138423e8558b6efce8169b74a82a0b8cabfcc b/dotscope/tests/samples/fuzz-regressions/crash-291138423e8558b6efce8169b74a82a0b8cabfcc new file mode 100644 index 00000000..c5835eaa Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-291138423e8558b6efce8169b74a82a0b8cabfcc differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-2ac2ad93bd33102f2178c81b4df2eb486a577054 b/dotscope/tests/samples/fuzz-regressions/crash-2ac2ad93bd33102f2178c81b4df2eb486a577054 new file mode 100644 index 00000000..dc3ff0b9 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-2ac2ad93bd33102f2178c81b4df2eb486a577054 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-2d515faa49c111a9235c6bd2e8d4fb4101a3e961 b/dotscope/tests/samples/fuzz-regressions/crash-2d515faa49c111a9235c6bd2e8d4fb4101a3e961 new file mode 100644 index 00000000..2b97722c Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-2d515faa49c111a9235c6bd2e8d4fb4101a3e961 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-2e514684c4426d239b3311528e7b288e9628f80c b/dotscope/tests/samples/fuzz-regressions/crash-2e514684c4426d239b3311528e7b288e9628f80c new file mode 100644 index 00000000..c0dd695f Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-2e514684c4426d239b3311528e7b288e9628f80c differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-36dc9407db12eb3dc0c8416cf2e05702860afa13 b/dotscope/tests/samples/fuzz-regressions/crash-36dc9407db12eb3dc0c8416cf2e05702860afa13 new file mode 100644 index 00000000..540dc5d4 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-36dc9407db12eb3dc0c8416cf2e05702860afa13 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-37862deb65bdd4829ef254d5aabea5612f6453cc b/dotscope/tests/samples/fuzz-regressions/crash-37862deb65bdd4829ef254d5aabea5612f6453cc new file mode 100644 index 00000000..db3c6d22 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-37862deb65bdd4829ef254d5aabea5612f6453cc differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-3ae512e97f3e229aca9592be6dff56336cb1e91b b/dotscope/tests/samples/fuzz-regressions/crash-3ae512e97f3e229aca9592be6dff56336cb1e91b new file mode 100644 index 00000000..06fa5661 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-3ae512e97f3e229aca9592be6dff56336cb1e91b differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-3fef7486fe48c8a62581600ec54c286fe0a8eac0 b/dotscope/tests/samples/fuzz-regressions/crash-3fef7486fe48c8a62581600ec54c286fe0a8eac0 new file mode 100644 index 00000000..f7f851a6 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-3fef7486fe48c8a62581600ec54c286fe0a8eac0 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-43dfaf4644a738c2588a689f8f7ee6e376fb2efa b/dotscope/tests/samples/fuzz-regressions/crash-43dfaf4644a738c2588a689f8f7ee6e376fb2efa new file mode 100644 index 00000000..447ffbc0 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-43dfaf4644a738c2588a689f8f7ee6e376fb2efa differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-47a66e432ec761640198d3085666a1469c9cd220 b/dotscope/tests/samples/fuzz-regressions/crash-47a66e432ec761640198d3085666a1469c9cd220 new file mode 100644 index 00000000..5d3424b1 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-47a66e432ec761640198d3085666a1469c9cd220 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-47d8ec47286d2cbd3b19de0fa1794680d415946f b/dotscope/tests/samples/fuzz-regressions/crash-47d8ec47286d2cbd3b19de0fa1794680d415946f new file mode 100644 index 00000000..a76a0c96 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-47d8ec47286d2cbd3b19de0fa1794680d415946f differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-481e4a026aa58f360df82d0468166082f7952209 b/dotscope/tests/samples/fuzz-regressions/crash-481e4a026aa58f360df82d0468166082f7952209 new file mode 100644 index 00000000..35777b57 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-481e4a026aa58f360df82d0468166082f7952209 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-492bc73afc8bf8152c4a0709cae6e4027a994546 b/dotscope/tests/samples/fuzz-regressions/crash-492bc73afc8bf8152c4a0709cae6e4027a994546 new file mode 100644 index 00000000..4ec40b66 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-492bc73afc8bf8152c4a0709cae6e4027a994546 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-4cd7b2fbf6ea790bb758ded03016d66bb16ce806 b/dotscope/tests/samples/fuzz-regressions/crash-4cd7b2fbf6ea790bb758ded03016d66bb16ce806 new file mode 100644 index 00000000..30609be6 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-4cd7b2fbf6ea790bb758ded03016d66bb16ce806 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-57b8789c285f2b0780cf595b1d61d1f35bc58498 b/dotscope/tests/samples/fuzz-regressions/crash-57b8789c285f2b0780cf595b1d61d1f35bc58498 new file mode 100644 index 00000000..6c8aca2d Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-57b8789c285f2b0780cf595b1d61d1f35bc58498 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-58f1bf8fadb1dd341c2329edfd03a8c25ef0cdea b/dotscope/tests/samples/fuzz-regressions/crash-58f1bf8fadb1dd341c2329edfd03a8c25ef0cdea new file mode 100644 index 00000000..c3f4f56d Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-58f1bf8fadb1dd341c2329edfd03a8c25ef0cdea differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-5b1b3c4c5502b086760d7a35cfe6068aa59cd7d8 b/dotscope/tests/samples/fuzz-regressions/crash-5b1b3c4c5502b086760d7a35cfe6068aa59cd7d8 new file mode 100644 index 00000000..530a617e Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-5b1b3c4c5502b086760d7a35cfe6068aa59cd7d8 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-5dd5b021b2fb26393da9782983b4885bcf44cd03 b/dotscope/tests/samples/fuzz-regressions/crash-5dd5b021b2fb26393da9782983b4885bcf44cd03 new file mode 100644 index 00000000..eeb6044f Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-5dd5b021b2fb26393da9782983b4885bcf44cd03 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-5f21342c1a5babd06401940c3e0c7ba697b51eb3 b/dotscope/tests/samples/fuzz-regressions/crash-5f21342c1a5babd06401940c3e0c7ba697b51eb3 new file mode 100644 index 00000000..479007a1 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-5f21342c1a5babd06401940c3e0c7ba697b51eb3 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-609a2f11943a71d2662c76520d5a34034c009aee b/dotscope/tests/samples/fuzz-regressions/crash-609a2f11943a71d2662c76520d5a34034c009aee new file mode 100644 index 00000000..d97d4295 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-609a2f11943a71d2662c76520d5a34034c009aee differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-6302bd3c8e304b6158cc34d058b52e13137914c5 b/dotscope/tests/samples/fuzz-regressions/crash-6302bd3c8e304b6158cc34d058b52e13137914c5 new file mode 100644 index 00000000..b69e1dbd Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-6302bd3c8e304b6158cc34d058b52e13137914c5 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-6b996e5a1942c4b4d0156dac6616e087aff8f761 b/dotscope/tests/samples/fuzz-regressions/crash-6b996e5a1942c4b4d0156dac6616e087aff8f761 new file mode 100644 index 00000000..0acaa1c9 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-6b996e5a1942c4b4d0156dac6616e087aff8f761 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-70f956eb1d3a16cc15785e0a891c560c92c7e510 b/dotscope/tests/samples/fuzz-regressions/crash-70f956eb1d3a16cc15785e0a891c560c92c7e510 new file mode 100644 index 00000000..0080b5ee Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-70f956eb1d3a16cc15785e0a891c560c92c7e510 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-72354448461017f512594a24b711e881977c1113 b/dotscope/tests/samples/fuzz-regressions/crash-72354448461017f512594a24b711e881977c1113 new file mode 100644 index 00000000..d1335f3f Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-72354448461017f512594a24b711e881977c1113 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-724a402d6aa00572bcee97f9fdb338bcdc6eb23a b/dotscope/tests/samples/fuzz-regressions/crash-724a402d6aa00572bcee97f9fdb338bcdc6eb23a new file mode 100644 index 00000000..891bbe5a Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-724a402d6aa00572bcee97f9fdb338bcdc6eb23a differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-734666e7bc4dfe6002caeb33c6601723247f240a b/dotscope/tests/samples/fuzz-regressions/crash-734666e7bc4dfe6002caeb33c6601723247f240a new file mode 100644 index 00000000..e70fcfb8 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-734666e7bc4dfe6002caeb33c6601723247f240a differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-775cadba90b566b8446fd4aaecb6467ea597f573 b/dotscope/tests/samples/fuzz-regressions/crash-775cadba90b566b8446fd4aaecb6467ea597f573 new file mode 100644 index 00000000..53cd3edd Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-775cadba90b566b8446fd4aaecb6467ea597f573 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-7c620960bc72cb2167598cdca54b3d28888c7a86 b/dotscope/tests/samples/fuzz-regressions/crash-7c620960bc72cb2167598cdca54b3d28888c7a86 new file mode 100644 index 00000000..7f2d810d Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-7c620960bc72cb2167598cdca54b3d28888c7a86 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-8ad3ddcf41569559112e56d4ca908642cb0a75aa b/dotscope/tests/samples/fuzz-regressions/crash-8ad3ddcf41569559112e56d4ca908642cb0a75aa new file mode 100644 index 00000000..ba1fdfc2 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-8ad3ddcf41569559112e56d4ca908642cb0a75aa differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-8dad878902939c3ad4373e17c97a74445ab9e049 b/dotscope/tests/samples/fuzz-regressions/crash-8dad878902939c3ad4373e17c97a74445ab9e049 new file mode 100644 index 00000000..dbf6b79c Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-8dad878902939c3ad4373e17c97a74445ab9e049 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-8f4840e448bccf840fea0fdd76317319a7aaadeb b/dotscope/tests/samples/fuzz-regressions/crash-8f4840e448bccf840fea0fdd76317319a7aaadeb new file mode 100644 index 00000000..db74fcf5 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-8f4840e448bccf840fea0fdd76317319a7aaadeb differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-95896fa8229590047497f1c02b31a794911a36a3 b/dotscope/tests/samples/fuzz-regressions/crash-95896fa8229590047497f1c02b31a794911a36a3 new file mode 100644 index 00000000..c06a1650 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-95896fa8229590047497f1c02b31a794911a36a3 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-96923484c1ec0e71328d1f786bf51b175a125357 b/dotscope/tests/samples/fuzz-regressions/crash-96923484c1ec0e71328d1f786bf51b175a125357 new file mode 100644 index 00000000..a8c43753 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-96923484c1ec0e71328d1f786bf51b175a125357 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-96c8c4e1b54ccef51eadda2752884c782cc87022 b/dotscope/tests/samples/fuzz-regressions/crash-96c8c4e1b54ccef51eadda2752884c782cc87022 new file mode 100644 index 00000000..8e7aa35e Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-96c8c4e1b54ccef51eadda2752884c782cc87022 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-9a60e737bf46c3f61b120c3b643557382404b910 b/dotscope/tests/samples/fuzz-regressions/crash-9a60e737bf46c3f61b120c3b643557382404b910 new file mode 100644 index 00000000..56a805c9 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-9a60e737bf46c3f61b120c3b643557382404b910 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-9ae6356f0d91cc40154facd72ec1dfd5aeedc91d b/dotscope/tests/samples/fuzz-regressions/crash-9ae6356f0d91cc40154facd72ec1dfd5aeedc91d new file mode 100644 index 00000000..e6bb34c0 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-9ae6356f0d91cc40154facd72ec1dfd5aeedc91d differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-a34eb04384b3ab29899faf88935bbabeb92bc749 b/dotscope/tests/samples/fuzz-regressions/crash-a34eb04384b3ab29899faf88935bbabeb92bc749 new file mode 100644 index 00000000..f4725583 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-a34eb04384b3ab29899faf88935bbabeb92bc749 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-a445a250d9cc0b18fd88b1e5cb777e685f9aebda b/dotscope/tests/samples/fuzz-regressions/crash-a445a250d9cc0b18fd88b1e5cb777e685f9aebda new file mode 100644 index 00000000..618ca9c1 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-a445a250d9cc0b18fd88b1e5cb777e685f9aebda differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-aa39a9253467829964c361ec8ccc7be6f938a9d2 b/dotscope/tests/samples/fuzz-regressions/crash-aa39a9253467829964c361ec8ccc7be6f938a9d2 new file mode 100644 index 00000000..b4bbb8d3 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-aa39a9253467829964c361ec8ccc7be6f938a9d2 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-acf1e2a5700b44a77491e91d174d4a5bae96d541 b/dotscope/tests/samples/fuzz-regressions/crash-acf1e2a5700b44a77491e91d174d4a5bae96d541 new file mode 100644 index 00000000..96938cb9 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-acf1e2a5700b44a77491e91d174d4a5bae96d541 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-ae5c3c6988a583b075d555a22fd34b4a45cbbad4 b/dotscope/tests/samples/fuzz-regressions/crash-ae5c3c6988a583b075d555a22fd34b4a45cbbad4 new file mode 100644 index 00000000..d1c48342 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-ae5c3c6988a583b075d555a22fd34b4a45cbbad4 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-b07c102a966113251346ee9faa80908b3078baaf b/dotscope/tests/samples/fuzz-regressions/crash-b07c102a966113251346ee9faa80908b3078baaf new file mode 100644 index 00000000..9eb8567e Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-b07c102a966113251346ee9faa80908b3078baaf differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-b66bfd4e8929b2967b5c20285b29af9d8097a642 b/dotscope/tests/samples/fuzz-regressions/crash-b66bfd4e8929b2967b5c20285b29af9d8097a642 new file mode 100644 index 00000000..ec938265 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-b66bfd4e8929b2967b5c20285b29af9d8097a642 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-bc64895e6f5519e64b5fcf7a5c2d6bb8d7c5a577 b/dotscope/tests/samples/fuzz-regressions/crash-bc64895e6f5519e64b5fcf7a5c2d6bb8d7c5a577 new file mode 100644 index 00000000..14c050c0 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-bc64895e6f5519e64b5fcf7a5c2d6bb8d7c5a577 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-be7edd4580c57b558efc804f9639901595d33500 b/dotscope/tests/samples/fuzz-regressions/crash-be7edd4580c57b558efc804f9639901595d33500 new file mode 100644 index 00000000..fc9f7491 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-be7edd4580c57b558efc804f9639901595d33500 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-c6a28afe5b093f432e61ff0cbc2eb632b25c191a b/dotscope/tests/samples/fuzz-regressions/crash-c6a28afe5b093f432e61ff0cbc2eb632b25c191a new file mode 100644 index 00000000..1eb2a797 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-c6a28afe5b093f432e61ff0cbc2eb632b25c191a differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-d386eb3edca5260c4693324b253918a1bc5d07dd b/dotscope/tests/samples/fuzz-regressions/crash-d386eb3edca5260c4693324b253918a1bc5d07dd new file mode 100644 index 00000000..63cddc0c Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-d386eb3edca5260c4693324b253918a1bc5d07dd differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-d3a715870e5eeaa6676173d8d80193760eebc568 b/dotscope/tests/samples/fuzz-regressions/crash-d3a715870e5eeaa6676173d8d80193760eebc568 new file mode 100644 index 00000000..ea622d70 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-d3a715870e5eeaa6676173d8d80193760eebc568 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-d4868fd455a55809b1afd691a5c6a18180c93eed b/dotscope/tests/samples/fuzz-regressions/crash-d4868fd455a55809b1afd691a5c6a18180c93eed new file mode 100644 index 00000000..4f74162b Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-d4868fd455a55809b1afd691a5c6a18180c93eed differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-d55961c78aaa23f88c5ea032f1f63d752940c885 b/dotscope/tests/samples/fuzz-regressions/crash-d55961c78aaa23f88c5ea032f1f63d752940c885 new file mode 100644 index 00000000..d7814956 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-d55961c78aaa23f88c5ea032f1f63d752940c885 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-d5a7ab596aee472c294cf8c88a079820140c2bd7 b/dotscope/tests/samples/fuzz-regressions/crash-d5a7ab596aee472c294cf8c88a079820140c2bd7 new file mode 100644 index 00000000..7ebdbfe1 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-d5a7ab596aee472c294cf8c88a079820140c2bd7 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-d90f43812c1a89f24f21e0a3c864b654d017f060 b/dotscope/tests/samples/fuzz-regressions/crash-d90f43812c1a89f24f21e0a3c864b654d017f060 new file mode 100644 index 00000000..9f53021b Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-d90f43812c1a89f24f21e0a3c864b654d017f060 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-dc1a9e462f45d6ce50f96b443f5940d414dfb3bf b/dotscope/tests/samples/fuzz-regressions/crash-dc1a9e462f45d6ce50f96b443f5940d414dfb3bf new file mode 100644 index 00000000..ea7928be Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-dc1a9e462f45d6ce50f96b443f5940d414dfb3bf differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-e49977f29a211ed3f4fbea7b2b9f8e30e08683df b/dotscope/tests/samples/fuzz-regressions/crash-e49977f29a211ed3f4fbea7b2b9f8e30e08683df new file mode 100644 index 00000000..3c3e9f0c Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-e49977f29a211ed3f4fbea7b2b9f8e30e08683df differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-e59d32d275d742ea4643d6d53a248a92f9c772e9 b/dotscope/tests/samples/fuzz-regressions/crash-e59d32d275d742ea4643d6d53a248a92f9c772e9 new file mode 100644 index 00000000..286845ea Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-e59d32d275d742ea4643d6d53a248a92f9c772e9 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-e5f06d1b74cdefbee7d3fac44ccf747ef2dc8385 b/dotscope/tests/samples/fuzz-regressions/crash-e5f06d1b74cdefbee7d3fac44ccf747ef2dc8385 new file mode 100644 index 00000000..53e83448 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-e5f06d1b74cdefbee7d3fac44ccf747ef2dc8385 differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-e9567af33b16b520b7de944e87ca5ab22ac541ea b/dotscope/tests/samples/fuzz-regressions/crash-e9567af33b16b520b7de944e87ca5ab22ac541ea new file mode 100644 index 00000000..d7715436 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-e9567af33b16b520b7de944e87ca5ab22ac541ea differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-ee0e9763c779daefdaae371ebe356938b39c9f9a b/dotscope/tests/samples/fuzz-regressions/crash-ee0e9763c779daefdaae371ebe356938b39c9f9a new file mode 100644 index 00000000..6891c455 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-ee0e9763c779daefdaae371ebe356938b39c9f9a differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-f286192c226474e7c760aca590f93351a7b109df b/dotscope/tests/samples/fuzz-regressions/crash-f286192c226474e7c760aca590f93351a7b109df new file mode 100644 index 00000000..19a0da07 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-f286192c226474e7c760aca590f93351a7b109df differ diff --git a/dotscope/tests/samples/fuzz-regressions/crash-fb9b72e6b88bc3866a5306bee5245f8e2b975cfd b/dotscope/tests/samples/fuzz-regressions/crash-fb9b72e6b88bc3866a5306bee5245f8e2b975cfd new file mode 100644 index 00000000..7d723d4b Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/crash-fb9b72e6b88bc3866a5306bee5245f8e2b975cfd differ diff --git a/dotscope/tests/samples/fuzz-regressions/leak-b96355f49d7918e856039d2c998de850e47ffdf3 b/dotscope/tests/samples/fuzz-regressions/leak-b96355f49d7918e856039d2c998de850e47ffdf3 new file mode 100644 index 00000000..21f878a8 Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/leak-b96355f49d7918e856039d2c998de850e47ffdf3 differ diff --git a/dotscope/tests/samples/fuzz-regressions/oom-21104105496866b7eea2ba90393fbfd44692b6a7 b/dotscope/tests/samples/fuzz-regressions/oom-21104105496866b7eea2ba90393fbfd44692b6a7 new file mode 100644 index 00000000..109f903f Binary files /dev/null and b/dotscope/tests/samples/fuzz-regressions/oom-21104105496866b7eea2ba90393fbfd44692b6a7 differ