diff --git a/.github/scripts/release_artifacts.py b/.github/scripts/release_artifacts.py new file mode 100644 index 0000000..b6e0319 --- /dev/null +++ b/.github/scripts/release_artifacts.py @@ -0,0 +1,172 @@ +"""Copy only the expected binary bytes from this run's untrusted artifacts.""" + +import hashlib +import io +import json +import os +from pathlib import Path +import re +import stat +from urllib.error import HTTPError +from urllib.parse import urlsplit +from urllib.request import HTTPRedirectHandler, Request, build_opener +import zipfile + + +MAX_BYTES = 64 * 1024 * 1024 +ARTIFACTS = {"Win32 Executable": "pd32.exe", "x64 Executable": "pd64.exe"} + + +class NoRedirect(HTTPRedirectHandler): + def redirect_request(self, request, fp, code, message, headers, new_url): + return None + + +def bounded_read(response, limit): + data = response.read(limit + 1) + if len(data) > limit: + raise ValueError("Response exceeds the release size limit") + return data + + +class GitHub: + def __init__(self): + self.base = os.environ["GITHUB_API_URL"].rstrip("/") + self.token = os.environ["GH_TOKEN"] + self.opener = build_opener(NoRedirect()) + if urlsplit(self.base).scheme != "https": + raise ValueError("The GitHub API must use HTTPS") + + def request(self, endpoint): + return Request(self.base + endpoint, headers={ + "Authorization": "Bearer " + self.token, + "Accept": "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28", + }) + + def json(self, endpoint, allow_missing=False): + try: + with self.opener.open(self.request(endpoint), timeout=60) as response: + return json.loads(bounded_read(response, 4 * 1024 * 1024)) + except HTTPError as error: + if error.code == 404 and allow_missing: + return None + raise + + def archive(self, repository, artifact_id): + endpoint = f"/repos/{repository}/actions/artifacts/{artifact_id}/zip" + try: + with self.opener.open(self.request(endpoint), timeout=60) as response: + return bounded_read(response, MAX_BYTES) + except HTTPError as error: + if error.code != 302: + raise + location = error.headers["Location"] + parsed = urlsplit(location) + if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password: + raise ValueError("Invalid artifact download redirect") + # Never forward the repository token to the signed blob URL or follow another redirect. + with self.opener.open(Request(location), timeout=60) as response: + return bounded_read(response, MAX_BYTES) + + +def select_artifacts(api, repository, run_id): + selected = {} + for page in range(1, 101): + listing = api.json(f"/repos/{repository}/actions/runs/{run_id}/artifacts?per_page=100&page={page}") + entries = listing["artifacts"] + for item in entries: + name = item["name"] + if name not in ARTIFACTS: + continue + if name in selected or item["expired"]: + raise ValueError("Duplicate or expired release artifact") + if type(item["id"]) is not int or item["id"] <= 0: + raise ValueError("Invalid release artifact ID") + if not 0 < item["size_in_bytes"] <= MAX_BYTES: + raise ValueError("Invalid release artifact size") + selected[name] = item + if len(entries) < 100: + break + else: + raise ValueError("Too many artifacts in the release run") + if selected.keys() != ARTIFACTS.keys(): + raise ValueError("Both architecture artifacts are required") + return selected + + +def validate_entry(entries): + if len(entries) != 1: + raise ValueError("A release artifact must contain exactly one entry") + entry = entries[0] + mode = entry.external_attr >> 16 + if (entry.filename != "pd.exe" or entry.orig_filename != "pd.exe" or entry.is_dir() + or stat.S_IFMT(mode) not in (0, stat.S_IFREG) + or entry.external_attr & 0x10 or entry.flag_bits & 1 + or entry.compress_type not in (zipfile.ZIP_STORED, zipfile.ZIP_DEFLATED) + or not 0 < entry.file_size <= MAX_BYTES + or not 0 <= entry.compress_size <= MAX_BYTES): + raise ValueError("Unexpected release archive entry") + return entry + + +def copy_binary(archive, destination): + if len(archive) > MAX_BYTES: + raise ValueError("Release archive exceeds the size limit") + with zipfile.ZipFile(io.BytesIO(archive)) as zipped: + entry = validate_entry(zipped.infolist()) + # Validate data/CRC before writing; archive-controlled paths are never opened. + with zipped.open(entry) as source: + data = bounded_read(source, MAX_BYTES) + if len(data) != entry.file_size: + raise ValueError("Release binary length mismatch") + with destination.open("xb") as output: + output.write(data) + return {"size": len(data), "sha256": hashlib.sha256(data).hexdigest()} + + +def verify_tag(api, repository, version, commit): + ref = api.json(f"/repos/{repository}/git/ref/tags/{version}", allow_missing=True) + if ref is None: + return + target = ref["object"] + for _ in range(8): + if target["type"] == "commit": + if target["sha"].lower() != commit.lower(): + raise ValueError("Existing version tag does not identify the tested commit") + return + if target["type"] != "tag" or not re.fullmatch(r"[0-9a-fA-F]{40}", target["sha"]): + break + target = api.json(f"/repos/{repository}/git/tags/{target['sha']}")["object"] + raise ValueError("Unsupported release tag target") + + +def main(): + repository = os.environ["GITHUB_REPOSITORY"] + run_id = os.environ["GITHUB_RUN_ID"] + commit = os.environ["RELEASE_COMMIT"] + version = os.environ["RELEASE_VERSION"] + if (not re.fullmatch(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", repository) + or not re.fullmatch(r"[0-9]+", run_id) + or not re.fullmatch(r"[0-9a-fA-F]{40}", commit) + or not (version == "Develop" or re.fullmatch(r"v[0-9]+\.[0-9]+(\.[0-9]+)?", version))): + raise ValueError("Invalid release identity") + api = GitHub() + if version != "Develop": + verify_tag(api, repository, version, commit) + selected = select_artifacts(api, repository, run_id) + output = Path("release-binaries") + output.mkdir() + manifest = {"repository": repository, "commit": commit, "run_id": run_id, "files": {}} + for name, filename in ARTIFACTS.items(): + artifact = selected[name] + result = copy_binary(api.archive(repository, artifact["id"]), output / filename) + result["artifact_id"] = artifact["id"] + manifest["files"][filename] = result + with (output / "provenance.json").open("x", encoding="utf-8") as destination: + json.dump(manifest, destination, indent=2) + destination.write("\n") + + +if __name__ == "__main__": + main() diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index f0a4fd7..03cd58a 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -22,17 +22,22 @@ jobs: platform: [Win32, x64] configuration: [Debug, Release] steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 with: ref: ${{ inputs.ref || github.sha }} - - uses: microsoft/setup-msbuild@v2 + persist-credentials: false + - uses: microsoft/setup-msbuild@6fb02220983dee41ce7ae257b6f4d8f9bf5ed4ce # v2 - name: Build application and unit tests run: msbuild pd.sln /m /v:minimal /p:Configuration=${{ matrix.configuration }} /p:Platform=${{ matrix.platform }} - name: Run unit tests run: | & ".\${{ matrix.platform }}\${{ matrix.configuration }}\pd_tests.exe" if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } - - uses: actions/upload-artifact@v4 + - name: Test CLI documentation, benchmark analysis and release artifact handling + env: + PD_TEST_EXE: .\${{ matrix.platform }}\${{ matrix.configuration }}\pd.exe + run: python -m unittest discover -s tests -p "test_*.py" + - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4 if: matrix.configuration == 'Release' with: name: ${{ matrix.platform }} Executable diff --git a/.github/workflows/deploy-create-releases.yml b/.github/workflows/deploy-create-releases.yml index 0897512..c6bca63 100644 --- a/.github/workflows/deploy-create-releases.yml +++ b/.github/workflows/deploy-create-releases.yml @@ -9,7 +9,7 @@ on: required: true version: type: string - description: Release version (for example v2.2) + description: Release version (for example v3.0.0) required: true permissions: @@ -35,16 +35,23 @@ jobs: needs: build runs-on: windows-2022 permissions: + actions: read contents: write + concurrency: + group: publish-${{ inputs.version }} + cancel-in-progress: false steps: - - uses: actions/download-artifact@v4 + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 with: - name: Win32 Executable - path: binaries\Win32 - - uses: actions/download-artifact@v4 - with: - name: x64 Executable - path: binaries\x64 + ref: ${{ github.workflow_sha }} + persist-credentials: false + sparse-checkout: .github/scripts + - name: Prepare allowlisted artifact bytes + env: + GH_TOKEN: ${{ github.token }} + RELEASE_VERSION: ${{ inputs.version }} + RELEASE_COMMIT: ${{ inputs.commit }} + run: python -I .github\scripts\release_artifacts.py - name: Publish tested binaries env: GH_TOKEN: ${{ github.token }} @@ -52,7 +59,5 @@ jobs: RELEASE_VERSION: ${{ inputs.version }} RELEASE_COMMIT: ${{ inputs.commit }} run: | - Move-Item binaries\Win32\pd.exe pd32.exe - Move-Item binaries\x64\pd.exe pd64.exe - gh release create $env:RELEASE_VERSION pd32.exe pd64.exe --target $env:RELEASE_COMMIT --title $env:RELEASE_VERSION --generate-notes + gh release create $env:RELEASE_VERSION release-binaries\pd32.exe release-binaries\pd64.exe release-binaries\provenance.json --target $env:RELEASE_COMMIT --title $env:RELEASE_VERSION --generate-notes if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } diff --git a/.github/workflows/deploy-develop.yml b/.github/workflows/deploy-develop.yml index d16a9f4..1c4604b 100644 --- a/.github/workflows/deploy-develop.yml +++ b/.github/workflows/deploy-develop.yml @@ -15,28 +15,45 @@ jobs: needs: build runs-on: windows-2022 permissions: + actions: read contents: write + concurrency: + group: publish-Develop + cancel-in-progress: false steps: - - uses: actions/download-artifact@v4 + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 with: - name: Win32 Executable - path: binaries\Win32 - - uses: actions/download-artifact@v4 - with: - name: x64 Executable - path: binaries\x64 + ref: ${{ github.workflow_sha }} + persist-credentials: false + sparse-checkout: .github/scripts + - name: Prepare allowlisted artifact bytes + env: + GH_TOKEN: ${{ github.token }} + RELEASE_VERSION: Develop + RELEASE_COMMIT: ${{ github.sha }} + run: python -I .github\scripts\release_artifacts.py - name: Publish tested development binaries env: GH_TOKEN: ${{ github.token }} GH_REPO: ${{ github.repository }} RELEASE_COMMIT: ${{ github.sha }} run: | - Move-Item binaries\Win32\pd.exe pd32.exe - Move-Item binaries\x64\pd.exe pd64.exe - gh release view Develop - if ($LASTEXITCODE -eq 0) { - gh release upload Develop pd32.exe pd64.exe --clobber + if ($env:GITHUB_EVENT_NAME -eq "push") { + $current = gh api "repos/$env:GH_REPO/git/ref/heads/main" --jq .object.sha + if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } + if ($current -ne $env:RELEASE_COMMIT) { + Write-Host "Skipping superseded main build $env:RELEASE_COMMIT" + exit 0 + } + } + $release = gh api --paginate "repos/$env:GH_REPO/releases?per_page=100" --jq '.[] | select(.tag_name == "Develop") | .id' + if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } + $notes = "Development binaries from $env:RELEASE_COMMIT. Verify both assets against provenance.json; the historical Develop tag is not moved." + if ($release) { + gh release upload Develop release-binaries\pd32.exe release-binaries\pd64.exe release-binaries\provenance.json --clobber + if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } + gh release edit Develop --notes $notes } else { - gh release create Develop pd32.exe pd64.exe --target $env:RELEASE_COMMIT --title Develop --prerelease --notes "Development binaries" + gh release create Develop release-binaries\pd32.exe release-binaries\pd64.exe release-binaries\provenance.json --target $env:RELEASE_COMMIT --title Develop --prerelease --notes $notes } if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } diff --git a/.github/workflows/deploy-releases.yml b/.github/workflows/deploy-releases.yml index 297d488..5c5ba4c 100644 --- a/.github/workflows/deploy-releases.yml +++ b/.github/workflows/deploy-releases.yml @@ -14,23 +14,28 @@ jobs: needs: build runs-on: windows-2022 permissions: + actions: read contents: write + concurrency: + group: publish-${{ github.ref_name }} + cancel-in-progress: false steps: - - uses: actions/download-artifact@v4 + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 with: - name: Win32 Executable - path: binaries\Win32 - - uses: actions/download-artifact@v4 - with: - name: x64 Executable - path: binaries\x64 + ref: ${{ github.workflow_sha }} + persist-credentials: false + sparse-checkout: .github/scripts + - name: Prepare allowlisted artifact bytes + env: + GH_TOKEN: ${{ github.token }} + RELEASE_VERSION: ${{ github.ref_name }} + RELEASE_COMMIT: ${{ github.sha }} + run: python -I .github\scripts\release_artifacts.py - name: Publish tested release binaries env: GH_TOKEN: ${{ github.token }} GH_REPO: ${{ github.repository }} RELEASE_VERSION: ${{ github.ref_name }} run: | - Move-Item binaries\Win32\pd.exe pd32.exe - Move-Item binaries\x64\pd.exe pd64.exe - gh release create $env:RELEASE_VERSION pd32.exe pd64.exe --verify-tag --title $env:RELEASE_VERSION --generate-notes + gh release create $env:RELEASE_VERSION release-binaries\pd32.exe release-binaries\pd64.exe release-binaries\provenance.json --verify-tag --title $env:RELEASE_VERSION --generate-notes if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } diff --git a/.gitignore b/.gitignore index 5cc55c5..fa0fab4 100644 --- a/.gitignore +++ b/.gitignore @@ -394,3 +394,6 @@ FakesAssemblies/ # Visual Studio 6 workspace options file *.opt >>>>>>> d26ace875aaef086a6992374a3169ca5c6db031d + +# Python benchmark analysis +__pycache__ diff --git a/README.md b/README.md index f32264e..5a9e777 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,19 @@ -# Process Dump +# Process Dump 3.0.0 + +The source version is **3.0.0** (Windows file/product version **3.0.0.0**). The console banner and executable resources share `pd\version.h`. See the [3.0.0 changes](#version-300-unreleased) and [command-line reference](#command-line-arguments). This source version does not imply that a tagged release has been published. + Process Dump is a Windows reverse-engineering command-line tool to dump malware memory components back to disk for analysis. Often malware files are packed and obfuscated before they are executed in order to avoid AV scanners, however when these files are executed they will often unpack or inject a clean version of the malware code in memory. A common task for malware researchers when analyzing malware is to dump this unpacked code back from memory to disk for scanning with AV products or for analysis with static analysis tools such as IDA. -Process Dump works for Windows 32 and 64 bit operating systems and can dump memory components from specific processes or from all processes currently running. Process Dump supports creation and use of a clean-hash database, so that dumping of all the clean files such as kernel32.dll can be skipped. It's main features include: +Process Dump works for Windows 32 and 64 bit operating systems and can dump memory components from specific processes or from all processes currently running. Process Dump supports creation and use of a clean-hash database, so that dumping of known modules such as kernel32.dll can be skipped. A hash match is an exclusion, not a trust verdict. Its main features include: * Dumps code from a specific process or all processes. * Finds and dumps hidden modules that are not properly loaded in processes. * Finds and dumps loose code chunks even if they aren't associated with a PE file. It builds a PE header and import table for the chunks. * Reconstructs imports using an aggressive approach. * Can run in close dump monitor mode ('-closemon'), where processes will be paused and dumped just before they terminate. -* Multi-threaded, so when you are dumping all running processes it will go pretty quickly. -* Can generate a clean hash database. Generate this before a machine is infected with malware so Process Dump will only dump the new malicious malware components. +* Shares persistent workers across modules for system dumping and live hashing; directory database generation also hashes files concurrently. +* Recovers missing/invalid entrypoints using full opcode signatures and a conservative x64 runtime-function fallback. +* Offers opt-in `-reexec` preparation for fresh-launch experiments, separate from normal forensic capture. +* Can generate clean-module and entrypoint databases. Generate a baseline on a known-clean system to reduce repeated output; unknown code is not necessarily malicious. I'm maintaining an official compiled release on my website here: https://split-code.com/processdump.html @@ -29,96 +34,202 @@ msbuild pd.sln /m /p:Configuration=Release /p:Platform=x64 Use `Platform=Win32` and `.\Win32\Release\pd_tests.exe` for 32-bit builds, or `Configuration=Debug` and the corresponding `Debug` directory. Tests use generated benign PE fixtures and this test process only; no administrator privileges, malware, or system-wide dumping are required. Failures exit nonzero and CRT assertions are reported to stderr instead of opening a blocking dialog. -The `--baseline` selector runs the compatibility suite. Its expected PE dump sizes, CRC32s, import-table bytes, and PE hashes were captured on the original implementation before security fixes, on both architectures. These expectations are checked on every run; security cases intentionally require safer behavior for malformed inputs. A single case can also be selected by its printed name. +The `--baseline` selector runs the compatibility suite. Non-reconstructed PE dumps and import-table serialization retain the fingerprints captured before the security fixes. Import-reconstructed fixtures now expect a separate `.pdimp` section rather than extension of the last section; their CRC32s and structural hashes were intentionally updated for that layout change. Expectations are checked on every run. A single case can also be selected by its printed name. CI builds and runs tests for all four Debug/Release and Win32/x64 combinations on Windows. Development, tagged, and manually selected-commit releases reuse this tested build; release publication is separate from read-only PR builds. Compiler stack checks, Control Flow Guard, ASLR, and DEP are explicitly enabled. +Release jobs use SHA-pinned actions and a preparation script from the workflow's own commit, never the manually selected build commit. Artifacts are read only from the current run, with exactly one ordinary `pd.exe` entry per architecture and 64 MiB archive/binary limits. The publisher validates ZIP metadata and CRC, then copies bytes to fixed, exclusively created filenames; it never extracts archive paths or executes downloaded binaries. Signed download redirects do not receive the repository token. Run `python -m unittest discover -s tests -p "test_*.py"` for the portable benchmark, documentation and release-preparation tests. To also check the built application's help, no-side-effect help exit and Windows version metadata, set `$env:PD_TEST_EXE = (Resolve-Path .\x64\Release\pd.exe).Path` before running that command. CI sets this for each build configuration; CLI checks use a private copy and never dump or hook processes. + +Release assets include `provenance.json` with source commit, run/artifact IDs, sizes, and SHA-256 hashes. Existing version tags must identify the tested commit. Publication is serialized per release; superseded automatic main builds do not replace development assets. Manual development dispatch still deliberately permits a maintainer-selected branch. `Develop` remains a rolling asset set with a historical, unmoved tag: use the manifest, not that tag, to identify its binaries. Multi-asset uploads are not transactional; a failed upload can leave a partial set, so verify both binary hashes against the manifest. Selecting a commit authorizes publication of its bytes, not execution of its build code in the write-enabled publisher. + +## Live reconstruction benchmarks + +`tests\benchmark_dumps.py` is an **opt-in** Windows benchmark using Python 3's standard library and Windows PowerShell. It snapshots running processes and dumps each accessible **main executable**, not every DLL or private allocation. It does not elevate, inject, suspend, hook, or terminate target processes. Unavailable/protected processes, exited/replaced PIDs, missing original files, and images over the 256 MiB limit are reported separately, not counted as successful tests. + +Keep a baseline build, then run both builds against the same inventory: + +```powershell +python .\tests\benchmark_dumps.py --exe C:\bench\before\pd.exe --report C:\bench\before.json +python .\tests\benchmark_dumps.py --exe .\x64\Release\pd.exe --snapshot C:\bench\before.json --compare C:\bench\before.json --report C:\bench\after.json +``` + +Repeat with `--imports` on **both** commands and distinct report names to exercise aggressive import reconstruction. Use `--pid 1234` (repeatable) or `--name spotify` to narrow a run. Each dumper child has a 30-second timeout, configurable with `--timeout`; a timeout terminates only that child. Dump files are deleted after each comparison, including failed runs. A free-space floor prevents further work below 512 MiB. Reports are never overwritten. + +Security software may quarantine reconstructed executables. The benchmark distinguishes unreadable output (`analysis_unavailable`) from invalid PE data (`invalid_dump`); neither counts as success. Do not disable protection to make a benchmark pass. Use `--skip-pid 1234` to record an explicit exclusion when repeating a known-blocked target. + +Reports record file sizes, zero-byte fractions, SHA-256 identities, section layout validity, 4 KiB executable-code block matches/losses, elapsed dumping time, and diagnostics. Paired comparisons require the same PID, creation time, path, original-file hash, and import mode. These are live, non-atomic snapshots: relocations, resolved imports, runtime patches, discarded pages, file overlays/signatures, and process exits prevent exact disk/memory equality. Timings include setup/export scanning and are diagnostic measurements, not controlled CPU-performance claims. Suspect flags are investigation leads, not proof of corruption. + +For **actual section-content fidelity**, add `--quality --memory-witness --unique-originals` to both commands. `--unique-originals` selects one process per executable path from the inventory; it does not claim to cover every process instance. `--quality` compares bytes at corresponding section RVAs, computes normalized per-section hashes, counts nonzero-to-zero changes, compares resource payload hashes (not just directory headers), and inspects import identities and IAT contents. HIGHLOW/DIR64 base relocations and intentionally rewritten debug file offsets are normalized; IAT changes are measured separately. Other relocation types are reported, not guessed. Sections without a matching name/RVA are reported as unmatched, never synthesized as all-zero comparisons. Differences in loaded versus on-disk layout make that reference inconclusive, for example after an application update or runtime image transformation. Paired IAT summaries exclude mismatched layouts; memory summaries report how many pairs have independent witnesses on both sides. Analyzer limits/errors produce `quality_unavailable`, not a claim that the dump is corrupt. + +`--memory-witness` additionally reads section contents independently, one page at a time, before and after the dump. It reports stable bytes that differ from the reconstruction separately from bytes that changed between those two reads. It excludes the IAT/debug fields intentionally rewritten by reconstruction and reports unreadable pages. A byte that changes and changes back between snapshots cannot be detected, so this is evidence of fidelity, not an atomic snapshot guarantee. Witness reads cover original section virtual extents, including initialized data beyond original raw file sizes; they can increase working set and analysis time. No memory bytes are included in reports. + +Repeated imports are intentional: separate references to the same function retain independent `FirstThunk` fixups so the loader can repair every location. The analyzer's repeated-slot counts are informational, not a deduplication request or failure criterion. Native tests exercise loader-style rebinding of adjacent repeated references alongside an original IAT. Existing import tables are restored using full PE32/PE64 pointer widths, including terminators; missing lookup tables retain captured bytes, and truncated/overlapping tables are diagnosed without partial rewriting. + +Reports contain local process names/paths and diagnostics, but no memory payloads. Keep them private and outside the repository. The analyzer's own portable tests run in CI via `python -m unittest discover -s tests -p test_benchmark_dumps.py`; CI never runs the live sweep. + +Reconstruction preserves virtual section sizes and image extent while omitting zero-filled raw tails; initialized runtime data is retained even beyond the original disk section size. Raw sizes and offsets honor file alignment. Reconstructed imports use `.pdimp` when unused section-header space exists; otherwise the last section is explicitly extended with a warning, which can still materialize a large virtual gap. Unmapped certificate references and stale checksums are cleared, and debug payload file offsets are relocated. These dumps are for analysis, not guaranteed runnable or signed copies of the originals. + +Section parsing honors extended optional-header sizes rather than assuming that the section table immediately follows the standard optional header. Export discovery includes ordinal-only and page-aligned exports, excludes holes and forwarder-string addresses, and rejects malformed table metadata without publishing partial results. Loaded export addresses redirected outside the module are preserved (as used by WOW64 USER32). Forwarded calls can still be recognized through the destination module's exports; forwarder strings themselves are not resolved as code. Consequently, aggressive-import counts may change without changing the captured section bytes or removing intentional per-location fixups. + +Address-specific dumping of a region without a PE header (or with `-g`) generates both PE32 and PE64 analysis images using the same bounded reconstruction path as loose code chunks. Generated images default to EXE; the internal DLL option now sets the DLL flag correctly. Import summaries include complete pointers at the end of the captured range, count their libraries, and distinguish ordinal identities. Ordinary PE/import serialization fingerprints remain unchanged, but affected loose-code import hashes can change; regenerate a clean database when relying on those hashes. Module enumeration resizes its buffer for large processes and reports failure if a changing list cannot be captured within eight attempts (at most 1,048,576 module handles). Concurrent clean-database lookups and updates are synchronized, and database-generation options have deterministic defaults. + +Native regression cases enforce a maximum 12 KiB dump (16 KiB with imports) for a fixture containing a 4 MiB mostly-zero virtual data section. They also cover valid sections beyond the former 60,000 KiB truncation threshold, sparse RVAs, unaligned image ends, high-address/terminal imports, occupied section-header space, and byte-preserving reloads of 32 varied PE32/PE64 layouts. + +## System scheduling and database generation + +`-system` now shares a persistent worker pool between per-process scanning, export parsing, individual module reconstruction/hashing, and loose-code processing. Idle workers can help with the last process instead of waiting for one process-owned thread. A producer waiting for its modules helps execute only module jobs: it never recursively starts another process or creates another thread pool. At most `-t` workers execute work (`-nt` uses one). Each process retains its handle, module snapshot and immutable export table until its jobs finish. Export aliases are merged in snapshot order, not completion order. Only small job metadata is queued; images are allocated when workers execute them, with at most 256 queued leaf jobs per active producer. + +The full region scan still discovers hidden/headerless code and retains the existing 1,000-page-per-region and loose-chunk limits; this is not a loaded-module-only shortcut. One region walk collects both PE candidates and executable heaps. Module exclusions finish before loose heaps are processed, retaining adjacent heaps. System dumping still updates its in-memory clean set between completed processes and performs one final new-PID scan after all initial work drains. Cross-process deduplication remains scheduling-dependent and snapshots remain non-atomic. The close-monitor and direct PID paths retain their synchronous lifecycle. + +`-db genquick` uses the same module pool for live hashing. `-db gen` and `-db add` also hash files concurrently using `-t`, with bounded batches of 256 file jobs; recursion, filename filtering and reparse-directory exclusions are retained. Database hashing no longer changes the caller's import-reconstruction setting. The CLI continues to require generation and dumping as separate commands. The three existing database formats/signatures are unchanged; compare their **sets**, not serialization order. Parallelism increases simultaneous image memory and is not a guarantee that 16 workers beat four on an I/O-bound machine. + +Two opt-in runners make these comparisons repeatable: + +```powershell +python .\tests\benchmark_system.py --probe .\x64\Release\pd_tests.exe --report C:\bench\system.json --threads 1 4 16 --repeats 3 --acknowledge-execution +python .\tests\benchmark_database.py --baseline C:\bench\before\pd.exe --exe .\x64\Release\pd.exe --manifest C:\bench\corpus.txt --report C:\bench\database.json +``` + +The system runner launches only its own benign test host containing 64 inert, non-executable mapped PE buffers plus its ordinary loaded modules. It uses the production system scheduler with an explicit one-PID inventory, modelling the final-process imbalance without dumping unrelated processes. It compares serial execution with 1/4/16 workers, alternates order, requires identical filenames and SHA-256 bytes for every dump and identical module/full-entrypoint/short-entrypoint hash sets, and removes captures afterward. `--flags` selects combinations of `i` (imports), `c` (loose chunks), `g` (generated headers) and `r` (reexecution); ordinary defaults test imports/chunks on and off. Synthetic images are never executed. The database runner stages the line-separated manifest's files privately, executes only the supplied dumper builds, and rejects missing, truncated or changed database sets. Reports are never overwritten; keep their local paths/hashes private. Run timing benchmarks separately from other builds/profilers. + +On the local 64-module tail workload, three alternating isolated trials of the final build gave median dump-plus-hash times of **523.02 ms serial / 342.58 ms with four workers / 320.87 ms with 16** for imports and loose chunks enabled. With both disabled, the corresponding times were **373.07 / 183.16 / 165.31 ms**. Live hashing with chunks and imports enabled improved **119.37 / 70.08 / 67.29 ms**. All 96 final timed runs preserved exact dump bytes and hash sets. An earlier 96-run measurement, another 24 runs covering generated-header/reexecution combinations, and a six-run Win32 smoke comparison also retained equivalence. All eight option combinations retained byte-identical output against the saved pre-change dumper (78 normal or 156 generated images). This measures a deliberately imbalanced tail, not a universal full-system speedup. + +For 621 identical staged PE files, three alternating trials of database generation produced median elapsed times of **500.65 ms before / 478.35 ms with one worker / 240.28 ms with four / 221.62 ms with 16**, with all three hash sets identical. The first baseline trial was cold and took 10.94 seconds; it is retained in the report rather than misrepresented as a 50-fold CPU improvement. Filesystem and security-scanner caching still influence these short runs. + +A separate real-machine `-db genquick -t 4 -nh` run completed with 3,237 clean hashes, 484 full entrypoint hashes and 121 short hashes. A subsequent default `-system -t 4` run using those private databases completed with 507 outputs (about 231 MiB), retaining imports and loose chunks. Captures were removed after recording their hashes. Existing access restrictions and the 256 MiB per-image bound still apply. Full `-db gen` traversal over all Windows/user/program folders was not timed end-to-end: its live-hashing component, recursive file worker, filters and all three output sets were exercised separately. + +## Entrypoint recovery + +Automatic recovery preserves an existing entrypoint in executable section storage, including legitimate RVA `0x2000` and low-alignment images. `-eprec` still forces a search; `-nep` disables it. Recovery scans executable section ranges rather than every image byte, requires a full opcode-hash match, and preserves the earliest exact-prefix/full-hash match regardless of verbosity. An eight-byte prefix alone is no longer enough to rewrite the entrypoint. If exact-prefix matching fails on x64, validated exception/unwind-table function starts provide a bounded fallback: only a unique full opcode-hash match is selected. Malformed or ambiguous metadata is not guessed. No original file, PDB or external symbol service is required. + +The read-only `pd_tests.exe --entrypoint-corpus C:\bench\corpus.txt C:\bench\recovery.csv` benchmark evaluates known on-disk entrypoints without executing the images. It compares seven policies: legacy weak/strong selection, executable-only strong selection, unique strong matches, function-boundary fallback, unique weak fallback, exact-match precedence, and x64 runtime-function-only fallback. It also invokes the production recovery method with the header entrypoint erased. Both an oracle database containing the image and a leave-one-image-out database are measured; the latter retains signatures shared by other files and is not a leave-one-family-out generalization claim. + +A SHA-256-deduplicated local corpus selected 621 system/application images (about 551 MiB); 577 had a nonzero executable entrypoint and were measured (263 PE32, 314 PE64). With each image's own signature contribution withheld, legacy recovery returned **323 correct / 252 wrong / 2 unresolved** results. The selected production policy returned **357 correct / 11 wrong / 209 unresolved**. With the image represented in the database, both recovered all **577/577**. Exact-only recovery produced 316 correct / 7 wrong / 254 unresolved held-out results; the selected runtime-function fallback adds 41 correct recoveries at the cost of four additional wrong guesses. Requiring uniqueness for every exact match unnecessarily lost seven known-image recoveries; scanning x86 call targets added no correct recoveries, so that extra work was omitted. Retaining weak-only guesses sharply increased false positives. These measurements support conservative abstention, not a claim of perfect recovery: opcode fingerprints can still collide and unusual/packed code or absent training signatures can remain unresolved. Unsuccessful forced recovery leaves the prior RVA unchanged and reports it. + +## Experimental re-execution + +`-reexec` prepares an analysis dump for a **fresh launch**, rather than preserving all captured state. It is off by default and never executes a dump itself. It restores the compiler's GS-cookie bootstrap value through validated load-config metadata, resets writable/non-executable zero-fill tails described by the captured section headers, and limits **new speculative** imports to originally file-backed writable/non-executable data outside delay-IAT sections. Existing import reconstruction and independent repeated eligible fixups remain supported. ASLR, DEP, CFG and delay-IAT protections are not disabled. + +These changes deliberately discard runtime data. They assume the section headers still describe the original on-disk layout; do not use this mode when faithful analysis of unpacked or runtime-generated data is the goal. Initialized file-backed globals, heaps, handles, TLS, dependencies, resources and external state are not restored. This is not a process checkpoint/restart facility or a guarantee that a dumped application will run. Normal dumping continues to preserve initialized virtual tails and aggressive import coverage. + +`tests\benchmark_reexecution.py` measures explicitly selected **trusted, benign** applications using a same-architecture `pd_tests.exe` debugger probe. It launches the original, an untouched staged disk-copy control, and the reconstruction with the same arguments and working directory. Only its newly created children are debugged and assigned to a kill-on-close job; unrelated running applications are untouched. A temporary entrypoint breakpoint is restored, including the instruction pointer, before capture. `--phase entry` captures at that breakpoint, `--phase ready` at the test host's `PD_REEXEC_READY` marker, and `--phase idle` after GUI input-idle. Entry/ready captures occur while a debug event suspends the child; idle capture is live and non-atomic. Dumping time is excluded from the observation duration. + +For example, retain a baseline `pd.exe`, create a private `C:\bench` directory, then compare Notepad startup on a machine with this matching language resource: + +```powershell +python .\tests\benchmark_reexecution.py --app C:\Windows\System32\notepad.exe --probe .\x64\Release\pd_tests.exe --dumper C:\bench\before\pd.exe --mui C:\Windows\System32\en-US\notepad.exe.mui --imports --report C:\bench\notepad-before.json --acknowledge-execution +python .\tests\benchmark_reexecution.py --app C:\Windows\System32\notepad.exe --probe .\x64\Release\pd_tests.exe --dumper .\x64\Release\pd.exe --mui C:\Windows\System32\en-US\notepad.exe.mui --imports --prepare --compare C:\bench\notepad-before.json --report C:\bench\notepad-after.json --acknowledge-execution +``` + +Omit `--imports` on both commands to disable additional import discovery. `--prepare` passes `-reexec`; `--arg=VALUE` is repeatable. Use `--arg=--reexecution-fixture --phase ready` with a frozen copy of `pd_tests.exe` as the application to exercise initialized CRT state, a worker thread and normal completion. `--milliseconds` controls the 100..30,000 ms observation window (default 3,000); each dumper has a separate 30-second deadline. Optional `--loader-snaps` collects bounded loader diagnostics only in the owned child, without changing system-wide settings. Both reports must use the same probe build, application, arguments, resources, capture phase, duration and OS; changed or failed controls make a comparison inconclusive. Controls must reach their entrypoint and either GUI idle or exit zero without timing out; intentional nonzero console exits also require manual interpretation rather than an automatic comparison. + +The optional `--mui` copies only the explicitly supplied matching `locale\app.exe.mui` beside both staged launches; no original executable code is substituted into the dump. Missing resource/dependency files and changed application directories can prevent an otherwise valid image from starting, which is why the staged original is a separate control. Reports distinguish loader, entrypoint, GUI input-idle, fixture thread/completion and clean-exit milestones, and compare captured console-output hashes. GUI input-idle proves startup progress, not editor functionality; exit zero alone does not prove GUI or fixture completion. Application hashes, module-relative exception locations and exit/exception codes are recorded. + +On Windows 11 build 22631, two trials of six frozen/system application images, both import modes and the relevant capture phases produced **40 matched baseline/prepared pairs**: 34 gained execution milestones, six retained their previous progress, and none lost milestones. Both x86/x64 test hosts completed their threaded fixture from entrypoint and initialized-CRT captures (16 runs). `sort.exe` produced the expected sorted bytes and `findstr.exe /?` the same help output as both controls (eight runs). Both x86/x64 Notepads captured at entrypoint reached GUI input-idle with the matching MUI (eight runs). Notepad captured at GUI idle progressed past its former loader rejection but still failed after entrypoint with `0xc0000409`, parameter 5 (eight runs). These are bounded startup observations, not general restart guarantees; the remaining failure was not worked around by resetting arbitrary initialized globals or weakening mitigations. + +This harness is **not a security sandbox**: children run with the caller's permissions and can affect files, the network and external applications. Never select malware or unknown images, bypass security software, or supply destructive application arguments. Output redirection is not size-limited; select bounded workloads. The normal unit suite never executes its synthetic PE fixtures, and CI never runs this opt-in launch benchmark. Temporary captures/staged copies are removed on normal/error unwinding; forced termination can leave files behind. Reports contain local paths, addresses and diagnostics, but no captured memory payloads; keep reports and any investigation dumps private and outside the repository. + +## Controlled performance benchmarks + +`pd_tests.exe --benchmark imports 4 8` runs eight reconstruction jobs on four workers. The opt-in workloads are `dense` (4 MiB), `sparse` (32 MiB with a 4 KiB initialized prefix), `imports` (65,536 independent fixups), `entrypoint` (a nonmatching entrypoint database), `entrypoint-empty`, `database` (known-hit/miss queries), and `exports` (discovery in the test process itself). Fixtures are never executed. No other process is inspected or modified. + +Save a reference Release test executable **with the same benchmark harness**, then alternate reference/candidate trials: + +```powershell +python .\tests\benchmark_performance.py --baseline C:\bench\before\pd_tests.exe --exe .\x64\Release\pd_tests.exe --report C:\bench\performance.json --threads 1 2 4 8 --jobs 16 --repeats 5 +``` + +Reports retain individual samples, median phase timings, ratios, executable hashes, worker CPU cycles, coarse thread CPU times, and process peak working set. `capture_ms`, `reconstruct_ms`, `write_ms`, `cleanup_ms`, and `service_ms` sum work across jobs; they are **not** parallel elapsed time. `worker_ms` ends when the last worker finishes processing; `wall_ms` also includes thread teardown/join. Fixture setup, output-name reservation, thread creation, output readback/CRC, and file deletion are outside these intervals. A warmup precedes every trial. Peak working set includes warmup but is sampled before final readback. Timing still reflects filesystem caching, security scanning, CPU scheduling/frequency, and memory-bandwidth contention; worker throughput and maximum thread count are not universal speedup guarantees. + +Every serialized job is checked after timing against the warmup's output size and CRC32, normalizing only the fixture's ASLR-dependent PE64 image base. The runner rejects changed fingerprints or mismatched configurations across trials/builds. Database queries verify their expected hit count; export discovery has no serialized output, so neither reports output-fingerprint equivalence. Use the native semantic tests and the separate live-content benchmark for those paths. Temporary fixture outputs are removed on normal/error unwinding; forcibly killing the benchmark can leave files behind. + +For allocation attribution, build `tests\pd_tests.vcxproj` with `/p:ProfileAllocations=true` and separate `OutDir`/`IntDir` directories. This replaces C++ `new`/`delete` in the **test executable only**, counting thread-local calls/requested bytes during each job, not every CRT/Windows allocation. Compare instrumented builds only to other instrumented builds, and confirm timings separately with ordinary Release builds. Allocation instrumentation cannot be combined with AddressSanitizer. Reports are local-only, never overwritten, and not uploaded automatically. + +The optimized pipeline keeps captured bytes read-only during packing, scans zero tails a machine word at a time, and avoids a whole-image copy when imports are disabled. Import records are contiguous; reconstruction borrows names from its immutable export list instead of allocating copies for each fixup. Standalone string-based import APIs still own their strings. Export-list transfers retain first-wins alias precedence and stable record addresses. Process capture reuses the dumper's handle while preserving module names and ownership. When entrypoint recovery is needed, each image acquires a consistent immutable snapshot of both entrypoint sets, skips empty databases, and scans without contending on the shared database lock. Workers share the cached snapshot until entrypoint data changes, rather than copying a potentially large database per image. Updates become visible to the next snapshot; ordinary clean-hash lookups remain synchronized. Repeated per-location imports and PE/import serialization semantics are unchanged. + ## Handling untrusted input PE headers, sections, imports, and exports are treated as untrusted data. Invalid ranges and reconstruction sizes above 256 MiB per image are rejected with a diagnostic instead of attempting unsafe allocations. This limit also applies to generated headers and reconstructed disk images. -Process and module names are encoded as single filename components. Dump files are **created only if the path does not exist**: existing dumps, hard links, and symbolic links are never overwritten. Use a fresh output directory when repeating a dump. Output directories containing reparse points (including junctions), unsupported device paths, and names exceeding supported Windows path/component lengths are rejected. Use an ordinary directory you control, outside the target's writable directories. +Process and module names are encoded as single filename components. Dump files are **created only if the path does not exist**: existing dumps, hard links, and symbolic links are never overwritten. Use a fresh output directory when repeating a dump. Output directories containing reparse points (including junctions), unsupported device paths, and names exceeding supported Windows path/component lengths are rejected. Creation uses `NtCreateFile` with `FILE_CREATE` and `OBJ_DONT_REPARSE`, relative to the held root directory, so intermediate reparses are rejected during creation even if directory metadata changes after validation. Directory handles also prevent rename/replacement during writing. Failure of the safe creation operation is reported without retrying through a weaker path API. Use an ordinary directory you control, outside the target's writable directories. Close monitoring uses a private release event rather than accepting a target-supplied thread ID. Stopping monitoring restores the original entry bytes and protection, then releases callbacks even if they have not started waiting yet. A published callback's two-page allocation and target-side event handle remain until that target exits; reclaiming them earlier could free code beneath an in-flight thread. Unpublished allocations and handles are cleaned up on failure. `pd_tests.exe --security` runs bounded validation and mock-based lifecycle tests; the tests contain no exploit PoCs and do not hook other processes. # Command-line arguments -Process dump can be used to dump all unknown code from memory ('-system' flag), dump specific processes, or run in a monitoring mode that dumps all processes just before they terminate. - -Before first usage of this tool, when on the clean workstation the clean excluding hash database can be generated by either: -* pd -db genquick -* pd -db gen +Use `pd.exe [options]`. Local builds are named `pd.exe`; release assets are `pd64.exe` and `pd32.exe`. Prefer the 64-bit build on 64-bit Windows to access both architectures. Administrator rights improve coverage but do not bypass protected processes. -Example Usage: -* pd -system -* pd -pid 419 -* pd -pid 0x1a3 -* pd -pid 0x1a3 -a 0x401000 -o c:\dump\ -c c:\dump\test\clean.db -* pd -p chrome.exe -* pd -p "(?i).\*chrome.\*" -* pd -closemon +Choose one dump or database-maintenance command per invocation; generate the baseline and dump in **separate commands**. `-db ignore` is a modifier that can accompany dumping. Flags are case-sensitive. No arguments, `--help`, `-help`, `-h` or `--h` print help and exit without running a dump or database command. -The command-line arguments can be grouped as follows: +Defaults: aggressive import reconstruction, loose-code discovery, and entrypoint recovery are enabled; `-reexec`, generated headers and verbose output are disabled. Dumps go to the current working directory; database files live beside the executable. The default worker count is 16. -**General Dumping Options** +**Dump commands and output** | Option | Description | |--------|-------------| -| -system | Dumps all modules not matching the clean hash database from all accessible processes into the working directory. | -| -pid \ | Dumps all modules not matching the clean hash database from the specified PID into the current working directory. Use a '0x' prefix to specify a hex PID. | -| -closemon | Runs in monitor mode. When any processes are terminating, process dump will first dump the process. | -| -p \ | Dumps all modules not matching the clean hash database from the process name found to match the filter into specified PID into the current working directory. | -| -a \ | Dumps a module at the specified base address from the process. | -| -o \ | Sets the default output root folder for dumped components. | +| `-system` | Dump unknown modules, hidden modules and eligible loose chunks from all accessible processes. Workers share module jobs, including the last process; one final scan discovers new PIDs after initial work drains. | +| `-pid ` | Dump one process. Use decimal or a `0x`-prefixed hexadecimal PID. | +| `-p ` | Match the **entire** process name using a case-sensitive ECMAScript regular expression. Quote the expression; use `.*` for a substring match. Multiple matches prompt for approval. Inline `(?i)` is not supported; use explicit character classes when needed. | +| `-closemon` | Hook termination of accessible processes and dump them before exit. This modifies target processes; use only in a controlled environment. Press Ctrl+C to stop and clean up hooks. It cannot guarantee capture of every termination path or protected process. | +| `-a
` | Dump at a decimal or `0x`-prefixed base address; **requires `-pid`**. Reuse a valid PE header. If absent, or with `-g`, generate both PE32 and PE64 analysis images. Clean filtering still applies. | +| `-o ` | Output directory, defaulting to the current working directory. Create an ordinary directory you control first. Existing files are never overwritten and reparse paths are rejected; use fresh directories for repeat captures. | -**Clean Hash Database Options** +**Database commands and paths** + +There are three databases: clean-module hashes (`clean.hashes`), full entrypoint-opcode hashes (`entrypoints.hashes`), and short entrypoint-prefix hashes (`shortentrypoints.hashes`). Defaults are beside the executable, not in the working/output directory. Maintenance commands save to these paths; use all three path overrides for an isolated baseline. Generation/addition extends existing sets rather than replacing them. Generate only on a known-clean system; skipped matches do not prove that all remaining code is malicious. | Option | Description | |--------|-------------| -| -db gen | Automatically processes a few common folders as well as all the currently running processes and adds the found module hashes to the clean hash database. It will add all files recursively in: `%WINDIR%`, `%HOMEPATH%`, `C:\Program Files\`, `C:\Program Files (x86)\`, as well as all modules in all running processes. These clean hashes will be added to the file `clean.hashes` in the application directory. During future process dumping commands, these known modules will not be dumped. It is recommended to run this command one time on a clean system prior to using the tool that way not too many modules will be dumped from memory.| -| -db genquick | Same as above, but only adds the hashes from all modules in all processes to the clean hash database. This is a much faster way to build the clean hash database, but it will be less complete. | -| -db add \ | Adds all the files in the specified directory recursively to the clean hash database. | -| -db rem \ | Removes all the files in the specified directory recursively from the clean hash database. | -| -nr | Disable recursion on hash database directory add or remove commands. | -| -db clean | Clears the clean hash database. | -| -db ignore | Ignores the clean hash database when dumping a process this time. All modules will be dumped even if a match is found. | -| -cdb \ | Full filepath to the clean hash database to use for this run if you'd like to override the default of `clean.hashes`. | -| -edb \ | Full filepath to the entrypoint hash database to use for this run. | -| -esdb \ | Full filepath to the entrypoint short hash database to use for this run. | - -**Output Options** +| `-db gen` | Add live module/chunk and entrypoint hashes, then recursively scan PE files in `%WINDIR%`, `%USERPROFILE%`, `C:\Program Files` and `C:\Program Files (x86)`. Uses `-t` for live and file work. `-nr` does not change this full recursive scan. | +| `-db genquick` | Add live module/chunk and entrypoint hashes only, without scanning disk folders. Uses the shared module pool; `-nc` excludes loose chunks. Faster but less complete than `gen`. | +| `-db add ` | Add PE file hashes from a directory to all three databases, recursively by default. Uses `-t`; reparse directories are excluded. | +| `-db remove `, `-db rem ` | Remove matching **clean-module** hashes for PE files in the directory, recursively by default. Does not remove entrypoint signatures. Removal remains serial. | +| `-nr` | Disable recursion for `-db add`, `-db remove` and `-db rem` only. | +| `-db clean` | Clear and save **all three databases**, including both entrypoint sets. | +| `-db ignore` | Ignore **all three loaded sets** in memory for this dump invocation, including entrypoint signatures; does not clear database files. In-run cross-process deduplication can still skip repeated modules. | +| `-cdb `, `-c ` | Override the clean-module database path. | +| `-edb ` | Override the full entrypoint-opcode database path. | +| `-esdb ` | Override the short entrypoint-prefix database path. | + +**Reconstruction and workers** | Option | Description | |--------|-------------| -| -v | Verbose mode where more details will be printed for debugging. | -| -nh | No header is printed in the output. | +| `-g` | Force generated PE32 and PE64 analysis headers instead of existing headers. Both architectures are emitted because raw code does not reliably identify its architecture. | +| `-ni` | Disable aggressive import reconstruction (enabled by default). | +| `-nc` | Disable loose-code dumping and hashing (enabled by default). Loaded and hidden PE modules are still scanned. | +| `-nep` | Disable entrypoint recovery and live entrypoint-signature collection; takes precedence over `-eprec`. File database add/gen still collect entrypoint signatures using their own parsing options. | +| `-eprec` | Force entrypoint recovery even when a valid executable entrypoint exists. Normally valid entrypoints are preserved. Requires full opcode matches; x64 runtime-function metadata provides a fallback. Weak prefixes alone are not accepted, and failure retains the previous RVA with a warning. Recovery is heuristic, not guaranteed. | +| `-reexec` | Experimental fresh-launch preparation, off by default. Reset the GS-cookie bootstrap value and writable/non-executable zero-fill tails; restrict new imports to file-backed writable data outside delay-IAT sections. Assumes original section layout and **discards captured state**. Does not execute dumps, restore a process checkpoint or disable ASLR/DEP/CFG. See [re-execution caveats](#experimental-re-execution). | +| `-t ` | Set worker count to at least 1 (default 16; decimal or `0x`-prefixed hex). Applies to system dumping, live/file database generation and close-monitor dump workers. Direct `-pid`/`-p` dumps remain synchronous. More workers can increase memory use and are not always faster. | +| `-nt` | Use one worker, equivalent to `-t 1`; the last `-t`/`-nt` wins. Close monitoring still has a separate monitoring thread. | -**Advanced Options** +**Help and diagnostics** | Option | Description | |--------|-------------| -| -g | Forces generation of PE headers from scratch, ignoring existing headers. | -| -eprec | Force the entry point to be reconstructed, even if a valid one appears to exist. | -| -ni | Disable import reconstruction. | -| -nc | Disable dumping of loose code regions. | -| -nt | Disable multithreading. | -| -nep | Disable entry point hashing. | -| -t \ | Sets the number of threads to use (default 16). | +| `--help`, `-help`, `-h`, `--h` | Print help and exit. No arguments also prints help. | +| `-v` | Enable verbose diagnostics (off by default). | +| `-nh` | Suppress the version/copyright banner, not other output; help text still prints. | + +Images and reconstructions above 256 MiB are rejected. Captures are live and non-atomic, can contain sensitive data, and are not guaranteed runnable or signed copies. Keep output private and review the [input/output safety notes](#handling-untrusted-input). # Usage Examples | Command | Description | | ------- | ----------- | -| `pd64.exe -db genquick` | Quickly build clean module database based on currently running processes. Process Dump in later tasks will only dump unrecognized modules. | -| `pd64.exe -system` | Dump all modules and hidden chunks from all processes while ignoring clean modules. | -| `pd64.exe -closemon` | Run in terminate monitor mode. This will dump all processes when they attempt to terminate. | -| `pd64.exe -pid 0x18A` | Dump modules and hidden chunks from a specific process ID. | -| `pd64.exe -p .\*chrome.\*` | Dump modules and hidden chunks by process name. | -| `pd64.exe -db gen` | Build a clean-hash database of known modules. This is used to avoid dumping known good modules in later tasks. | -| `pd64.exe -pid 0x1a3 -a 0xffb4000` | Dump code from a specific address in PID. This will generate two files for analysis, with reconstructed 32bit and 64bit PE headers: `notepad_exe_x64_hidden_FFB40000.exe` and `notepad_exe_x86_hidden_FFB40000.exe`. | - -Sure, here's a more streamlined version of the information: +| `pd64.exe -db genquick -t 4` | Build a live baseline on a known-clean system. Run separately from dumping. | +| `pd64.exe -system -t 4 -o C:\dumps` | Dump unknown modules/chunks across accessible processes into a pre-created output directory. | +| `pd64.exe -closemon -o C:\dumps` | Monitor hooked termination paths in a controlled environment; Ctrl+C stops monitoring. | +| `pd64.exe -pid 0x18A` | Dump modules and eligible loose chunks from one process. | +| `pd64.exe -p ".*chrome.*" -ni -nc` | Match process names containing lowercase `chrome`, without aggressive imports or loose chunks. | +| `pd64.exe -p "chrome[.]exe"` | Match exactly `chrome.exe`. | +| `pd64.exe -db gen` | Extend the live and disk baseline using the documented folders. | +| `pd64.exe -db add C:\baseline -nr -t 4` | Add PE files in just this directory with four workers. | +| `pd64.exe -pid 0x1a3 -a 0x401000 -o C:\dumps` | Dump at a chosen base address. Generates both architectures only when a header is absent or `-g` is supplied. | +| `pd64.exe -pid 419 -reexec -o C:\dumps` | Prepare dumps for fresh-launch experiments; does not launch them. | +| `pd64.exe -pid 419 -eprec -cdb C:\baseline\clean.hashes -edb C:\baseline\entrypoints.hashes -esdb C:\baseline\shortentrypoints.hashes` | Force recovery using an isolated set of database paths. | ## Sandbox Usage @@ -126,16 +237,16 @@ When using Process Dump in an automated sandbox or for manual anti-malware resea - **Build the Clean Hash Database:** Run `pd64.exe -db gen` or for a faster less complete process, use `pd64.exe -db genquick`. Depending on your situation, you may want to snapshot your VM after creating this clean hash database that way it doesn't need to be repeated each time. -- **Start the Process Dump Terminate Monitor:** Keep `pd64.exe -closemon` running in the background. It will dump all intermediate processes used by the malware. +- **Start the Process Dump Terminate Monitor:** Keep `pd64.exe -closemon` running in the background. It attempts to capture accessible processes using hooked termination paths; capture of every intermediate process is not guaranteed. -- **Execute the Malware File:** Monitor the malware installation. `pd64.exe` will automatically dump any process that tries to close. +- **Observe the Sample in an Isolated Lab:** Monitor the sample's behavior while `pd64.exe` captures hooked processes that terminate. Never use close monitoring as a containment or safety mechanism. - **Dump the Running Malware from Memory:** When ready, use `pd64.exe -system` to dump all processes. The dumped components will be found in the working directory of `pd64.exe`. To change the output path, use the `-o` flag. # Notes on the naming convention of dumped modules: -* 'hiddemodule' in the filename instead of the module name indicates the module was not properly registered in the process. +* 'hiddenmodule' in the filename instead of the module name indicates the module was not properly registered in the process. * 'codechunk' in the filename means that it is a reconstructed dump from a loose executable region. This can be for example injected code that did not have a PE header. Codechunks will be dumped twice, once with a reconstructed x86 and again with a reconstructed x64 header. Example filenames of dumped files @@ -147,6 +258,15 @@ Example filenames of dumped files # Version history +## Version 3.0.0 (unreleased) +* Hardened untrusted PE parsing, reconstruction bounds, output naming/creation, close-monitor cleanup and release artifact handling. +* Improved sparse section packing and content fidelity, PE32/PE64 imports and exports, raw-address reconstruction, and handling of extended/malformed headers. Intentional per-location IAT fixups are preserved. +* Added persistent module-level system scheduling, parallel live/file database hashing and lower-allocation reconstruction, with repeatable output-equivalence benchmarks. +* Improved entrypoint selection: preserve valid executable entrypoints, reject weak-only guesses, and use validated x64 runtime-function metadata as a fallback. Accuracy and remaining false recoveries are documented above. +* Added opt-in `-reexec` preparation and a controlled execution-progress benchmark; ordinary forensic dumps continue to preserve captured state. +* Added native regression tests, portable benchmark/release/CLI checks and Win32/x64 Debug/Release CI with modern Visual Studio toolsets. +* Unified the 3.0.0 banner and Windows version resources, refreshed all CLI flags/defaults/aliases, and made help exit without running a requested dump or database command. + ## Version 2.1 (February 12th, 2017) * Fixed a bug where the last section in some cases would instead be filled with zeros. Thanks to megastupidmonkey for reporting this issue. * Fixed a bug where 64-bit base addresses would be truncated to a 32-bit address. It now properly keeps the full 64-bit module base address. Thanks to megastupidmonkey for reporting this issue. diff --git a/pd/dump_output.h b/pd/dump_output.h index 491b5d7..c9e4374 100644 --- a/pd/dump_output.h +++ b/pd/dump_output.h @@ -1,26 +1,18 @@ #pragma once #include +#include #include #include #include #include +#include class dump_directory_locks { std::vector handles; -public: - ~dump_directory_locks() + bool accept(HANDLE handle) { - for (HANDLE handle : handles) - CloseHandle(handle); - } - - bool hold(const std::string& path) - { - HANDLE handle = CreateFileA(path.c_str(), FILE_READ_ATTRIBUTES, - FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, - FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT, NULL); if (handle == INVALID_HANDLE_VALUE) return false; BY_HANDLE_FILE_INFORMATION info = {}; @@ -35,8 +27,77 @@ class dump_directory_locks handles.push_back(handle); return true; } +public: + HANDLE root() const { return handles.empty() ? INVALID_HANDLE_VALUE : handles.front(); } + + ~dump_directory_locks() + { + for (HANDLE handle : handles) + CloseHandle(handle); + } + + bool hold(const std::string& path) + { + return accept(CreateFileA(path.c_str(), FILE_READ_ATTRIBUTES, FILE_SHARE_READ | FILE_SHARE_WRITE, + NULL, OPEN_EXISTING, FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT, NULL)); + } + bool hold_relative(const std::string& path); }; +namespace dump_output +{ + inline HANDLE open_no_reparse(HANDLE root, const char* relative_path, bool directory, + decltype(&NtCreateFile) create_file = NtCreateFile) + { + if (root == NULL || root == INVALID_HANDLE_VALUE || relative_path == NULL || + relative_path[0] == 0 || relative_path[0] == '\\' || relative_path[0] == '/' || + strchr(relative_path, ':') != NULL) + { + SetLastError(ERROR_INVALID_NAME); + return INVALID_HANDLE_VALUE; + } + WCHAR name_buffer[MAX_PATH]; + const int length = MultiByteToWideChar(AreFileApisANSI() ? CP_ACP : CP_OEMCP, MB_ERR_INVALID_CHARS, + relative_path, -1, name_buffer, MAX_PATH); + if (length == 0) + return INVALID_HANDLE_VALUE; + UNICODE_STRING name = {}; + name.Buffer = name_buffer; + name.Length = static_cast((length - 1) * sizeof(WCHAR)); + name.MaximumLength = static_cast(length * sizeof(WCHAR)); + OBJECT_ATTRIBUTES attributes = {}; + attributes.Length = sizeof(attributes); + attributes.RootDirectory = root; + attributes.ObjectName = &name; + attributes.Attributes = OBJ_CASE_INSENSITIVE | OBJ_DONT_REPARSE; + IO_STATUS_BLOCK io = {}; + HANDLE file = INVALID_HANDLE_VALUE; + const NTSTATUS status = create_file(&file, + (directory ? FILE_READ_ATTRIBUTES : GENERIC_WRITE | DELETE) | SYNCHRONIZE, &attributes, &io, + NULL, FILE_ATTRIBUTE_NORMAL, directory ? FILE_SHARE_READ | FILE_SHARE_WRITE : 0, + directory ? FILE_OPEN : FILE_CREATE, + (directory ? FILE_DIRECTORY_FILE : FILE_NON_DIRECTORY_FILE) | + FILE_SYNCHRONOUS_IO_NONALERT | FILE_OPEN_REPARSE_POINT, NULL, 0); + if (status < 0) + { + SetLastError(RtlNtStatusToDosError(status)); + return INVALID_HANDLE_VALUE; + } + return file; + } + + inline HANDLE create_no_reparse(HANDLE root, const char* relative_path, + decltype(&NtCreateFile) create_file = NtCreateFile) + { + return open_no_reparse(root, relative_path, false, create_file); + } +} + +inline bool dump_directory_locks::hold_relative(const std::string& path) +{ + return accept(dump_output::open_no_reparse(root(), path.c_str(), true)); +} + inline bool write_new_dump(const char* filename, const unsigned char* data, SIZE_T size) { if (filename == NULL || data == NULL || size == 0) @@ -71,7 +132,7 @@ inline bool write_new_dump(const char* filename, const unsigned char* data, SIZE } // Keep every parent directory open without delete sharing until creation and writing finish. - // This prevents a target from replacing a parent with a junction between validation and use. + // These handles prevent renames, but cannot prevent in-place reparse-point changes. dump_directory_locks locks; if (!locks.hold(path.substr(0, root_end))) { @@ -80,15 +141,15 @@ inline bool write_new_dump(const char* filename, const unsigned char* data, SIZE } for (SIZE_T end = path.find('\\', root_end); end != std::string::npos; end = path.find('\\', end + 1)) { - if (!locks.hold(path.substr(0, end))) + if (!locks.hold_relative(path.substr(root_end, end - root_end))) { fprintf(stderr, "ERROR: Dump directory is unavailable or contains a reparse point: '%s' (Windows error %lu).\n", filename, GetLastError()); return false; } } - HANDLE file = CreateFileA(path.c_str(), GENERIC_WRITE | DELETE, 0, NULL, CREATE_NEW, - FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OPEN_REPARSE_POINT, NULL); + // Reject reparses during the actual create, not only during the earlier directory checks. + HANDLE file = dump_output::create_no_reparse(locks.root(), path.c_str() + root_end); if (file == INVALID_HANDLE_VALUE) { fprintf(stderr, "ERROR: Cannot create dump '%s'; existing files are not overwritten (Windows error %lu).\n", diff --git a/pd/dump_process.cpp b/pd/dump_process.cpp index 7544119..4183fdf 100644 --- a/pd/dump_process.cpp +++ b/pd/dump_process.cpp @@ -1,6 +1,7 @@ #include "StdAfx.h" #include "dump_process.h" #include "dump_path.h" +#include "module_work.h" @@ -146,7 +147,7 @@ MBI_BASIC_INFO dump_process::get_mbi_info(unsigned __int64 address) result.base = mbi.BaseAddress; result.end = mbi.BaseAddress + mbi.RegionSize; result.protect = mbi.Protect; - result.valid = mbi.State != MEM_FREE && !(mbi.Protect & (PAGE_NOACCESS | PAGE_GUARD)); + result.valid = mbi.State == MEM_COMMIT && !(mbi.Protect & (PAGE_NOACCESS | PAGE_GUARD)); result.executable = (mbi.Protect & (PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY)) > 0; } else if (blockSize == sizeof(_MEMORY_BASIC_INFORMATION32)) @@ -156,7 +157,7 @@ MBI_BASIC_INFO dump_process::get_mbi_info(unsigned __int64 address) result.base = mbi32->BaseAddress; result.end = mbi32->BaseAddress + mbi32->RegionSize; result.protect = mbi32->Protect; - result.valid = mbi32->State != MEM_FREE && !(mbi32->Protect & (PAGE_NOACCESS | PAGE_GUARD)); + result.valid = mbi32->State == MEM_COMMIT && !(mbi32->Protect & (PAGE_NOACCESS | PAGE_GUARD)); result.executable = (mbi32->Protect & (PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY)) > 0; } @@ -164,146 +165,63 @@ MBI_BASIC_INFO dump_process::get_mbi_info(unsigned __int64 address) } -int dump_process::get_all_hashes(unordered_set* output_hashes, unordered_set* output_hashes_eps, unordered_set* output_hashes_ep_shorts) +std::vector dump_process::scan_regions(set& executable_heaps) +{ + std::vector bases; + for (unsigned __int64 address = 0;;) + { + const auto region = get_mbi_info(address); + if (region.base != 0 && region.end > region.base && region.valid) + { + if (_options->DumpChunks && region.executable) + executable_heaps.insert(region.base); + if (_options->Verbose) + printf("INFO: Scanning region 0x%llX to 0x%llX for MZ headers.\n", region.base, region.end); + unsigned __int64 base = region.base - region.base % PAGE_SIZE; + for (int page = 0; page < 1000 && base < region.end && region.end - base > 0x300; ++page, base += PAGE_SIZE) + { + WORD signature = 0; + SIZE_T read = 0; + if (ReadProcessMemory(_ph, reinterpret_cast(base), &signature, sizeof(signature), &read) && + read == sizeof(signature) && signature == IMAGE_DOS_SIGNATURE) + bases.push_back(base); + } + } + if (region.end <= address) + break; + address = region.end; + } + return bases; +} + +int dump_process::get_all_hashes(unordered_set* output_hashes, unordered_set* output_hashes_eps, unordered_set* output_hashes_ep_shorts, work_pool* pool) { // Adds all the modules in the process to the output array if( _ph != NULL ) { - if( !_options->DumpChunks || build_export_list() ) // Only build export list if getting hashes for code chunks + if( !_options->DumpChunks || build_export_list(pool) ) // Only build export list if getting hashes for code chunks { // First build a list of the modules - module_list* modules = new module_list( _pid ); - - // Set the max address of the target process - unsigned __int64 maxAddress = 0; - maxAddress = 0xffffffffffffffff; // Not a problem for 32bit targets + module_list module_snapshot(_pid); + auto* modules = &module_snapshot; - // Walk the process heaps - __int64 address = 0; - - // First loop to build a list of executable heaps for later use in locating loose executable heaps not associated with any modules set executable_heaps; - if( _options->DumpChunks ) - { - while (address < maxAddress) - { - MBI_BASIC_INFO mbi_info = get_mbi_info(address); - - // Check if this is a loose executable heap - if( mbi_info.base > 0 && mbi_info.end > 0 && mbi_info.valid && mbi_info.executable) - { - executable_heaps.insert(mbi_info.base); - } - - if( mbi_info.end + 1 <= address ) - break; - address = mbi_info.end + 1; - } - } - - // Load all the PE files in the process - address = 0; - - while (address < maxAddress) - { - // Load this heap information - MBI_BASIC_INFO mbi_info = get_mbi_info(address); - - if (mbi_info.base > 0 && mbi_info.end > 0 && mbi_info.valid ) - { - if( _options->Verbose ) - fprintf( stdout, "INFO: Scanning from region from 0x%llX to 0x%llX for MZ headers.\n", mbi_info.base, mbi_info.end ); - - // This heap may have a PE file, check all page alignments for a "MZ". - unsigned __int64 base = mbi_info.base - (mbi_info.base % PAGE_SIZE); // shouldn't be required. - char output[2]; - SIZE_T out_read; - int count = 0; - while(mbi_info.base + 0x300 < mbi_info.end && count < 1000 ) // Skip the rest of the section if we have looped over 1000 pages. - { - if( ReadProcessMemory( _ph, (LPCVOID) ((unsigned char*) base), output, 2, &out_read) && out_read == 2 ) - { - if( output[0] == 'M' && output[1] == 'Z' ) - { - if( _options->Verbose ) - fprintf( stdout, "INFO: Found MZ header at %llX.\n", base ); - - // Bingo, possible MZ file - pe_header* header = new pe_header( _pid, (void*) base, modules, _options ); - - header->process_pe_header(); - header->process_sections(); - - if( header->somewhat_parsed() ) - { - // Exclude all executable regions from this PE module - unsigned __int64 end_address = header->get_virtual_size() + base; - for (set::iterator it=executable_heaps.begin(); it!=executable_heaps.end(); ) - { - if( *it <= end_address && *it >= base ) - { - // We've accounted for this executable heap, remove it from the loose heap list - it = executable_heaps.erase(it); - } - else - { - it++; - } - } - - // Check hash - unsigned __int64 hash = header->get_hash(); - if( hash != 0 && !_db_clean->contains(hash) && output_hashes->count( hash ) == 0 ) - { - // Add this to the output hash array - output_hashes->insert( hash ); - } - - if (output_hashes_ep_shorts != NULL) - { - // Also get the entrypoint hash - hash = header->get_hash_ep_short(); - - if (hash != 0) - { - if( output_hashes_ep_shorts->count(hash) == 0 ) - { - // Add this to the output hash array - output_hashes_ep_shorts->insert(hash); - } - } - - if (output_hashes_eps != NULL) - { - // Also get the entrypoint hash - hash = header->get_hash_ep(); - - if (hash != 0 && output_hashes_eps->count(hash) == 0) - { - // Add this to the output hash array - output_hashes_eps->insert(hash); - } - } - } - - - - - - } - delete header; - } - } - - base += PAGE_SIZE; - count++; - } - } - - if (mbi_info.end + 1 <= address) - break; - address = mbi_info.end + 1; - } + const auto bases = scan_regions(executable_heaps); + std::mutex results; + module_work(pool, bases, [&](unsigned __int64 base) { + pe_header header(_ph, reinterpret_cast(base), modules, _options); + if (!header.process_pe_header() || !header.process_sections() || !header.somewhat_parsed()) + return; + const auto hash = header.get_hash(); + const auto short_hash = output_hashes_ep_shorts ? header.get_hash_ep_short() : 0; + const auto full_hash = output_hashes_ep_shorts && output_hashes_eps ? header.get_hash_ep() : 0; + const bool known = _db_clean->contains(hash); + std::lock_guard lock(results); + exclude_module_heaps(executable_heaps, base, header.get_virtual_size()); + if (hash && !known) output_hashes->insert(hash); + if (short_hash) output_hashes_ep_shorts->insert(short_hash); + if (full_hash) output_hashes_eps->insert(full_hash); + }); if( _options->DumpChunks ) @@ -313,6 +231,7 @@ int dump_process::get_all_hashes(unordered_set* output_hashes, fprintf( stdout, "INFO: Looking at unattached executable heaps...\n" ); int count_new_header_hashes = 0; + std::vector chunks; for (set::iterator it=executable_heaps.begin(); it!=executable_heaps.end(); it++) { // Unattached executable page. First check hash of crc32 of first 2kb of memory since import reconstruction hashing is @@ -336,31 +255,24 @@ int dump_process::get_all_hashes(unordered_set* output_hashes, // Add this header hash output_hashes->insert( chunk_header_hash ); - - // Calculate the generic import reference hash as well - pe_header* header = new pe_header( _pid, (void*) *it, modules, _options ); - header->build_pe_header( 0x1000, true, 1 ); // 64bit, only build it with the 1 executable section for performance reasons - header->process_sections(); - - // Get the import attributes of this header - IMPORT_SUMMARY import_summary = header->get_imports_information(&this->_export_list); - - // Check hash - if( import_summary.HASH_GENERIC != 0 && !_db_clean->contains(import_summary.HASH_GENERIC) && output_hashes->count( import_summary.HASH_GENERIC ) == 0 ) - { - if( _options->Verbose ) - fprintf( stdout, "INFO: Adding hash from unattached heap at 0x%llX to process hash list: Hash=0x%llX\n", *it, import_summary.HASH_GENERIC ); - - output_hashes->insert( import_summary.HASH_GENERIC ); - } - delete header; + chunks.push_back(*it); } } + module_work(pool, chunks, [&](unsigned __int64 base) { + pe_header header(_ph, reinterpret_cast(base), modules, _options); + if (!header.build_pe_header(0x1000, true, 1) || !header.process_sections()) + return; + const auto summary = header.get_imports_information(&_export_list); + if (summary.HASH_GENERIC && !_db_clean->contains(summary.HASH_GENERIC)) + { + std::lock_guard lock(results); + output_hashes->insert(summary.HASH_GENERIC); + } + }); if( _options->Verbose ) fprintf( stdout, "INFO: Done looking at unattached executable heaps...\n" ); } - delete modules; } } else if( _options->Verbose ) @@ -390,7 +302,7 @@ unsigned __int64 dump_process::hash_codechunk_header(__int64 base) return 0; } -bool dump_process::build_export_list() +bool dump_process::build_export_list(work_pool* pool) { // Walk through each module, building the export list for this process. This will be used for import reconstruction // Returns: True if there are any modules to dump, False if there is nothing to dump. @@ -403,23 +315,23 @@ bool dump_process::build_export_list() if (_ph != NULL) { // First build a list of the modules - module_list* modules = new module_list(_pid); - - // Loop through each of these modules, grabbing their exports - for (unordered_map::const_iterator item = modules->_modules.begin(); item != modules->_modules.end(); ++item) - { - pe_header* header = new pe_header(_pid, (void*) item->first, modules, _options); - if (header->process_pe_header() && header->process_sections() && header->process_export_directory()) + module_list modules(_pid); + std::vector bases; + for (const auto& item : modules._modules) bases.push_back(item.first); + std::vector indices; + for (size_t i = 0; i < bases.size(); ++i) indices.push_back(i); + std::vector> exports(bases.size()); + module_work(pool, indices, [&](size_t i) { + pe_header header(_ph, reinterpret_cast(bases[i]), &modules, _options); + if (header.process_pe_header() && header.process_sections() && header.process_export_directory()) { - // Load it's exports - this->_export_list.add_exports(header->get_exports()); + exports[i].reset(new export_list); + exports[i]->take_exports(*header.get_exports()); } - - // Cleanup - delete header; - } - - delete modules; + }); + // Merge in snapshot order, not completion order, to preserve alias precedence. + for (auto& entry : exports) + if (entry) _export_list.take_exports(*entry); } _export_list_built = true; } @@ -438,11 +350,11 @@ bool dump_process::build_export_list(export_list* result, char* library, module_ { if (strcmpi(item->second->short_name, library) == 0) { - pe_header* header = new pe_header(_pid, (void*)item->first, modules, _options); + pe_header* header = new pe_header(_ph, (void*)item->first, modules, _options); if (header->process_pe_header() && header->process_sections() && header->process_export_directory()) { // Load its exports - result->add_exports(header->get_exports()); + result->take_exports(*header->get_exports()); } // Cleanup @@ -534,23 +446,23 @@ void dump_process::dump_region(__int64 base) if( !_options->ImportRec || build_export_list() ) { module_list* modules = new module_list( _pid ); - pe_header* header = new pe_header( _pid, (void*) base, modules, _options ); + pe_header* header = new pe_header( _ph, (void*) base, modules, _options ); if( _options->ForceGenHeader || !header->process_pe_header() ) { if( _options->Verbose ) - printf( "Generating 32-bit PE header for module at %llX.\n", base ); + printf( "Generating 64-bit PE header for module at %llX.\n", base ); // Build the pe header as 32 and 64 bit since it could be either - header->build_pe_header(0x1000ffff, true ); - dump_header(header, base, _pid); + if (header->build_pe_header(0x1000, true)) + dump_header(header, base, _pid); delete header; if( _options->Verbose ) - printf( "Generating 64-bit PE header for module at %llX.\n", base ); - header = new pe_header( _pid, (void*) base, modules, _options ); - header->build_pe_header(0x1000ffff, false ); - dump_header(header, base, _pid); + printf( "Generating 32-bit PE header for module at %llX.\n", base ); + header = new pe_header( _ph, (void*) base, modules, _options ); + if (header->build_pe_header(0x1000, false)) + dump_header(header, base, _pid); } else { @@ -645,136 +557,49 @@ bool dump_process::monitor_close_dump_and_resume() return false; // not hooked } -void dump_process::dump_all() +void dump_process::dump_all(work_pool* pool) { // Walk through the pages while dumping all MZ files that do not match our good hash database. printf( "dumping process %s with pid 0x%x...\n", this->_process_name, this->_pid ); if( _ph != NULL ) { // First build the export list for this process - if( build_export_list() ) + if ((!_options->ImportRec && !_options->DumpChunks) || build_export_list(pool)) { // First build a list of the modules - module_list* modules = new module_list( _pid ); + module_list module_snapshot(_pid); + auto* modules = &module_snapshot; - // Set the max address of the target process - unsigned __int64 maxAddress = 0; - maxAddress = 0xffffffffffffffff; // Not a problem for 32bit targets - - // Walk the process heaps - __int64 address = 0; - - // First loop to build a list of executable heaps for later use in locating loose executable heaps not associated with any modules set executable_heaps; - if (_options->DumpChunks) - { - while (address < maxAddress) + const auto bases = scan_regions(executable_heaps); + std::mutex results; + module_work(pool, bases, [&](unsigned __int64 base) { + pe_header header(_ph, reinterpret_cast(base), modules, _options); + if (header.process_pe_header() && header.process_sections() && + header.somewhat_parsed() && header.process_import_directory()) { - MBI_BASIC_INFO mbi_info = get_mbi_info(address); - - // Check if this is a loose executable heap (this check causes it to fail.) - if (mbi_info.base > 0 && mbi_info.end > 0 && mbi_info.valid && mbi_info.executable) { - executable_heaps.insert(mbi_info.base); + std::lock_guard lock(results); + exclude_module_heaps(executable_heaps, base, header.get_virtual_size()); } - - if (mbi_info.end + 1 <= address) - break; - address = mbi_info.end + 1; - } - } - - // Load all the PE files in the process - address = 0; - while (address < maxAddress) - { - MBI_BASIC_INFO mbi_info = get_mbi_info(address); - - // Check if this is a loose executable heap - if ( mbi_info.base > 0 && mbi_info.end > 0 && mbi_info.valid ) - { - if (_options->Verbose) - fprintf(stdout, "INFO: Scanning from region from 0x%llX to 0x%llX for MZ headers.\n", mbi_info.base, mbi_info.end); - - - // This heap may have a PE file, check all page alignments for a "MZ". - unsigned __int64 base = mbi_info.base - (mbi_info.base % PAGE_SIZE); // shouldn't be required. - char output[2]; - SIZE_T out_read; - int count = 0; - while (base + 0x300 < mbi_info.end && count < 1000) // Skip the rest of the section if we have looped over 1000 pages. + const auto hash = header.get_hash(); + if (hash != 0 && !_db_clean->contains(hash)) { - if (ReadProcessMemory(_ph, (LPCVOID)((unsigned char*)base), output, 2, &out_read) && out_read == 2) + if (_options->ForceGenHeader) { - if (output[0] == 'M' && output[1] == 'Z') + printf("Dumping a module but ignoring existing PE Header for module at 0x%llX.\n", base); + for (bool win64 : {true, false}) { - // Bingo, possible MZ file - pe_header* header = new pe_header(_pid, (void*)base, modules, _options); - - // Use the existing PE header for the dumping - if (header->process_pe_header()) - { - if (header->process_sections() && header->somewhat_parsed() && header->process_import_directory()) - { - // Exclude all executable regions from this PE module - unsigned __int64 end_address = header->get_virtual_size() + base; - for (set::iterator it = executable_heaps.begin(); it != executable_heaps.end(); ) - { - if ( *it <= end_address && *it >= base) - { - // We've accounted for this executable heap, remove it from the loose heap list - it = executable_heaps.erase(it); - } - else - { - it++; - } - } - - // Check hash - unsigned __int64 hash = header->get_hash(); - if (hash != 0 && !_db_clean->contains(hash)) - { - if (_options->ForceGenHeader) - { - // Use the existing PE header only to get the hash, then generate a PE header for the dumping. - fprintf(stdout, "Dumping a module but ignoring existing PE Header for module at 0x%llX.\n", base); - pe_header* header_dump = new pe_header(_pid, (void*)base, modules, _options); - header_dump->build_pe_header(0x1000, true); // 64bit - dump_header(header_dump, base, _pid); - delete header_dump; - - header_dump = new pe_header(_pid, (void*)base, modules, _options); - header_dump->build_pe_header(0x1000, false); // 32bit - dump_header(header_dump, base, _pid); - delete header_dump; - } - else if (header->process_disk_image(&this->_export_list, this->_db_clean)) - { - write_dump(header, base, _pid); - } - } - } - } - - // Cleanup - delete header; + pe_header generated(_ph, reinterpret_cast(base), modules, _options); + if (generated.build_pe_header(0x1000, win64)) + dump_header(&generated, base, _pid); } } - else - { - - } - - base += PAGE_SIZE; - count++; + else if (header.process_disk_image(&_export_list, _db_clean)) + write_dump(&header, base, _pid); } } - - if (mbi_info.end + 1 <= address) - break; - address = mbi_info.end + 1; - } + }); if( _options->DumpChunks ) { @@ -783,6 +608,7 @@ void dump_process::dump_all() fprintf( stdout, "INFO: Looking at unattached executable heaps...\n" ); int count_new_header_hashes = 0; + std::vector chunks; for (set::iterator it=executable_heaps.begin(); it!=executable_heaps.end(); it++) { // Unattached executable page. First check hash of crc32 of first 2kb of memory since import reconstruction hashing is @@ -804,44 +630,33 @@ void dump_process::dump_all() break; // Too many new code chunks } - // Calculate the generic import reference hash as well - pe_header* header = new pe_header( _pid, (void*) *it, modules, _options ); - header->build_pe_header( 0x1000, true, 1 ); // 64bit, only build it with the 1 executable section for performance reasons - header->process_sections(); - - // Get the import attributes of this header - IMPORT_SUMMARY import_summary = header->get_imports_information(&this->_export_list); - - // Check hash - if( import_summary.HASH_GENERIC != 0 && !_db_clean->contains(import_summary.HASH_GENERIC) ) + chunks.push_back(*it); + } + } + module_work(pool, chunks, [&](unsigned __int64 base) { + pe_header header(_ph, reinterpret_cast(base), modules, _options); + if (!header.build_pe_header(0x1000, true, 1) || !header.process_sections()) + return; + const auto summary = header.get_imports_information(&_export_list); + if (summary.HASH_GENERIC && !_db_clean->contains(summary.HASH_GENERIC) && + header.somewhat_parsed() && summary.COUNT_UNIQUE_IMPORT_ADDRESSES >= 2) + { + printf("Dumping unattached executable code chunk from 0x%llX.\n", base); + for (bool win64 : {true, false}) { - if( _options->Verbose ) - fprintf( stdout, "INFO: Unattached executable heap at 0x%llX found with %i imports matched.\n", *it, import_summary.COUNT_UNIQUE_IMPORT_ADDRESSES ); - - if( header->somewhat_parsed() && import_summary.COUNT_UNIQUE_IMPORT_ADDRESSES >= 2 ) // Require at least 5 imports for dumping + pe_header generated(_ph, reinterpret_cast(base), modules, _options); + if (generated.build_pe_header(0x1000, win64)) { - fprintf( stdout, "Dumping unattached executable code chunk from 0x%llX.\n", *it ); - pe_header* header_dump = new pe_header( _pid, (void*) *it, modules, _options ); - header_dump->build_pe_header( 0x1000, true ); // 64bit - header_dump->set_name("codechunk"); - dump_header(header_dump, *it, _pid); - delete header_dump; - - header_dump = new pe_header( _pid, (void*) *it, modules, _options ); - header_dump->build_pe_header( 0x1000, false ); // 32bit - header_dump->set_name("codechunk"); - dump_header(header_dump, *it, _pid); - delete header_dump; + generated.set_name("codechunk"); + dump_header(&generated, base, _pid); } } - delete header; } - } + }); if( _options->Verbose ) fprintf( stdout, "INFO: Done looking at unattached executable heaps...\n" ); } - delete modules; } } } diff --git a/pd/dump_process.h b/pd/dump_process.h index 6c7ada1..05d573e 100644 --- a/pd/dump_process.h +++ b/pd/dump_process.h @@ -12,6 +12,7 @@ #include "hash.h" #include #include "terminate_monitor_hook.h" +#include "work_pool.h" #define PAGE_SIZE 0x1000 #define CODECHUNK_HEADER_HASH_SIZE 0x200 // First X bytes are CRC32'd of each loose code chunk. Only unique CRC32s are processed deeply. @@ -50,16 +51,17 @@ class dump_process MBI_BASIC_INFO get_mbi_info(unsigned __int64 address); bool write_dump(pe_header* header, unsigned __int64 base, DWORD pid); + std::vector scan_regions(set& executable_heaps); public: dump_process(DWORD pid, pe_hash_database* db, PD_OPTIONS* options, bool quieter); - void dump_all(); + void dump_all(work_pool* pool = NULL); void dump_region(__int64 base); void dump_header(pe_header* header, __int64 base, DWORD pid); DWORD get_pid() { return _pid; }; - bool build_export_list(); + bool build_export_list(work_pool* pool = NULL); bool build_export_list(export_list* result, char* library, module_list* modules); - int get_all_hashes(unordered_set* output_hashes, unordered_set* output_hashes_eps, unordered_set* output_hashes_ep_shorts); + int get_all_hashes(unordered_set* output_hashes, unordered_set* output_hashes_eps, unordered_set* output_hashes_ep_shorts, work_pool* pool = NULL); unsigned __int64 hash_codechunk_header(__int64 base); bool is64(); bool get_process_name(char* process_name, SIZE_T byte_length); diff --git a/pd/export_list.cpp b/pd/export_list.cpp index 8dcacf0..4ce752a 100644 --- a/pd/export_list.cpp +++ b/pd/export_list.cpp @@ -62,41 +62,25 @@ export_list::export_list() bool export_list::contains(unsigned __int64 address) { - // Look up a 64-bit value - if ( address <= UINT_MAX ) - return contains((unsigned __int32)address); - - if (address > _max64 || address < _min64 || (address & ~_bits64) > 0) - { - // We know there is no match by this quick filtering. This improves performance hugely. - return false; - } - - // Lookup the address - unordered_set::const_iterator got = _addresses.find(address); - if (got != _addresses.end()) - { - return true; - } - return false; + return lookup(address) != NULL; } bool export_list::contains(unsigned __int32 address) { - // Look up a 32-bit value - if (address > _max32 || address < _min32 || (address & ~_bits32) > 0) - { - // We know there is no match by this quick filtering. This improves performance hugely. - return false; - } + return lookup(address) != NULL; +} - // Lookup the address - unordered_set::const_iterator got = _addresses.find(address); - if (got != _addresses.end()) +const export_entry* export_list::lookup(unsigned __int64 address) const +{ + if (address <= UINT_MAX) { - return true; + if (address > _max32 || address < _min32 || (address & ~_bits32) != 0) + return NULL; } - return false; + else if (address > _max64 || address < _min64 || (address & ~_bits64) != 0) + return NULL; + const auto found = _address_to_exports.find(address); + return found == _address_to_exports.end() ? NULL : found->second; } unsigned __int64 export_list::find_export(char* library, char* name, bool is64) @@ -137,34 +121,52 @@ void export_list::add_export(unsigned __int64 address, export_entry* entry) { // Register this export address for quick lookups later if (_address_to_exports.count(address) == 0) + _add_owned(address, std::unique_ptr(new export_entry(entry))); +} + +void export_list::_add_owned(unsigned __int64 address, std::unique_ptr entry) +{ + if (_address_to_exports.emplace(address, entry.get()).second) { - _address_to_exports.insert(std::pair(address, new export_entry(entry))); + entry.release(); + _update_filters(address); + } +} - if (_addresses.count(address) == 0) - { - _addresses.insert(address); +void export_list::_update_filters(unsigned __int64 address) +{ + if (address > UINT_MAX) + { + _max64 = (std::max)(_max64, address); + _min64 = (std::min)(_min64, address); + _bits64 |= address; + } + else + { + _max32 = (std::max)(_max32, static_cast(address)); + _min32 = (std::min)(_min32, static_cast(address)); + _bits32 |= static_cast(address); + } +} - // Update our quick-lookup values - if ( address > UINT_MAX ) - { - // 64bit value - if (address > _max64) - _max64 = address; - if (address < _min64) - _min64 = address; - _bits64 = _bits64 | address; - } - else - { - // 32bit value - if (address > _max32) - _max32 = address; - if (address < _min32) - _min32 = address; - _bits32 = _bits32 | address; - } - } +void export_list::take_exports(export_list& other) +{ + if (&other == this) + return; + for (auto entry = other._address_to_exports.begin(); entry != other._address_to_exports.end();) + { + if (_address_to_exports.emplace(entry->first, entry->second).second) + _update_filters(entry->first); + else + delete entry->second; + entry = other._address_to_exports.erase(entry); } + other._min64 = _UI64_MAX; + other._max64 = UINT_MAX; + other._min32 = UINT_MAX; + other._max32 = 0; + other._bits32 = 0; + other._bits64 = 0; } bool export_list::add_exports(export_list* other) @@ -200,7 +202,7 @@ namespace } } -bool export_list::add_exports(unsigned char* image, SIZE_T image_size, unsigned __int64 image_base, IMAGE_EXPORT_DIRECTORY* export_directory, bool is64) +bool export_list::add_exports(unsigned char* image, SIZE_T image_size, unsigned __int64 image_base, IMAGE_EXPORT_DIRECTORY* export_directory, bool is64, DWORD directory_size) { if (image == NULL || export_directory == NULL || !test_read(image, image_size, reinterpret_cast(export_directory), sizeof(*export_directory))) @@ -208,6 +210,9 @@ bool export_list::add_exports(unsigned char* image, SIZE_T image_size, unsigned IMAGE_EXPORT_DIRECTORY directory; memcpy(&directory, export_directory, sizeof(directory)); + const SIZE_T directory_rva = reinterpret_cast(export_directory) - image; + if (directory_size != 0 && (directory_size < sizeof(directory) || !range_fits(image_size, directory_rva, directory_size))) + return false; if (!export_array_fits(image_size, directory.AddressOfFunctions, directory.NumberOfFunctions, sizeof(DWORD)) || !export_array_fits(image_size, directory.AddressOfNames, directory.NumberOfNames, sizeof(DWORD)) || !export_array_fits(image_size, directory.AddressOfNameOrdinals, directory.NumberOfNames, sizeof(WORD))) @@ -225,6 +230,24 @@ bool export_list::add_exports(unsigned char* image, SIZE_T image_size, unsigned return false; } + export_list parsed; + const auto add_function = [&](DWORD index, const char* name) { + DWORD rva; + memcpy(&rva, image + directory.AddressOfFunctions + static_cast(index) * sizeof(DWORD), sizeof(rva)); + if (rva == 0) + return true; + // EAT entries inside the export directory name a forwarder, not executable/data addresses. + if (directory_size != 0 && rva >= directory_rva && rva - directory_rva < directory_size) + return true; + // Loaded export tables can redirect functions outside their own module (for example WOW64 USER32). + if (image_base > _UI64_MAX - rva || index > MAXWORD || directory.Base > MAXWORD - index) + return false; + const auto address = image_base + rva; + if (!parsed.contains(address)) + parsed._add_owned(address, std::unique_ptr(new export_entry(library_name.c_str(), + name, static_cast(directory.Base + index), rva, address, is64))); + return true; + }; for (DWORD i = 0; i < directory.NumberOfNames; ++i) { WORD ordinal_relative; @@ -238,19 +261,14 @@ bool export_list::add_exports(unsigned char* image, SIZE_T image_size, unsigned if (!export_string(image, image_size, name_offset, name)) return false; - DWORD rva; - memcpy(&rva, image + directory.AddressOfFunctions + static_cast(ordinal_relative) * sizeof(DWORD), sizeof(rva)); - if (image_base > _UI64_MAX - rva) + if (!add_function(ordinal_relative, name.c_str())) return false; - - // Keep the existing filter used to avoid false matches during import repair. - if (rva % 0x1000 != 0) - { - const unsigned __int64 address = image_base + rva; - export_entry entry(library_name.c_str(), name.c_str(), static_cast(directory.Base + ordinal_relative), rva, address, is64); - add_export(address, &entry); - } } + // Named entries win aliases at the same address; all remaining EAT entries are ordinal imports. + for (DWORD i = 0; i < directory.NumberOfFunctions; ++i) + if (!add_function(i, NULL)) + return false; + take_exports(parsed); return true; } diff --git a/pd/export_list.h b/pd/export_list.h index 5560d03..22ca27b 100644 --- a/pd/export_list.h +++ b/pd/export_list.h @@ -7,6 +7,7 @@ #include #include #include +#include #include "utils.h" #include "limits.h" @@ -39,14 +40,16 @@ class export_list unsigned __int64 _bits64; unordered_map _address_to_exports; // List of export addresses in this export list - unordered_set _addresses; // List of export addresses + void _update_filters(unsigned __int64 address); + void _add_owned(unsigned __int64 address, std::unique_ptr entry); public: export_list(); - bool add_exports(unsigned char* image, SIZE_T image_size, unsigned __int64 image_base, IMAGE_EXPORT_DIRECTORY* header_export_directory, bool is64); + bool add_exports(unsigned char* image, SIZE_T image_size, unsigned __int64 image_base, IMAGE_EXPORT_DIRECTORY* header_export_directory, bool is64, DWORD directory_size = 0); bool add_exports(export_list* other); + void take_exports(export_list& other); void add_export(unsigned __int64 address, export_entry* entry); // Find export addresses in a process @@ -56,6 +59,7 @@ class export_list bool contains(unsigned __int64 address); bool contains(unsigned __int32 address); export_entry find(unsigned __int64 address); + const export_entry* lookup(unsigned __int64 address) const; unsigned __int64 get_min64() { return _min64; }; unsigned __int64 get_max64() { return _max64; }; unsigned __int32 get_min32() { return _min32; }; diff --git a/pd/module_list.cpp b/pd/module_list.cpp index 332e49d..00e2367 100644 --- a/pd/module_list.cpp +++ b/pd/module_list.cpp @@ -11,20 +11,18 @@ module_list::module_list( DWORD pid ) { #if defined(_WIN64) // List modules on a 64 bit machine. A 64 bit machine is assumed to be Windows Vista+ - HMODULE hMods[2048]; - DWORD cbNeeded; - unsigned int i; + std::vector hMods; HANDLE ph = OpenProcess( PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid ); if( ph != NULL ) { - if( EnumProcessModulesEx(ph, hMods, sizeof(hMods), &cbNeeded, LIST_MODULES_ALL)) + if (module_snapshot::read(ph, hMods)) { - for ( i = 0; i < (cbNeeded / sizeof(HMODULE)); i++ ) + for (HMODULE handle : hMods) { // Query the module basic information MODULEINFO info; - if( GetModuleInformation( ph, hMods[i], &info, sizeof(MODULEINFO) ) ) + if( GetModuleInformation( ph, handle, &info, sizeof(MODULEINFO) ) ) { // Check if this base address is already occupied unordered_map::const_iterator item = _modules.find( (unsigned __int64) info.lpBaseOfDll ); @@ -32,7 +30,7 @@ module_list::module_list( DWORD pid ) if( item == _modules.end() ) { // Add this module - _modules[(unsigned __int64) info.lpBaseOfDll ] = new module( ph, hMods[i], info ); + _modules[(unsigned __int64) info.lpBaseOfDll ] = new module( ph, handle, info ); } } else @@ -42,6 +40,9 @@ module_list::module_list( DWORD pid ) } } + else + PrintLastError(L"module_list EnumProcessModulesEx"); + CloseHandle( ph ); } else diff --git a/pd/module_list.h b/pd/module_list.h index e4897b6..61f08aa 100644 --- a/pd/module_list.h +++ b/pd/module_list.h @@ -10,12 +10,43 @@ #include #include #include +#include using namespace std::tr1; extern bool global_flag_verbose; +namespace module_snapshot +{ + using enumerator = BOOL (WINAPI *)(HANDLE, HMODULE*, DWORD, LPDWORD, DWORD); + + inline bool read(HANDLE process, std::vector& modules, enumerator enumerate = EnumProcessModulesEx) + { + modules.resize(2048); + for (int attempt = 0; attempt < 8; ++attempt) + { + DWORD needed = 0; + if (!enumerate(process, modules.data(), static_cast(modules.size() * sizeof(HMODULE)), &needed, LIST_MODULES_ALL)) + { + modules.clear(); + return false; + } + if (needed % sizeof(HMODULE) != 0 || needed / sizeof(HMODULE) > 1024 * 1024) + break; + if (needed / sizeof(HMODULE) <= modules.size()) + { + modules.resize(needed / sizeof(HMODULE)); + return true; + } + modules.resize(needed / sizeof(HMODULE)); + } + modules.clear(); + SetLastError(ERROR_BAD_LENGTH); + return false; + } +} + namespace module_names { inline bool to_ansi(const wchar_t* source, SIZE_T source_count, std::string& result, const char* description) diff --git a/pd/module_work.h b/pd/module_work.h new file mode 100644 index 0000000..c4aa826 --- /dev/null +++ b/pd/module_work.h @@ -0,0 +1,20 @@ +#pragma once + +#include "work_pool.h" +#include + +template +void module_work(work_pool* pool, const std::vector& items, F action) +{ + if (pool) + pool->parallel_for(items, action); + else + for (const auto& item : items) action(item); +} + +inline void exclude_module_heaps(std::set& heaps, unsigned __int64 base, unsigned __int64 size) +{ + if (size > _UI64_MAX - base) + return; + heaps.erase(heaps.lower_bound(base), heaps.lower_bound(base + size)); +} diff --git a/pd/pd.cpp b/pd/pd.cpp index e7c0db3..97576d0 100644 --- a/pd/pd.cpp +++ b/pd/pd.cpp @@ -13,6 +13,8 @@ #include #include "pd.h" #include "close_watcher.h" +#include "system_work.h" +#include "version.h" #define NMD_ASSEMBLY_IMPLEMENTATION #include "nmd_assembly.h" @@ -130,244 +132,29 @@ void add_process_hashes( DWORD pid, pe_hash_database* db, PD_OPTIONS* options ) db->add_hashes_eps( new_hashes_eps, new_hashes_ep_shorts); } -void add_process_hashes_worker(Queue* work_queue, pe_hash_database* db, PD_OPTIONS* options) -{ - // Gather hashes from the work queue process list - while (!work_queue->empty()) - { - // Process the hashes for this process - PROCESSENTRY32 entry; - if (work_queue->pop(entry)) - { - add_process_hashes(entry.th32ProcessID, db, options); - } - } -} - - - void add_system_hashes( pe_hash_database* db, PD_OPTIONS* options ) { - // Add clean hashes from all processes on the system right now - - // Build a list of the hashes from all processes - options->ImportRec = false; // Force no import reconstruciton - - // Build the queue of work from the currently running processes - Queue work_queue; - int total_work_count = 0; - - PROCESSENTRY32 entry; - entry.dwSize = sizeof(PROCESSENTRY32); - HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, NULL); - - if( snapshot != INVALID_HANDLE_VALUE ) - { - if (Process32First(snapshot, &entry) == TRUE) - { - while (Process32Next(snapshot, &entry) == TRUE) - { - printf("...adding process to work queue: pid 0x%x,%S\n", entry.th32ProcessID, entry.szExeFile); - work_queue.push(entry); - total_work_count++; - } - } - CloseHandle(snapshot); - } - - // Create threads processing the work queue - thread** threads = new thread*[options->NumberOfThreads]; - for (int i = 0; i < options->NumberOfThreads; i++) - { - threads[i] = new thread(add_process_hashes_worker, &work_queue, db, options); - } - - // Wait for queue depletion - int count = 0; - bool still_working = false; - int running_count = 1; - while (!work_queue.empty() || running_count > 0) - { - // Check threads state - running_count = 0; - for (int i = 0; i < options->NumberOfThreads; i++) - { - if( WaitForSingleObject(threads[i]->native_handle(), 1) == WAIT_TIMEOUT ) - running_count++; - } - - if (count % 10 == 0) - { - // Print the status - int waiting_count = work_queue.count(); - printf("Hash Queue -> Waiting: %i\tRunning: %i\tComplete: %i\n", waiting_count, running_count, total_work_count - (waiting_count + running_count)); - } - - // Wait - Sleep(50); - count++; - } - - // Wait for thread completions - for (int i = 0; i < options->NumberOfThreads; i++) - { - threads[i]->join(); // blocks until each thread is finished - } - - // Cleanup - printf("...cleaning up system memory hashes factory\n"); - for (int i = 0; i < options->NumberOfThreads; i++) - { - delete threads[i]; - } - delete []threads; -} - - -void dump_process_worker(Queue* work_queue, pe_hash_database* db, PD_OPTIONS* options) -{ - // Gather hashes from the work queue process list - unordered_set new_hashes; - while (!work_queue->empty()) - { - // Process the hashes for this process - PROCESSENTRY32 entry; - if (work_queue->pop(entry)) - { - // Process this process - - // Dump - dump_process* dumper = new dump_process(entry.th32ProcessID, db, options, true); - dumper->dump_all(); - - // Exclude these hashes from the next dumps - dumper->get_all_hashes(&new_hashes, NULL, NULL); - db->add_hashes(new_hashes); - new_hashes.clear(); - - delete dumper; - } - } + work_pool pool(options->NumberOfThreads); + system_work(pool, false, system_processes, [=](DWORD pid, work_pool& workers) { + unordered_set hashes, full, prefixes; + dump_process dumper(pid, db, options, true); + dumper.get_all_hashes(&hashes, &full, &prefixes, &workers); + db->add_hashes(hashes); + db->add_hashes_eps(full, prefixes); + }); } void dump_system(pe_hash_database* db, PD_OPTIONS* options) { - // Dump modules from all running processes - - // Build the queue of work from the currently running processes - Queue work_queue; - unordered_set dumping_pids; - - PROCESSENTRY32 entry; - entry.dwSize = sizeof(PROCESSENTRY32); - HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, NULL); - int total_work_count = 0; - - if (snapshot != INVALID_HANDLE_VALUE) - { - if (Process32First(snapshot, &entry) == TRUE) - { - while (Process32Next(snapshot, &entry) == TRUE) - { - printf("...adding process to work queue: pid 0x%x,%S\n", entry.th32ProcessID, entry.szExeFile); - work_queue.push(entry); - dumping_pids.insert(entry.th32ProcessID); - total_work_count++; - } - } - CloseHandle(snapshot); - } - - // Create threads processing the work queue - thread** threads = new thread*[options->NumberOfThreads]; - for (int i = 0; i < options->NumberOfThreads; i++) - { - threads[i] = new thread(dump_process_worker, &work_queue, db, options); - } - - // Wait for queue depletion - int count = 0; - bool still_working = false; - int running_count = 1; - bool added_new_processes = false; - - while (!work_queue.empty() || running_count > 0 || !added_new_processes) - { - // Check threads state - running_count = 0; - for (int i = 0; i < options->NumberOfThreads; i++) - { - if (WaitForSingleObject(threads[i]->native_handle(), 1) == WAIT_TIMEOUT) - running_count++; - } - - // Add any newly started processes at the very end - if (!added_new_processes && work_queue.empty() && running_count ) - { - printf("...adding new processes since we started this job\n"); - HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, NULL); - if (snapshot != INVALID_HANDLE_VALUE) - { - if (Process32First(snapshot, &entry) == TRUE) - { - while (Process32Next(snapshot, &entry) == TRUE) - { - if (dumping_pids.count(entry.th32ProcessID) == 0) - { - printf("...adding new process to work queue: pid 0x%x,%S\n", entry.th32ProcessID, entry.szExeFile); - work_queue.push(entry); - dumping_pids.insert(entry.th32ProcessID); - total_work_count++; - - // Add a new worker thread if needed - if (running_count < options->NumberOfThreads) - { - // Add a new worker thread - for (int i = 0; i < options->NumberOfThreads; i++) - { - if (WaitForSingleObject(threads[i]->native_handle(), 1) != WAIT_TIMEOUT) - { - threads[i]->join(); // blocks until each thread is finished - delete threads[i]; - threads[i] = new thread(dump_process_worker, &work_queue, db, options); - break; - } - } - } - } - } - } - CloseHandle(snapshot); - } - added_new_processes = true; - } - - if (count % 10 == 0) - { - // Print the status - int waiting_count = work_queue.count(); - printf("Dump Queue -> Waiting: %i\tRunning: %i\tComplete: %i\n", waiting_count, running_count, total_work_count - (waiting_count + running_count)); - } - - // Wait - Sleep(50); - count++; - } - - // Wait for thread completions - for (int i = 0; i < options->NumberOfThreads; i++) - { - threads[i]->join(); // blocks until each thread is finished - } - - // Cleanup - printf("...cleaning up system dump factory\n"); - for (int i = 0; i < options->NumberOfThreads; i++) - { - delete threads[i]; - } - delete []threads; + work_pool pool(options->NumberOfThreads); + system_work(pool, true, system_processes, [=](DWORD pid, work_pool& workers) { + dump_process dumper(pid, db, options, true); + dumper.dump_all(&workers); + unordered_set hashes; + dumper.get_all_hashes(&hashes, NULL, NULL, &workers); + db->add_hashes(hashes); + }); } @@ -377,11 +164,9 @@ void dump_system(pe_hash_database* db, PD_OPTIONS* options) bool global_flag_verbose = false; -int _tmain(int argc, _TCHAR* argv[]) +int _tmain(int argc, _TCHAR* argv[]) try { - get_privileges( GetCurrentProcess() ); - // Process the flags WCHAR* filter = NULL; char* processNameFilter = NULL; @@ -445,6 +230,8 @@ int _tmain(int argc, _TCHAR* argv[]) flagRecursion = false; else if( lstrcmp(argv[i],L"-ni") == 0 ) options.ImportRec = false; + else if (lstrcmp(argv[i], L"-reexec") == 0) + options.Reexecution = true; else if( lstrcmp(argv[i],L"-nc") == 0 ) options.DumpChunks = false; else if (lstrcmp(argv[i], L"-nep") == 0) @@ -796,7 +583,7 @@ int _tmain(int argc, _TCHAR* argv[]) if( flagHeader ) { - printf("Process Dump v2.2 (dev)\n"); + printf("Process Dump v%s\n", PD_VERSION_STRING); printf(" Copyright © 2017, Geoff McDonald\n"); printf(" http://www.split-code.com/\n"); printf(" https://github.com/glmcdona/Process-Dump\n\n"); @@ -804,47 +591,95 @@ int _tmain(int argc, _TCHAR* argv[]) if( flagHelp ) { - // Print help page - printf("Process Dump (pd.exe) is a tool used to dump both 32 and 64 bit executable modules back to disk from memory within a process address space. This tool is able to find and dump hidden modules as well as loose executable code chunks, and it uses a clean hash database to exclude dumping of known clean files. This tool uses an aggressive import reconstruction approach that links all DWORD/QWORDs that point to an export in the process to the corresponding export function. Process dump can be used to dump all unknown code from memory ('-system' flag), dump specific processes, or run in a monitoring mode that dumps all processes just before they terminate.\n\n"); - printf("Before first usage of this tool, when on the clean workstation the clean exclusing hash database can be generated by either:\n"); - printf("\tpd -db gen\n"); - printf("\tpd -db genquick\n\n"); - printf("Example Usage:\n"); - printf("\tpd -system\n"); - printf("\tpd -pid 419\n"); - printf("\tpd -pid 0x1a3\n"); - printf("\tpd -pid 0x1a3 -a 0x401000 -o c:\\dump\\ -c c:\\dump\\test\\clean.db\n"); - printf("\tpd -p chrome.exe\n"); - printf("\tpd -p \"(?i).*chrome.*\"\n"); - printf("\tpd -closemon\n\n"); - - printf("Options:\n"); - printf("\t-system\t\tDumps all modules not matching the clean hash database\n\t\t\tfrom all accessible processes into the working\n\t\t\tdirectory.\n\n"); - printf("\t-pid \tDumps all modules not matching the clean hash database\n\t\t\tfrom the specified pid into the current working\n\t\t\tdirectory. Use a '0x' prefix to specify a hex PID.\n\n"); - printf("\t-closemon\t\tRuns in monitor mode. When any processes are terminating\n\t\t\tprocess dump will first dump the process.\n\n"); - printf("\t-p \tDumps all modules not matching the clean hash database\n\t\t\tfrom the process name found to match the filter into\n\t\t\tspecified pid into the current working directory.\n\n"); - printf("\t-g\t\tForces generation of PE headers from scratch, ignoring existing headers.\n\n"); - printf("\t-o \tSets the default output root folder for dumped components.\n\n"); - //printf("\t-log\t\tRuns in log generation mode. No files are dumped, logfiles are generated for analysis. Specify a network path as the output and batch run to generate a report of multiple workstations.\n\n"); - //printf("\t-netcollect\t\tRuns in network sample collection mode. Specify a network path as the output, dumped files will be organized by hash.\n\n"); - printf("\t-v\t\tVerbose.\n\n"); - printf("\t-nh\t\tNo header is printed in the output.\n\n"); - printf("\t-nr\t\tDisable recursion on hash database directory add or\n\t\t\tremove commands.\n\n"); - printf("\t-ni\t\tDisable import reconstruction.\n\n"); - printf("\t-nc\t\tDisable dumping of loose code regions.\n\n"); - printf("\t-nt\t\tDisable multithreading.\n\n"); - printf("\t-nep\t\tDisable entry point hashing.\n\n"); - printf("\t-eprec\t\tForce the entry point to be reconstructed, even if a valid one appears to exist.\n\n"); - printf("\t-t \t\tSets the number of threads to use (default 16).\n\n"); - printf("\t-cdb \t\tFull filepath to the clean hash database to use for this run.\n\n"); - printf("\t-edb \t\tFull filepath to the entrypoint hash database to use for this run.\n\n"); - printf("\t-esdb \t\tFull filepath to the entrypoint short hash database to use for this run.\n\n"); - printf("\t-db gen\t\tAutomatically processes a few common folders as well as\n\t\t\tall the currently running processes and adds the found\n\t\t\tmodule hashes to the clean hash database. It will add\n\t\t\tall files recursively in: \n\t\t\t\t%%WINDIR%% \n\t\t\t\t%%HOMEPATH%% \n\t\t\t\tC:\\Program Files\\ \n\t\t\t\tC:\\Program Files (x86)\\ \n\t\t\tAs well as all modules in all running processes \n\n"); - printf("\t-db genquick\tAdds the hashes from all modules in all processes to\n\t\t\tthe clean hash database. Run this on a clean system.\n\n"); - printf("\t-db add \tAdds all the files in the specified directory\n\t\t\trecursively to the clean hash database. \n\n"); - printf("\t-db rem \tRemoves all the files in the specified directory\n\t\t\trecursively from the clean hash database. \n\n"); - printf("\t-db clean\tClears the clean hash database.\n\n"); - printf("\t-db ignore\tIgnores the clean hash database when dumping a process\n\t\t\tthis time. All modules will be dumped even if a match\n\t\t\tis found.\n\n"); + fputs(R"help(Usage: pd.exe [options] + +Reconstruct PE32/PE64 modules, hidden modules and loose executable chunks from +process memory for analysis. Dumps are not guaranteed runnable or signed. +Build a baseline database on a known-clean system before collecting unknown code. +Use pd64.exe on 64-bit Windows, pd32.exe on 32-bit Windows; local builds use pd.exe. +Administrator rights improve access but do not bypass protected processes. + +Commands (choose one; run database maintenance separately from dumping): + -system Dump unknown modules/chunks from all accessible processes. + Workers share module jobs, including the last process; + one final scan checks for newly discovered PIDs. + -pid Dump one process. PID may be decimal or 0x-prefixed hex. + -p Match the entire process name, case-sensitive (ECMAScript). + Quote the expression; multiple matches prompt for approval. + -closemon Hook termination of accessible processes and dump before + they exit. Intrusive; use only in a controlled environment. + Press Ctrl+C to stop and clean up hooks. + -db gen Add live hashes and PE files recursively from %WINDIR%, + %USERPROFILE%, C:\Program Files, C:\Program Files (x86). + -db genquick Add live module/chunk and entrypoint hashes only. + -db add Add PE file hashes recursively from a directory. + -db remove Remove matching clean-module hashes from a directory; + -db rem alias of remove. Entrypoint signatures are retained. + -db clean Clear and save ALL THREE databases, including entrypoints. + +Dump/reconstruction options: + -a
Dump at a decimal or 0x-prefixed base address; requires -pid. + Reuse a valid PE header; absent headers or -g generate + both PE32 and PE64 analysis images. Clean filtering applies. + -o Output directory (default: current working directory). + Use an existing ordinary directory you control. Existing + files are never overwritten; reparse paths are rejected. + -g Force generated PE32 and PE64 headers instead of originals. + -ni Disable aggressive import reconstruction (default: on). + -nc Disable loose-code dumping/hashing (default: on). + -nep Disable entrypoint recovery and live entrypoint hashing. + File database add/gen still collect entrypoint signatures. + -eprec Force recovery even for an existing executable entrypoint. + Default: preserve valid entrypoints. Full opcode matches + are required; x64 runtime-function metadata is a fallback. + No weak-prefix-only guesses; failure retains the old RVA. + -nep takes precedence. Recovery remains heuristic. + -reexec Experimental fresh-launch preparation (default: off). + Reset the GS cookie and writable/non-executable zero-fill + tails; restrict new imports to file-backed writable data + outside delay-IAT sections. Assumes original section layout + and discards captured state. Does not run the dump, restore + a checkpoint, or disable ASLR/DEP/CFG. + +Database options: + -db ignore Ignore ALL THREE loaded databases for this dump invocation, + including entrypoint signatures; files are not cleared. + In-run cross-process deduplication can still skip repeats. + -nr Disable recursion for -db add/remove/rem (not -db gen). + -cdb Clean-module database (alias: -c ). + -edb Full entrypoint-opcode database. + -esdb Short entrypoint-prefix database. + Defaults beside pd.exe: clean.hashes, entrypoints.hashes, + shortentrypoints.hashes. Maintenance updates these files. + Hash matches are exclusions, not a trust verdict. + +Workers and display: + -t Worker count, at least 1 (default: 16). Applies to -system, + live/file database generation and close-monitor workers. + Direct -pid/-p dumps remain synchronous. + -nt Use one worker (same as -t 1). Last -t/-nt wins. + Close monitoring still has a separate monitoring thread. + -v Verbose diagnostics (default: off). + -nh Suppress the version/copyright banner, not other output. + --help Print help and exit; aliases: -help, -h, --h. + No arguments also prints help. + +Examples (create the output directory first): + pd.exe -db genquick -t 4 + pd.exe -system -t 4 -o C:\dumps + pd.exe -pid 419 + pd.exe -pid 0x1a3 -a 0x401000 -o C:\dumps -c C:\baseline\clean.hashes + pd.exe -p "chrome[.]exe" + pd.exe -p ".*chrome.*" -ni -nc + pd.exe -pid 419 -reexec -o C:\dumps + pd.exe -db add C:\baseline -nr -t 4 + pd.exe -closemon -o C:\dumps + +Images/reconstructions above 256 MiB are rejected. Captures are live, non-atomic, +and may contain sensitive data; keep output private. See README.md for benchmark +commands, recovery accuracy, reexecution caveats and output safety details. +)help", stdout); + return 0; } // Sanity check on flags @@ -902,7 +737,7 @@ int _tmain(int argc, _TCHAR* argv[]) printf("Adding all files in folder '%s' to clean hash and entrypoint database...\n", add_directory); int count_before = db->count(); - db->add_folder(add_directory, L"*", flagRecursion); + db->add_folder(add_directory, L"*", flagRecursion, options.NumberOfThreads); printf("Added %i new hashes to the database. It now has %i hashes.\n", db->count() - count_before, db->count()); db->save(); }else if( flagDB_remove ) @@ -931,25 +766,25 @@ int _tmain(int argc, _TCHAR* argv[]) // Add a bunch of folders to the database count_before = db->count(); printf("Adding files in %%WINDIR%% to clean hash database...\n"); - db->add_folder("%WINDIR%", L"*", true); + db->add_folder("%WINDIR%", L"*", true, options.NumberOfThreads); printf("...added %i new hashes from %%WINDIR%%.\n", db->count() - count_before); db->save(); count_before = db->count(); printf("Adding files in %%USERPROFILE%% to clean hash database...\n"); - db->add_folder("%USERPROFILE%", L"*", true); + db->add_folder("%USERPROFILE%", L"*", true, options.NumberOfThreads); printf("...added %i new hashes from %%USERPROFILE%%.\n", db->count() - count_before); db->save(); count_before = db->count(); printf("Adding files in 'C:\\Program Files\\' to clean hash database...\n"); - db->add_folder("C:\\Program Files\\", L"*", true); + db->add_folder("C:\\Program Files\\", L"*", true, options.NumberOfThreads); printf("...added %i new hashes from 'C:\\Program Files\\'.\n", db->count() - count_before); db->save(); count_before = db->count(); printf("Adding files in C:\\Program Files (x86)\\ to clean hash database...\n"); - db->add_folder("C:\\Program Files (x86)\\", L"*", true); + db->add_folder("C:\\Program Files (x86)\\", L"*", true, options.NumberOfThreads); printf("...added %i new hashes from 'C:\\Program Files (x86)\\'.\n", db->count() - count_before); db->save(); @@ -1083,5 +918,8 @@ int _tmain(int argc, _TCHAR* argv[]) return 0; } - - +catch (const std::exception& error) +{ + fprintf(stderr, "ERROR: Process Dump failed: %s.\n", error.what()); + return 1; +} diff --git a/pd/pd.rc b/pd/pd.rc index 4c48f40..51accac 100644 --- a/pd/pd.rc +++ b/pd/pd.rc @@ -1,6 +1,7 @@ // Microsoft Visual C++ generated resource script. // #include "resource.h" +#include "version.h" #define APSTUDIO_READONLY_SYMBOLS ///////////////////////////////////////////////////////////////////////////// @@ -51,8 +52,8 @@ END // VS_VERSION_INFO VERSIONINFO - FILEVERSION 2,2,0,0 - PRODUCTVERSION 2,2,0,0 + FILEVERSION PD_VERSION_RESOURCE + PRODUCTVERSION PD_VERSION_RESOURCE FILEFLAGSMASK 0x3fL #ifdef _DEBUG FILEFLAGS 0x1L @@ -69,12 +70,12 @@ BEGIN BEGIN VALUE "CompanyName", "Geoff McDonald" VALUE "FileDescription", "Process Dump. https://github.com/glmcdona/Process-Dump" - VALUE "FileVersion", "2.2.0.0" + VALUE "FileVersion", PD_VERSION_STRING ".0" VALUE "InternalName", "pd.exe" VALUE "LegalCopyright", "Copyright (C) 2022" VALUE "OriginalFilename", "pd.exe" VALUE "ProductName", "Process Dump" - VALUE "ProductVersion", "2.2.0.0" + VALUE "ProductVersion", PD_VERSION_STRING ".0" END END BLOCK "VarFileInfo" @@ -97,4 +98,3 @@ END ///////////////////////////////////////////////////////////////////////////// #endif // not APSTUDIO_INVOKED - diff --git a/pd/pd.vcxproj b/pd/pd.vcxproj index 1201d21..e5dcd51 100644 --- a/pd/pd.vcxproj +++ b/pd/pd.vcxproj @@ -248,7 +248,7 @@ Guard - Psapi.lib;Advapi32.lib;%(AdditionalDependencies) + Psapi.lib;Advapi32.lib;ntdll.lib;%(AdditionalDependencies) true true Guard @@ -295,6 +295,7 @@ + @@ -305,6 +306,9 @@ + + + diff --git a/pd/pd.vcxproj.filters b/pd/pd.vcxproj.filters index 97c9427..9e0f468 100644 --- a/pd/pd.vcxproj.filters +++ b/pd/pd.vcxproj.filters @@ -116,12 +116,24 @@ Header Files + + Header Files + + + Header Files + + + Header Files + Header Files Header Files + + Header Files + diff --git a/pd/pe_hash_database.cpp b/pd/pe_hash_database.cpp index 9ed8502..9407041 100644 --- a/pd/pe_hash_database.cpp +++ b/pd/pe_hash_database.cpp @@ -160,15 +160,35 @@ bool pe_hash_database::clear_database() _clean_hashes.clear(); _ep_hashes.clear(); _epshort_hashes.clear(); + _entrypoint_snapshot.reset(); LeaveCriticalSection( &_lock ); return true; } -bool pe_hash_database::add_hashes(unordered_set hashes) +std::shared_ptr pe_hash_database::snapshot_entrypoints() +{ + EnterCriticalSection(&_lock); + std::unique_ptr release(&_lock, LeaveCriticalSection); + if (!_entrypoint_snapshot) + { + auto result = std::make_shared(); + result->full.insert(_ep_hashes.begin(), _ep_hashes.end()); + result->short_hashes.insert(_epshort_hashes.begin(), _epshort_hashes.end()); + for (auto hash : result->short_hashes) + { + result->minimum = (std::min)(result->minimum, hash); + result->maximum = (std::max)(result->maximum, hash); + } + _entrypoint_snapshot = std::move(result); + } + return _entrypoint_snapshot; +} + +bool pe_hash_database::add_hashes(const unordered_set& hashes) { EnterCriticalSection( &_lock ); - for (unordered_set::iterator it=hashes.begin(); it!=hashes.end(); it++) + for (auto it=hashes.begin(); it!=hashes.end(); it++) { if( *it != 0 && _clean_hashes.count( *it ) == 0 ) { @@ -182,23 +202,25 @@ bool pe_hash_database::add_hashes(unordered_set hashes) } -bool pe_hash_database::add_hashes_eps(unordered_set hashes, unordered_set hashes_short) +bool pe_hash_database::add_hashes_eps(const unordered_set& hashes, const unordered_set& hashes_short) { EnterCriticalSection(&_lock); - for (unordered_set::iterator it = hashes.begin(); it != hashes.end(); it++) + for (auto it = hashes.begin(); it != hashes.end(); it++) { if (*it != 0 && _ep_hashes.count(*it) == 0) { _ep_hashes.insert(*it); + _entrypoint_snapshot.reset(); } } - for (unordered_set::iterator it = hashes_short.begin(); it != hashes_short.end(); it++) + for (auto it = hashes_short.begin(); it != hashes_short.end(); it++) { if (*it != 0 && _epshort_hashes.count(*it) == 0) { _epshort_hashes.insert(*it); + _entrypoint_snapshot.reset(); } } @@ -208,7 +230,16 @@ bool pe_hash_database::add_hashes_eps(unordered_set hashes, un } -bool pe_hash_database::add_folder( char* dir_name, WCHAR* filter, bool recursively ) +bool pe_hash_database::add_folder(char* dir_name, WCHAR* filter, bool recursively, int threads) +{ + work_pool pool(threads); + size_t queued = 0; + const bool result = _add_folder(dir_name, filter, recursively, pool, queued); + pool.wait(); + return result; +} + +bool pe_hash_database::_add_folder(char* dir_name, WCHAR* filter, bool recursively, work_pool& pool, size_t& queued) { // Expand the environment names in the directory char dir_name_expanded[PATH_MAX + 1]; @@ -242,7 +273,7 @@ bool pe_hash_database::add_folder( char* dir_name, WCHAR* filter, bool recursive char* directory = new char[strlen(dir_name_expanded) + strlen(ent->d_name) + 2]; sprintf_s(directory, strlen(dir_name_expanded) + strlen(ent->d_name) + 2, "%s\\%s", dir_name_expanded, ent->d_name); - add_folder(directory, filter, recursively); + _add_folder(directory, filter, recursively, pool, queued); // Cleanup delete[] directory; @@ -258,23 +289,24 @@ bool pe_hash_database::add_folder( char* dir_name, WCHAR* filter, bool recursive filename[length] = 0; sprintf(filename, "%s\\%S", dir_name_expanded, result); - // Processes the specified file - FILE* fh = fopen(filename, "rb"); - if (fh != NULL) - { - if (_is_mz(fh)) + const std::string path(filename); + pool.submit([this, path] { + FILE* fh = fopen(path.c_str(), "rb"); + if (fh != NULL) { + const bool pe = _is_mz(fh); fclose(fh); - - add_file(filename); + if (pe) + { + auto mutable_path = path; + add_file(&mutable_path[0]); + } } else - fclose(fh); - } - else { - // Error - fprintf(stderr, "Error opening file %s: %s.\n", filename, strerror(errno)); - } + fprintf(stderr, "Error opening file %s: %s.\n", path.c_str(), strerror(errno)); + }); + if (++queued % 256 == 0) + pool.wait(); delete[] filename; } } @@ -381,17 +413,26 @@ bool pe_hash_database::remove_folder( char* dir_name, WCHAR* filter, bool recurs bool pe_hash_database::contains(unsigned __int64 hash) { - return _clean_hashes.count( hash ) != 0; + EnterCriticalSection(&_lock); + const bool found = _clean_hashes.count(hash) != 0; + LeaveCriticalSection(&_lock); + return found; } bool pe_hash_database::contains_ep(unsigned __int64 hash) { - return _ep_hashes.count(hash) != 0; + EnterCriticalSection(&_lock); + const bool found = _ep_hashes.count(hash) != 0; + LeaveCriticalSection(&_lock); + return found; } bool pe_hash_database::contains_epshort(unsigned __int64 hash) { - return _epshort_hashes.count(hash) != 0; + EnterCriticalSection(&_lock); + const bool found = _epshort_hashes.count(hash) != 0; + LeaveCriticalSection(&_lock); + return found; } bool pe_hash_database::add_file(char* file) @@ -430,6 +471,7 @@ bool pe_hash_database::add_file(char* file) if (_ep_hashes.count(hash_ep) == 0) { _ep_hashes.insert(hash_ep); + _entrypoint_snapshot.reset(); printf("...new entrypoint hash %s,0x%llX\n", file, hash_ep); } LeaveCriticalSection(&_lock); @@ -442,6 +484,7 @@ bool pe_hash_database::add_file(char* file) if (_epshort_hashes.count(hash_ep_short) == 0) { _epshort_hashes.insert(hash_ep_short); + _entrypoint_snapshot.reset(); printf("...new entrypoint short hash %s,0x%llX\n", file, hash_ep_short); } LeaveCriticalSection(&_lock); @@ -576,6 +619,7 @@ bool pe_hash_database::save() pe_hash_database::~pe_hash_database(void) { delete[] _clean_database_path; + delete[] _ep_database_path; + delete[] _epshort_database_path; DeleteCriticalSection(&_lock); } - diff --git a/pd/pe_hash_database.h b/pd/pe_hash_database.h index 243aa85..7b14062 100644 --- a/pd/pe_hash_database.h +++ b/pd/pe_hash_database.h @@ -2,6 +2,7 @@ #include +#include #include #include #include @@ -11,6 +12,7 @@ #include "pe_header.h" #include "simple.h" #include // std::mutex +#include "work_pool.h" @@ -22,6 +24,15 @@ using namespace std::tr1; class pe_hash_database { +public: + struct entrypoint_hashes + { + unordered_set full; + unordered_set short_hashes; + unsigned __int64 minimum = _UI64_MAX; + unsigned __int64 maximum = 0; + }; +private: char* _clean_database_path; // Known clean modules char* _ep_database_path; // EntryPoint hashes char* _epshort_database_path; // EntryPoint short hashes @@ -29,11 +40,14 @@ class pe_hash_database unordered_set _clean_hashes; unordered_set _ep_hashes; unordered_set _epshort_hashes; + std::shared_ptr _entrypoint_snapshot; bool _is_mz(FILE* stream); + bool _add_folder(char* dir_name, WCHAR* filter, bool recursively, work_pool& pool, size_t& queued); CRITICAL_SECTION _lock; public: + std::shared_ptr snapshot_entrypoints(); pe_hash_database(char* clean_database_name, char* ep_database_name, char* epshort_database_name); int count(); @@ -42,10 +56,10 @@ class pe_hash_database bool clear_database(); - bool add_hashes(unordered_set hashes); - bool add_hashes_eps(unordered_set hashes, unordered_set hashes_short); + bool add_hashes(const unordered_set& hashes); + bool add_hashes_eps(const unordered_set& hashes, const unordered_set& hashes_short); - bool add_folder( char* dir_name, WCHAR* filter, bool recursively ); + bool add_folder(char* dir_name, WCHAR* filter, bool recursively, int threads = 1); bool remove_folder( char* dir_name, WCHAR* filter, bool recursively ); bool contains(unsigned __int64 hash); bool contains_epshort(unsigned __int64 hash); diff --git a/pd/pe_header.cpp b/pd/pe_header.cpp index 1a74f10..e1b516d 100644 --- a/pd/pe_header.cpp +++ b/pd/pe_header.cpp @@ -219,8 +219,15 @@ pe_header::pe_header( HANDLE ph, void* base, module_list* modules, PD_OPTIONS* o this->_disk_image_size = 0; this->_unique_hash = 0; - this->_stream = (stream_wrapper*) new process_stream( ph, base ); + this->_stream = (stream_wrapper*) new process_stream( ph, base, modules ); _original_base = base; + if (modules != NULL) + { + _name_filepath_long = new char[FILEPATH_SIZE]; + _name_filepath_long_size = _stream->get_long_name(_name_filepath_long, FILEPATH_SIZE); + _name_filepath_short = new char[FILEPATH_SIZE]; + _name_filepath_short_size = _stream->get_short_name(_name_filepath_short, FILEPATH_SIZE); + } if( _options->Verbose ) fprintf( stdout, "INFO: Initialized header for module name %s.\n", this->get_name() ); @@ -604,6 +611,7 @@ IMPORT_SUMMARY pe_header::get_imports_information( export_list* exports, __int64 // Gets the number of distinct import addresses that are imported. unordered_set import_addresses; + unordered_set import_libraries; if( _options->Verbose ) printf( "INFO: Building import information.\n" ); @@ -619,21 +627,22 @@ IMPORT_SUMMARY pe_header::get_imports_information( export_list* exports, __int64 hash hasher; - if( this->_parsed_sections && exports != NULL && _image_size >= 8 && size_limit >= 8 ) + if( this->_parsed_sections && exports != NULL && size_limit >= 4 ) { + const SIZE_T limit = static_cast((std::min)(_image_size, size_limit)); // Add matches to exports in this process unsigned __int32 cand32_last = 0; unsigned __int64 cand64_last = 0; - for(__int64 offset = 0; offset < _image_size - 8 && offset < size_limit - 8; offset+=4 ) + for(SIZE_T offset = 0; range_fits(limit, offset, sizeof(DWORD)); offset += sizeof(DWORD)) { // Check if this 4-gram or 8-gram points to an export unsigned __int32 cand32 = *((__int32*)(_image + offset)); if (cand32 != cand32_last) { - if (exports->contains(cand32)) + if (const export_entry* found = exports->lookup(cand32)) { - export_entry entry = exports->find(cand32); + const export_entry& entry = *found; // Found an import reference unordered_set::const_iterator gotImportAddress = import_addresses.find(cand32); @@ -650,8 +659,15 @@ IMPORT_SUMMARY pe_header::get_imports_information( export_list* exports, __int64 hash_generic = hash_generic ^ hasher(string(entry.name)); hash_specific = hash_specific ^ hasher(string(entry.name)); } + else + { + const size_t ordinal_hash = hasher("#" + std::to_string(entry.ord)); + hash_generic ^= ordinal_hash; + hash_specific ^= ordinal_hash; + } if (entry.library_name != NULL) { + import_libraries.insert(entry.library_name); hash_generic = hash_generic ^ (hasher(string(entry.library_name)) << 1); hash_specific = hash_specific ^ (hasher(string(entry.library_name)) << 1); } @@ -663,12 +679,14 @@ IMPORT_SUMMARY pe_header::get_imports_information( export_list* exports, __int64 } cand32_last = cand32; + if (!range_fits(limit, offset, sizeof(unsigned __int64))) + continue; unsigned __int64 cand64 = *((unsigned __int64*)(_image + offset)); if (cand64 != cand64_last && cand64 > 0xffffffff) { - if (exports->contains(cand64)) + if (const export_entry* found = exports->lookup(cand64)) { - export_entry entry = exports->find(cand64); + const export_entry& entry = *found; // Found an import reference unordered_set::const_iterator gotImportAddress = import_addresses.find(cand64); @@ -685,8 +703,15 @@ IMPORT_SUMMARY pe_header::get_imports_information( export_list* exports, __int64 hash_generic = hash_generic ^ hasher(string(entry.name)); hash_specific = hash_specific ^ hasher(string(entry.name)); } + else + { + const size_t ordinal_hash = hasher("#" + std::to_string(entry.ord)); + hash_generic ^= ordinal_hash; + hash_specific ^= ordinal_hash; + } if (entry.library_name != NULL) { + import_libraries.insert(entry.library_name); hash_generic = hash_generic ^ (hasher(string(entry.library_name)) << 1); hash_specific = hash_specific ^ (hasher(string(entry.library_name)) << 1); } @@ -702,6 +727,7 @@ IMPORT_SUMMARY pe_header::get_imports_information( export_list* exports, __int64 result.HASH_GENERIC = hash_generic; result.HASH_SPECIFIC = hash_specific; + result.COUNT_UNIQUE_IMPORT_LIBRARIES = import_libraries.size(); if( _options->Verbose ) { @@ -792,9 +818,9 @@ bool pe_header::build_pe_header( __int64 size, bool amd64, int num_sections_limi _header_pe32->FileHeader.PointerToSymbolTable = 0; _header_pe32->FileHeader.SizeOfOptionalHeader = sizeof(IMAGE_OPTIONAL_HEADER32); if( _options->ReconstructHeaderAsDll ) - _header_pe32->FileHeader.Characteristics = 0x0002; // Exe: 0x0002 + _header_pe32->FileHeader.Characteristics = IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_DLL | IMAGE_FILE_32BIT_MACHINE; else - _header_pe32->FileHeader.Characteristics = 0x2000; // Dll: 0x2000 + _header_pe32->FileHeader.Characteristics = IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_32BIT_MACHINE; _header_pe32->OptionalHeader.Magic=0x10b; _header_pe32->OptionalHeader.MajorLinkerVersion=0x08; _header_pe32->OptionalHeader.MinorLinkerVersion=0x00; @@ -841,9 +867,9 @@ bool pe_header::build_pe_header( __int64 size, bool amd64, int num_sections_limi _header_pe64->FileHeader.PointerToSymbolTable = 0; _header_pe64->FileHeader.SizeOfOptionalHeader = sizeof(IMAGE_OPTIONAL_HEADER64); if( _options->ReconstructHeaderAsDll ) - _header_pe64->FileHeader.Characteristics = 0x0002; // Exe: 0x0002 + _header_pe64->FileHeader.Characteristics = IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_DLL; else - _header_pe64->FileHeader.Characteristics = 0x2000; // Dll: 0x2000 + _header_pe64->FileHeader.Characteristics = IMAGE_FILE_EXECUTABLE_IMAGE; _header_pe64->OptionalHeader.Magic=0x020b; _header_pe64->OptionalHeader.MajorLinkerVersion=0x08; _header_pe64->OptionalHeader.MinorLinkerVersion=0x00; @@ -1021,8 +1047,11 @@ bool pe_header::process_sections( ) _clear_images(); if (!_parsed_pe_32 && !_parsed_pe_64) return false; - const SIZE_T nt_size = _parsed_pe_32 ? sizeof(IMAGE_NT_HEADERS32) : sizeof(IMAGE_NT_HEADERS64); - const SIZE_T section_offset = static_cast(_header_dos->e_lfanew) + nt_size; + const WORD optional_size = _parsed_pe_32 ? _header_pe32->FileHeader.SizeOfOptionalHeader : _header_pe64->FileHeader.SizeOfOptionalHeader; + const SIZE_T minimum_optional_size = _parsed_pe_32 ? sizeof(IMAGE_OPTIONAL_HEADER32) : sizeof(IMAGE_OPTIONAL_HEADER64); + if (optional_size < minimum_optional_size) + return _reject_size(); + const SIZE_T section_offset = static_cast(_header_dos->e_lfanew) + sizeof(DWORD) + sizeof(IMAGE_FILE_HEADER) + optional_size; WORD section_count = _parsed_pe_32 ? _header_pe32->FileHeader.NumberOfSections : _header_pe64->FileHeader.NumberOfSections; const DWORD declared_image_size = _parsed_pe_32 ? _header_pe32->OptionalHeader.SizeOfImage : _header_pe64->OptionalHeader.SizeOfImage; const DWORD headers_size = _parsed_pe_32 ? _header_pe32->OptionalHeader.SizeOfHeaders : _header_pe64->OptionalHeader.SizeOfHeaders; @@ -1039,12 +1068,6 @@ bool pe_header::process_sections( ) if (!range_fits(MAX_PE_IMAGE_SIZE, last.VirtualAddress, last.Misc.VirtualSize)) return _reject_size(); SIZE_T checked_image_size = last.VirtualAddress + last.Misc.VirtualSize; - if (last.Misc.VirtualSize > MAX_SECTION_SIZE) - { - const SIZE_T maximum_end = static_cast(last.VirtualAddress) + MAX_SECTION_SIZE; - checked_image_size = declared_image_size > last.VirtualAddress && declared_image_size < maximum_end ? - declared_image_size : maximum_end; - } checked_image_size = max(checked_image_size, declared_image_size); if (checked_image_size == 0 || checked_image_size > MAX_PE_IMAGE_SIZE || !range_fits(checked_image_size, 0, headers_size)) @@ -1066,8 +1089,7 @@ bool pe_header::process_sections( ) if( this->_parsed_pe_32 ) { // Attempt to parse the sections - unsigned char* base_pe = _header_dos->e_lfanew + _raw_header; - unsigned char* base_sections = base_pe + sizeof(*_header_pe32); + unsigned char* base_sections = _raw_header + section_offset; if( _header_pe32->FileHeader.NumberOfSections > 0x100 ) { char* location = new char[FILEPATH_SIZE + 1]; @@ -1112,42 +1134,7 @@ bool pe_header::process_sections( ) } - // Calculate the total size of the virtual image by inspecting the last section - DWORD image_size = 0; - if( this->_num_sections > 0 ) - { - if( _header_sections[_num_sections - 1].Misc.VirtualSize > MAX_SECTION_SIZE ) - { - // Smartly choose _header_pe32->OptionalHeader.SizeOfImage or last section plus max section size - if( _header_pe32->OptionalHeader.SizeOfImage > _header_sections[_num_sections - 1].VirtualAddress && - _header_pe32->OptionalHeader.SizeOfImage < _header_sections[_num_sections - 1].VirtualAddress + MAX_SECTION_SIZE ) - { - // Use the _header_pe32->OptionalHeader.SizeOfImage, since it seems valid - char* location = new char[FILEPATH_SIZE + 1]; - _stream->get_location(location, FILEPATH_SIZE + 1); - fprintf( stderr, "WARNING: module '%s' at %s. Image size of last section appears incorrect, using image size specified by optional header instead since it appears valid. This could be as a result of a custom code to load a library by means other than LoadLibrary().\n", - this->get_name(), location); - delete[] location; - image_size = _header_pe32->OptionalHeader.SizeOfImage; - } - else - { - // Assume a really large last section since _header_pe32->OptionalHeader.SizeOfImage appears invalid. - char* location = new char[FILEPATH_SIZE + 1]; - _stream->get_location(location, FILEPATH_SIZE + 1); - fprintf( stderr, "WARNING: module '%s' at %s. Image size of last section appears incorrect, using built-in max section size of 0x%x instead. This could be as a result of a custom code to load a library by means other than LoadLibrary().\n", - this->get_name(), location, - MAX_SECTION_SIZE); - delete[] location; - image_size = _header_sections[_num_sections - 1].VirtualAddress + MAX_SECTION_SIZE; - } - } - else - image_size = _header_sections[_num_sections - 1].VirtualAddress + - _header_sections[_num_sections - 1].Misc.VirtualSize; - } - if( _header_pe32->OptionalHeader.SizeOfImage > image_size ) - image_size = _header_pe32->OptionalHeader.SizeOfImage; + const SIZE_T image_size = checked_image_size; // Perform a sanity check on the resulting image size if( image_size > MAX_PE_IMAGE_SIZE ) @@ -1233,8 +1220,7 @@ bool pe_header::process_sections( ) else if( this->_parsed_pe_64 ) { // Attempt to parse the sections - unsigned char* base_pe = _header_dos->e_lfanew + _raw_header; - unsigned char* base_sections = base_pe + sizeof(*_header_pe64); + unsigned char* base_sections = _raw_header + section_offset; if( _header_pe64->FileHeader.NumberOfSections > 0x100 ) { char* location = new char[FILEPATH_SIZE + 1]; @@ -1275,42 +1261,7 @@ bool pe_header::process_sections( ) } } - // Calculate the total size of the virtual image by inspecting the last section - DWORD image_size = 0; - if( this->_num_sections > 0 ) - { - if( _header_sections[_num_sections - 1].Misc.VirtualSize > MAX_SECTION_SIZE ) - { - // Smartly choose _header_pe64->OptionalHeader.SizeOfImage or last section plus max section size - if( _header_pe64->OptionalHeader.SizeOfImage > _header_sections[_num_sections - 1].VirtualAddress && - _header_pe64->OptionalHeader.SizeOfImage < _header_sections[_num_sections - 1].VirtualAddress + MAX_SECTION_SIZE ) - { - // Use the _header_pe64->OptionalHeader.SizeOfImage, since it seems valid - char* location = new char[FILEPATH_SIZE + 1]; - _stream->get_location(location, FILEPATH_SIZE + 1); - fprintf( stderr, "WARNING: module '%s' at %s. Image size of last section appears incorrect, using image size specified by optional header instead since it appears valid. This could be as a result of a custom code to load a library by means other than LoadLibrary().\n", - this->get_name(), location); - delete[] location; - image_size = _header_pe64->OptionalHeader.SizeOfImage; - } - else - { - // Assume a really large last section since _header_pe64->OptionalHeader.SizeOfImage appears invalid. - char* location = new char[FILEPATH_SIZE + 1]; - _stream->get_location(location, FILEPATH_SIZE + 1); - fprintf( stderr, "WARNING: module '%s' at %s. Image size of last section appears incorrect, using built-in max section size of 0x%x instead. This could be as a result of a custom code to load a library by means other than LoadLibrary().\n", - this->get_name(), location, - MAX_SECTION_SIZE); - delete[] location; - image_size = _header_sections[_num_sections - 1].VirtualAddress + MAX_SECTION_SIZE; - } - } - else - image_size = _header_sections[_num_sections - 1].VirtualAddress + - _header_sections[_num_sections - 1].Misc.VirtualSize; - } - if( _header_pe64->OptionalHeader.SizeOfImage > image_size ) - image_size = _header_pe64->OptionalHeader.SizeOfImage; + const SIZE_T image_size = checked_image_size; // Perform a sanity check on the resulting image size if( image_size > MAX_PE_IMAGE_SIZE ) @@ -1400,6 +1351,343 @@ bool pe_header::process_sections( ) return false; } +bool pe_header::_append_import_section(DWORD rva, DWORD size) +{ + if (!range_fits(MAX_PE_IMAGE_SIZE, rva, size)) + return _reject_size(); + const SIZE_T next_header = reinterpret_cast(_header_sections + _num_sections) - _raw_header; + const DWORD headers_size = _parsed_pe_32 ? _header_pe32->OptionalHeader.SizeOfHeaders : _header_pe64->OptionalHeader.SizeOfHeaders; + if (_num_sections < 0x100 && range_fits(_raw_header_size, next_header, sizeof(IMAGE_SECTION_HEADER)) && + range_fits(headers_size, next_header, sizeof(IMAGE_SECTION_HEADER)) && + std::all_of(_raw_header + next_header, _raw_header + next_header + sizeof(IMAGE_SECTION_HEADER), + [](unsigned char byte) { return byte == 0; })) + { + // Keep a preceding BSS tail virtual instead of materializing its zeros. + IMAGE_SECTION_HEADER& section = _header_sections[_num_sections++]; + memset(§ion, 0, sizeof(section)); + memcpy(section.Name, ".pdimp", 6); + section.VirtualAddress = rva; + section.Misc.VirtualSize = size; + section.SizeOfRawData = size; + section.Characteristics = IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE; + if (_parsed_pe_32) + _header_pe32->FileHeader.NumberOfSections = static_cast(_num_sections); + else + _header_pe64->FileHeader.NumberOfSections = static_cast(_num_sections); + return true; + } + IMAGE_SECTION_HEADER& last = _header_sections[_num_sections - 1]; + if (last.VirtualAddress > rva) + return _reject_size(); + fprintf(stderr, "WARNING: No section-header space for reconstructed imports in '%s'; extending the last section.\n", get_name()); + last.Misc.VirtualSize = rva + size - last.VirtualAddress; + last.SizeOfRawData = last.Misc.VirtualSize; + return true; +} + +bool pe_header::_writable_data_range(SIZE_T rva, SIZE_T width) const +{ + for (int i = 0; i < _num_sections; ++i) + { + const auto& section = _header_sections[i]; + const SIZE_T span = (std::max)(section.Misc.VirtualSize, section.SizeOfRawData); + if (rva >= section.VirtualAddress && range_fits(span, rva - section.VirtualAddress, width)) + return (section.Characteristics & IMAGE_SCN_MEM_WRITE) != 0 && + (section.Characteristics & IMAGE_SCN_MEM_EXECUTE) == 0; + } + return false; +} + +bool pe_header::_reexecution_import_ranges(std::vector>& ranges) const +{ + const auto& directory = _parsed_pe_32 ? _header_pe32->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT] : + _header_pe64->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT]; + std::vector delay_iats; + if (directory.VirtualAddress != 0) + { + if (!range_fits(_image_size, directory.VirtualAddress, directory.Size)) + return false; + bool terminated = false; + for (SIZE_T offset = 0; range_fits(directory.Size, offset, 8 * sizeof(DWORD)); offset += 8 * sizeof(DWORD)) + { + DWORD fields[8]; + memcpy(fields, _image + directory.VirtualAddress + offset, sizeof(fields)); + if (std::all_of(fields, fields + 8, [](DWORD value) { return value == 0; })) + { + terminated = true; + break; + } + ULONGLONG iat = fields[3]; + if ((fields[0] & 1) == 0) + { + const ULONGLONG base = _original_base != NULL ? reinterpret_cast(_original_base) : + (_parsed_pe_32 ? _header_pe32->OptionalHeader.ImageBase : _header_pe64->OptionalHeader.ImageBase); + if (iat < base) + return false; + iat -= base; + } + const SIZE_T width = _parsed_pe_32 ? sizeof(DWORD) : sizeof(ULONGLONG); + if (iat == 0 || iat > MAX_PE_IMAGE_SIZE || !range_fits(_image_size, static_cast(iat), width)) + return false; + delay_iats.push_back(static_cast(iat)); + } + if (!terminated) + return false; + std::sort(delay_iats.begin(), delay_iats.end()); + } + for (int i = 0; i < _num_sections; ++i) + { + const auto& section = _header_sections[i]; + const SIZE_T span = (std::max)(section.Misc.VirtualSize, section.SizeOfRawData); + const auto delay = std::lower_bound(delay_iats.begin(), delay_iats.end(), section.VirtualAddress); + // The loader can protect delay-IAT sections before binding normal imports. + if ((section.Characteristics & IMAGE_SCN_MEM_WRITE) && !(section.Characteristics & IMAGE_SCN_MEM_EXECUTE) && + (delay == delay_iats.end() || *delay - section.VirtualAddress >= span)) + ranges.emplace_back(section.VirtualAddress, section.SizeOfRawData); + } + return true; +} + +bool pe_header::_reexecution_cookie_rva(SIZE_T& rva) const +{ + rva = _image_size; + const auto& config = _parsed_pe_32 ? _header_pe32->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG] : + _header_pe64->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG]; + const SIZE_T width = _parsed_pe_32 ? sizeof(DWORD) : sizeof(ULONGLONG); + const SIZE_T field = _parsed_pe_32 ? offsetof(IMAGE_LOAD_CONFIG_DIRECTORY32, SecurityCookie) : + offsetof(IMAGE_LOAD_CONFIG_DIRECTORY64, SecurityCookie); + if (config.VirtualAddress == 0 || config.Size < field + width) + return true; + if (!range_fits(_image_size, config.VirtualAddress, config.Size)) + return false; + DWORD declared_size = 0; + memcpy(&declared_size, _image + config.VirtualAddress, sizeof(declared_size)); + if (declared_size < field + width) + return true; + ULONGLONG cookie = 0; + memcpy(&cookie, _image + config.VirtualAddress + field, width); + if (cookie == 0) + return true; + const ULONGLONG base = _original_base != NULL ? + (_parsed_pe_32 ? static_cast(reinterpret_cast(_original_base)) : reinterpret_cast(_original_base)) : + (_parsed_pe_32 ? _header_pe32->OptionalHeader.ImageBase : _header_pe64->OptionalHeader.ImageBase); + if (cookie < base || cookie - base > MAX_PE_IMAGE_SIZE || + !range_fits(_image_size, static_cast(cookie - base), width) || + !_writable_data_range(static_cast(cookie - base), width)) + return false; + rva = static_cast(cookie - base); + return true; +} + +bool pe_header::_pack_disk_image(const unsigned char* image, SIZE_T size, SIZE_T cookie_rva) +{ + const DWORD alignment = _parsed_pe_32 ? _header_pe32->OptionalHeader.SectionAlignment : _header_pe64->OptionalHeader.SectionAlignment; + DWORD& headers_size = _parsed_pe_32 ? _header_pe32->OptionalHeader.SizeOfHeaders : _header_pe64->OptionalHeader.SizeOfHeaders; + DWORD& file_alignment = _parsed_pe_32 ? _header_pe32->OptionalHeader.FileAlignment : _header_pe64->OptionalHeader.FileAlignment; + DWORD& image_size = _parsed_pe_32 ? _header_pe32->OptionalHeader.SizeOfImage : _header_pe64->OptionalHeader.SizeOfImage; + IMAGE_DATA_DIRECTORY* directories = _parsed_pe_32 ? + _header_pe32->OptionalHeader.DataDirectory : _header_pe64->OptionalHeader.DataDirectory; + const SIZE_T table_end = reinterpret_cast(_header_sections + _num_sections) - _raw_header; + const __int64 aligned_headers = _section_align(static_cast<__int64>((std::max)(headers_size, table_end)), alignment); + const __int64 aligned_image = _section_align(static_cast<__int64>(size), alignment); + if (!image_size_fits(aligned_headers) || !image_size_fits(aligned_image) || + aligned_headers > static_cast<__int64>(size)) + return _reject_size(); + headers_size = static_cast(aligned_headers); + file_alignment = alignment; + image_size = static_cast(aligned_image); + SIZE_T required_space = headers_size; + std::vector copy_sizes; + for (int i = 0; i < _num_sections; ++i) + { + IMAGE_SECTION_HEADER& section = _header_sections[i]; + const SIZE_T span = (std::max)(section.Misc.VirtualSize, section.SizeOfRawData); + if (!range_fits(size, section.VirtualAddress, span) || + (span != 0 && section.VirtualAddress < headers_size)) + return _reject_size(); + SIZE_T initialized = span; + while (initialized >= sizeof(SIZE_T)) + { + SIZE_T word; + memcpy(&word, image + section.VirtualAddress + initialized - sizeof(word), sizeof(word)); + if (word != 0) + break; + initialized -= sizeof(word); + } + while (initialized != 0 && image[section.VirtualAddress + initialized - 1] == 0) + --initialized; + for (int directory = 0; directory < IMAGE_NUMBEROF_DIRECTORY_ENTRIES; ++directory) + { + if (directory == IMAGE_DIRECTORY_ENTRY_SECURITY) + continue; + const auto& entry = directories[directory]; + if (entry.VirtualAddress >= section.VirtualAddress && + range_fits(span, entry.VirtualAddress - section.VirtualAddress, entry.Size)) + initialized = (std::max)(initialized, entry.VirtualAddress - section.VirtualAddress + entry.Size); + } + const SIZE_T cookie_width = _parsed_pe_32 ? sizeof(DWORD) : sizeof(ULONGLONG); + if (cookie_rva < _image_size && cookie_rva >= section.VirtualAddress && + range_fits(span, cookie_rva - section.VirtualAddress, cookie_width)) + initialized = (std::max)(initialized, cookie_rva - section.VirtualAddress + cookie_width); + const __int64 raw_size = _section_align(static_cast<__int64>(initialized), alignment); + if (!image_size_fits(raw_size) || !range_fits(MAX_PE_IMAGE_SIZE, required_space, static_cast(raw_size))) + return _reject_size(); + section.SizeOfRawData = static_cast(raw_size); + section.PointerToRawData = raw_size == 0 ? 0 : static_cast(required_space); + if (raw_size != 0 && (section.Characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA)) + section.Characteristics = (section.Characteristics & ~IMAGE_SCN_CNT_UNINITIALIZED_DATA) | IMAGE_SCN_CNT_INITIALIZED_DATA; + copy_sizes.push_back(initialized); + required_space += static_cast(raw_size); + } + // Certificates use file offsets and are not mapped into the process. + directories[IMAGE_DIRECTORY_ENTRY_SECURITY] = {}; + if (_parsed_pe_32) + _header_pe32->OptionalHeader.CheckSum = 0; + else + _header_pe64->OptionalHeader.CheckSum = 0; + _disk_image = new unsigned char[required_space](); + _disk_image_size = required_space; + memcpy(_disk_image, _raw_header, (std::min)(headers_size, _raw_header_size)); + for (int i = 0; i < _num_sections; ++i) + { + if (copy_sizes[i] != 0) + memcpy(_disk_image + _header_sections[i].PointerToRawData, + image + _header_sections[i].VirtualAddress, copy_sizes[i]); + } + const auto file_offset = [this, headers_size](DWORD rva, SIZE_T count, SIZE_T& offset) { + if (rva < headers_size && range_fits(headers_size, rva, count)) + { + offset = rva; + return true; + } + for (int i = 0; i < _num_sections; ++i) + { + const auto& section = _header_sections[i]; + if (rva >= section.VirtualAddress && range_fits(section.SizeOfRawData, rva - section.VirtualAddress, count)) + { + offset = section.PointerToRawData + rva - section.VirtualAddress; + return true; + } + } + return false; + }; + const auto& debug = directories[IMAGE_DIRECTORY_ENTRY_DEBUG]; + SIZE_T debug_offset = 0; + if (debug.Size != 0 && file_offset(debug.VirtualAddress, debug.Size, debug_offset)) + { + for (SIZE_T i = 0; i < debug.Size / sizeof(IMAGE_DEBUG_DIRECTORY); ++i) + { + auto& entry = *reinterpret_cast(_disk_image + debug_offset + i * sizeof(IMAGE_DEBUG_DIRECTORY)); + SIZE_T data_offset = 0; + entry.PointerToRawData = entry.AddressOfRawData != 0 && + file_offset(entry.AddressOfRawData, entry.SizeOfData, data_offset) ? static_cast(data_offset) : 0; + } + } + if (cookie_rva < _image_size) + { + const SIZE_T width = _parsed_pe_32 ? sizeof(DWORD) : sizeof(ULONGLONG); + SIZE_T offset = 0; + if (!file_offset(static_cast(cookie_rva), width, offset)) + { + fprintf(stderr, "WARNING: Cannot prepare '%s' for re-execution: GS cookie has no raw storage.\n", get_name()); + return false; + } + // The loader/CRT replaces this ABI sentinel with a fresh process-specific cookie. + const ULONGLONG bootstrap = _parsed_pe_32 ? 0xbb40e64eULL : 0x2b992ddfa232ULL; + memcpy(_disk_image + offset, &bootstrap, width); + if (_options->Verbose) + printf("INFO: Restored GS cookie bootstrap value for re-execution of '%s'.\n", get_name()); + } + return true; +} + +bool pe_header::_executable_range(SIZE_T rva, SIZE_T length) const +{ + if (!range_fits(_image_size, rva, length)) + return false; + for (int i = 0; i < _num_sections; ++i) + { + const auto& section = _header_sections[i]; + if ((section.Characteristics & IMAGE_SCN_MEM_EXECUTE) && rva >= section.VirtualAddress && + range_fits((std::max)(section.Misc.VirtualSize, section.SizeOfRawData), rva - section.VirtualAddress, length)) + return true; + } + return false; +} + +void pe_header::_recover_entrypoint(pe_hash_database* database) +{ + if (!database || !_options->EntryPointHash) + return; + DWORD& original = _parsed_pe_32 ? _header_pe32->OptionalHeader.AddressOfEntryPoint : + _header_pe64->OptionalHeader.AddressOfEntryPoint; + if (!_options->ForceReconstructEntryPoint && original != 0 && _executable_range(original, 1)) + return; + const auto hashes = database->snapshot_entrypoints(); + if (hashes->full.empty()) + return; + DWORD best = 0; + for (int i = 0; i < _num_sections; ++i) + { + const auto& section = _header_sections[i]; + const SIZE_T size = (std::max)(section.Misc.VirtualSize, section.SizeOfRawData); + if (!(section.Characteristics & IMAGE_SCN_MEM_EXECUTE) || !range_fits(_image_size, section.VirtualAddress, size)) + continue; + for (SIZE_T offset = 0; range_fits(size, offset, sizeof(ULONGLONG)); ++offset) + { + const SIZE_T rva = section.VirtualAddress + offset; + if (best && rva >= best) + break; + ULONGLONG prefix = 0; + memcpy(&prefix, _image + rva, sizeof(prefix)); + if (prefix >= hashes->minimum && prefix <= hashes->maximum && hashes->short_hashes.count(prefix) && + rva != 0) + { + const auto full = _hash_asm(rva); + if (full && hashes->full.count(full)) best = static_cast(rva); + } + } + } + if (!best && _parsed_pe_64) + { + // x64 unwind metadata gives bounded function starts even when relocated + // operands or build changes prevent the raw eight-byte prefilter matching. + const auto& directory = _header_pe64->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION]; + bool valid = directory.VirtualAddress != 0 && directory.Size % sizeof(IMAGE_RUNTIME_FUNCTION_ENTRY) == 0 && + range_fits(_image_size, directory.VirtualAddress, directory.Size); + DWORD previous = 0, candidate = 0; + bool ambiguous = false; + for (SIZE_T offset = 0; valid && offset < directory.Size; offset += sizeof(IMAGE_RUNTIME_FUNCTION_ENTRY)) + { + IMAGE_RUNTIME_FUNCTION_ENTRY function; + memcpy(&function, _image + directory.VirtualAddress + offset, sizeof(function)); + if (function.BeginAddress >= function.EndAddress || (offset && function.BeginAddress < previous) || + !_executable_range(function.BeginAddress, function.EndAddress - function.BeginAddress) || + !range_fits(_image_size, function.UnwindData & ~1u, sizeof(DWORD))) + { + valid = false; + break; + } + previous = function.EndAddress; + const auto full = _hash_asm(function.BeginAddress); + if (full && hashes->full.count(full)) + { + if (candidate && candidate != function.BeginAddress) ambiguous = true; + candidate = function.BeginAddress; + } + } + if (valid && !ambiguous) + best = candidate; + } + if (best) + { + printf("INFO: Recovered entrypoint for '%s': %x -> %x (opcode hash match).\n", get_name(), original, best); + original = best; + } + else + fprintf(stderr, "WARNING: No unambiguous supported entrypoint match for '%s'; retaining RVA %x.\n", get_name(), original); +} + bool pe_header::process_disk_image( export_list* exports, pe_hash_database* hash_database) { delete[] _disk_image; @@ -1422,57 +1710,49 @@ bool pe_header::process_disk_image( export_list* exports, pe_hash_database* hash _header_sections[i].SizeOfRawData > MAX_PE_IMAGE_SIZE) return _reject_size(); } + std::vector> import_ranges; + std::vector> zero_fill; + if (_options->Reexecution) + for (int i = 0; i < _num_sections; ++i) + { + const auto& section = _header_sections[i]; + if ((section.Characteristics & IMAGE_SCN_MEM_WRITE) && !(section.Characteristics & IMAGE_SCN_MEM_EXECUTE) && + section.Misc.VirtualSize > section.SizeOfRawData) + { + if (!range_fits(_image_size, section.VirtualAddress, section.Misc.VirtualSize)) + return _reject_size(); + zero_fill.emplace_back(section.VirtualAddress + section.SizeOfRawData, + section.Misc.VirtualSize - section.SizeOfRawData); + } + } + const auto restore_zero_fill = [&](unsigned char* image) { + for (const auto& range : zero_fill) + memset(image + range.first, 0, range.second); + }; + SIZE_T cookie_rva = _image_size; + if (_options->Reexecution && !_reexecution_cookie_rva(cookie_rva)) + { + fprintf(stderr, "WARNING: Cannot prepare '%s' for re-execution: invalid load configuration or GS cookie location.\n", get_name()); + return false; + } + if (_options->Reexecution && _options->ImportRec && !_reexecution_import_ranges(import_ranges)) + { + fprintf(stderr, "WARNING: Cannot prepare '%s' for re-execution: invalid delay-import metadata.\n", get_name()); + return false; + } + const auto allow_import = [&](SIZE_T rva, SIZE_T width) { + if (cookie_rva < _image_size && rva < cookie_rva + width && cookie_rva < rva + width) + return false; + return !_options->Reexecution || std::any_of(import_ranges.begin(), import_ranges.end(), + [&](const std::pair& range) { + return rva >= range.first && range_fits(range.second, rva - range.first, width); + }); + }; + _recover_entrypoint(hash_database); if( this->_parsed_sections ) { if( this->_parsed_pe_32 ) { - // Re-build the Original Entry Point OEP if it looks to be not valid - if (hash_database != NULL && (_header_pe32->OptionalHeader.AddressOfEntryPoint == 0 || - _header_pe32->OptionalHeader.AddressOfEntryPoint == 0x2000 || - !range_fits(_image_size, _header_pe32->OptionalHeader.AddressOfEntryPoint, 20) || - _options->ForceReconstructEntryPoint)) - { - printf("INFO: Re-building entrypoint. Original entrypoint invalid: %x\n", _header_pe32->OptionalHeader.AddressOfEntryPoint); - - // The entry-point looks invalid, search for candidates to reconstruct it - unsigned __int64 best_entrypoint = 0; - - for (__int64 offset = 0x1000; _image_size >= 8 && offset < _image_size - 8; offset += 1) - { - // Check if this is a possible entrypoint - unsigned __int64 cand = *((__int64*)(_image + offset)); - - // Lookup the address - if (hash_database->contains_epshort(cand)) - { - // This is a possible entrypoint, this is a weak correlation but we'll use it if it's all we have - if (best_entrypoint == 0) - { - best_entrypoint = offset; - } - if( _options->Verbose ) - printf("INFO: Possible entrypoint found (weak): %x\n", offset); - - // Validate that the full hash matches a known entrypoint - cand = _hash_asm(offset); - if (hash_database->contains_ep(cand)) - { - best_entrypoint = offset; - printf("INFO: Possible entrypoint found (strong): %x\n", offset); - if (!_options->Verbose) - break; - } - } - } - - // Update the entrypoint - if (best_entrypoint != 0) - { - _header_pe32->OptionalHeader.AddressOfEntryPoint = best_entrypoint; - printf("INFO: Updated entrypoint to: %x\n", best_entrypoint); - } - } - // Reconstruct PE imports aggressively using our knowledge of all the exports addresses in this process // Technique: // 1. 'exports' defines all valid export addresses in this process @@ -1493,25 +1773,21 @@ bool pe_header::process_disk_image( export_list* exports, pe_hash_database* hash // Add matches to exports in this process int count = 0; unsigned __int64 cand_last = 0; - for(__int64 offset = 0; _image_size >= 8 && offset < _image_size - 8; offset+=4 ) + for (SIZE_T offset = 0; range_fits(_image_size, offset, sizeof(DWORD)); offset += sizeof(DWORD)) { - // Check if this 4-gram or 8-gram points to an export - unsigned __int64 cand = *((__int32*)(_image + offset)); + unsigned __int32 cand = 0; + memcpy(&cand, _image + offset, sizeof(cand)); - if ( cand_last != cand && exports->contains( cand ) ) + const export_entry* entry = cand_last != cand ? exports->lookup(cand) : NULL; + if (entry != NULL && allow_import(offset, sizeof(DWORD))) { - export_entry entry = exports->find(cand); - // Add this to be reconstructed as an import - if (entry.name != NULL) - peimp->add_fixup(entry.library_name, entry.name, offset, this->_parsed_pe_64); - else - peimp->add_fixup(entry.library_name, entry.ord, offset, this->_parsed_pe_64); + peimp->add_fixup(*entry, offset, this->_parsed_pe_64); if (!peimp->valid()) return _reject_size(); count++; } - else + else if (entry == NULL) { cand_last = cand; } @@ -1528,26 +1804,24 @@ bool pe_header::process_disk_image( export_list* exports, pe_hash_database* hash return _reject_size(); const __int64 new_section_size = _section_align(data_size + descriptor_size, section_alignment); const __int64 section_start = _section_align(static_cast<__int64>(_image_size), section_alignment); - const __int64 last_size = _section_align(static_cast<__int64>(_header_sections[_num_sections - 1].Misc.VirtualSize), section_alignment); larger_image_size = section_start + new_section_size; - if (new_section_size <= 0 || section_start < static_cast<__int64>(_image_size) || last_size < 0 || - !image_size_fits(larger_image_size) || !image_size_fits(last_size + new_section_size) || - !image_size_fits(_header_sections[_num_sections - 1].VirtualAddress + last_size + new_section_size)) + if (new_section_size <= 0 || section_start < static_cast<__int64>(_image_size) || + !image_size_fits(larger_image_size)) return _reject_size(); larger_image_storage.reset(new unsigned char[static_cast(larger_image_size)]); larger_image = larger_image_storage.get(); - memset(larger_image, 0, larger_image_size); + memset(larger_image + _image_size, 0, larger_image_size - _image_size); memcpy(larger_image, _image, _image_size); + restore_zero_fill(larger_image); if( _options->Verbose ) printf( "INFO: Writing added import table.\n" ); // Write to the new section if (!peimp->build_table(larger_image + static_cast(section_start), - new_section_size, static_cast<__int64>(_image_size), 0, descriptor_size)) + new_section_size, section_start, 0, descriptor_size) || + !_append_import_section(static_cast(section_start), static_cast(new_section_size))) return _reject_size(); - _header_sections[_num_sections - 1].Misc.VirtualSize = static_cast(last_size + new_section_size); - _header_sections[_num_sections - 1].SizeOfRawData = _header_sections[_num_sections - 1].Misc.VirtualSize; if( _options->Verbose ) printf( "INFO: Updating import data directory.\n" ); @@ -1560,10 +1834,14 @@ bool pe_header::process_disk_image( export_list* exports, pe_hash_database* hash else { larger_image_size = _image_size; - larger_image_storage.reset(new unsigned char[static_cast(larger_image_size)]); - larger_image = larger_image_storage.get(); - memset(larger_image, 0, larger_image_size); - memcpy(larger_image, _image, _image_size); + larger_image = _image; + if (!zero_fill.empty()) + { + larger_image_storage.reset(new unsigned char[_image_size]); + larger_image = larger_image_storage.get(); + memcpy(larger_image, _image, _image_size); + restore_zero_fill(larger_image); + } } if( _original_base != 0 ) @@ -1572,188 +1850,10 @@ bool pe_header::process_disk_image( export_list* exports, pe_hash_database* hash _header_pe32->OptionalHeader.ImageBase = (DWORD) _original_base; } - // Change the physical alignment to use the virtual alignment - if( _options->Verbose ) - printf( "INFO: Adjusting file alignment to %x.\n", _header_pe32->OptionalHeader.SectionAlignment); - _header_pe32->OptionalHeader.FileAlignment = _header_pe32->OptionalHeader.SectionAlignment; - - // Adjust the physical size of each section to use the virtual size - const __int64 aligned_headers = _section_align(static_cast<__int64>(_header_pe32->OptionalHeader.SizeOfHeaders), section_alignment); - if (!image_size_fits(aligned_headers)) - return _reject_size(); - _header_pe32->OptionalHeader.SizeOfHeaders = static_cast(aligned_headers); - SIZE_T required_space = static_cast(aligned_headers); - - for( int i = 0; i < _num_sections; i++ ) - { - // Correct the VirtualSize of the section if it is too large - if( this->_header_sections[i].Misc.VirtualSize > MAX_SECTION_SIZE ) - { - if( _options->Verbose ) - printf( "INFO: Calculating required space for section %i.\n", i); - - if( i + 1 < _num_sections && - this->_header_sections[i+1].VirtualAddress > this->_header_sections[i].VirtualAddress && - this->_header_sections[i+1].VirtualAddress < this->_header_sections[i].VirtualAddress + MAX_SECTION_SIZE ) - { - // Calculate the virtual size manually - char* location = new char[FILEPATH_SIZE + 1]; - _stream->get_location(location, FILEPATH_SIZE + 1); - fprintf( stderr, "WARNING: module '%s' at %s. Large section size for section %i of 0x%x changed to 0x%x based on image size as part of sanity check. This could be as a result of a custom code to load a library by means other than LoadLibrary().\n", - this->get_name(), location, i, this->_header_sections[i].Misc.VirtualSize, this->_header_sections[i+1].VirtualAddress - this->_header_sections[i].VirtualAddress ); - delete[] location; - this->_header_sections[i].Misc.VirtualSize = this->_header_sections[i+1].VirtualAddress - this->_header_sections[i].VirtualAddress; - } - else - { - // Use MAX_SECTION_SIZE - char* location = new char[FILEPATH_SIZE + 1]; - _stream->get_location(location, FILEPATH_SIZE + 1); - fprintf( stderr, "WARNING: module '%s' at %s. Large section size for section %i of 0x%x changed to 0x%x based on maximum section size as part of sanity check. This could be as a result of a custom code to load a library by means other than LoadLibrary().\n", - this->get_name(), location, i, this->_header_sections[i].Misc.VirtualSize, MAX_SECTION_SIZE ); - delete[] location; - this->_header_sections[i].Misc.VirtualSize = MAX_SECTION_SIZE; - } - } - - // Truncate VirtualSize to fit inside image size - if( this->_header_sections[i].VirtualAddress > static_cast(larger_image_size) ) - return _reject_size(); - if( this->_header_sections[i].Misc.VirtualSize > static_cast(larger_image_size) - this->_header_sections[i].VirtualAddress ) - { - char* location = new char[FILEPATH_SIZE + 1]; - _stream->get_location(location, FILEPATH_SIZE + 1); - DWORD new_size = larger_image_size - this->_header_sections[i].VirtualAddress; - fprintf( stderr, "WARNING: module '%s' at %s. Large section size for section %i of 0x%x being truncated to 0x%x to fit within the image size. This could be as a result of a custom code to load a library by means other than LoadLibrary().\n", - this->get_name(), location, i, this->_header_sections[i].Misc.VirtualSize, new_size ); - delete[] location; - this->_header_sections[i].Misc.VirtualSize = new_size; - } - - // Adjust the physical size to be at least the same size as the virtual size - if( this->_header_sections[i].Misc.VirtualSize > _header_sections[i].SizeOfRawData ) - { - _header_sections[i].SizeOfRawData = this->_header_sections[i].Misc.VirtualSize; - } - - // Update the pointer to raw data to be correct - _header_sections[i].PointerToRawData = required_space; - - const __int64 next_space = _section_align(static_cast<__int64>(required_space) + _header_sections[i].SizeOfRawData, section_alignment); - if (!image_size_fits(next_space) || - !range_fits(static_cast(larger_image_size), _header_sections[i].VirtualAddress, _header_sections[i].SizeOfRawData)) - return _reject_size(); - required_space = static_cast(next_space); - } - - // Set the size of image - _header_pe32->OptionalHeader.SizeOfImage = required_space; - - if( _options->Verbose ) - printf( "INFO: Copying the corrected memory PE header into file PE header format.\n"); - - // Copy over the modified PE header into the imaged version - if( _test_read( larger_image, larger_image_size, larger_image, _header_pe32->OptionalHeader.SizeOfHeaders ) && - _test_read( _raw_header, _raw_header_size, _raw_header, _header_pe32->OptionalHeader.SizeOfHeaders ) ) - { - memcpy( larger_image, _raw_header, _header_pe32->OptionalHeader.SizeOfHeaders ); - } - else if( _test_read( larger_image, larger_image_size, larger_image, _raw_header_size ) && - _test_read( _raw_header, _raw_header_size, _raw_header, _raw_header_size ) ) - { - memcpy( larger_image, _raw_header, _raw_header_size ); - } - - if( _header_pe32->OptionalHeader.SectionAlignment >= _header_pe32->OptionalHeader.FileAlignment ) - { - // Pack it down into a disk image of the file - if( _options->Verbose ) - printf( "INFO: Packing down memory sections into the file.\n"); - - // Allocate the necessary space for the physical image and initialize it to zero - _disk_image_size = required_space; - _disk_image = new unsigned char[_disk_image_size]; - memset(_disk_image, 0, _disk_image_size); - - // Copy the header - if( _test_read( _disk_image, _disk_image_size, _disk_image, _section_align(_header_pe32->OptionalHeader.SizeOfHeaders, _header_pe32->OptionalHeader.FileAlignment) ) && - _test_read( larger_image, larger_image_size, larger_image, _section_align(_header_pe32->OptionalHeader.SizeOfHeaders, _header_pe32->OptionalHeader.FileAlignment) ) ) - { - memcpy( _disk_image, larger_image, _header_pe32->OptionalHeader.SizeOfHeaders ); - } - - if( _parsed_sections ) - { - // Copy the sections one-by-one - for( int i = 0; i < _num_sections; i++ ) - { - if( _options->Verbose ) - printf( "INFO: Packing down section %i.\n", i); - - // Copy this section if the source and destination are both within acceptable bounds - if( range_fits(_disk_image_size, _header_sections[i].PointerToRawData, _header_sections[i].SizeOfRawData) && - range_fits(static_cast(larger_image_size), _header_sections[i].VirtualAddress, _header_sections[i].SizeOfRawData) ) - { - memcpy( _disk_image + _header_sections[i].PointerToRawData, larger_image + _header_sections[i].VirtualAddress, _header_sections[i].SizeOfRawData ); - } - } - } - - if( _options->Verbose ) - printf( "INFO: Done processing disk image.\n"); - - return true; - } + return _pack_disk_image(larger_image, static_cast(larger_image_size), cookie_rva); } else if( this->_parsed_pe_64 ) { - // Re-build the Original Entry Point OEP if it looks to be not valid - if (hash_database != NULL && (_header_pe64->OptionalHeader.AddressOfEntryPoint == 0 || - _header_pe64->OptionalHeader.AddressOfEntryPoint == 0x2000 || - !range_fits(_image_size, _header_pe64->OptionalHeader.AddressOfEntryPoint, 20) || - _options->ForceReconstructEntryPoint)) - { - printf("INFO: Re-building entrypoint. Original entrypoint invalid: %x\n", _header_pe64->OptionalHeader.AddressOfEntryPoint); - - // The entry-point looks invalid, search for candidates to reconstruct it - unsigned __int64 best_entrypoint = 0; - - for (__int64 offset = 0x1000; _image_size >= 8 && offset < _image_size - 8; offset += 1) - { - // Check if this is a possible entrypoint - unsigned __int64 cand = *((__int64*)(_image + offset)); - - // Lookup the address - if (hash_database->contains_epshort(cand)) - { - // This is a possible entrypoint, this is a weak correlation but we'll use it if it's all we have - if (best_entrypoint == 0) - { - best_entrypoint = offset; - } - if (_options->Verbose) - printf("INFO: Possible entrypoint found (weak): %x\n", offset); - - // Validate that the full hash matches a known entrypoint - cand = _hash_asm(offset); - if (hash_database->contains_ep(cand)) - { - best_entrypoint = offset; - printf("INFO: Possible entrypoint found (strong): %x\n", offset); - if (!_options->Verbose) - break; - } - } - } - - // Update the entrypoint - if (best_entrypoint != 0) - { - _header_pe64->OptionalHeader.AddressOfEntryPoint = best_entrypoint; - printf("INFO: Updated entrypoint to: %x\n", best_entrypoint); - } - } - // Reconstruct PE imports aggressively using our knowledge of all the exports addresses in this process // Technique: // 1. 'exports' defines all valid export addresses in this process @@ -1774,25 +1874,21 @@ bool pe_header::process_disk_image( export_list* exports, pe_hash_database* hash // Add matches to exports in this process int count = 0; unsigned __int64 cand_last = 0; - for(__int64 offset = 0; _image_size >= 8 && offset < _image_size - 8; offset+=4 ) + for (SIZE_T offset = 0; range_fits(_image_size, offset, sizeof(unsigned __int64)); offset += sizeof(DWORD)) { - // Check if this 4-gram or 8-gram points to an export - unsigned __int64 cand = *((unsigned __int64*)(_image + offset)); + unsigned __int64 cand = 0; + memcpy(&cand, _image + offset, sizeof(cand)); - if (cand_last != cand && exports->contains(cand)) + const export_entry* entry = cand_last != cand ? exports->lookup(cand) : NULL; + if (entry != NULL && allow_import(offset, sizeof(ULONGLONG))) { - export_entry entry = exports->find(cand); - // Add this to be reconstructed as an import - if (entry.name != NULL) - peimp->add_fixup(entry.library_name, entry.name, offset, this->_parsed_pe_64); - else - peimp->add_fixup(entry.library_name, entry.ord, offset, this->_parsed_pe_64); + peimp->add_fixup(*entry, offset, this->_parsed_pe_64); if (!peimp->valid()) return _reject_size(); count++; } - else + else if (entry == NULL) { cand_last = cand; } @@ -1809,26 +1905,24 @@ bool pe_header::process_disk_image( export_list* exports, pe_hash_database* hash return _reject_size(); const __int64 new_section_size = _section_align(data_size + descriptor_size, section_alignment); const __int64 section_start = _section_align(static_cast<__int64>(_image_size), section_alignment); - const __int64 last_size = _section_align(static_cast<__int64>(_header_sections[_num_sections - 1].Misc.VirtualSize), section_alignment); larger_image_size = section_start + new_section_size; - if (new_section_size <= 0 || section_start < static_cast<__int64>(_image_size) || last_size < 0 || - !image_size_fits(larger_image_size) || !image_size_fits(last_size + new_section_size) || - !image_size_fits(_header_sections[_num_sections - 1].VirtualAddress + last_size + new_section_size)) + if (new_section_size <= 0 || section_start < static_cast<__int64>(_image_size) || + !image_size_fits(larger_image_size)) return _reject_size(); larger_image_storage.reset(new unsigned char[static_cast(larger_image_size)]); larger_image = larger_image_storage.get(); - memset(larger_image, 0, larger_image_size); + memset(larger_image + _image_size, 0, larger_image_size - _image_size); memcpy(larger_image, _image, _image_size); + restore_zero_fill(larger_image); if( _options->Verbose ) printf( "INFO: Writing added import table.\n" ); // Write to the new section if (!peimp->build_table(larger_image + static_cast(section_start), - new_section_size, static_cast<__int64>(_image_size), 0, descriptor_size)) + new_section_size, section_start, 0, descriptor_size) || + !_append_import_section(static_cast(section_start), static_cast(new_section_size))) return _reject_size(); - _header_sections[_num_sections - 1].Misc.VirtualSize = static_cast(last_size + new_section_size); - _header_sections[_num_sections - 1].SizeOfRawData = _header_sections[_num_sections - 1].Misc.VirtualSize; if( _options->Verbose ) printf( "INFO: Updating import data directory.\n" ); @@ -1841,10 +1935,14 @@ bool pe_header::process_disk_image( export_list* exports, pe_hash_database* hash else { larger_image_size = _image_size; - larger_image_storage.reset(new unsigned char[static_cast(larger_image_size)]); - larger_image = larger_image_storage.get(); - memset(larger_image, 0, larger_image_size); - memcpy(larger_image, _image, _image_size); + larger_image = _image; + if (!zero_fill.empty()) + { + larger_image_storage.reset(new unsigned char[_image_size]); + larger_image = larger_image_storage.get(); + memcpy(larger_image, _image, _image_size); + restore_zero_fill(larger_image); + } } @@ -1855,139 +1953,7 @@ bool pe_header::process_disk_image( export_list* exports, pe_hash_database* hash _header_pe64->OptionalHeader.ImageBase = reinterpret_cast<__int64> (_original_base); } - // Change the physical alignment to use the virtual alignment - if( _options->Verbose ) - printf( "INFO: Adjusting file alignment to %x.\n", _header_pe64->OptionalHeader.SectionAlignment); - _header_pe64->OptionalHeader.FileAlignment = _header_pe64->OptionalHeader.SectionAlignment; - - // Adjust the physical size of each section to use the virtual size - const __int64 aligned_headers = _section_align(static_cast<__int64>(_header_pe64->OptionalHeader.SizeOfHeaders), section_alignment); - if (!image_size_fits(aligned_headers)) - return _reject_size(); - _header_pe64->OptionalHeader.SizeOfHeaders = static_cast(aligned_headers); - SIZE_T required_space = static_cast(aligned_headers); - - for( int i = 0; i < _num_sections; i++ ) - { - // Correct the VirtualSize of the section if it is too large - if( this->_header_sections[i].Misc.VirtualSize > MAX_SECTION_SIZE ) - { - if( _options->Verbose ) - printf( "INFO: Calculating required space for section %i.\n", i); - - if( i + 1 < _num_sections && - this->_header_sections[i+1].VirtualAddress > this->_header_sections[i].VirtualAddress && - this->_header_sections[i+1].VirtualAddress < this->_header_sections[i].VirtualAddress + MAX_SECTION_SIZE ) - { - // Calculate the virtual size manually - char* location = new char[FILEPATH_SIZE + 1]; - _stream->get_location(location, FILEPATH_SIZE + 1); - fprintf( stderr, "WARNING: module '%s' at %s. Large section size for section %i of 0x%x changed to 0x%x based on image virtual size as part of sanity check. This could be as a result of a custom code to load a library by means other than LoadLibrary().\n", - this->get_name(), location, i, this->_header_sections[i].Misc.VirtualSize, this->_header_sections[i+1].VirtualAddress - this->_header_sections[i].VirtualAddress ); - delete[] location; - this->_header_sections[i].Misc.VirtualSize = this->_header_sections[i+1].VirtualAddress - this->_header_sections[i].VirtualAddress; - } - else - { - // Use MAX_SECTION_SIZE - char* location = new char[FILEPATH_SIZE + 1]; - _stream->get_location(location, FILEPATH_SIZE + 1); - fprintf( stderr, "WARNING: module '%s' at %s. Large section size for section %i of 0x%x changed to 0x%x based on maximum section size as part of sanity check. This could be as a result of a custom code to load a library by means other than LoadLibrary().\n", - this->get_name(), location, i, this->_header_sections[i].Misc.VirtualSize, MAX_SECTION_SIZE ); - delete[] location; - this->_header_sections[i].Misc.VirtualSize = MAX_SECTION_SIZE; - } - } - - // Truncate VirtualSize to fit inside image size - if( this->_header_sections[i].VirtualAddress > static_cast(larger_image_size) ) - return _reject_size(); - if( this->_header_sections[i].Misc.VirtualSize > static_cast(larger_image_size) - this->_header_sections[i].VirtualAddress ) - { - char* location = new char[FILEPATH_SIZE + 1]; - _stream->get_location(location, FILEPATH_SIZE + 1); - DWORD new_size = larger_image_size - this->_header_sections[i].VirtualAddress; - fprintf( stderr, "WARNING: module '%s' at %s. Large section size for section %i of 0x%x being truncated to 0x%x to fit within the image size. This could be as a result of a custom code to load a library by means other than LoadLibrary().\n", - this->get_name(), location, i, this->_header_sections[i].Misc.VirtualSize, new_size ); - delete[] location; - this->_header_sections[i].Misc.VirtualSize = new_size; - } - - // Adjust the physical size to be at least the same size as the virtual size - if( this->_header_sections[i].Misc.VirtualSize > _header_sections[i].SizeOfRawData ) - { - _header_sections[i].SizeOfRawData = this->_header_sections[i].Misc.VirtualSize; - } - - // Update the pointer to raw data to be correct - _header_sections[i].PointerToRawData = required_space; - - const __int64 next_space = _section_align(static_cast<__int64>(required_space) + _header_sections[i].SizeOfRawData, section_alignment); - if (!image_size_fits(next_space) || - !range_fits(static_cast(larger_image_size), _header_sections[i].VirtualAddress, _header_sections[i].SizeOfRawData)) - return _reject_size(); - required_space = static_cast(next_space); - } - - // Set the size of image - _header_pe64->OptionalHeader.SizeOfImage = required_space; - - if( _options->Verbose ) - printf( "INFO: Copying the corrected memory PE header into file PE header format.\n"); - - // Copy over the modified PE header into the imaged version - if( _test_read( larger_image, larger_image_size, larger_image, _header_pe64->OptionalHeader.SizeOfHeaders ) && - _test_read( _raw_header, _raw_header_size, _raw_header, _header_pe64->OptionalHeader.SizeOfHeaders ) ) - { - memcpy( larger_image, _raw_header, _header_pe64->OptionalHeader.SizeOfHeaders ); - } - else if( _test_read( larger_image, larger_image_size, larger_image, _raw_header_size ) && - _test_read( _raw_header, _raw_header_size, _raw_header, _raw_header_size ) ) - { - memcpy( larger_image, _raw_header, _raw_header_size ); - } - - if( _header_pe64->OptionalHeader.SectionAlignment >= _header_pe64->OptionalHeader.FileAlignment ) - { - // Pack it down into a disk image of the file - if( _options->Verbose ) - printf( "INFO: Packing down memory sections into the file.\n"); - - // Allocate the necessary space for the physical image and initialize it to zero - _disk_image_size = required_space; - _disk_image = new unsigned char[_disk_image_size]; - memset(_disk_image, 0, _disk_image_size); - - // Copy the header - if( _test_read( _disk_image, _disk_image_size, _disk_image, _section_align(_header_pe64->OptionalHeader.SizeOfHeaders, _header_pe64->OptionalHeader.FileAlignment) ) && - _test_read( larger_image, larger_image_size, larger_image, _section_align(_header_pe64->OptionalHeader.SizeOfHeaders, _header_pe64->OptionalHeader.FileAlignment) ) ) - { - memcpy( _disk_image, larger_image, _header_pe64->OptionalHeader.SizeOfHeaders ); - } - - if( _parsed_sections ) - { - // Copy the sections one-by-one - for( int i = 0; i < _num_sections; i++ ) - { - if( _options->Verbose ) - printf( "INFO: Packing down section %i.\n", i); - - // Copy this section if the source and destination are both within acceptable bounds - if( range_fits(_disk_image_size, _header_sections[i].PointerToRawData, _header_sections[i].SizeOfRawData) && - range_fits(static_cast(larger_image_size), _header_sections[i].VirtualAddress, _header_sections[i].SizeOfRawData) ) - { - memcpy( _disk_image + _header_sections[i].PointerToRawData, larger_image + _header_sections[i].VirtualAddress, _header_sections[i].SizeOfRawData ); - } - } - } - - - if( _options->Verbose ) - printf( "INFO: Done processing disk image.\n"); - - return true; - } + return _pack_disk_image(larger_image, static_cast(larger_image_size), cookie_rva); } } return false; @@ -2030,17 +1996,31 @@ bool pe_header::process_import_directory( ) { SIZE_T destination = _header_import_descriptors[i].FirstThunk; SIZE_T source = _header_import_descriptors[i].OriginalFirstThunk; - // Preserve the existing DWORD-at-a-time reconstruction for both formats. - while (range_fits(_image_size, destination, sizeof(IMAGE_THUNK_DATA32)) && - range_fits(_image_size, source, sizeof(IMAGE_THUNK_DATA32))) + if (source == 0 || destination == 0 || source == destination) + continue; + const SIZE_T width = _parsed_pe_64 ? sizeof(IMAGE_THUNK_DATA64) : sizeof(IMAGE_THUNK_DATA32); + SIZE_T length = 0; + bool terminated = false; + while (range_fits(_image_size, source, length) && + range_fits(_image_size - source, length, width)) { - if (*reinterpret_cast(_image + destination) == 0 || - *reinterpret_cast(_image + source) == 0) + unsigned __int64 value = 0; + memcpy(&value, _image + source + length, width); + length += width; + if (value == 0) + { + terminated = true; break; - memmove(_image + destination, _image + source, sizeof(IMAGE_THUNK_DATA32)); - destination += sizeof(IMAGE_THUNK_DATA32); - source += sizeof(IMAGE_THUNK_DATA32); + } + } + if (!terminated || !range_fits(_image_size, destination, length) || + (source < destination + length && destination < source + length)) + { + fprintf(stderr, "WARNING: Invalid or overlapping import thunks in '%s'; preserving captured IAT bytes.\n", get_name()); + continue; } + // Copy complete pointer-sized entries, including the terminator, only after validating both tables. + memcpy(_image + destination, _image + source, length); } return true; } @@ -2053,11 +2033,12 @@ bool pe_header::process_export_directory( ) if( (this->_parsed_pe_32 || this->_parsed_pe_64) && _image != NULL ) { - SIZE_T exports_rva; + IMAGE_DATA_DIRECTORY directory; if( _parsed_pe_32 ) - exports_rva = _header_pe32->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress; + directory = _header_pe32->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT]; else - exports_rva = _header_pe64->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress; + directory = _header_pe64->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT]; + const SIZE_T exports_rva = directory.VirtualAddress; if( exports_rva != 0 && range_fits(_image_size, exports_rva, sizeof(IMAGE_EXPORT_DIRECTORY)) ) { @@ -2065,8 +2046,13 @@ bool pe_header::process_export_directory( ) _header_export_directory = ((IMAGE_EXPORT_DIRECTORY*) base_exports); // Parse this export directory - _export_list = new export_list(); - _export_list->add_exports( _image, _image_size, (__int64) _original_base, _header_export_directory, this->_parsed_pe_64); + std::unique_ptr exports(new export_list()); + if (!exports->add_exports(_image, _image_size, (__int64)_original_base, _header_export_directory, _parsed_pe_64, directory.Size)) + { + fprintf(stderr, "WARNING: Invalid export table in '%s'; excluding it from import reconstruction.\n", get_name()); + return false; + } + _export_list = exports.release(); return true; } diff --git a/pd/pe_header.h b/pd/pe_header.h index f5496a8..1092f73 100644 --- a/pd/pe_header.h +++ b/pd/pe_header.h @@ -24,9 +24,6 @@ using namespace std::tr1; #define EP_HASH_OPCODES_MIN 30 #define EP_HASH_OPCODES_MAX 100 -// 10MB -#define MAX_SECTION_SIZE (1024 * 1000) * 60 - class pe_hash_database; static bool static_zero_init = false; @@ -48,6 +45,7 @@ struct IMPORT_SUMMARY class pe_header { + friend struct entrypoint_benchmark; unsigned __int64 _unique_hash; unsigned __int64 _unique_hash_ep; @@ -106,6 +104,13 @@ class pe_header void _clear_images(); void _clear_header(); bool _reject_size(); + bool _append_import_section(DWORD rva, DWORD size); + bool _pack_disk_image(const unsigned char* image, SIZE_T size, SIZE_T cookie_rva); + bool _writable_data_range(SIZE_T rva, SIZE_T width) const; + bool _reexecution_import_ranges(std::vector>& ranges) const; + bool _reexecution_cookie_rva(SIZE_T& rva) const; + bool _executable_range(SIZE_T rva, SIZE_T length) const; + void _recover_entrypoint(pe_hash_database* database); unsigned __int64 _hash_asm(SIZE_T offset); unsigned __int64 _hash_short_asm(SIZE_T offset); diff --git a/pd/pe_imports.cpp b/pd/pe_imports.cpp index e11463a..1307367 100644 --- a/pd/pe_imports.cpp +++ b/pd/pe_imports.cpp @@ -30,12 +30,12 @@ namespace } } -void pe_imports::_add(std::unique_ptr library) +void pe_imports::_add(import_library library) { __int64 descriptor_size = 0, extra_size = 0; - library->get_table_size(descriptor_size, extra_size); - const SIZE_T overhead = sizeof(import_library) + 2 * sizeof(import_library*); - if (!_valid || !library->valid() || !bounded_size(descriptor_size) || !bounded_size(extra_size) || + library.get_table_size(descriptor_size, extra_size); + const SIZE_T overhead = 2 * sizeof(import_library); + if (!_valid || !library.valid() || !bounded_size(descriptor_size) || !bounded_size(extra_size) || !table_range(MAX_PE_IMAGE_SIZE, descriptor_size, extra_size) || !table_range(MAX_PE_IMAGE_SIZE, descriptor_size + extra_size, overhead) || !table_range(MAX_PE_IMAGE_SIZE, _owned_size, descriptor_size + extra_size + overhead)) @@ -44,19 +44,27 @@ void pe_imports::_add(std::unique_ptr library) return; } _owned_size += static_cast(descriptor_size + extra_size) + overhead; + if (_libraries.size() == _libraries.capacity()) + _libraries.reserve(_libraries.empty() ? 1 : 2 * _libraries.size()); _libraries.push_back(std::move(library)); } void pe_imports::add_fixup(char* library_name, int ordinal, __int64 rva, bool win64) { if (_valid) - _add(std::unique_ptr(new import_library(library_name, ordinal, rva, win64))); + _add(import_library(library_name, ordinal, rva, win64)); } void pe_imports::add_fixup(char* library_name, char* proc_name, __int64 rva, bool win64) { if (_valid) - _add(std::unique_ptr(new import_library(library_name, proc_name, rva, win64))); + _add(import_library(library_name, proc_name, rva, win64)); +} + +void pe_imports::add_fixup(const export_entry& entry, __int64 rva, bool win64) +{ + if (_valid) + _add(import_library(entry, rva, win64)); } void pe_imports::get_table_size(__int64& descriptor_size, __int64& extra_size) @@ -66,8 +74,8 @@ void pe_imports::get_table_size(__int64& descriptor_size, __int64& extra_size) descriptor_size = extra_size = -1; return; } - for (const auto& library : _libraries) - library->get_table_size(descriptor_size, extra_size); + for (auto& library : _libraries) + library.get_table_size(descriptor_size, extra_size); add_size(descriptor_size, sizeof(IMAGE_IMPORT_DESCRIPTOR)); } @@ -83,9 +91,9 @@ bool pe_imports::build_table(unsigned char* section, __int64 section_size, __int !table_rva(section_rva, 0, section_size)) return false; - for (const auto& library : _libraries) + for (auto& library : _libraries) { - if (!library->build_table(section, section_size, section_rva, descriptor_offset, extra_offset)) + if (!library.build_table(section, section_size, section_rva, descriptor_offset, extra_offset)) return false; } memset(section + static_cast(descriptor_offset), 0, sizeof(IMAGE_IMPORT_DESCRIPTOR)); @@ -126,60 +134,76 @@ pe_imports::pe_imports(unsigned char* image, __int64 image_size, IMAGE_IMPORT_DE void pe_imports::add_descriptor(IMAGE_IMPORT_DESCRIPTOR* descriptor) { if (_valid) - _add(std::unique_ptr(new import_library(descriptor, _win64))); + _add(import_library(descriptor, _win64)); } pe_imports::~pe_imports(void) = default; -import_library::~import_library(void) -{ - delete _descriptor; - delete[] reinterpret_cast(_import_by_name); - delete[] _library_name; - delete _thunk_entry; -} - import_library::import_library(IMAGE_IMPORT_DESCRIPTOR* descriptor, bool win64) { if (descriptor != NULL) - _descriptor = new IMAGE_IMPORT_DESCRIPTOR(*descriptor); + { + _descriptor = *descriptor; + _valid = true; + } } import_library::import_library(char* library_name, int ordinal, __int64 rva, bool win64) +{ + _initialize(library_name, NULL, ordinal, rva, win64); + _copy_names(); +} + +import_library::import_library(char* library_name, char* proc_name, __int64 rva, bool win64) +{ + if (proc_name != NULL) + _initialize(library_name, proc_name, 0, rva, win64); + _copy_names(); +} + +import_library::import_library(const export_entry& entry, __int64 rva, bool win64) +{ + _initialize(entry.library_name, entry.name, entry.ord, rva, win64); +} + +void import_library::_initialize(const char* library_name, const char* proc_name, int ordinal, __int64 rva, bool win64) { if (library_name == NULL || rva < 0 || rva > MAXDWORD) return; const SIZE_T length = strnlen_s(library_name, MAX_PE_IMAGE_SIZE); if (length >= MAX_PE_IMAGE_SIZE) return; - _library_name_len = length + 1; - _library_name = new char[_library_name_len]; - memcpy(_library_name, library_name, _library_name_len); - _descriptor = new IMAGE_IMPORT_DESCRIPTOR(); - _descriptor->TimeDateStamp = MAXDWORD; - _descriptor->ForwarderChain = MAXDWORD; - _descriptor->FirstThunk = static_cast(rva); - _thunk_entry = new IMAGE_THUNK_DATA64(); - _thunk_entry->u1.Ordinal = (win64 ? IMAGE_ORDINAL_FLAG64 : IMAGE_ORDINAL_FLAG32) | (ordinal & 0xffff); + _library_name_len = static_cast(length + 1); + _library_name = library_name; + if (proc_name != NULL) + { + const SIZE_T proc_length = strnlen_s(proc_name, MAX_PE_IMAGE_SIZE); + if (proc_length > MAX_PE_IMAGE_SIZE - sizeof(WORD) - 1 || + !range_fits(MAX_PE_IMAGE_SIZE, _library_name_len, proc_length + 1)) + return; + _proc_name_len = static_cast(proc_length + 1); + _proc_name = proc_name; + } + _descriptor.TimeDateStamp = MAXDWORD; + _descriptor.ForwarderChain = MAXDWORD; + _descriptor.FirstThunk = static_cast(rva); + _thunk_entry.u1.Ordinal = proc_name == NULL ? + (win64 ? IMAGE_ORDINAL_FLAG64 : IMAGE_ORDINAL_FLAG32) | (ordinal & 0xffff) : 0; + _valid = true; } -import_library::import_library(char* library_name, char* proc_name, __int64 rva, bool win64) - : import_library(library_name, 0, rva, win64) +void import_library::_copy_names() { - if (!valid()) + if (!_valid) return; - const SIZE_T length = proc_name == NULL ? MAX_PE_IMAGE_SIZE : strnlen_s(proc_name, MAX_PE_IMAGE_SIZE); - if (length > MAX_PE_IMAGE_SIZE - sizeof(WORD) - 1) + _owned_names.reset(new char[static_cast(_library_name_len) + _proc_name_len]); + memcpy(_owned_names.get(), _library_name, _library_name_len); + _library_name = _owned_names.get(); + if (_proc_name != NULL) { - delete _descriptor; - _descriptor = NULL; - return; + memcpy(_owned_names.get() + _library_name_len, _proc_name, _proc_name_len); + _proc_name = _owned_names.get() + _library_name_len; } - _import_by_name_len = length + 1 + sizeof(WORD); - _import_by_name = reinterpret_cast(new char[_import_by_name_len]); - _import_by_name->Hint = 0; - memcpy(_import_by_name->Name, proc_name, length + 1); - _thunk_entry->u1.AddressOfData = 0; } void import_library::get_table_size(__int64& descriptor_size, __int64& extra_size) @@ -189,8 +213,9 @@ void import_library::get_table_size(__int64& descriptor_size, __int64& extra_siz descriptor_size = extra_size = -1; return; } - add_size(extra_size, _import_by_name_len); - if (_thunk_entry != NULL) + if (_proc_name != NULL) + add_size(extra_size, sizeof(WORD) + _proc_name_len); + if (_library_name != NULL) add_size(extra_size, 2 * sizeof(IMAGE_THUNK_DATA64)); add_size(extra_size, _library_name_len); add_size(descriptor_size, sizeof(IMAGE_IMPORT_DESCRIPTOR)); @@ -208,17 +233,18 @@ bool import_library::build_table(unsigned char* section, __int64 section_size, _ !table_rva(section_rva, 0, section_size)) return false; - IMAGE_IMPORT_DESCRIPTOR descriptor = *_descriptor; + IMAGE_IMPORT_DESCRIPTOR descriptor = _descriptor; __int64 import_name_rva = 0; - if (_import_by_name != NULL) + if (_proc_name != NULL) { - memcpy(section + static_cast(extra_offset), _import_by_name, _import_by_name_len); + memset(section + static_cast(extra_offset), 0, sizeof(WORD)); + memcpy(section + static_cast(extra_offset) + sizeof(WORD), _proc_name, _proc_name_len); import_name_rva = section_rva + extra_offset; - extra_offset += _import_by_name_len; + extra_offset += sizeof(WORD) + _proc_name_len; } - if (_thunk_entry != NULL) + if (_library_name != NULL) { - IMAGE_THUNK_DATA64 thunk = *_thunk_entry; + IMAGE_THUNK_DATA64 thunk = _thunk_entry; if (import_name_rva != 0) thunk.u1.AddressOfData = import_name_rva; memcpy(section + static_cast(extra_offset), &thunk, sizeof(thunk)); diff --git a/pd/pe_imports.h b/pd/pe_imports.h index e71fc81..cc4ac8b 100644 --- a/pd/pe_imports.h +++ b/pd/pe_imports.h @@ -1,6 +1,7 @@ #pragma once #include "windows.h" #include "utils.h" +#include "export_list.h" #include #include @@ -9,24 +10,29 @@ class import_library { - char* _library_name = NULL; - - IMAGE_IMPORT_DESCRIPTOR* _descriptor = NULL; - IMAGE_THUNK_DATA64* _thunk_entry = NULL; // Use 64 bit definition for both 32 and 64 bit modules. - IMAGE_IMPORT_BY_NAME* _import_by_name = NULL; - SIZE_T _import_by_name_len = 0; - SIZE_T _library_name_len = 0; + IMAGE_THUNK_DATA64 _thunk_entry = {}; + IMAGE_IMPORT_DESCRIPTOR _descriptor = {}; + DWORD _library_name_len = 0; + DWORD _proc_name_len = 0; + bool _valid = false; + const char* _library_name = NULL; + const char* _proc_name = NULL; + std::unique_ptr _owned_names; + void _initialize(const char* library_name, const char* proc_name, int ordinal, __int64 rva, bool win64); + void _copy_names(); public: import_library(IMAGE_IMPORT_DESCRIPTOR* descriptor, bool win64); import_library(char* library_name, int ordinal, __int64 rva, bool win64); import_library(char* library_name, char* proc_name, __int64 rva, bool win64); + // Borrowed names remain valid while the immutable export list is alive. + import_library(const export_entry& entry, __int64 rva, bool win64); bool build_table(unsigned char* section, __int64 section_size, __int64 section_rva, __int64 &descriptor_offset, __int64 &extra_offset); void get_table_size(__int64 &descriptor_size, __int64 &extra_size); - bool valid() const { return _descriptor != NULL; } - char* GetName(); - ~import_library(void); + bool valid() const { return _valid; } + import_library(import_library&&) noexcept = default; + import_library& operator=(import_library&&) noexcept = default; import_library(const import_library&) = delete; import_library& operator=(const import_library&) = delete; }; @@ -36,8 +42,8 @@ class pe_imports bool _win64; bool _valid = true; SIZE_T _owned_size = sizeof(IMAGE_IMPORT_DESCRIPTOR); - std::vector> _libraries; - void _add(std::unique_ptr library); + std::vector _libraries; + void _add(import_library library); public: pe_imports(unsigned char* image, __int64 image_size, IMAGE_IMPORT_DESCRIPTOR* imports, bool win64); void add_descriptor(IMAGE_IMPORT_DESCRIPTOR* descriptor); @@ -47,7 +53,7 @@ class pe_imports void add_fixup(char* library_name, int ordinal, __int64 rva, bool win64); void add_fixup(char* library_name, char* proc_name, __int64 rva, bool win64); + void add_fixup(const export_entry& entry, __int64 rva, bool win64); //char* build_table(); ~pe_imports(void); }; - diff --git a/pd/simple.h b/pd/simple.h index fedf6a5..4f19b90 100644 --- a/pd/simple.h +++ b/pd/simple.h @@ -23,12 +23,15 @@ class PD_OPTIONS bool DumpChunks; // Dump loose code chunks bool EntryPointHash; bool ForceReconstructEntryPoint; + bool Reexecution; int NumberOfThreads; __int64 EntryPointOverride; char* output_path; - PD_OPTIONS() + PD_OPTIONS() : ImportRec(true), ForceGenHeader(false), Verbose(false), ReconstructHeaderAsDll(false), + DumpChunks(true), EntryPointHash(true), ForceReconstructEntryPoint(false), Reexecution(false), NumberOfThreads(16), + EntryPointOverride(-1) { output_path = new char[1]; strcpy(output_path,""); diff --git a/pd/system_work.h b/pd/system_work.h new file mode 100644 index 0000000..348044c --- /dev/null +++ b/pd/system_work.h @@ -0,0 +1,42 @@ +#pragma once + +#include "work_pool.h" +#include +#include +#include +#include +#include + +inline std::vector system_processes() +{ + HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0); + if (snapshot == INVALID_HANDLE_VALUE) + throw std::runtime_error("Unable to enumerate processes: Windows error " + std::to_string(GetLastError())); + std::unique_ptr close(snapshot, CloseHandle); + PROCESSENTRY32 entry = {}; + entry.dwSize = sizeof(entry); + std::vector result; + BOOL more = Process32First(snapshot, &entry); + while (more) + { + result.push_back(entry.th32ProcessID); + more = Process32Next(snapshot, &entry); + } + if (GetLastError() != ERROR_NO_MORE_FILES) + throw std::runtime_error("Incomplete process enumeration: Windows error " + std::to_string(GetLastError())); + return result; +} + +template +void system_work(work_pool& pool, bool rescan, Enumerate enumerate, Action action) +{ + std::unordered_set seen; + // Drain all producer and module work before the one final discovery pass. + for (int pass = 0; pass < (rescan ? 2 : 1); ++pass) + { + for (DWORD pid : enumerate()) + if (seen.insert(pid).second) + pool.submit([=, &pool] { action(pid, pool); }); + pool.wait(); + } +} diff --git a/pd/version.h b/pd/version.h new file mode 100644 index 0000000..1112674 --- /dev/null +++ b/pd/version.h @@ -0,0 +1,4 @@ +#pragma once + +#define PD_VERSION_STRING "3.0.0" +#define PD_VERSION_RESOURCE 3,0,0,0 diff --git a/pd/work_pool.h b/pd/work_pool.h new file mode 100644 index 0000000..7ef617f --- /dev/null +++ b/pd/work_pool.h @@ -0,0 +1,157 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Producers may wait for module work, but help execute only leaf jobs while waiting. +// This avoids both pool starvation and recursively nesting process lifetimes. +class work_pool +{ + struct group + { + size_t pending = 0; + std::exception_ptr error; + }; + struct job + { + std::function run; + group* completion; + }; + std::mutex mutex_; + std::condition_variable changed_; + std::deque producers_, leaves_; + std::vector workers_; + group root_; + bool stopping_ = false; + + void execute(job task) + { + std::exception_ptr error; + try { task.run(); } + catch (...) { error = std::current_exception(); } + { + std::lock_guard lock(mutex_); + if (error && !task.completion->error) + task.completion->error = error; + --task.completion->pending; + } + changed_.notify_all(); + } + + void worker() + { + for (;;) + { + job task; + { + std::unique_lock lock(mutex_); + changed_.wait(lock, [&] { return stopping_ || !leaves_.empty() || !producers_.empty(); }); + if (stopping_ && leaves_.empty() && producers_.empty()) + return; + auto& queue = leaves_.empty() ? producers_ : leaves_; + task = std::move(queue.front()); + queue.pop_front(); + } + execute(std::move(task)); + } + } + + void wait(group& completion, bool help) + { + std::unique_lock lock(mutex_); + while (completion.pending != 0) + { + if (help && !leaves_.empty()) + { + auto task = std::move(leaves_.front()); + leaves_.pop_front(); + lock.unlock(); + execute(std::move(task)); + lock.lock(); + } + else + changed_.wait(lock); + } + if (completion.error) + std::rethrow_exception(completion.error); + } + +public: + explicit work_pool(int threads) + { + if (threads < 1) + throw std::invalid_argument("worker count must be positive"); + try + { + for (int i = 0; i < threads; ++i) + workers_.emplace_back([this] { worker(); }); + } + catch (...) + { + { + std::lock_guard lock(mutex_); + stopping_ = true; + } + changed_.notify_all(); + for (auto& thread : workers_) thread.join(); + throw; + } + } + + ~work_pool() + { + { + std::lock_guard lock(mutex_); + stopping_ = true; + } + changed_.notify_all(); + for (auto& thread : workers_) thread.join(); + } + + void submit(std::function task) + { + { + std::lock_guard lock(mutex_); + producers_.push_back({std::move(task), &root_}); + ++root_.pending; + } + changed_.notify_one(); + } + + void wait() { wait(root_, false); } + + template + void parallel_for(const std::vector& items, F action) + { + // Bound queued metadata and captured references even for huge process maps. + for (size_t first = 0; first < items.size(); first += 256) + { + group completion; + try + { + std::lock_guard lock(mutex_); + for (size_t i = first; i < (std::min)(items.size(), first + 256); ++i) + { + leaves_.push_back({[&, i] { action(items[i]); }, &completion}); + ++completion.pending; + } + } + catch (...) + { + const auto error = std::current_exception(); + changed_.notify_all(); + wait(completion, true); + std::rethrow_exception(error); + } + changed_.notify_all(); + wait(completion, true); + } + } +}; diff --git a/tests/baseline_tests.cpp b/tests/baseline_tests.cpp index e27a387..6387355 100644 --- a/tests/baseline_tests.cpp +++ b/tests/baseline_tests.cpp @@ -60,13 +60,13 @@ static void pe_roundtrip(bool win64, bool imports) std::vector bytes = output.read(); require(bytes.size() >= 0x2000, "dump truncated"); require(bytes[0] == 'M' && bytes[1] == 'Z' && bytes[0x1000] == 0x90, "dump data changed"); - const DWORD expected_crc = imports ? (win64 ? 0x11056f9c : 0xd6ca78d0) : (win64 ? 0x0f6a2ed3 : 0xfd941dbd); - require(bytes.size() == (imports ? 12288 : 8192), "baseline dump size changed"); - require(crc32buf(reinterpret_cast(bytes.data()), bytes.size()) == expected_crc, "baseline dump bytes changed"); - require(header.get_hash() == (imports ? 0xcd4757e79d476f71ULL : 0xcd8757e79d476f71ULL), "baseline PE hash changed"); printf("FINGERPRINT PE%s imports=%d size=%zu crc=%08lx hash=%016llx\n", win64 ? "64" : "32", imports, bytes.size(), crc32buf(reinterpret_cast(bytes.data()), bytes.size()), header.get_hash()); + const DWORD expected_crc = imports ? (win64 ? 0x64a456d2 : 0xd9e601a9) : (win64 ? 0x0f6a2ed3 : 0xfd941dbd); + require(bytes.size() == (imports ? 12288 : 8192), "baseline dump size changed"); + require(crc32buf(reinterpret_cast(bytes.data()), bytes.size()) == expected_crc, "baseline dump bytes changed"); + require(header.get_hash() == (imports ? 0x391c57a138fafe6cULL : 0xcd8757e79d476f71ULL), "baseline PE hash changed"); } static void import_table_roundtrip(bool win64) @@ -132,6 +132,18 @@ int main(int argc, char** argv) _CrtSetReportFile(_CRT_ASSERT, _CRTDBG_FILE_STDERR); #endif setvbuf(stdout, NULL, _IONBF, 0); + if (argc > 1 && strcmp(argv[1], "--benchmark") == 0) + return run_performance_benchmark(argc - 2, argv + 2); + if (argc > 1 && strcmp(argv[1], "--reexecute") == 0) + return run_reexecution_probe(argc - 2, argv + 2); + if (argc == 2 && strcmp(argv[1], "--reexecution-fixture") == 0) + return run_reexecution_fixture(); + if (argc > 1 && strcmp(argv[1], "--entrypoint-corpus") == 0) + return run_entrypoint_benchmark(argc - 2, argv + 2); + if (argc == 2 && strcmp(argv[1], "--system-fixture") == 0) + return run_system_fixture(); + if (argc > 1 && strcmp(argv[1], "--system-work") == 0) + return run_system_benchmark(argc - 2, argv + 2); std::vector tests = { {"crc-vectors", crc_vectors}, {"file-stream", file_stream_roundtrip}, {"process-stream", process_stream_roundtrip}, @@ -153,6 +165,12 @@ int main(int argc, char** argv) append_name_safety_tests(tests); append_hook_tests(tests); append_output_tests(tests); + append_reconstruction_tests(tests); + append_pipeline_tests(tests); + append_performance_tests(tests); + append_reexecution_tests(tests); + append_scheduling_tests(tests); + append_entrypoint_tests(tests); } int failures = 0, count = 0; for (const auto& test : tests) diff --git a/tests/benchmark_database.py b/tests/benchmark_database.py new file mode 100644 index 0000000..7077e74 --- /dev/null +++ b/tests/benchmark_database.py @@ -0,0 +1,90 @@ +"""Compare database-generation file traversal on a fixed, privately staged PE corpus.""" + +import argparse +import hashlib +import json +from pathlib import Path +import shutil +import statistics +import struct +import subprocess +import tempfile +import time + + +def database_sets(folder): + result = {} + for name in ("clean.hashes", "entrypoints.hashes", "shortentrypoints.hashes"): + data = (folder / name).read_bytes() + if len(data) % 8: + raise ValueError(f"Truncated database: {name}") + result[name] = sorted({item[0] for item in struct.iter_unpack(" len(data): + raise ValueError("truncated PE header") + return struct.unpack_from(fmt, data, offset) + + if data[:2] != b"MZ": + raise ValueError("missing DOS signature") + nt, = unpack(" len(data): + problems.append("headers_outside_file") + for i in range(count): + offset = optional + optional_size + i * 40 + name, virtual_size, rva, raw_size, raw, _, _, _, _, flags = unpack("<8sIIIIIIHHI", offset) + section = dict(name=name.rstrip(b"\0").decode("ascii", "replace"), + rva=rva, virtual_size=virtual_size, raw=raw, raw_size=raw_size, flags=flags) + sections.append(section) + if raw_size and raw + raw_size > len(data): + problems.append(f"section_{i}_outside_file") + if rva + virtual_size > image_size: + problems.append(f"section_{i}_outside_image") + if raw_size and file_alignment and (raw % file_alignment or raw_size % file_alignment): + problems.append(f"section_{i}_raw_alignment") + return dict(machine=machine, image_size=image_size, headers_size=headers_size, + section_alignment=section_alignment, file_alignment=file_alignment, + sections=sections, problems=problems) + + +def file_metrics(path): + with path.open("rb") as stream, mmap.mmap(stream.fileno(), 0, access=mmap.ACCESS_READ) as data: + info = pe_info(data) + digest, zeros = hashlib.sha256(), 0 + for offset in range(0, len(data), 1024 * 1024): + chunk = data[offset:offset + 1024 * 1024] + digest.update(chunk) + zeros += chunk.count(0) + info.update(bytes=len(data), zero_bytes=zeros, zero_fraction=zeros / len(data), + sha256=digest.hexdigest()) + return info + + +def compare_code(original, dump, source, reconstructed): + total, equal, lost = 0, 0, 0 + with original.open("rb") as first, dump.open("rb") as second: + for section in source["sections"]: + if not section["flags"] & 0x20000000: + continue + match = next((s for s in reconstructed["sections"] if s["rva"] == section["rva"]), None) + size = min(section["raw_size"], section["virtual_size"] or section["raw_size"]) + for offset in range(0, size, BLOCK): + count = min(BLOCK, size - offset) + first.seek(section["raw"] + offset) + expected = first.read(count) + actual = b"" + if match and offset + count <= match["raw_size"]: + second.seek(match["raw"] + offset) + actual = second.read(count) + total += 1 + equal += actual == expected + lost += bool(expected.strip(b"\0")) and not actual.strip(b"\0") + return dict(code_blocks=total, identical_code_blocks=equal, lost_code_blocks=lost, + code_block_match=equal / total if total else None) + + +def compare_reports(before, after): + if before["imports"] != after["imports"]: + raise ValueError("comparison requires the same import reconstruction setting") + def identity(row): + return row["pid"], row.get("started"), row.get("path") + previous = {identity(row): row for row in before["results"]} + pairs, regressions = [], [] + for row in after["results"]: + old = previous.get(identity(row)) + if not old or old["status"] != "ok": + continue + if row["status"] in ("dump_failed", "invalid_dump", "timeout"): + regressions.append(dict(pid=row["pid"], name=row["name"], status=row["status"])) + if row["status"] == "ok" and old["original"]["sha256"] == row["original"]["sha256"]: + pairs.append((old, row)) + result = dict( + comparable_processes=len(pairs), dump_regressions=regressions, + before_bytes=sum(old["dump"]["bytes"] for old, _ in pairs), + after_bytes=sum(new["dump"]["bytes"] for _, new in pairs), + before_structural_failures=sum(bool(old["dump"]["problems"]) for old, _ in pairs), + after_structural_failures=sum(bool(new["dump"]["problems"]) for _, new in pairs), + before_lost_code_blocks=sum(old["lost_code_blocks"] for old, _ in pairs), + after_lost_code_blocks=sum(new["lost_code_blocks"] for _, new in pairs), + before_median_seconds=statistics.median(old["seconds"] for old, _ in pairs) if pairs else None, + after_median_seconds=statistics.median(new["seconds"] for _, new in pairs) if pairs else None) + quality_pairs = [(a["quality"], b["quality"]) for a, b in pairs if "quality" in a and "quality" in b] + if quality_pairs: + matched = [(a, b) for a, b in quality_pairs if a["matching_file_layout"] and b["matching_file_layout"]] + witnessed = [(a["memory_witness"], b["memory_witness"]) for a, b in matched + if a["memory_witness"] is not None and b["memory_witness"] is not None] + result["quality"] = dict( + compared_images=len(quality_pairs), + matching_layout_images=len(matched), + before_iat_different_slots=sum(a["iat"]["different_slots"] for a, _ in matched), + after_iat_different_slots=sum(b["iat"]["different_slots"] for _, b in matched), + before_duplicate_iat_slots=sum(a["duplicate_iat_slots"] for a, _ in quality_pairs), + after_duplicate_iat_slots=sum(b["duplicate_iat_slots"] for _, b in quality_pairs), + witnessed_images=len(witnessed), + before_stable_memory_differences=sum(a["stable_different_bytes"] for a, _ in witnessed), + after_stable_memory_differences=sum(b["stable_different_bytes"] for _, b in witnessed)) + return result + + +def snapshot(): + command = r""" + $ErrorActionPreference = 'Stop' + [Console]::OutputEncoding = [Text.UTF8Encoding]::new($false) + $rows = foreach ($p in Get-Process) { + $row = [ordered]@{pid=$p.Id; name=$p.ProcessName; status='ready'} + try { + $row.started = $p.StartTime.ToUniversalTime().Ticks.ToString() + $module = $p.MainModule + if (!$module -or !$module.FileName) { throw 'No accessible main module' } + $row.path = $module.FileName + $row.base = $module.BaseAddress.ToInt64().ToString() + } catch { $row.status='unavailable'; $row.reason=$_.Exception.Message } + [pscustomobject]$row + } + ConvertTo-Json -InputObject @($rows) -Depth 4 -Compress + """ + result = subprocess.run(["powershell.exe", "-NoProfile", "-NonInteractive", "-Command", command], + capture_output=True, check=True, encoding="utf-8-sig") + return json.loads(result.stdout) + + +@contextmanager +def process_identity(entry): + kernel = ctypes.WinDLL("kernel32", use_last_error=True) + kernel.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel.OpenProcess.restype = wintypes.HANDLE + kernel.GetProcessTimes.argtypes = [wintypes.HANDLE] + [ctypes.POINTER(wintypes.FILETIME)] * 4 + kernel.GetProcessTimes.restype = wintypes.BOOL + kernel.GetExitCodeProcess.argtypes = [wintypes.HANDLE, ctypes.POINTER(wintypes.DWORD)] + kernel.GetExitCodeProcess.restype = wintypes.BOOL + kernel.CloseHandle.argtypes = [wintypes.HANDLE] + kernel.CloseHandle.restype = wintypes.BOOL + handle = kernel.OpenProcess(0x1000, False, entry["pid"]) # QUERY_LIMITED_INFORMATION + if not handle: + yield False + return + try: + times = [wintypes.FILETIME() for _ in range(4)] + code = wintypes.DWORD() + valid = kernel.GetProcessTimes(handle, *(ctypes.byref(t) for t in times)) + ticks = (times[0].dwHighDateTime << 32) + times[0].dwLowDateTime + 504911232000000000 + alive = kernel.GetExitCodeProcess(handle, ctypes.byref(code)) and code.value == 259 + # Retain the process object through the dump so its PID cannot be reused mid-run. + yield bool(valid and alive and str(ticks) == entry["started"]) + finally: + kernel.CloseHandle(handle) + + +def benchmark(entry, args, root): + row = dict(entry, imports=args.imports) + if entry["pid"] in args.skip_pid: + return dict(row, status="excluded", reason="explicit --skip-pid") + if entry["status"] != "ready": + return row + with process_identity(entry) as same: + if not same: + return dict(row, status="exited_or_replaced_or_inaccessible") + return benchmark_running(entry, args, root) + + +def benchmark_running(entry, args, root): + row = dict(entry, imports=args.imports) + original = Path(entry["path"]) + try: + source = file_metrics(original) + except (OSError, ValueError, struct.error) as error: + return dict(row, status="original_unavailable", reason=str(error)) + row["original"] = source + if source["image_size"] > MAX_IMAGE: + return dict(row, status="image_limit", reason="image exceeds the dumper's 256 MiB limit") + if shutil.disk_usage(root).free < 2 * MAX_IMAGE: + raise RuntimeError("less than 512 MiB free; refusing further dumps") + with tempfile.TemporaryDirectory(prefix=f"pid{entry['pid']}-", dir=root) as temporary: + folder = Path(temporary) + command = [str(args.exe), "-pid", str(entry["pid"]), "-a", hex(int(entry["base"])), + "-o", str(folder), "-db", "ignore", "-nep", "-nc", "-nt", "-nh", "-v"] + if not args.imports: + command.append("-ni") + try: + witness_before = snapshot_sections(entry, source) if args.quality and args.memory_witness else None + except (OSError, ValueError) as error: + return dict(row, status="witness_unavailable", reason=str(error)) + start = time.perf_counter() + try: + result = subprocess.run(command, cwd=folder, capture_output=True, + timeout=args.timeout, errors="replace") + except subprocess.TimeoutExpired: + # subprocess.run terminates only the dumper child, never the target. + return dict(row, status="timeout", seconds=round(time.perf_counter() - start, 4)) + row.update(seconds=round(time.perf_counter() - start, 4), exit_code=result.returncode) + row["diagnostics"] = [line for line in (result.stdout + result.stderr).splitlines() + if any(word in line.lower() for word in ("warning", "error", "failed", "invalid"))] + files = list(folder.glob("*.exe")) + list(folder.glob("*.dll")) + list(folder.glob("*.bin")) + if result.returncode or len(files) != 1: + return dict(row, status="dump_failed", outputs=len(files)) + try: + output = file_metrics(files[0]) + row.update(dump=output, size_ratio=output["bytes"] / source["bytes"], + **compare_code(original, files[0], source, output)) + if args.quality: + try: + with original.open("rb") as first, files[0].open("rb") as second, \ + mmap.mmap(first.fileno(), 0, access=mmap.ACCESS_READ) as source_bytes, \ + mmap.mmap(second.fileno(), 0, access=mmap.ACCESS_READ) as dump_bytes: + row["quality"] = compare_quality(source_bytes, dump_bytes, source, output) + if witness_before is not None and row["quality"]["matching_file_layout"]: + row["quality"] = compare_quality(source_bytes, dump_bytes, source, output, witness_before, + snapshot_sections(entry, source)) + except (ValueError, struct.error) as error: + return dict(row, status="quality_unavailable", reason=str(error)) + except OSError as error: + return dict(row, status="analysis_unavailable", reason=str(error)) + except (ValueError, struct.error) as error: + return dict(row, status="invalid_dump", reason=str(error)) + row["status"] = "ok" + row["suspect"] = bool(output["problems"] or row["lost_code_blocks"] or + (row["size_ratio"] > 4 and output["zero_fraction"] > 0.75)) + if "quality" in row: + quality = row["quality"] + row["suspect"] |= bool(quality["import_errors"] or quality["directory_errors"] or + (quality["matching_file_layout"] and quality["iat"]["different_slots"]) or + (quality["memory_witness"] or {}).get("stable_different_bytes")) + return row + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--exe", type=Path, required=True, help="built x64 Process Dump executable to exercise") + parser.add_argument("--report", type=Path, required=True, help="new local JSON report (never overwritten)") + parser.add_argument("--snapshot", type=Path, help="reuse a previous report's process inventory") + parser.add_argument("--compare", type=Path, help="include paired metrics against this earlier report") + parser.add_argument("--pid", type=int, action="append", help="restrict to these PIDs") + parser.add_argument("--skip-pid", type=int, action="append", default=[], + help="record an exclusion, e.g. a dump already blocked by security software") + parser.add_argument("--name", help="case-insensitive process-name substring") + parser.add_argument("--imports", action="store_true", help="enable aggressive import reconstruction") + parser.add_argument("--quality", action="store_true", help="compare RVA bytes and PE directories in detail") + parser.add_argument("--memory-witness", action="store_true", + help="independently read sections pagewise before/after dumping (requires --quality)") + parser.add_argument("--unique-originals", action="store_true", help="one process per original path for a bounded corpus") + parser.add_argument("--timeout", type=float, default=30, help="seconds per dumper child") + args = parser.parse_args() + args.exe = args.exe.resolve(strict=True) + args.report = args.report.resolve() + if args.report.exists(): + parser.error("report already exists; choose a new path") + if args.timeout <= 0: + parser.error("--timeout must be positive") + if args.memory_witness and not args.quality: + parser.error("--memory-witness requires --quality") + inventory = json.loads(args.snapshot.read_text(encoding="utf-8"))["inventory"] if args.snapshot else snapshot() + selected = [p for p in inventory if (not args.pid or p["pid"] in args.pid) and + (not args.name or args.name.lower() in p["name"].lower())] + if args.unique_originals: + seen, unique = set(), [] + for entry in selected: + identity = entry.get("path", "").lower() + if identity and identity not in seen and entry["pid"] not in args.skip_pid: + unique.append(entry) + seen.add(identity) + selected = unique + report = dict(schema=1, created=datetime.now(timezone.utc).isoformat(), + executable=str(args.exe), executable_sha256=hashlib.sha256(args.exe.read_bytes()).hexdigest(), + imports=args.imports, quality=args.quality, memory_witness=args.memory_witness, + unique_originals=args.unique_originals, + excluded_pids=args.skip_pid, inventory=inventory, results=[]) + previous = json.loads(args.compare.read_text(encoding="utf-8")) if args.compare else None + if previous and previous["imports"] != args.imports: + parser.error("--compare must have the same import reconstruction setting") + args.report.parent.mkdir(parents=True, exist_ok=True) + with args.report.open("x", encoding="utf-8") as report_file: + with tempfile.TemporaryDirectory(prefix="pd-bench-") as root: + try: + for number, entry in enumerate(selected, 1): + row = benchmark(entry, args, root) + report["results"].append(row) + print(f"{number}/{len(selected)} {entry['name']} PID={entry['pid']} " + f"{row['status']} ratio={row.get('size_ratio', 0):.2f} " + f"lost_code={row.get('lost_code_blocks', 0)}", flush=True) + finally: + report["summary"] = dict(Counter(row["status"] for row in report["results"])) + report["suspect_count"] = sum(row.get("suspect", False) for row in report["results"]) + if previous: + report["comparison"] = compare_reports(previous, report) + json.dump(report, report_file, indent=2) + print(json.dumps(dict(summary=report["summary"], suspect_count=report["suspect_count"]))) + if "comparison" in report: + print(json.dumps(report["comparison"])) + + +if __name__ == "__main__": + main() diff --git a/tests/benchmark_performance.py b/tests/benchmark_performance.py new file mode 100644 index 0000000..9d1f203 --- /dev/null +++ b/tests/benchmark_performance.py @@ -0,0 +1,103 @@ +"""Opt-in controlled native dumping benchmarks; never dumps another process.""" + +import argparse +from datetime import datetime, timezone +import hashlib +import json +import os +from pathlib import Path +import platform +import statistics +import subprocess + + +WORKLOADS = ("dense", "sparse", "imports", "entrypoint", "entrypoint-empty", "database", "exports") +METRICS = ("wall_ms", "worker_ms", "capture_ms", "reconstruct_ms", "write_ms", "cleanup_ms", "service_ms", "cpu_ms", + "cpu_cycles", "allocations", "allocated_bytes", "peak_working_set") + + +def read_result(output, workload, threads, jobs): + rows = [json.loads(line[5:]) for line in output.splitlines() if line.startswith("PERF ")] + if len(rows) != 1: + raise ValueError("benchmark did not return exactly one PERF result") + row = rows[0] + if (row["workload"], row["threads"], row["jobs"]) != (workload, threads, jobs): + raise ValueError("benchmark returned a different workload") + return row + + +def summarize(rows): + if not rows: + raise ValueError("benchmark has no samples") + identities = {(r["workload"], r["threads"], r["jobs"], r["profiled_allocations"], + r["output_bytes"], r["normalized_crc32"]) for r in rows} + if len(identities) != 1: + raise ValueError("workload configuration or reconstructed bytes differ") + return dict(workload=rows[0]["workload"], threads=rows[0]["threads"], jobs=rows[0]["jobs"], + output_bytes=rows[0]["output_bytes"], normalized_crc32=rows[0]["normalized_crc32"], + profiled_allocations=rows[0]["profiled_allocations"], samples=len(rows), + **{metric: statistics.median(r[metric] for r in rows) for metric in METRICS}) + + +def compare(before, after): + for key in ("workload", "threads", "jobs", "profiled_allocations", "output_bytes", "normalized_crc32"): + if before[key] != after[key]: + raise ValueError(f"before/after benchmark differs: {key}") + return dict(workload=after["workload"], threads=after["threads"], + equivalent_output=True if after["output_bytes"] else None, + **{metric + "_ratio": after[metric] / before[metric] if before[metric] else None + for metric in METRICS}) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--exe", type=Path, required=True, help="Release pd_tests.exe with benchmark mode") + parser.add_argument("--baseline", type=Path, help="saved reference pd_tests.exe built with the same benchmark harness") + parser.add_argument("--report", type=Path, required=True, help="new local JSON report (never overwritten)") + parser.add_argument("--workload", choices=WORKLOADS, action="append") + parser.add_argument("--threads", nargs="+", type=int, default=[1, 4]) + parser.add_argument("--jobs", type=int, default=8) + parser.add_argument("--repeats", type=int, default=5) + args = parser.parse_args() + if args.repeats < 1 or not 1 <= args.jobs <= 256 or any(not 1 <= t <= min(8, args.jobs) for t in args.threads): + parser.error("repeats must be positive; jobs 1..256; threads 1..min(8, jobs)") + executables = {"after": args.exe.resolve(strict=True)} + if args.baseline: + executables["before"] = args.baseline.resolve(strict=True) + report = dict(created=datetime.now(timezone.utc).isoformat(), platform=platform.platform(), + cpu_count=os.cpu_count(), executables={name: dict(path=str(path), + sha256=hashlib.sha256(path.read_bytes()).hexdigest()) for name, path in executables.items()}, + samples=[], summaries=[], comparisons=[]) + args.report.parent.mkdir(parents=True, exist_ok=True) + with args.report.open("x", encoding="utf-8") as stream: + try: + for workload in args.workload or WORKLOADS: + for threads in args.threads: + samples = {name: [] for name in executables} + for repetition in range(args.repeats): + order = list(executables) + if repetition % 2: + order.reverse() + for name in order: + completed = subprocess.run([str(executables[name]), "--benchmark", workload, + str(threads), str(args.jobs)], + capture_output=True, text=True, timeout=120) + if completed.returncode: + raise RuntimeError(f"{name} {workload}/{threads} failed:\n{completed.stdout}\n{completed.stderr}") + result = read_result(completed.stdout, workload, threads, args.jobs) + samples[name].append(result) + report["samples"].append(dict(version=name, repetition=repetition, **result)) + summaries = {name: summarize(rows) for name, rows in samples.items()} + report["summaries"].extend(dict(version=name, **row) for name, row in summaries.items()) + if "before" in summaries: + comparison = compare(summaries["before"], summaries["after"]) + report["comparisons"].append(comparison) + print(json.dumps(comparison), flush=True) + else: + print(json.dumps(summaries["after"]), flush=True) + finally: + json.dump(report, stream, indent=2) + + +if __name__ == "__main__": + main() diff --git a/tests/benchmark_reexecution.py b/tests/benchmark_reexecution.py new file mode 100644 index 0000000..1bcefcc --- /dev/null +++ b/tests/benchmark_reexecution.py @@ -0,0 +1,198 @@ +"""Opt-in execution of explicitly selected benign native programs and their dumps.""" + +import argparse +from datetime import datetime, timezone +import hashlib +import json +import os +import re +import platform +from pathlib import Path +import shutil +import subprocess +import tempfile + +from benchmark_dumps import pe_info + + +def milestones(row): + return { + "loader_breakpoint": bool(row.get("loader_breakpoint")), + "entrypoint": bool(row.get("entrypoint")), + "gui_input_idle": bool(row.get("input_idle")), + "fixture_thread": "PD_REEXEC_THREAD" in row.get("markers", []), + "fixture_complete": "PD_REEXEC_DONE" in row.get("markers", []), + "clean_exit": bool(row.get("entrypoint")) and not row.get("timed_out") and row.get("exit_code") == 0, + } + + +def progress(row): + if "PD_REEXEC_DONE" in row.get("markers", []): + return "fixture_complete" + if row.get("input_idle"): + return "gui_input_idle" + if row.get("entrypoint"): + return "entrypoint" + if row.get("loader_breakpoint"): + return "loader_breakpoint" + return "process_created" + + +def compare_reports(before, after): + identity = ("schema", "application", "application_sha256", "probe_sha256", "loader_snaps", + "arguments", "cwd", "imports", "phase", "milliseconds", "mui", "os") + if any(before.get(key) != after.get(key) for key in identity): + raise ValueError("Reexecution comparison requires the same application, arguments, capture stage, resources, OS and duration") + if before["status"] != "compared" or after["status"] != "compared": + return {"comparable": False, "reason": "capture unavailable"} + outcome = ("status", "entrypoint", "input_idle", "timed_out", "exit_code", "stdout_sha256", "markers") + for control in ("original", "disk_copy"): + if any(report.get(control, {}).get("status") != "observed" for report in (before, after)): + return {"comparable": False, "reason": f"{control} launch or probe failure"} + if any(not row.get("entrypoint") or not (row.get("input_idle") or milestones(row)["clean_exit"]) + for row in (before[control], after[control])): + return {"comparable": False, "reason": f"{control} did not reach GUI idle or exit cleanly"} + if any( + before[control].get(key) != after[control].get(key) for key in outcome): + return {"comparable": False, "reason": f"{control} outcome changed"} + old, new = before["reconstructed"], after["reconstructed"] + if old["status"] != "observed" or new["status"] != "observed": + return {"comparable": False, "reason": "launch or probe failure"} + first, second = milestones(old), milestones(new) + return {"comparable": True, "before": first, "after": second, + "gained": [key for key in first if second[key] and not first[key]], + "lost": [key for key in first if first[key] and not second[key]], + "before_exit": old["exit_code"], "after_exit": new["exit_code"], + "before_same_stdout": before["same_stdout"], "after_same_stdout": after["same_stdout"]} + + +def same_stdout(first, second): + return (first.get("status") == second.get("status") == "observed" and + bool(first.get("stdout_sha256")) and first["stdout_sha256"] == second.get("stdout_sha256")) + + +def run_probe(args, executable, report, phase, output): + command = [str(args.probe), "--reexecute", str(executable), str(report), str(args.milliseconds), + phase, str(args.dumper), str(output), str(int(args.imports)), *args.arg] + environment = os.environ.copy() + environment["PD_REEXEC_PREPARE"] = "1" if args.prepare else "0" + environment["PD_REEXEC_LOADER_SNAPS"] = "1" if args.loader_snaps else "0" + result = subprocess.run(command, capture_output=True, timeout=args.milliseconds / 1000 + 45, env=environment, + cwd=args.cwd, encoding="utf-8", errors="replace") + if result.returncode != 0 or not report.exists(): + return {"status": "probe_failed", "returncode": result.returncode, + "diagnostic": (result.stderr + result.stdout)[-8192:]} + row = json.loads(report.read_text(encoding="utf-8")) + stdout = Path(str(report) + ".stdout").read_bytes() + row.update(status="observed", progress=progress(row), stdout_bytes=len(stdout), + stdout_sha256=hashlib.sha256(stdout).hexdigest()) + dump_log = Path(str(report) + ".dump.log") + if dump_log.exists(): + row["dump_log"] = dump_log.read_text(encoding="utf-8", errors="replace")[-16384:] + return row + + +def benchmark(args): + original = args.app.read_bytes() + metadata = pe_info(original) + probe = pe_info(args.probe.read_bytes()) + if metadata["machine"] != probe["machine"]: + raise ValueError("The probe and application must have the same architecture") + result = { + "schema": 1, "timestamp": datetime.now(timezone.utc).isoformat(), + "application": str(args.app), "application_sha256": hashlib.sha256(original).hexdigest(), + "dumper_sha256": hashlib.sha256(args.dumper.read_bytes()).hexdigest(), + "probe_sha256": hashlib.sha256(args.probe.read_bytes()).hexdigest(), + "arguments": args.arg, "cwd": str(args.cwd), "imports": args.imports, + "phase": args.phase, "milliseconds": args.milliseconds, "prepare": args.prepare, + "loader_snaps": args.loader_snaps, + "os": platform.version(), + } + if args.mui: + if args.mui.name.lower() != args.app.name.lower() + ".mui" or not re.fullmatch(r"[A-Za-z0-9-]+", args.mui.parent.name): + raise ValueError("--mui must name the app's .mui file inside a locale directory") + result["mui"] = {"locale": args.mui.parent.name, "sha256": hashlib.sha256(args.mui.read_bytes()).hexdigest()} + with tempfile.TemporaryDirectory(prefix="pd-rerun-") as temporary: + root = Path(temporary) + output = root / "capture" + output.mkdir() + source = run_probe(args, args.app, root / "source.json", args.phase, output) + result["original"] = source + for name in ("control", "reconstructed"): + target = root / name + target.mkdir() + if args.mui: + locale = target / args.mui.parent.name + locale.mkdir() + shutil.copyfile(args.mui, locale / (args.app.name + ".mui")) + control = root / "control" / args.app.name + shutil.copyfile(args.app, control) + result["disk_copy"] = run_probe(args, control, root / "control.json", "none", output) + dumps = list(output.glob("*.exe")) + if source["status"] != "observed" or not source.get("captured") or len(dumps) != 1: + result["status"] = "capture_unavailable" + result["dump_count"] = len(dumps) + return result + # Preserve the original basename without changing any dumped bytes. + launch = root / "reconstructed" / args.app.name + dumps[0].rename(launch) + dumped = launch.read_bytes() + result["dump_sha256"] = hashlib.sha256(dumped).hexdigest() + result["dump_bytes"] = len(dumped) + result["dump_pe"] = pe_info(dumped) + result["reconstructed"] = run_probe(args, launch, root / "rerun.json", "none", output) + result["status"] = "compared" + result["same_stdout"] = same_stdout(source, result["reconstructed"]) + result["same_copy_stdout"] = same_stdout(result["disk_copy"], result["reconstructed"]) + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--app", type=Path, required=True) + parser.add_argument("--probe", type=Path, required=True, help="same-architecture pd_tests.exe") + parser.add_argument("--dumper", type=Path, required=True) + parser.add_argument("--report", type=Path, required=True) + parser.add_argument("--arg", action="append", default=[]) + parser.add_argument("--cwd", type=Path, default=Path.cwd()) + parser.add_argument("--phase", choices=["entry", "ready", "idle"], default="entry") + parser.add_argument("--milliseconds", type=int, default=3000) + parser.add_argument("--imports", action="store_true") + parser.add_argument("--prepare", action="store_true", help="pass experimental -reexec to the dumper") + parser.add_argument("--loader-snaps", action="store_true", help="collect per-child Windows loader diagnostics") + parser.add_argument("--mui", type=Path, help="copy this app's trusted locale\\app.exe.mui beside both staged launches") + parser.add_argument("--compare", type=Path, help="compare against a report using the same application and controls") + parser.add_argument("--acknowledge-execution", action="store_true", + help="confirm the selected executable and its dumps are benign and may execute") + args = parser.parse_args() + if not args.acknowledge_execution: + parser.error("--acknowledge-execution is required; never select malware or an unknown process image") + if not 100 <= args.milliseconds <= 30000: + parser.error("--milliseconds must be 100..30000") + for name in ("app", "probe", "dumper", "report", "cwd"): + setattr(args, name, getattr(args, name).resolve()) + if args.mui: + args.mui = args.mui.resolve() + if shutil.disk_usage(tempfile.gettempdir()).free < 512 * 1024 * 1024: + raise RuntimeError("Less than 512 MiB free; refusing execution benchmark") + before = json.loads(args.compare.read_text(encoding="utf-8")) if args.compare else None + with args.report.open("x", encoding="utf-8") as report: + try: + result = benchmark(args) + if before: + result["comparison"] = compare_reports(before, result) + except (OSError, ValueError, subprocess.TimeoutExpired) as error: + json.dump({"status": "benchmark_failed", "diagnostic": str(error)}, report, indent=2) + raise + json.dump(result, report, indent=2) + report.write("\n") + print(json.dumps({key: result[key] for key in ("status", "phase", "imports")})) + for name in ("original", "disk_copy", "reconstructed"): + if name in result: + row = result[name] + print(name, row.get("progress", row["status"]), "exit", hex(row.get("exit_code", 0)), + "exceptions", [hex(e["code"]) for e in row.get("exceptions", [])]) + + +if __name__ == "__main__": + main() diff --git a/tests/benchmark_system.py b/tests/benchmark_system.py new file mode 100644 index 0000000..eca8536 --- /dev/null +++ b/tests/benchmark_system.py @@ -0,0 +1,96 @@ +"""Opt-in bounded system-scheduler benchmark against one owned benign fixture process.""" + +import argparse +import hashlib +import json +from pathlib import Path +import statistics +import subprocess +import tempfile +import time + + +def fingerprints(folder): + return {p.name: hashlib.sha256(p.read_bytes()).hexdigest() + for p in sorted(folder.iterdir()) if p.is_file()} + + +def run(args): + samples = [] + with tempfile.TemporaryDirectory(prefix="pd-system-") as temporary: + root = Path(temporary) + with (root / "host.log").open("w") as log: + child = subprocess.Popen([str(args.probe), "--system-fixture"], stdout=log, stderr=subprocess.STDOUT) + try: + deadline = time.monotonic() + 10 + while "PD_SYSTEM_READY" not in (root / "host.log").read_text(): + if child.poll() is not None or time.monotonic() > deadline: + raise RuntimeError("Owned system fixture did not become ready") + time.sleep(0.05) + references = {} + for trial in range(args.repeats): + workers = [0, *args.threads] if trial % 2 == 0 else [*reversed(args.threads), 0] + for flags in args.flags: + for mode in ("dump", "hash"): + for threads in workers: + tag = f"{trial}-{flags or 'none'}-{mode}-{threads}" + folder = root / tag + folder.mkdir() + report = root / (tag + ".json") + with (root / (tag + ".log")).open("w") as log: + subprocess.run([str(args.probe), "--system-work", str(child.pid), mode, str(threads), + str(folder), str(report), flags], stdout=log, stderr=subprocess.STDOUT, + check=True, timeout=90) + row = json.loads(report.read_text()) + row.update(threads=threads, trial=trial, flags=flags, mode=mode, files=fingerprints(folder)) + if len(row["hashes"]) < 64 or (mode == "dump" and len(row["files"]) < 64): + raise RuntimeError(f"Owned fixture coverage missing for {tag}") + identity = {k: row[k] for k in ("hashes", "full", "prefixes", "files")} + key = (flags, mode) + if key not in references: + references[key] = identity + if identity != references[key]: + raise RuntimeError(f"Output/database mismatch for {tag}") + samples.append(row) + for path in folder.iterdir(): + path.unlink() + folder.rmdir() + print(tag, row["milliseconds"], flush=True) + return {"schema": 1, "probe_sha256": hashlib.sha256(args.probe.read_bytes()).hexdigest(), + "samples": samples, "equivalent": True, + "medians_ms": {f"{flags or 'none'}-{mode}-{threads}": + statistics.median(r["milliseconds"] for r in samples + if (r["flags"], r["mode"], r["threads"]) == (flags, mode, threads)) + for flags in args.flags for mode in ("dump", "hash") for threads in [0, *args.threads]}} + finally: + if child.poll() is None: + child.terminate() + child.wait(timeout=10) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--probe", type=Path, required=True) + parser.add_argument("--report", type=Path, required=True) + parser.add_argument("--threads", type=int, nargs="+", default=[1, 4, 16]) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--flags", nargs="+", default=["", "i", "c", "ic"]) + parser.add_argument("--acknowledge-execution", action="store_true") + args = parser.parse_args() + if not args.acknowledge_execution: + parser.error("--acknowledge-execution is required for the benign test-host launch") + if not 1 <= args.repeats <= 5 or any(not 1 <= t <= 64 for t in args.threads): + parser.error("use 1..5 repeats and 1..64 workers") + if any(set(flags) - set("icgr") for flags in args.flags): + parser.error("flags must contain only i(imports), c(chunks), g(generated headers), r(reexecution)") + args.probe = args.probe.resolve() + with args.report.open("x") as output: + try: + json.dump(run(args), output, indent=2) + except (OSError, ValueError, RuntimeError, subprocess.SubprocessError) as error: + json.dump({"status": "failed", "diagnostic": str(error)}, output) + raise + + +if __name__ == "__main__": + main() diff --git a/tests/entrypoint_benchmark.cpp b/tests/entrypoint_benchmark.cpp new file mode 100644 index 0000000..605070f --- /dev/null +++ b/tests/entrypoint_benchmark.cpp @@ -0,0 +1,272 @@ +#include "test_support.h" +#include "pe_hash_database.h" +#include +#include +#include +#include + +struct entrypoint_benchmark +{ + struct sample + { + std::string path; + DWORD truth; + unsigned __int64 full, prefix; + }; + + static bool executable(pe_header& image, SIZE_T rva, SIZE_T length = 1) + { + for (int i = 0; i < image._num_sections; ++i) + { + const auto& section = image._header_sections[i]; + if ((section.Characteristics & IMAGE_SCN_MEM_EXECUTE) && rva >= section.VirtualAddress && + range_fits((std::max)(section.Misc.VirtualSize, section.SizeOfRawData), rva - section.VirtualAddress, length)) + return true; + } + return false; + } + + static DWORD unique(const std::set& candidates) + { + return candidates.size() == 1 ? *candidates.begin() : 0; + } + + static int run(const char* manifest, const char* report) + { + test_options options; + options.EntryPointHash = true; + std::ifstream input(manifest); + require(input.good(), "cannot read entrypoint corpus manifest"); + std::vector samples; + std::map full_counts, prefix_counts; + std::string path; + while (std::getline(input, path)) + { + if (!path.empty() && path.back() == '\r') path.pop_back(); + pe_header image(&path[0], &options); + if (!image.process_pe_header() || !image.process_sections()) continue; + const DWORD truth = image._parsed_pe_32 ? image._header_pe32->OptionalHeader.AddressOfEntryPoint : + image._header_pe64->OptionalHeader.AddressOfEntryPoint; + if (!truth || !executable(image, truth)) continue; + const auto full = image.get_hash_ep(), prefix = image.get_hash_ep_short(); + samples.push_back({path, truth, full, prefix}); + if (full) ++full_counts[full]; + if (prefix) ++prefix_counts[prefix]; + } + std::ofstream output(report); + require(output.good(), "cannot create entrypoint corpus report"); + output << "index,architecture,truth,heldout,legacy,code_first,unique_full,boundary_fallback,weak_unique,full_candidates,fallback_candidates,prefer_exact,runtime_only,actual,ms,path\n"; + for (size_t index = 0; index < samples.size(); ++index) + { + auto& source = samples[index]; + pe_header image(&source.path[0], &options); + require(image.process_pe_header() && image.process_sections(), "corpus changed during measurement"); + for (bool heldout : {false, true}) + { + const auto started = std::chrono::steady_clock::now(); + const auto known_full = [&](unsigned __int64 hash) { + auto found = full_counts.find(hash); + return hash && found != full_counts.end() && found->second > (heldout && hash == source.full ? 1 : 0); + }; + const auto known_prefix = [&](unsigned __int64 hash) { + auto found = prefix_counts.find(hash); + return hash && found != prefix_counts.end() && found->second > (heldout && hash == source.prefix ? 1 : 0); + }; + DWORD first = 0, strong = 0, first_code = 0; + std::set full, weak, fallback, boundaries, runtime; + for (SIZE_T rva = 0; range_fits(image._image_size, rva, 8); ++rva) + { + unsigned __int64 prefix = 0; + memcpy(&prefix, image._image + rva, 8); + if (known_prefix(prefix)) + { + const bool matched = known_full(image._hash_asm(rva)); + if (rva >= 0x1000 && rva < image._image_size - 8) + { + if (!first) first = static_cast(rva); + if (matched && !strong) strong = static_cast(rva); + } + if (executable(image, rva, 8)) + { + weak.insert(static_cast(rva)); + if (matched) + { + full.insert(static_cast(rva)); + if (!first_code) first_code = static_cast(rva); + } + } + } + if (executable(image, rva, 5) && image._image[rva] == 0xe8) + { + INT32 displacement = 0; + memcpy(&displacement, image._image + rva + 1, 4); + const auto target = static_cast<__int64>(rva) + 5 + displacement; + if (target > 0 && target < static_cast<__int64>(image._image_size) && executable(image, target, 20)) + boundaries.insert(static_cast(target)); + } + } + if (image._parsed_pe_64) + { + const auto& directory = image._header_pe64->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION]; + if (range_fits(image._image_size, directory.VirtualAddress, directory.Size)) + for (SIZE_T i = 0; range_fits(directory.Size, i, sizeof(IMAGE_RUNTIME_FUNCTION_ENTRY)); i += sizeof(IMAGE_RUNTIME_FUNCTION_ENTRY)) + { + IMAGE_RUNTIME_FUNCTION_ENTRY function; + memcpy(&function, image._image + directory.VirtualAddress + i, sizeof(function)); + if (function.BeginAddress < function.EndAddress && executable(image, function.BeginAddress, 20)) + { + boundaries.insert(function.BeginAddress); + if (known_full(image._hash_asm(function.BeginAddress))) + runtime.insert(function.BeginAddress); + } + } + } + fallback = full; + for (DWORD rva : boundaries) + if (known_full(image._hash_asm(rva))) fallback.insert(rva); + temporary_file clean, ep, short_ep; + pe_hash_database database(clean.path, ep.path, short_ep.path); + unordered_set full_set, prefix_set; + for (const auto& item : full_counts) if (known_full(item.first)) full_set.insert(item.first); + for (const auto& item : prefix_counts) if (known_prefix(item.first)) prefix_set.insert(item.first); + database.add_hashes_eps(full_set, prefix_set); + DWORD& entry = image._parsed_pe_32 ? image._header_pe32->OptionalHeader.AddressOfEntryPoint : + image._header_pe64->OptionalHeader.AddressOfEntryPoint; + entry = 0; + image._recover_entrypoint(&database); + const double ms = std::chrono::duration(std::chrono::steady_clock::now() - started).count(); + output << index << ',' << (image.is_64() ? 64 : 32) << ',' << source.truth << ',' << heldout << ',' + << (strong ? strong : first) << ',' << first_code << ',' << unique(full) << ',' << unique(fallback) << ',' + << (unique(full) ? unique(full) : unique(weak)) << ',' << full.size() << ',' << fallback.size() << ',' + << (first_code ? first_code : unique(fallback)) << ',' << (first_code ? first_code : unique(runtime)) + << ',' << entry << ',' << ms << ",\"" << source.path << "\"\n"; + output.flush(); + } + } + require(output.good(), "entrypoint report write failed"); + printf("Measured %zu corpus images.\n", samples.size()); + return 0; + } + + static void recovery_rules(bool win64) + { + test_options options; + options.EntryPointHash = true; + auto bytes = pe_fixture(win64); + memset(bytes.data() + 0x200, 0xcc, 0x1000); + auto* code = bytes.data() + 0x240; + memset(code, 0x90, 128); + code[0] = 0xe8; + code[1] = 0x34; code[2] = 0x12; code[3] = code[4] = 0; + code[40] = 0xc3; + temporary_file file, clean, ep, prefix; + file.write(bytes); + pe_hash_database database(clean.path, ep.path, prefix.path); + pe_header image(file.path, &options); + require(image.process_pe_header() && image.process_sections(), "recovery rules parse failed"); + const auto full = image._hash_asm(0x1040), short_hash = image._hash_short_asm(0x1040); + require(full && short_hash, "recovery rules signatures missing"); + database.add_hashes_eps({full}, {short_hash}); + DWORD& entry = win64 ? image._header_pe64->OptionalHeader.AddressOfEntryPoint : + image._header_pe32->OptionalHeader.AddressOfEntryPoint; + entry = 0; + image._recover_entrypoint(&database); + require(entry == 0x1040, "exact recovery failed"); + entry = 0x1000; + image._recover_entrypoint(&database); + require(entry == 0x1000, "valid entrypoint was overwritten"); + options.ForceReconstructEntryPoint = true; + options.Verbose = true; + image._recover_entrypoint(&database); + require(entry == 0x1040, "forced/verbose recovery changed selection"); + options.EntryPointHash = false; + entry = 0; + image._recover_entrypoint(&database); + require(entry == 0, "disabled recovery ran"); + options.EntryPointHash = true; + image._header_sections[0].Characteristics &= ~IMAGE_SCN_MEM_EXECUTE; + image._recover_entrypoint(&database); + require(entry == 0, "nonexecutable signature was selected"); + image._header_sections[0].Characteristics |= IMAGE_SCN_MEM_EXECUTE; + if (win64) + { + auto& directory = image._header_pe64->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION]; + directory = {0x1e00, sizeof(IMAGE_RUNTIME_FUNCTION_ENTRY)}; + IMAGE_RUNTIME_FUNCTION_ENTRY function = {0x1040, 0x1080, 0x1f00}; + memcpy(image._image + 0x1e00, &function, sizeof(function)); + image._image[0x1041] ^= 0x10; + image._recover_entrypoint(&database); + require(entry == 0x1040, "runtime metadata did not recover a changed raw prefix"); + entry = 0; + function.EndAddress = 0x4000; + memcpy(image._image + 0x1e00, &function, sizeof(function)); + image._recover_entrypoint(&database); + require(entry == 0, "invalid runtime extent was accepted"); + function.EndAddress = 0x1080; + memcpy(image._image + 0x1e00, &function, sizeof(function)); + memcpy(image._image + 0x1100, image._image + 0x1040, 128); + function = {0x1100, 0x1140, 0x1f00}; + memcpy(image._image + 0x1e00 + sizeof(function), &function, sizeof(function)); + directory.Size *= 2; + image._recover_entrypoint(&database); + require(entry == 0, "ambiguous metadata fallback guessed an entrypoint"); + directory.Size = sizeof(function) - 1; + image._recover_entrypoint(&database); + require(entry == 0, "partial runtime record was accepted"); + } + options.ForceReconstructEntryPoint = false; + options.Verbose = false; + for (DWORD start : {0x200u, 0x2000u}) + { + auto shifted = bytes; + IMAGE_SECTION_HEADER* section = NULL; + if (win64) + { + auto* header = reinterpret_cast(shifted.data() + 0x80); + header->OptionalHeader.AddressOfEntryPoint = start; + header->OptionalHeader.SectionAlignment = start == 0x200 ? 0x200 : 0x1000; + header->OptionalHeader.SizeOfImage = start + 0x1000; + section = reinterpret_cast(header + 1); + } + else + { + auto* header = reinterpret_cast(shifted.data() + 0x80); + header->OptionalHeader.AddressOfEntryPoint = start; + header->OptionalHeader.SectionAlignment = start == 0x200 ? 0x200 : 0x1000; + header->OptionalHeader.SizeOfImage = start + 0x1000; + section = reinterpret_cast(header + 1); + } + section->VirtualAddress = start; + file.write(shifted); + pe_header low(file.path, &options); + require(low.process_pe_header() && low.process_sections(), "shifted entrypoint fixture parse failed"); + DWORD& found = win64 ? low._header_pe64->OptionalHeader.AddressOfEntryPoint : + low._header_pe32->OptionalHeader.AddressOfEntryPoint; + low._recover_entrypoint(&database); + require(found == start, "valid low-alignment or RVA 0x2000 entrypoint was overwritten"); + found = 0; + low._recover_entrypoint(&database); + require(found == start + 0x40, "executable scan missed low-alignment entrypoint"); + } + } +}; + +void append_entrypoint_tests(std::vector& tests) +{ + tests.push_back({"pe32-entrypoint-selection-rules", [] { entrypoint_benchmark::recovery_rules(false); }}); + tests.push_back({"pe64-entrypoint-selection-rules", [] { entrypoint_benchmark::recovery_rules(true); }}); +} + +int run_entrypoint_benchmark(int argc, char** argv) +{ + try + { + require(argc == 2, "usage: --entrypoint-corpus manifest.txt report.csv"); + return entrypoint_benchmark::run(argv[0], argv[1]); + } + catch (const std::exception& error) + { + fprintf(stderr, "Entrypoint benchmark failed: %s\n", error.what()); + return 1; + } +} diff --git a/tests/name_safety_tests.cpp b/tests/name_safety_tests.cpp index be03c6f..a9589b3 100644 --- a/tests/name_safety_tests.cpp +++ b/tests/name_safety_tests.cpp @@ -21,7 +21,7 @@ namespace entry->AddressOfFunctions = 80; entry->AddressOfNames = 84; entry->AddressOfNameOrdinals = 88; - const DWORD rva = 0x180, name_offset = 248; + const DWORD rva = 0x78, name_offset = 248; const WORD ordinal = 0; memcpy(bytes.data() + 80, &rva, sizeof(rva)); memcpy(bytes.data() + 84, &name_offset, sizeof(name_offset)); @@ -47,16 +47,16 @@ namespace export_list exports; require(fixture.parse(exports, fixture.bytes.size()), "terminated export at image boundary rejected"); char library[] = "ordinary.dll", name[] = "Example"; - require(exports.find_export(library, name, false) == export_fixture::image_base + 0x180, + require(exports.find_export(library, name, false) == export_fixture::image_base + 0x78, "bounded export lookup changed"); - export_entry entry = exports.find(export_fixture::image_base + 0x180); + export_entry entry = exports.find(export_fixture::image_base + 0x78); require(strcmp(entry.library_name, library) == 0 && strcmp(entry.name, name) == 0, "export strings were not copied exactly"); export_list missing_name_terminator; require(!fixture.parse(missing_name_terminator, fixture.bytes.size() - 1), "export name without an in-range terminator accepted"); - require(!missing_name_terminator.contains(export_fixture::image_base + 0x180), + require(!missing_name_terminator.contains(export_fixture::image_base + 0x78), "unterminated export name registered"); export_list missing_library_terminator; diff --git a/tests/output_tests.cpp b/tests/output_tests.cpp index d66c214..d3a7d18 100644 --- a/tests/output_tests.cpp +++ b/tests/output_tests.cpp @@ -27,9 +27,103 @@ static void invalid_output_rejected() require(!write_new_dump(NULL, &data, 1), "null output accepted"); } +static unsigned int native_create_calls = 0; +static NTSTATUS native_failure = 0; +static bool native_directory = false; + +static NTSTATUS NTAPI checked_native_create(PHANDLE file, ACCESS_MASK access, POBJECT_ATTRIBUTES attributes, + PIO_STATUS_BLOCK io, PLARGE_INTEGER allocation, ULONG file_attributes, ULONG sharing, + ULONG disposition, ULONG options, PVOID ea, ULONG ea_size) +{ + ++native_create_calls; + require(file != NULL && io != NULL && allocation == NULL, "native creation output parameters changed"); + require(attributes->RootDirectory == reinterpret_cast(static_cast(42)), + "creation did not retain the verified root handle"); + require(attributes->Attributes == (OBJ_CASE_INSENSITIVE | OBJ_DONT_REPARSE), + "creation does not reject reparses atomically"); + const std::wstring name(attributes->ObjectName->Buffer, attributes->ObjectName->Length / sizeof(WCHAR)); + require(name == L"nested directory\\file.exe" && + attributes->ObjectName->MaximumLength == attributes->ObjectName->Length + sizeof(WCHAR), + "root-relative Unicode name changed"); + require(access == ((native_directory ? FILE_READ_ATTRIBUTES : GENERIC_WRITE | DELETE) | SYNCHRONIZE) && + sharing == (native_directory ? FILE_SHARE_READ | FILE_SHARE_WRITE : 0) && + disposition == (native_directory ? FILE_OPEN : FILE_CREATE) && + options == ((native_directory ? FILE_DIRECTORY_FILE : FILE_NON_DIRECTORY_FILE) | + FILE_SYNCHRONOUS_IO_NONALERT | FILE_OPEN_REPARSE_POINT) && + file_attributes == FILE_ATTRIBUTE_NORMAL && ea == NULL && ea_size == 0, + "native creation weakened exclusive, synchronous file creation"); + return native_failure; +} + +static void native_output_fail_closed() +{ + const HANDLE root = reinterpret_cast(static_cast(42)); + // Exercise API failures, not an attacker-controlled filesystem or a race. + for (bool directory : {false, true}) + { + native_directory = directory; + for (NTSTATUS status : {static_cast(0xc000050bL), static_cast(0xc0000022L), + static_cast(0xc000000dL)}) + { + native_failure = status; + native_create_calls = 0; + const HANDLE result = directory ? + dump_output::open_no_reparse(root, "nested directory\\file.exe", true, checked_native_create) : + dump_output::create_no_reparse(root, "nested directory\\file.exe", checked_native_create); + require(result == INVALID_HANDLE_VALUE, "native create failure was ignored"); + require(native_create_calls == 1 && GetLastError() == RtlNtStatusToDosError(status), + "native failure was retried or not propagated"); + } + } +} + +static void native_output_requires_relative_path() +{ + native_create_calls = 0; + const HANDLE root = reinterpret_cast(static_cast(42)); + for (const char* path : {"", "\\absolute", "/absolute", "C:\\absolute", "file:stream"}) + require(dump_output::create_no_reparse(root, path, checked_native_create) == INVALID_HANDLE_VALUE, + "non-relative native output path accepted"); + require(dump_output::create_no_reparse(NULL, "file", checked_native_create) == INVALID_HANDLE_VALUE && + dump_output::create_no_reparse(INVALID_HANDLE_VALUE, "file", checked_native_create) == INVALID_HANDLE_VALUE && + dump_output::create_no_reparse(root, NULL, checked_native_create) == INVALID_HANDLE_VALUE, + "invalid native creation arguments accepted"); + require(native_create_calls == 0, "invalid native name reached the filesystem"); +} + +static void nested_output_roundtrip() +{ + struct directory_fixture + { + temporary_file root; + std::string nested, output; + directory_fixture() : nested(std::string(root.path) + "\\nested directory"), output(nested + "\\file.bin") + { + require(DeleteFileA(root.path) && CreateDirectoryA(root.path, NULL), "temporary directory creation failed"); + } + ~directory_fixture() + { + DeleteFileA(output.c_str()); + RemoveDirectoryA(nested.c_str()); + RemoveDirectoryA(root.path); + } + } directory; + require(CreateDirectoryA(directory.nested.c_str(), NULL) != 0, "nested directory creation failed"); + const unsigned char data[] = {1, 2, 3}; + require(write_new_dump(directory.output.c_str(), data, sizeof(data)), "nested dump creation failed"); + file_stream input(const_cast(directory.output.c_str())); + unsigned char read[sizeof(data)] = {}; + SIZE_T count = 0; + require(input.read(0, sizeof(read), read, &count) && count == sizeof(data) && + memcmp(data, read, sizeof(data)) == 0, "nested output contents changed"); +} + void append_output_tests(std::vector& tests) { tests.emplace_back("output-new-file", new_output_roundtrip); tests.emplace_back("output-preserve-existing", existing_output_unchanged); tests.emplace_back("output-invalid-path", invalid_output_rejected); + tests.emplace_back("output-native-fail-closed", native_output_fail_closed); + tests.emplace_back("output-native-relative-path", native_output_requires_relative_path); + tests.emplace_back("output-nested-file", nested_output_roundtrip); } diff --git a/tests/pd_tests.vcxproj b/tests/pd_tests.vcxproj index af9177d..7ae8915 100644 --- a/tests/pd_tests.vcxproj +++ b/tests/pd_tests.vcxproj @@ -27,6 +27,7 @@ $(MSBuildThisFileDirectory)..\pd;%(AdditionalIncludeDirectories) WIN32;_CONSOLE;UNICODE;_UNICODE;%(PreprocessorDefinitions) + PD_PROFILE_ALLOCATIONS;WIN32;_CONSOLE;UNICODE;_UNICODE;%(PreprocessorDefinitions) NotUsing Level3 ProgramDatabase @@ -38,7 +39,7 @@ Guard - Shlwapi.lib;Psapi.lib;Advapi32.lib;%(AdditionalDependencies) + Shlwapi.lib;Psapi.lib;Advapi32.lib;ntdll.lib;%(AdditionalDependencies) Console true true @@ -51,4 +52,7 @@ + + + diff --git a/tests/pe_quality.py b/tests/pe_quality.py new file mode 100644 index 0000000..9ec52e9 --- /dev/null +++ b/tests/pe_quality.py @@ -0,0 +1,381 @@ +"""Bounded PE fidelity analysis; results contain counts/hashes, not captured bytes.""" + +from collections import Counter +import ctypes +from ctypes import wintypes +import hashlib +import struct + + +MAX_IMAGE = 256 * 1024 * 1024 +DIRECTORIES = ("exports", "imports", "resources", "exceptions", "security", "relocations", + "debug", "architecture", "globalptr", "tls", "load_config", "bound_imports", + "iat", "delay_imports", "clr", "reserved") + + +def snapshot_sections(entry, info): + size = info["image_size"] + if not 0 < size <= MAX_IMAGE: + raise ValueError("memory witness image exceeds limit") + kernel = ctypes.WinDLL("kernel32", use_last_error=True) + kernel.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel.OpenProcess.restype = wintypes.HANDLE + kernel.ReadProcessMemory.argtypes = [wintypes.HANDLE, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t)] + kernel.ReadProcessMemory.restype = wintypes.BOOL + kernel.CloseHandle.argtypes = [wintypes.HANDLE] + kernel.CloseHandle.restype = wintypes.BOOL + handle = kernel.OpenProcess(0x410, False, entry["pid"]) + if not handle: + raise ctypes.WinError(ctypes.get_last_error()) + image, readable = bytearray(size), set() + buffer = (ctypes.c_char * size).from_buffer(image) + try: + for section in info["sections"]: + start = section["rva"] + end = start + max(section["raw_size"], section["virtual_size"]) + if end > size: + raise ValueError("memory witness section exceeds image") + for rva in range(start, end, 4096): + count, got = min(4096, end - rva), ctypes.c_size_t() + ok = kernel.ReadProcessMemory(handle, int(entry["base"]) + rva, + ctypes.byref(buffer, rva), count, ctypes.byref(got)) + if ok and got.value == count: + readable.add((rva, count)) + finally: + kernel.CloseHandle(handle) + return image, readable + + +class PE: + def __init__(self, data, info): + self.info = info + nt = struct.unpack_from(" optional + optional_size: + raise ValueError("directories exceed optional header") + for i in range(16): + self.directories.append(struct.unpack_from(" size or raw + count > len(data): + raise ValueError("section outside file or image") + self.image[rva:rva + count] = data[raw:raw + count] + + def read(self, rva, count): + if rva < 0 or count < 0 or rva + count > len(self.image): + raise ValueError("RVA outside image") + return self.image[rva:rva + count] + + def number(self, rva, width=4): + return int.from_bytes(self.read(rva, width), "little") + + def string(self, rva): + if rva == 0: + raise ValueError("null string RVA") + data = self.read(rva, min(4096, len(self.image) - rva)) + end = data.find(0) + if end < 0: + raise ValueError("unterminated string") + return bytes(data[:end]).decode("ascii", "strict") + + def raw_offset(self, rva, size): + if rva < self.info["headers_size"] and rva + size <= self.info["headers_size"]: + return rva + for s in self.info["sections"]: + if rva >= s["rva"] and rva + size <= s["rva"] + s["raw_size"]: + return s["raw"] + rva - s["rva"] + return None + + def imports(self, delay=False): + rva, size = self.directories[13 if delay else 1] + entries, errors, descriptors, slots = [], [], 0, [] + if not rva: + return entries, errors, descriptors, slots + stride = 32 if delay else 20 + limit = min(size // stride, MAX_IMAGE // stride) + terminated = False + for index in range(limit): + try: + values = struct.unpack("<8I" if delay else "<5I", self.read(rva + index * stride, stride)) + if not any(values): + terminated = True + break + if delay: + attributes, name, _, iat, lookup, _, _, _ = values + if not attributes & 1: + name, iat, lookup = (value - self.base if value else 0 for value in (name, iat, lookup)) + else: + lookup, _, _, name, iat = values + library = self.string(name).lower() + lookup = lookup or iat + if not lookup or not iat: + raise ValueError("missing lookup or IAT") + descriptors += 1 + for item in range(65536): + value = self.number(lookup + item * self.width, self.width) + if value == 0: + break + slot = iat + item * self.width + self.read(slot, self.width) + symbol = ("#" + str(value & 0xffff)) if value & (1 << (self.width * 8 - 1)) else self.string(value + 2) + if not symbol: + raise ValueError("empty import name") + entries.append((library, symbol, slot)) + slots.append(slot) + else: + raise ValueError("unterminated thunk table") + except (ValueError, UnicodeError, struct.error) as error: + errors.append(f"descriptor {index}: {error}") + if len(errors) >= 20: + break + if not terminated: + errors.append("descriptor terminator not reached") + return entries, errors, descriptors, slots + + def relocate(self, base): + rva, size = self.directories[5] + end, cursor, count = rva + size, rva, 0 + errors = [] + if not rva: + return count, errors + delta = base - self.base + while cursor < end: + try: + page, block_size = struct.unpack(" end: + raise ValueError("invalid relocation block") + block = self.read(cursor + 8, block_size - 8) + for (entry,) in struct.iter_unpack("> 12, entry & 4095 + if kind == 0: + continue + if kind not in (3, 10): + raise ValueError(f"unsupported relocation type {kind}") + width = 4 if kind == 3 else 8 + target = page + offset + value = (self.number(target, width) + delta) & ((1 << (width * 8)) - 1) + self.image[target:target + width] = value.to_bytes(width, "little") + count += 1 + cursor += block_size + except (ValueError, struct.error) as error: + errors.append(str(error)) + break + return count, errors + + def resources(self): + root, size = self.directories[2] + leaves, seen, hashes = {}, set(), {} + hashed_bytes = 0 + if not root: + return leaves + def relative(offset, count): + if offset < 0 or offset + count > size: + raise ValueError("resource entry outside directory") + return self.read(root + offset, count) + def walk(offset, path): + nonlocal hashed_bytes + if len(path) > 16 or len(seen) > 65536 or offset in seen: + raise ValueError("invalid resource tree") + seen.add(offset) + named, ids = struct.unpack_from("= 65536: + raise ValueError("resource leaf limit exceeded") + data_rva, count, codepage, _ = struct.unpack(" MAX_IMAGE: + raise ValueError("resource payload analysis limit exceeded") + hashes[key] = hashlib.sha256(self.read(data_rva, count)).hexdigest() + hashed_bytes += count + leaves[identity] = count, codepage, hashes[key] + walk(0, ()) + return leaves + + def directory_health(self): + errors, hashes = [], {} + for i, (rva, size) in enumerate(self.directories): + if not rva or not size or i == 4: + continue + try: + content = self.read(rva, size) + hashes[DIRECTORIES[i]] = hashlib.sha256(content).hexdigest() + except ValueError as error: + errors.append(f"{DIRECTORIES[i]}: {error}") + rva, size = self.directories[6] + if rva: + try: + for i in range(size // 28): + entry = struct.unpack(" len(dump.image): + continue + first, second = old[rva:rva + count], new[rva:rva + count] + reconstructed = dump.read(rva, count) + if first == second == reconstructed: + compared += count - rewrite_mask[rva:rva + count].count(1) + continue + for i, (a, b, value) in enumerate(zip(first, second, reconstructed)): + if rewrite_mask[rva + i]: + continue + compared += 1 + if a != b: + changed += 1 + elif value != a: + stable_difference += 1 + missing += a != 0 and value == 0 + witness = dict(compared_bytes=compared, changing_live_bytes=changed, + stable_different_bytes=stable_difference, stable_nonzero_to_zero_bytes=missing, + before_unreadable_bytes=sum(max(s["raw_size"], s["virtual_size"]) for s in source_info["sections"]) - + sum(count for _, count in old_readable), + after_unreadable_bytes=sum(max(s["raw_size"], s["virtual_size"]) for s in source_info["sections"]) - + sum(count for _, count in new_readable)) + return dict( + matching_file_layout=source.timestamp == dump.timestamp and all( + any(a["name"] == b["name"] and a["rva"] == b["rva"] and a["virtual_size"] == b["virtual_size"] + for b in dump_info["sections"]) for a in source_info["sections"]), + original_timestamp=source.timestamp, dump_timestamp=dump.timestamp, + import_descriptors=descriptors, import_entries=len(imports), + original_import_entries=len(original_imports), + missing_original_imports=len(set(original_imports) - set(imports)), + duplicate_iat_slots=sum(count - 1 for count in Counter(slots).values() if count > 1), + import_errors=import_errors, original_import_errors=original_errors, + iat=dump.iat_health(source), + delay_import_errors=delay_errors, original_delay_import_errors=source_delay_errors, + missing_delay_imports=len(set(source_delay) - set(delay)), + relocation_entries=relocation_count, relocation_errors=relocation_errors, + directory_errors=directory_errors, resource_errors=resource_errors, + resources_compared=resource_count, changed_resources=missing_resources, + memory_witness=witness, + unchanged_directories=[name for name, value in source_hashes.items() if hashes.get(name) == value], + sections=sections, + immutable_unexplained_bytes=sum(s["unexplained_bytes"] for s in sections if not s["writable"] and not s["discardable"])) diff --git a/tests/performance_benchmarks.cpp b/tests/performance_benchmarks.cpp new file mode 100644 index 0000000..6ebd19e --- /dev/null +++ b/tests/performance_benchmarks.cpp @@ -0,0 +1,321 @@ +#include "test_support.h" +#include "hash.h" +#include "dump_process.h" +#include +#include +#include +#include +#include + +namespace +{ + thread_local bool count_allocations = false; + thread_local unsigned long long allocation_count = 0, allocation_bytes = 0; +} + +#ifdef PD_PROFILE_ALLOCATIONS +void* operator new(size_t size) +{ + void* memory = malloc(size ? size : 1); + if (memory == NULL) + throw std::bad_alloc(); + if (count_allocations) + { + ++allocation_count; + allocation_bytes += size; + } + return memory; +} +void* operator new[](size_t size) { return ::operator new(size); } +void operator delete(void* memory) noexcept { free(memory); } +void operator delete[](void* memory) noexcept { free(memory); } +void operator delete(void* memory, size_t) noexcept { free(memory); } +void operator delete[](void* memory, size_t) noexcept { free(memory); } +#endif + +namespace +{ + long long ticks() + { + LARGE_INTEGER value; + QueryPerformanceCounter(&value); + return value.QuadPart; + } + + unsigned long long cpu_ticks() + { + FILETIME created, exited, kernel, user; + require(GetThreadTimes(GetCurrentThread(), &created, &exited, &kernel, &user) != 0, "thread CPU timing failed"); + return (static_cast(kernel.dwHighDateTime) << 32) + kernel.dwLowDateTime + + (static_cast(user.dwHighDateTime) << 32) + user.dwLowDateTime; + } + + unsigned long long cpu_cycles() + { + ULONG64 cycles = 0; + require(QueryThreadCycleTime(GetCurrentThread(), &cycles) != 0, "thread cycle timing failed"); + return cycles; + } + + struct result + { + long long capture = 0, reconstruct = 0, write = 0, cleanup = 0, total = 0; + unsigned long long cpu = 0, cycles = 0, allocations = 0, allocated_bytes = 0; + SIZE_T output_bytes = 0; + DWORD crc = 0; + std::unique_ptr output; + }; + + void verify_output(result& measured) + { + if (!measured.output) + return; + auto bytes = measured.output->read(); + require(bytes.size() >= 0x80 + sizeof(IMAGE_NT_HEADERS64), "benchmark output header truncated"); + reinterpret_cast(bytes.data() + 0x80)->OptionalHeader.ImageBase = 0; + measured.output_bytes = bytes.size(); + measured.crc = crc32buf(reinterpret_cast(bytes.data()), bytes.size()); + } + + class fixture + { + public: + unsigned char* memory = NULL; + SIZE_T size; + explicit fixture(const std::string& workload) + { + const SIZE_T payload = workload == "sparse" ? 32 * 1024 * 1024 : 4 * 1024 * 1024; + auto header = pe_fixture(true); + size = 0x1000 + payload; + memory = static_cast(VirtualAlloc(NULL, size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE)); + require(memory != NULL, "benchmark allocation failed"); + memcpy(memory, header.data(), 0x200); + auto* nt = reinterpret_cast(memory + 0x80); + nt->OptionalHeader.SizeOfImage = static_cast(size); + if (workload == "entrypoint" || workload == "entrypoint-empty") + nt->OptionalHeader.AddressOfEntryPoint = 0; + auto* section = reinterpret_cast(nt + 1); + section->Misc.VirtualSize = static_cast(payload); + section->SizeOfRawData = static_cast(payload); + memset(memory + 0x1000, 0x90, workload == "sparse" ? 4096 : payload); + if (workload == "imports") + for (SIZE_T offset = 0x1000; offset + 8 <= size; offset += 64) + { + const unsigned __int64 address = 0x7fff12340000ULL + ((offset / 64) % 256) * 16; + memcpy(memory + offset, &address, sizeof(address)); + } + DWORD previous; + if (VirtualProtect(memory, 0x1000, PAGE_READONLY, &previous) == 0) + { + VirtualFree(memory, 0, MEM_RELEASE); + memory = NULL; + throw std::runtime_error("benchmark header protection failed"); + } + } + ~fixture() { if (memory != NULL) VirtualFree(memory, 0, MEM_RELEASE); } + fixture(const fixture&) = delete; + fixture& operator=(const fixture&) = delete; + }; + + std::unique_ptr prepare_output(const std::string& workload) + { + std::unique_ptr output; + if (workload != "database" && workload != "exports") + { + output.reset(new temporary_file()); + require(DeleteFileA(output->path) != 0, "benchmark output reservation failed"); + } + return output; + } + + result measure(const std::string& workload, fixture& input, export_list& exports, pe_hash_database& database, + bool verify, std::unique_ptr output = {}) + { + result measured; + test_options options; + options.ImportRec = workload == "imports"; + options.EntryPointHash = true; + measured.output = output ? std::move(output) : prepare_output(workload); + allocation_count = allocation_bytes = 0; + const auto cpu_start = cpu_ticks(); + const auto cycles_start = cpu_cycles(); + const auto start = ticks(); + count_allocations = true; + try + { + if (workload == "database") + { + unsigned int found = 0; + for (unsigned int i = 0; i < 200000; ++i) + found += database.contains(1 + i % 65536); + require(found == 101696, "database benchmark result changed"); + measured.reconstruct = ticks() - start; + } + else if (workload == "exports") + { + dump_process process(GetCurrentProcessId(), &database, &options, true); + require(process.build_export_list(), "export discovery benchmark failed"); + measured.reconstruct = ticks() - start; + } + else + { + pe_header image(GetCurrentProcess(), input.memory, NULL, &options); + require(image.process_pe_header() && image.process_sections() && image.process_import_directory(), + "benchmark capture failed"); + const auto captured = ticks(); + measured.capture = captured - start; + require(image.process_disk_image(&exports, &database), "benchmark reconstruction failed"); + const auto reconstructed = ticks(); + measured.reconstruct = reconstructed - captured; + require(image.write_image(measured.output->path), "benchmark write failed"); + measured.write = ticks() - reconstructed; + } + measured.total = ticks() - start; + measured.cleanup = measured.total - measured.capture - measured.reconstruct - measured.write; + count_allocations = false; + } + catch (...) + { + count_allocations = false; + throw; + } + measured.cpu = cpu_ticks() - cpu_start; + measured.cycles = cpu_cycles() - cycles_start; + measured.allocations = allocation_count; + measured.allocated_bytes = allocation_bytes; + if (verify) + verify_output(measured); + return measured; + } + + unsigned int argument(const char* text, unsigned int maximum) + { + char* end; + const unsigned long value = strtoul(text, &end, 10); + require(text[0] != 0 && *end == 0 && value > 0 && value <= maximum, "invalid benchmark count"); + return static_cast(value); + } +} + +int run_performance_benchmark(int argc, char** argv) +{ + try + { + require(argc == 3, "usage: --benchmark dense|sparse|imports|entrypoint|entrypoint-empty|database|exports threads jobs"); + const std::string workload = argv[0]; + require(workload == "dense" || workload == "sparse" || workload == "imports" || workload == "entrypoint" || + workload == "entrypoint-empty" || workload == "database" || workload == "exports", "unknown benchmark workload"); + const unsigned int thread_count = argument(argv[1], 8), jobs = argument(argv[2], 256); + require(jobs >= thread_count, "jobs must cover all benchmark workers"); + fixture input(workload); + export_list exports; + for (unsigned int i = 0; i < 256; ++i) + { + const auto address = 0x7fff12340000ULL + i * 16; + const std::string name = "ExampleFunction" + std::to_string(i); + export_entry entry("example.dll", name.c_str(), static_cast(i + 1), 0x1000, address, true); + exports.add_export(address, &entry); + } + temporary_file clean, ep, short_ep; + pe_hash_database database(clean.path, ep.path, short_ep.path); + if (workload == "database") + { + unordered_set values; + for (unsigned int i = 1; i <= 32768; ++i) + values.insert(i); + database.add_hashes(values); + } + if (workload == "entrypoint") + database.add_hashes_eps({0x1234}, {0x1122334455667788ULL}); + const auto reference = measure(workload, input, exports, database, true); + std::vector results(jobs); + for (auto& item : results) + item.output = prepare_output(workload); + std::vector workers; + workers.reserve(thread_count); + std::vector worker_finished(thread_count); + std::atomic next(0), ready(0); + std::atomic start(false), failed(false); + std::mutex error_lock; + std::string error; + try + { + for (unsigned int i = 0; i < thread_count; ++i) + workers.emplace_back([&, i] { + ++ready; + while (!start) + SwitchToThread(); + try + { + for (unsigned int job; !failed && (job = next.fetch_add(1)) < jobs;) + results[job] = measure(workload, input, exports, database, false, std::move(results[job].output)); + } + catch (const std::exception& exception) + { + std::lock_guard lock(error_lock); + error = exception.what(); + failed = true; + } + worker_finished[i] = ticks(); + }); + } + catch (...) + { + failed = true; + start = true; + for (auto& worker : workers) + worker.join(); + throw; + } + while (ready != thread_count) + SwitchToThread(); + const auto begin = ticks(); + start = true; + for (auto& worker : workers) + worker.join(); + const auto elapsed = ticks() - begin; + const auto worker_elapsed = *std::max_element(worker_finished.begin(), worker_finished.end()) - begin; + if (failed) + throw std::runtime_error(error); + PROCESS_MEMORY_COUNTERS memory = {}; + memory.cb = sizeof(memory); + require(GetProcessMemoryInfo(GetCurrentProcess(), &memory, sizeof(memory)) != 0, "memory counters failed"); + result sum; + for (auto& item : results) + { + verify_output(item); + require(reference.output_bytes == item.output_bytes && reference.crc == item.crc, "parallel output differs"); + sum.capture += item.capture; + sum.reconstruct += item.reconstruct; + sum.write += item.write; + sum.cleanup += item.cleanup; + sum.total += item.total; + sum.cpu += item.cpu; + sum.cycles += item.cycles; + sum.allocations += item.allocations; + sum.allocated_bytes += item.allocated_bytes; + } + LARGE_INTEGER frequency; + QueryPerformanceFrequency(&frequency); + const double milliseconds = 1000.0 / frequency.QuadPart; +#ifdef PD_PROFILE_ALLOCATIONS + const char* profiled = "true"; +#else + const char* profiled = "false"; +#endif + printf("PERF {\"workload\":\"%s\",\"threads\":%u,\"jobs\":%u,\"wall_ms\":%.3f,\"worker_ms\":%.3f,\"capture_ms\":%.3f," + "\"reconstruct_ms\":%.3f,\"write_ms\":%.3f,\"cleanup_ms\":%.3f,\"service_ms\":%.3f,\"cpu_ms\":%.3f,\"cpu_cycles\":%llu," + "\"profiled_allocations\":%s,\"allocations\":%llu,\"allocated_bytes\":%llu," + "\"peak_working_set\":%zu,\"output_bytes\":%zu,\"normalized_crc32\":\"%08lx\"}\n", + workload.c_str(), thread_count, jobs, elapsed * milliseconds, worker_elapsed * milliseconds, sum.capture * milliseconds, + sum.reconstruct * milliseconds, sum.write * milliseconds, sum.cleanup * milliseconds, sum.total * milliseconds, + sum.cpu / 10000.0, sum.cycles, profiled, sum.allocations, sum.allocated_bytes, memory.PeakWorkingSetSize, + reference.output_bytes, reference.crc); + return 0; + } + catch (const std::exception& error) + { + fprintf(stderr, "BENCHMARK FAILED: %s\n", error.what()); + return 1; + } +} diff --git a/tests/performance_tests.cpp b/tests/performance_tests.cpp new file mode 100644 index 0000000..944878b --- /dev/null +++ b/tests/performance_tests.cpp @@ -0,0 +1,198 @@ +#include "test_support.h" +#include "pe_hash_database.h" +#include +#include + +namespace +{ + std::vector serialize(pe_imports& imports) + { + __int64 descriptors = 0, extra = 0; + imports.get_table_size(descriptors, extra); + require(imports.valid() && descriptors >= 0 && extra >= 0, "import records invalid"); + std::vector bytes(static_cast(descriptors + extra), 0); + require(imports.build_table(bytes.data(), bytes.size(), 0x5000, 0, descriptors), "import serialization failed"); + return bytes; + } + + void import_storage_equivalence(bool win64) + { + export_list exports; + pe_imports owned(NULL, 0, NULL, win64), borrowed(NULL, 0, NULL, win64); + IMAGE_IMPORT_DESCRIPTOR original = {}; + original.FirstThunk = 0x1400; + original.OriginalFirstThunk = 0x1800; + original.Name = 0x1200; + owned.add_descriptor(&original); + borrowed.add_descriptor(&original); + for (unsigned int i = 0; i < 257; ++i) + { + char library[] = "example.dll", name[] = "ExampleFunction"; + const unsigned __int64 address = win64 ? 0x7fff12340000ULL + i : 0x12340000ULL + i; + export_entry entry(library, i % 2 ? name : NULL, static_cast(i + 1), 0x1000, address, win64); + exports.add_export(address, &entry); + const auto* found = exports.lookup(address); + require(found != NULL, "borrowed export lookup failed"); + borrowed.add_fixup(*found, 0x2000 + i * 8, win64); + if (i % 2) + owned.add_fixup(library, name, 0x2000 + i * 8, win64); + else + owned.add_fixup(library, i + 1, 0x2000 + i * 8, win64); + memset(library, 'x', sizeof(library) - 1); + memset(name, 'x', sizeof(name) - 1); + } + original.FirstThunk = 0; + require(serialize(owned) == serialize(borrowed), "moved/borrowed import storage changed serialized bytes"); + } + + void export_transfer() + { + export_list destination, source; + export_entry existing("first.dll", "First", 1, 0x1000, 0x12340000, false); + export_entry duplicate("second.dll", "Second", 2, 0x1000, 0x12340000, false); + export_entry ordinal("ordinal.dll", NULL, 65535, 0x1000, 0x7fff12340000ULL, true); + destination.add_export(existing.address, &existing); + source.add_export(duplicate.address, &duplicate); + source.add_export(ordinal.address, &ordinal); + const auto* stable = source.lookup(ordinal.address); + destination.take_exports(source); + require(destination.lookup(ordinal.address) == stable, "export ownership transfer copied or lost the entry"); + require(strcmp(destination.lookup(existing.address)->name, "First") == 0, "duplicate export precedence changed"); + require(source.lookup(existing.address) == NULL && source.lookup(ordinal.address) == NULL, + "transferred exports retained stale source addresses"); + source.add_export(duplicate.address, &duplicate); + require(source.lookup(duplicate.address) != NULL, "transferred-from export list cannot be reused"); + destination.take_exports(destination); + require(destination.lookup(ordinal.address) == stable, "self-transfer invalidated exports"); + export_list copied; + copied.add_exports(&destination); + require(copied.lookup(ordinal.address) != stable && copied.lookup(ordinal.address)->ord == 65535, + "copying export API lost independent ownership"); + } + + void entrypoint_snapshots() + { + temporary_file clean, ep, short_ep; + std::shared_ptr saved; + { + pe_hash_database database(clean.path, ep.path, short_ep.path); + require(database.snapshot_entrypoints()->short_hashes.empty(), "empty snapshot is not empty"); + database.add_hashes_eps({1, _UI64_MAX}, {1, _UI64_MAX}); + saved = database.snapshot_entrypoints(); + require(database.snapshot_entrypoints() == saved, "unchanged entrypoint sets were copied again"); + database.add_hashes_eps({1}, {1}); + database.add_hashes({3}); + require(database.snapshot_entrypoints() == saved, "unrelated/duplicate update invalidated the entrypoint snapshot"); + database.clear_database(); + require(database.snapshot_entrypoints()->short_hashes.empty(), "clear left a stale entrypoint snapshot"); + require(saved->minimum == 1 && saved->maximum == _UI64_MAX && saved->full.size() == 2 && + saved->short_hashes.size() == 2, "entrypoint snapshot does not own its hashes"); + std::atomic done(false); + std::thread writer([&] { + for (unsigned int i = 1; i <= 1000; ++i) + database.add_hashes_eps({i}, {i}); + done = true; + }); + bool valid = true; + do + { + const auto snapshot = database.snapshot_entrypoints(); + valid = valid && snapshot->full == snapshot->short_hashes; + } while (!done); + writer.join(); + require(valid, "snapshot observed a partially published entrypoint pair"); + } + require(saved->short_hashes.count(_UI64_MAX) == 1, "snapshot depended on database lifetime"); + } + + void snapshot_file_update() + { + temporary_file source, clean, ep, short_ep; + auto bytes = pe_fixture(true); + bytes[0x200] = 0x31; + bytes[0x201] = 0xc0; + source.write(bytes); + pe_hash_database database(clean.path, ep.path, short_ep.path); + const auto before = database.snapshot_entrypoints(); + require(database.add_file(source.path), "fixture database insertion failed"); + const auto after = database.snapshot_entrypoints(); + require(before != after && before->full.empty() && before->short_hashes.empty() && + after->full.size() == 1 && after->short_hashes.size() == 1, "file insertion left stale entrypoint hashes"); + } + + template + void entrypoint_recovery(bool win64, bool strong) + { + auto bytes = pe_fixture(win64); + auto* nt = reinterpret_cast(bytes.data() + 0x80); + memset(bytes.data() + 0x200, 0xcc, 0x1000); + memset(bytes.data() + 0x240, 0x90, 8); + bytes[0x240] = 0x31; + bytes[0x241] = 0xc0; + bytes[0x248] = 0xc3; + memset(bytes.data() + 0x280, 0x90, 8); + bytes[0x280] = 0x31; + bytes[0x281] = 0xc0; + bytes[0x288] = 0x31; + bytes[0x289] = 0xc0; + bytes[0x28a] = 0xc3; + nt->OptionalHeader.AddressOfEntryPoint = 0x1080; + temporary_file source, clean, ep, short_ep, output; + source.write(bytes); + test_options options; + options.EntryPointHash = true; + unsigned __int64 full, prefix; + { + pe_header header(source.path, &options); + require(header.process_pe_header() && header.process_sections(), "entrypoint fixture failed"); + full = header.get_hash_ep(); + prefix = header.get_hash_ep_short(); + require(full != 0 && prefix != 0, "entrypoint fixture hashes are empty"); + } + pe_hash_database database(clean.path, ep.path, short_ep.path); + database.add_hashes_eps(strong ? unordered_set{full} : unordered_set{}, {prefix}); + nt->OptionalHeader.AddressOfEntryPoint = 0; + source.write(bytes); + pe_header header(source.path, &options); + require(header.process_pe_header() && header.process_sections() && header.process_disk_image(NULL, &database), + "entrypoint recovery failed"); + require(DeleteFileA(output.path) && header.write_image(output.path), "entrypoint output failed"); + const auto dumped = output.read(); + require(reinterpret_cast(dumped.data() + 0x80)->OptionalHeader.AddressOfEntryPoint == + (strong ? 0x1080 : 0), "entrypoint strong recovery or weak-only abstention changed"); + } + + void borrowed_process_identity() + { + std::unique_ptr process( + OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, GetCurrentProcessId()), CloseHandle); + require(process != NULL, "self process handle failed"); + module_list modules(GetCurrentProcessId()); + void* base = GetModuleHandleW(NULL); + test_options options; + { + pe_header owning(GetCurrentProcessId(), base, &modules, &options); + pe_header borrowed(process.get(), base, &modules, &options); + require(strcmp(owning.get_name(), borrowed.get_name()) == 0 && + strcmp(borrowed.get_name(), "hiddenmodule") != 0, "borrowed handle changed module identity"); + } + WORD signature = 0; + SIZE_T read = 0; + require(ReadProcessMemory(process.get(), base, &signature, sizeof(signature), &read) && + read == sizeof(signature) && signature == IMAGE_DOS_SIGNATURE, "borrowed process handle was closed"); + } +} + +void append_performance_tests(std::vector& tests) +{ + tests.push_back({"pe32-import-storage-equivalence", [] { import_storage_equivalence(false); }}); + tests.push_back({"pe64-import-storage-equivalence", [] { import_storage_equivalence(true); }}); + tests.push_back({"export-ownership-transfer", export_transfer}); + tests.push_back({"entrypoint-snapshot-consistency", entrypoint_snapshots}); + tests.push_back({"entrypoint-snapshot-file-update", snapshot_file_update}); + tests.push_back({"pe32-weak-entrypoint-recovery", [] { entrypoint_recovery(false, false); }}); + tests.push_back({"pe64-weak-entrypoint-recovery", [] { entrypoint_recovery(true, false); }}); + tests.push_back({"pe32-strong-entrypoint-recovery", [] { entrypoint_recovery(false, true); }}); + tests.push_back({"pe64-strong-entrypoint-recovery", [] { entrypoint_recovery(true, true); }}); + tests.push_back({"borrowed-process-handle-identity", borrowed_process_identity}); +} diff --git a/tests/pipeline_tests.cpp b/tests/pipeline_tests.cpp new file mode 100644 index 0000000..5d9b779 --- /dev/null +++ b/tests/pipeline_tests.cpp @@ -0,0 +1,305 @@ +#include "test_support.h" +#include "dump_process.h" +#include "dump_path.h" +#include +#include +#include + +namespace +{ + template + NT& nt(std::vector& data) + { + return *reinterpret_cast(data.data() + 0x80); + } + + class raw_pages + { + public: + unsigned char* bytes; + raw_pages() + { + bytes = static_cast(VirtualAlloc(NULL, 0x4000, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE)); + require(bytes != NULL, "raw allocation failed"); + DWORD previous; + require(VirtualProtect(bytes + 0x3000, 0x1000, PAGE_NOACCESS, &previous) != 0, "raw boundary setup failed"); + for (SIZE_T i = 0; i < 0x1000; ++i) + bytes[0x2000 + i] = static_cast(i % 251 + 1); + } + ~raw_pages() { VirtualFree(bytes, 0, MEM_RELEASE); } + }; + + std::vector write_reconstruction(pe_header& header) + { + temporary_file output; + require(DeleteFileA(output.path) != 0, "output reservation failed"); + require(header.write_image(output.path), "reconstruction output failed"); + return output.read(); + } + + template + void extended_optional_header(bool win64) + { + for (bool imports : {false, true}) + { + auto bytes = pe_fixture(win64); + auto& header = nt(bytes); + auto* old_section = reinterpret_cast(&header + 1); + memmove(old_section + 16, old_section, sizeof(IMAGE_SECTION_HEADER)); + memset(old_section, 0, 16); + header.FileHeader.SizeOfOptionalHeader += 16; + temporary_file input; + input.write(bytes); + test_options options; + options.ImportRec = imports; + pe_header image(input.path, &options); + require(image.process_pe_header() && image.process_sections(), "extended optional header was not parsed"); + require(image.process_import_directory(), "extended-header imports rejected"); + export_list exports; + require(image.process_disk_image(&exports, NULL), "extended-header reconstruction failed"); + auto dump = write_reconstruction(image); + const auto* section = reinterpret_cast(dump.data() + 0x80 + sizeof(NT) + 16); + require(section->VirtualAddress == 0x1000 && section->SizeOfRawData == 0x1000, + "extended optional header changed the section layout"); + require(memcmp(dump.data() + section->PointerToRawData, bytes.data() + 0x200, 0x1000) == 0, + "extended optional header lost section content"); + } + } + + void export_discovery(DWORD function_rva, bool named, bool valid, bool forwarded, bool invalid_tail = false) + { + auto bytes = pe_fixture(true); + nt(bytes).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT] = {0x1800, 0x100}; + memset(bytes.data() + 0xa00, 0, 0x200); + auto& directory = *reinterpret_cast(bytes.data() + 0xa00); + directory.Name = 0x1960; + directory.Base = 7; + directory.NumberOfFunctions = invalid_tail ? 2 : 1; + directory.NumberOfNames = named ? (invalid_tail ? 2 : 1) : 0; + directory.AddressOfFunctions = 0x1900; + directory.AddressOfNames = named ? 0x1940 : 0; + directory.AddressOfNameOrdinals = named ? 0x1950 : 0; + memcpy(bytes.data() + 0xb00, &function_rva, sizeof(function_rva)); + if (invalid_tail) + { + const DWORD outside = 0x3000; + memcpy(bytes.data() + 0xb04, &outside, sizeof(outside)); + } + const DWORD name = 0x1980; + memcpy(bytes.data() + 0xb40, &name, sizeof(name)); + if (invalid_tail) + { + const DWORD invalid_name = 0x3000; + const WORD second_ordinal = 1; + memcpy(bytes.data() + 0xb44, &invalid_name, sizeof(invalid_name)); + memcpy(bytes.data() + 0xb52, &second_ordinal, sizeof(second_ordinal)); + } + memcpy(bytes.data() + 0xb60, "example.dll", sizeof("example.dll")); + memcpy(bytes.data() + 0xb80, "Example", sizeof("Example")); + memcpy(bytes.data() + 0xa80, "other.Example", sizeof("other.Example")); + temporary_file input; + input.write(bytes); + test_options options; + pe_header image(input.path, &options); + require(image.process_pe_header() && image.process_sections(), "export fixture rejected"); + require(image.process_export_directory() == valid, "invalid export result was hidden"); + if (!valid) + { + require(image.get_exports() == NULL, "invalid export table published a partial result"); + return; + } + require(image.get_exports()->contains(static_cast(function_rva)) == !forwarded, + "real export omitted or forwarder string treated as a callable address"); + if (!forwarded) + { + auto entry = image.get_exports()->find(function_rva); + require(entry.ord == 7 && (entry.name != NULL) == named, "export identity lost"); + } + } + + template + void terminal_import_summary(bool win64) + { + auto bytes = pe_fixture(win64); + const unsigned __int64 address = win64 ? 0x7fff12345678ULL : 0xf2345678ULL; + const SIZE_T width = win64 ? 8 : 4; + memcpy(bytes.data() + bytes.size() - width, &address, width); + temporary_file input; + input.write(bytes); + test_options options; + pe_header image(input.path, &options); + require(image.process_pe_header() && image.process_sections(), "summary fixture rejected"); + export_list exports; + export_entry entry("example.dll", "Example", 1, 0x1000, address, win64); + exports.add_export(address, &entry); + const auto summary = image.get_imports_information(&exports); + require(summary.COUNT_UNIQUE_IMPORT_ADDRESSES == 1, "summary omitted a terminal import pointer"); + require(summary.COUNT_UNIQUE_IMPORT_LIBRARIES == 1, "summary omitted the import library count"); + require(image.get_imports_information(&exports, 0x2000 - width).COUNT_UNIQUE_IMPORT_ADDRESSES == 0, + "summary read beyond its requested size limit"); + unsigned __int64 ordinal_hashes[2]; + for (WORD ordinal = 7; ordinal <= 8; ++ordinal) + { + export_list ordinal_exports; + export_entry ordinal_entry("example.dll", NULL, ordinal, 0x1000, address, win64); + ordinal_exports.add_export(address, &ordinal_entry); + ordinal_hashes[ordinal - 7] = image.get_imports_information(&ordinal_exports).HASH_GENERIC; + pe_imports imports(NULL, 0, NULL, win64); + imports.add_fixup(ordinal_entry.library_name, ordinal_entry.ord, 0x1000, win64); + __int64 descriptors = 0, extra = 0; + imports.get_table_size(descriptors, extra); + std::vector table(static_cast(descriptors + extra)); + require(imports.build_table(table.data(), table.size(), 0x3000, 0, descriptors), "ordinal import table failed"); + const auto& descriptor = *reinterpret_cast(table.data()); + unsigned __int64 thunk = 0; + memcpy(&thunk, table.data() + descriptor.OriginalFirstThunk - 0x3000, width); + require(thunk == ((win64 ? IMAGE_ORDINAL_FLAG64 : IMAGE_ORDINAL_FLAG32) | ordinal), + "ordinal import encoding lost its identity"); + } + require(ordinal_hashes[0] != ordinal_hashes[1], "different ordinal imports have the same generic identity"); + } + + void generated_image_type(bool win64) + { + raw_pages memory; + for (bool dll : {false, true}) + { + test_options options; + options.ReconstructHeaderAsDll = dll; + pe_header image(GetCurrentProcess(), memory.bytes + 0x2000, NULL, &options); + require(image.build_pe_header(0x1000, win64, 1) && image.process_sections(), "generated image failed"); + require(image.is_dll() == dll && image.is_exe() != dll, "generated image type flag is reversed"); + require(image.process_import_directory() && image.process_disk_image(NULL, NULL), "raw image reconstruction failed"); + auto dump = write_reconstruction(image); + const auto* section = reinterpret_cast(dump.data() + 0xe0 + + (win64 ? sizeof(IMAGE_NT_HEADERS64) : sizeof(IMAGE_NT_HEADERS32))); + require(section->SizeOfRawData == 0x1000 && + memcmp(dump.data() + section->PointerToRawData, memory.bytes + 0x2000, 0x1000) == 0, + "generated image lost raw bytes"); + } + } + + void raw_address_dump() + { + raw_pages memory; + temporary_file clean, ep, short_ep, directory; + require(DeleteFileA(directory.path) && CreateDirectoryA(directory.path, NULL), "raw output directory failed"); + std::vector paths; + auto cleanup = [&paths](char* path) { + for (const auto& file : paths) + DeleteFileA(file.c_str()); + RemoveDirectoryA(path); + }; + std::unique_ptr output(directory.path, cleanup); + pe_hash_database database(clean.path, ep.path, short_ep.path); + test_options options; + options.set_output_path(directory.path); + dump_process dumper(GetCurrentProcessId(), &database, &options, true); + char process_name[1024]; + require(dumper.get_process_name(process_name, sizeof(process_name)), "self process name missing"); + for (bool win64 : {false, true}) + { + std::string file, error; + require(dump_path::make_filename(directory.path, process_name, GetCurrentProcessId(), "hiddenmodule", + reinterpret_cast(memory.bytes + 0x2000), win64, "exe", file, error), "raw filename failed"); + paths.push_back(file); + } + dumper.dump_region(reinterpret_cast(memory.bytes + 0x2000)); + for (const auto& file : paths) + { + require(GetFileAttributesA(file.c_str()) != INVALID_FILE_ATTRIBUTES, "address-specific raw dumping produced no image"); + test_options check_options; + pe_header check(const_cast(file.c_str()), &check_options); + require(check.process_pe_header() && check.process_sections(), "raw output is not a readable PE"); + } + } + + BOOL WINAPI many_modules(HANDLE, HMODULE* modules, DWORD capacity, LPDWORD needed, DWORD) + { + const DWORD count = 4097; + *needed = count * sizeof(HMODULE); + for (DWORD i = 0; i < (std::min)(count, capacity / sizeof(HMODULE)); ++i) + modules[i] = reinterpret_cast(static_cast(i + 1)); + return TRUE; + } + + BOOL WINAPI growing_modules(HANDLE, HMODULE*, DWORD capacity, LPDWORD needed, DWORD) + { + *needed = capacity + sizeof(HMODULE); + return TRUE; + } + + BOOL WINAPI unavailable_modules(HANDLE, HMODULE*, DWORD, LPDWORD, DWORD) + { + SetLastError(ERROR_ACCESS_DENIED); + return FALSE; + } + + void module_snapshot_bounds() + { + std::vector modules; + require(module_snapshot::read(NULL, modules, many_modules), "large module snapshot failed"); + require(modules.size() == 4097, "large module list was truncated"); + for (SIZE_T i = 0; i < modules.size(); ++i) + require(reinterpret_cast(modules[i]) == i + 1, "unfilled module slot exposed"); + require(!module_snapshot::read(NULL, modules, growing_modules) && modules.empty(), + "perpetually growing module snapshot was accepted"); + require(!module_snapshot::read(NULL, modules, unavailable_modules) && modules.empty() && + GetLastError() == ERROR_ACCESS_DENIED, "module enumeration error was hidden"); + } + + void default_dump_options() + { + PD_OPTIONS options; + require(options.ImportRec && !options.ForceGenHeader && !options.Verbose && !options.ReconstructHeaderAsDll && + options.DumpChunks && options.EntryPointHash && !options.ForceReconstructEntryPoint && + !options.Reexecution && options.NumberOfThreads == 16 && options.EntryPointOverride == -1 && options.output_path[0] == 0, + "default dump/database options are not initialized consistently"); + } + + void concurrent_database_lookup() + { + temporary_file clean, ep, short_ep; + pe_hash_database database(clean.path, ep.path, short_ep.path); + std::atomic valid(true); + std::vector workers; + for (unsigned int worker = 0; worker < 4; ++worker) + workers.emplace_back([&, worker] { + for (unsigned int batch = 0; batch < 300; ++batch) + { + unordered_set values; + for (unsigned int i = 0; i < 64; ++i) + values.insert(1 + worker * 19200 + batch * 64 + i); + database.add_hashes(values); + database.add_hashes_eps(values, values); + for (auto value : values) + if (!database.contains(value) || !database.contains_ep(value) || !database.contains_epshort(value)) + valid = false; + } + }); + for (auto& worker : workers) + worker.join(); + require(valid && database.count() == 76800 && database.count_eps() == 76800 && + database.count_epshorts() == 76800, "concurrent database updates lost hashes"); + } +} + +void append_pipeline_tests(std::vector& tests) +{ + tests.push_back({"pe32-extended-optional-header", [] { extended_optional_header(false); }}); + tests.push_back({"pe64-extended-optional-header", [] { extended_optional_header(true); }}); + tests.push_back({"page-aligned-export", [] { export_discovery(0x1000, true, true, false); }}); + tests.push_back({"ordinal-only-export", [] { export_discovery(0x1100, false, true, false); }}); + tests.push_back({"forwarded-export-is-not-code", [] { export_discovery(0x1880, true, true, true); }}); + tests.push_back({"redirected-export-preserved", [] { export_discovery(0x3000, true, true, false); }}); + tests.push_back({"export-hole-is-not-an-address", [] { export_discovery(0, true, true, true); }}); + tests.push_back({"invalid-export-table-is-atomic", [] { export_discovery(0x1000, true, false, false, true); }}); + tests.push_back({"pe32-terminal-import-summary", [] { terminal_import_summary(false); }}); + tests.push_back({"pe64-terminal-import-summary", [] { terminal_import_summary(true); }}); + tests.push_back({"pe32-generated-image-type", [] { generated_image_type(false); }}); + tests.push_back({"pe64-generated-image-type", [] { generated_image_type(true); }}); + tests.push_back({"raw-address-dump", raw_address_dump}); + tests.push_back({"large-and-changing-module-snapshots", module_snapshot_bounds}); + tests.push_back({"default-dump-options", default_dump_options}); + tests.push_back({"concurrent-database-lookup", concurrent_database_lookup}); +} diff --git a/tests/reconstruction_tests.cpp b/tests/reconstruction_tests.cpp new file mode 100644 index 0000000..80416bb --- /dev/null +++ b/tests/reconstruction_tests.cpp @@ -0,0 +1,399 @@ +#include "test_support.h" + +namespace +{ + template + NT& nt(std::vector& bytes) + { + return *reinterpret_cast(bytes.data() + 0x80); + } + + template + IMAGE_SECTION_HEADER* sections(std::vector& bytes) + { + return reinterpret_cast(&nt(bytes) + 1); + } + + std::vector mapped_fixture(bool win64, SIZE_T size) + { + const auto file = pe_fixture(win64); + std::vector image(size); + memcpy(image.data(), file.data(), 0x200); + memcpy(image.data() + 0x1000, file.data() + 0x200, 0x1000); + return image; + } + + std::vector reconstruct(std::vector& image, bool imports, + unsigned __int64 export_address = 0x12345678) + { + test_options options; + options.ImportRec = imports; + pe_header header(GetCurrentProcess(), image.data(), NULL, &options); + require(header.process_pe_header() && header.process_sections(), "mapped fixture rejected"); + require(header.process_import_directory(), "fixture imports rejected"); + export_list exports; + char library[] = "example.dll", function[] = "ExampleFunction"; + export_entry entry(library, function, 1, 0x1000, export_address, header.is_64()); + exports.add_export(export_address, &entry); + require(header.process_disk_image(&exports, NULL), "reconstruction failed"); + temporary_file output; + require(DeleteFileA(output.path) != 0, "could not reserve output path"); + require(header.write_image(output.path), "reconstruction write failed"); + return output.read(); + } + + template + SIZE_T file_offset(std::vector& bytes, DWORD rva) + { + auto& header = nt(bytes); + for (WORD i = 0; i < header.FileHeader.NumberOfSections; ++i) + { + const auto& section = sections(bytes)[i]; + if (rva >= section.VirtualAddress && rva - section.VirtualAddress < section.SizeOfRawData) + { + const SIZE_T offset = section.PointerToRawData + rva - section.VirtualAddress; + require(offset < bytes.size(), "RVA maps outside the dump"); + return offset; + } + } + throw std::runtime_error("RVA is not backed by section data"); + } + + template + void sparse_data(bool win64, bool imports) + { + auto image = mapped_fixture(win64, 0x402000); + auto& header = nt(image); + header.FileHeader.NumberOfSections = 2; + header.OptionalHeader.SizeOfImage = static_cast(image.size()); + auto& data = sections(image)[1]; + memcpy(data.Name, ".bss", 4); + data.VirtualAddress = 0x2000; + data.Misc.VirtualSize = 0x400000; + data.Characteristics = IMAGE_SCN_CNT_UNINITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE; + image[0x2123] = 0x5a; // Preserve initialized runtime data in an otherwise zero-filled section. + auto dump = reconstruct(image, imports); + require(dump.size() <= (imports ? 0x4000u : 0x3000u), "zero-filled virtual tail expanded on disk"); + require(sections(dump)[1].Misc.VirtualSize == 0x400000, "virtual allocation was truncated"); + require((sections(dump)[1].Characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA) != 0, + "initialized runtime data is still marked uninitialized"); + require(dump[file_offset(dump, 0x2123)] == 0x5a, "runtime-initialized data was lost"); + require(nt(dump).OptionalHeader.SizeOfImage >= image.size(), "virtual image size was replaced by disk size"); + } + + template + void sparse_rvas(bool win64) + { + auto image = mapped_fixture(win64, 0x9000); + auto& header = nt(image); + header.FileHeader.NumberOfSections = 2; + header.OptionalHeader.SizeOfImage = static_cast(image.size()); + auto& data = sections(image)[1]; + memcpy(data.Name, ".data", 5); + data.VirtualAddress = 0x8000; + data.Misc.VirtualSize = data.SizeOfRawData = 0x1000; + data.Characteristics = IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE; + image[0x8fff] = 0x5a; + auto dump = reconstruct(image, false); + require(dump.size() == 0x3000, "RVA gap was materialized on disk"); + require(nt(dump).OptionalHeader.SizeOfImage == 0x9000, "SizeOfImage does not cover section RVAs"); + require(dump[file_offset(dump, 0x8fff)] == 0x5a, "sparse section content lost"); + } + + template + void unaligned_imports(bool win64) + { + auto image = mapped_fixture(win64, 0x2801); + auto& header = nt(image); + header.OptionalHeader.SizeOfImage = static_cast(image.size()); + sections(image)[0].Misc.VirtualSize = 0x1801; + image.back() = 0x5a; + const unsigned __int64 address = 0x12345678; + memcpy(image.data() + 0x1020, &address, win64 ? 8 : 4); + auto dump = reconstruct(image, true); + const auto directory = nt(dump).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT]; + const SIZE_T offset = file_offset(dump, directory.VirtualAddress); + require(range_fits(dump.size(), offset, sizeof(IMAGE_IMPORT_DESCRIPTOR)), "import descriptor outside dump"); + const auto& descriptor = *reinterpret_cast(dump.data() + offset); + const SIZE_T name = file_offset(dump, descriptor.Name); + require(range_fits(dump.size(), name, sizeof("example.dll")), "import name outside dump"); + require(memcmp(dump.data() + name, "example.dll", sizeof("example.dll")) == 0, + "import name RVA uses the unaligned image end"); + } + + template + void large_section(bool win64) + { + const DWORD size = 60 * 1000 * 1024 + 0x2000; + auto image = mapped_fixture(win64, size + 0x1000); + nt(image).OptionalHeader.SizeOfImage = static_cast(image.size()); + auto& section = sections(image)[0]; + section.Misc.VirtualSize = size; + image.back() = 0x5a; + auto dump = reconstruct(image, false); + require(sections(dump)[0].Misc.VirtualSize == size, "valid large section was arbitrarily truncated"); + require(dump[file_offset(dump, static_cast(image.size() - 1))] == 0x5a, + "large section tail lost"); + } + + template + void full_header(bool win64) + { + auto image = mapped_fixture(win64, 0x5000); + nt(image).OptionalHeader.SizeOfImage = static_cast(image.size()); + // Keep header-resident data where an extra section header would otherwise go. + unsigned char* reserved = reinterpret_cast(sections(image) + 1); + memset(reserved, 0x5a, sizeof(IMAGE_SECTION_HEADER)); + auto dump = reconstruct(image, true); + require(nt(dump).FileHeader.NumberOfSections == 1, "header data overwritten by import section"); + require(memcmp(reinterpret_cast(sections(dump) + 1), reserved, + sizeof(IMAGE_SECTION_HEADER)) == 0, "header-resident data lost"); + const auto directory = nt(dump).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT]; + require(directory.VirtualAddress == 0x5000, "import section start changed"); + file_offset(dump, directory.VirtualAddress + directory.Size - 1); + require(sections(dump)[0].Misc.VirtualSize == 0x5000, "fallback does not cover the import table"); + } + + template + void file_only_metadata(bool win64) + { + auto image = mapped_fixture(win64, 0x2000); + auto& header = nt(image); + header.OptionalHeader.CheckSum = 0x1234; + header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY] = {0x2200, 0x200}; + header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG] = {0x1100, sizeof(IMAGE_DEBUG_DIRECTORY)}; + auto& debug = *reinterpret_cast(image.data() + 0x1100); + debug = {}; + debug.SizeOfData = 32; + debug.AddressOfRawData = 0x1200; + debug.PointerToRawData = 0x400; + auto dump = reconstruct(image, false); + const auto& output_header = nt(dump); + const auto& certificate = output_header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY]; + require(certificate.VirtualAddress == 0 && certificate.Size == 0, "unmapped certificate offset retained"); + require(output_header.OptionalHeader.CheckSum == 0, "stale checksum retained"); + const auto& output_debug = *reinterpret_cast(dump.data() + file_offset(dump, 0x1100)); + require(output_debug.PointerToRawData == file_offset(dump, 0x1200), "debug payload file offset was not relocated"); + } + + template + void terminal_import(bool win64) + { + auto image = mapped_fixture(win64, 0x2000); + const unsigned __int64 address = win64 ? 0x7fff12345678ULL : 0xf2345678ULL; + const SIZE_T width = win64 ? 8 : 4; + memcpy(image.data() + image.size() - width, &address, width); + auto dump = reconstruct(image, true, address); + const auto& directory = nt(dump).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT]; + require(directory.Size == 2 * sizeof(IMAGE_IMPORT_DESCRIPTOR), "final pointer-sized import was omitted"); + const auto& descriptor = *reinterpret_cast(dump.data() + file_offset(dump, directory.VirtualAddress)); + require(descriptor.FirstThunk == image.size() - width, "terminal import references the wrong location"); + } + + template + void layout_matrix(bool win64) + { + for (int variant = 0; variant < 16; ++variant) + { + auto image = mapped_fixture(win64, 0x40000); + auto& header = nt(image); + header.OptionalHeader.SizeOfHeaders = 0x400; + header.FileHeader.NumberOfSections = 6; + DWORD rva = 0x1000; + for (int i = 0; i < 6; ++i) + { + auto& section = sections(image)[i]; + section = {}; + memcpy(section.Name, ".data", 5); + section.VirtualAddress = rva; + section.Misc.VirtualSize = 0x1000 + ((variant + i) % 4) * 0x1000 + i * 17; + section.SizeOfRawData = i % 2 ? 0 : 0x200; + section.Characteristics = IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE; + memset(image.data() + rva, 0, section.Misc.VirtualSize); + if (i % 3) + { + image[rva + 17] = 0x5a; + image[rva + section.Misc.VirtualSize - 1] = 0x6b; + } + rva = (rva + section.Misc.VirtualSize + 0xfff) & ~0xfff; + rva += (variant % 3) * 0x1000; + } + header.OptionalHeader.SizeOfImage = rva; + auto dump = reconstruct(image, variant % 2 != 0); + require(nt(dump).OptionalHeader.SizeOfImage >= rva, "matrix virtual extent shrank"); + temporary_file source; + source.write(dump); + test_options options; + pe_header reloaded(source.path, &options); + require(reloaded.process_pe_header() && reloaded.process_sections(), "reconstructed layout cannot be reloaded"); + for (int i = 0; i < 6; ++i) + { + const auto& before = sections(image)[i]; + const auto& after = sections(dump)[i]; + require(after.VirtualAddress == before.VirtualAddress && after.Misc.VirtualSize == before.Misc.VirtualSize, + "matrix section virtual layout changed"); + require(after.SizeOfRawData % 0x1000 == 0 && after.PointerToRawData % 0x1000 == 0, + "matrix raw alignment invalid"); + require(range_fits(dump.size(), after.PointerToRawData, after.SizeOfRawData), "matrix raw range invalid"); + for (DWORD j = 0; j < before.Misc.VirtualSize; ++j) + { + const unsigned char value = j < after.SizeOfRawData ? dump[after.PointerToRawData + j] : 0; + require(value == image[before.VirtualAddress + j], "matrix lost section content"); + } + } + } + } + + template + void import_section_contents(bool win64, bool has_lookup) + { + auto image = mapped_fixture(win64, 0x3000); + auto& header = nt(image); + header.OptionalHeader.SizeOfImage = static_cast(image.size()); + sections(image)[0].Misc.VirtualSize = 0x2000; + header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT] = {0x1800, 40}; + memset(image.data() + 0x1800, 0, 0x400); + auto& descriptor = *reinterpret_cast(image.data() + 0x1800); + descriptor.Name = 0x1900; + descriptor.FirstThunk = 0x1b00; + descriptor.OriginalFirstThunk = has_lookup ? 0x1a00 : 0; + memcpy(image.data() + 0x1900, "example.dll", sizeof("example.dll")); + memcpy(image.data() + 0x1922, "ExampleFunction", sizeof("ExampleFunction")); + const SIZE_T width = win64 ? 8 : 4; + const unsigned __int64 lookup[] = {0x1920, (win64 ? IMAGE_ORDINAL_FLAG64 : IMAGE_ORDINAL_FLAG32) | 7, 0}; + const unsigned __int64 resolved[] = {win64 ? 0x7fff12345678ULL : 0x72345678ULL, + win64 ? 0x7fff23456789ULL : 0x73456789ULL, has_lookup ? 0x76543210ULL : 0}; + for (int i = 0; i < 3; ++i) + { + memcpy(image.data() + 0x1a00 + i * width, &lookup[i], width); + memcpy(image.data() + 0x1b00 + i * width, &resolved[i], width); + } + auto expected = image; + if (has_lookup) + memcpy(expected.data() + 0x1b00, expected.data() + 0x1a00, 3 * width); + for (bool imports : {false, true}) + { + auto dump = reconstruct(image, imports); + const auto& section = sections(dump)[0]; + require(section.Misc.VirtualSize == 0x2000, "original section size changed"); + for (SIZE_T i = 0; i < 0x2000; ++i) + { + const unsigned char actual = i < section.SizeOfRawData ? dump[section.PointerToRawData + i] : 0; + require(actual == expected[0x1000 + i], + has_lookup ? "section content differs from complete pointer-width IAT restoration" : + "missing OriginalFirstThunk corrupted section bytes using DOS header data"); + } + } + } + + template + void truncated_lookup_preserves_content(bool win64) + { + const DWORD width = win64 ? 8 : 4; + const DWORD locations[][2] = {{0x3000 - width, 0x1b00}, {0x1a00, 0x3000 - width}, + {0x1a00, 0x1a00 + width}, {0x1a00, 0x1a00}}; + for (const auto& location : locations) + { + auto image = mapped_fixture(win64, 0x3000); + auto& header = nt(image); + header.OptionalHeader.SizeOfImage = static_cast(image.size()); + sections(image)[0].Misc.VirtualSize = 0x2000; + header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT] = {0x1800, 40}; + memset(image.data() + 0x1800, 0, 0x400); + auto& descriptor = *reinterpret_cast(image.data() + 0x1800); + descriptor.Name = 0x1900; + descriptor.OriginalFirstThunk = location[0]; + descriptor.FirstThunk = location[1]; + memcpy(image.data() + 0x1900, "example.dll", sizeof("example.dll")); + const unsigned __int64 lookup = 0x1920, resolved = 0x76543210; + memcpy(image.data() + descriptor.OriginalFirstThunk, &lookup, width); + memcpy(image.data() + descriptor.FirstThunk, &resolved, width); + auto dump = reconstruct(image, false); + const auto& section = sections(dump)[0]; + for (SIZE_T i = 0; i < 0x2000; ++i) + { + const unsigned char actual = i < section.SizeOfRawData ? dump[section.PointerToRawData + i] : 0; + require(actual == image[0x1000 + i], "truncated, overlapping or shared thunks rewrote section contents"); + } + } + } + + template + void repeated_import_fixups(bool win64) + { + auto image = mapped_fixture(win64, 0x3000); + auto& header = nt(image); + header.OptionalHeader.SizeOfImage = static_cast(image.size()); + sections(image)[0].Misc.VirtualSize = 0x2000; + header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT] = {0x1800, 40}; + memset(image.data() + 0x1800, 0, 0x400); + auto& descriptor = *reinterpret_cast(image.data() + 0x1800); + descriptor.Name = 0x1900; + descriptor.FirstThunk = 0x1b00; + descriptor.OriginalFirstThunk = 0x1a00; + memcpy(image.data() + 0x1900, "example.dll", sizeof("example.dll")); + memcpy(image.data() + 0x1922, "ExampleFunction", sizeof("ExampleFunction")); + const SIZE_T width = win64 ? 8 : 4; + const unsigned __int64 name = 0x1920, address = 0x12345678, relocated_address = 0x76543210; + memcpy(image.data() + 0x1a00, &name, width); + const DWORD fixups[] = {0x1100, static_cast(0x1100 + width), 0x1b00}; + for (DWORD rva : fixups) + memcpy(image.data() + rva, &address, width); + auto dump = reconstruct(image, true); + const auto directory = nt(dump).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT]; + require(directory.Size == 4 * sizeof(IMAGE_IMPORT_DESCRIPTOR), "per-location repeated imports were removed"); + unsigned int patched = 0; + for (SIZE_T i = 0; i < 3; ++i) + { + const auto& entry = *reinterpret_cast( + dump.data() + file_offset(dump, directory.VirtualAddress) + i * sizeof(IMAGE_IMPORT_DESCRIPTOR)); + const SIZE_T lookup_offset = file_offset(dump, entry.OriginalFirstThunk); + unsigned __int64 lookup = 0; + memcpy(&lookup, dump.data() + lookup_offset, width); + const SIZE_T symbol = file_offset(dump, static_cast(lookup) + 2); + require(memcmp(dump.data() + symbol, "ExampleFunction", sizeof("ExampleFunction")) == 0, + "repeated reference changed import identity"); + require(entry.FirstThunk == 0x1100 || entry.FirstThunk == 0x1100 + width || entry.FirstThunk == 0x1b00, + "import fixup targets the wrong location"); + memcpy(dump.data() + file_offset(dump, entry.FirstThunk), &relocated_address, width); + ++patched; + } + require(patched == 3, "original IAT or repeated-reference fixup missing"); + for (DWORD rva : fixups) + { + unsigned __int64 value = 0; + memcpy(&value, dump.data() + file_offset(dump, rva), width); + require(value == relocated_address, "loader-style rebinding did not repair every reference"); + } + } +} + +void append_reconstruction_tests(std::vector& tests) +{ + tests.push_back({"pe32-sparse-data", [] { sparse_data(false, false); }}); + tests.push_back({"pe64-sparse-data", [] { sparse_data(true, false); }}); + tests.push_back({"pe32-sparse-data-imports", [] { sparse_data(false, true); }}); + tests.push_back({"pe64-sparse-data-imports", [] { sparse_data(true, true); }}); + tests.push_back({"pe32-sparse-rvas", [] { sparse_rvas(false); }}); + tests.push_back({"pe64-sparse-rvas", [] { sparse_rvas(true); }}); + tests.push_back({"pe32-unaligned-imports", [] { unaligned_imports(false); }}); + tests.push_back({"pe64-unaligned-imports", [] { unaligned_imports(true); }}); + tests.push_back({"pe32-large-section", [] { large_section(false); }}); + tests.push_back({"pe64-large-section", [] { large_section(true); }}); + tests.push_back({"pe32-full-header-imports", [] { full_header(false); }}); + tests.push_back({"pe64-full-header-imports", [] { full_header(true); }}); + tests.push_back({"pe32-file-metadata", [] { file_only_metadata(false); }}); + tests.push_back({"pe64-file-metadata", [] { file_only_metadata(true); }}); + tests.push_back({"pe32-terminal-import", [] { terminal_import(false); }}); + tests.push_back({"pe64-terminal-import", [] { terminal_import(true); }}); + tests.push_back({"pe32-layout-matrix", [] { layout_matrix(false); }}); + tests.push_back({"pe64-layout-matrix", [] { layout_matrix(true); }}); + tests.push_back({"pe32-import-section-content", [] { import_section_contents(false, true); }}); + tests.push_back({"pe64-import-section-content", [] { import_section_contents(true, true); }}); + tests.push_back({"pe32-no-lookup-section-content", [] { import_section_contents(false, false); }}); + tests.push_back({"pe64-no-lookup-section-content", [] { import_section_contents(true, false); }}); + tests.push_back({"pe32-truncated-lookup-content", [] { truncated_lookup_preserves_content(false); }}); + tests.push_back({"pe64-truncated-lookup-content", [] { truncated_lookup_preserves_content(true); }}); + tests.push_back({"pe32-repeated-import-fixups", [] { repeated_import_fixups(false); }}); + tests.push_back({"pe64-repeated-import-fixups", [] { repeated_import_fixups(true); }}); +} diff --git a/tests/reexecution_probe.cpp b/tests/reexecution_probe.cpp new file mode 100644 index 0000000..f7a7f87 --- /dev/null +++ b/tests/reexecution_probe.cpp @@ -0,0 +1,434 @@ +#include "test_support.h" +#include +#include +#include +#include + +namespace +{ + class owned_handle + { + public: + HANDLE value; + explicit owned_handle(HANDLE handle = NULL) : value(handle) {} + ~owned_handle() { if (value != NULL && value != INVALID_HANDLE_VALUE) CloseHandle(value); } + owned_handle(const owned_handle&) = delete; + owned_handle& operator=(const owned_handle&) = delete; + }; + + void checked(bool success, const char* operation, DWORD error) + { + if (!success) + throw std::runtime_error(std::string(operation) + ": Windows error " + std::to_string(error)); + } + + void checked(bool success, const char* operation) + { + checked(success, operation, GetLastError()); + } + + std::string quote(const std::string& text) + { + std::string result = "\""; + size_t slashes = 0; + for (char ch : text) + { + if (ch == '\\') { ++slashes; continue; } + result.append(slashes * (ch == '"' ? 2 : 1), '\\'); + slashes = 0; + if (ch == '"') result += '\\'; + result += ch; + } + result.append(slashes * 2, '\\'); + return result + '"'; + } + + std::string json_string(const std::string& text) + { + std::ostringstream output; + output << '"'; + const char hex[] = "0123456789abcdef"; + for (unsigned char ch : text) + { + if (ch == '"' || ch == '\\') output << '\\' << ch; + else if (ch < 32 || ch >= 127) output << "\\u00" << hex[ch >> 4] << hex[ch & 15]; + else output << ch; + } + output << '"'; + return output.str(); + } + + void write_report(const char* path, const std::string& text) + { + owned_handle file(CreateFileA(path, GENERIC_WRITE, 0, NULL, CREATE_NEW, FILE_ATTRIBUTE_NORMAL, NULL)); + checked(file.value != INVALID_HANDLE_VALUE, "Create report"); + DWORD written = 0; + checked(WriteFile(file.value, text.data(), static_cast(text.size()), &written, NULL) && + written == text.size(), "Write report"); + } + + struct module + { + uintptr_t base; + DWORD size; + std::string name; + }; + + struct observation + { + bool loader = false, entry = false, idle = false, timeout = false; + bool capture_attempted = false, captured = false; + DWORD exit_code = STILL_ACTIVE, dump_exit = STILL_ACTIVE; + uintptr_t image_base = 0; + DWORD entry_rva = 0, dlls = 0, threads = 0, dropped_messages = 0; + std::vector markers, exceptions, loader_messages; + + std::string json() const + { + std::ostringstream out; + out << "{\"loader_breakpoint\":" << loader << ",\"entrypoint\":" << entry + << ",\"input_idle\":" << idle << ",\"timed_out\":" << timeout + << ",\"capture_attempted\":" << capture_attempted << ",\"captured\":" << captured + << ",\"exit_code\":" << exit_code << ",\"dumper_exit_code\":" << dump_exit + << ",\"image_base\":" << image_base << ",\"entry_rva\":" << entry_rva + << ",\"loaded_dlls\":" << dlls << ",\"created_threads\":" << threads << ",\"markers\":["; + for (size_t i = 0; i < markers.size(); ++i) out << (i ? "," : "") << json_string(markers[i]); + out << "],\"exceptions\":["; + for (size_t i = 0; i < exceptions.size(); ++i) out << (i ? "," : "") << exceptions[i]; + out << "],\"dropped_loader_messages\":" << dropped_messages << ",\"loader_messages\":["; + for (size_t i = 0; i < loader_messages.size(); ++i) out << (i ? "," : "") << json_string(loader_messages[i]); + out << "]}\n"; + return out.str(); + } + }; + + template bool read_remote(HANDLE process, uintptr_t address, T& value) + { + SIZE_T read = 0; + return ReadProcessMemory(process, reinterpret_cast(address), &value, sizeof(value), &read) && + read == sizeof(value); + } + + bool image_header(HANDLE process, uintptr_t base, IMAGE_NT_HEADERS& header) + { + IMAGE_DOS_HEADER dos = {}; + return read_remote(process, base, dos) && dos.e_magic == IMAGE_DOS_SIGNATURE && + dos.e_lfanew > 0 && dos.e_lfanew < 1024 * 1024 && + read_remote(process, base + dos.e_lfanew, header) && header.Signature == IMAGE_NT_SIGNATURE && + header.OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR_MAGIC; + } + + void breakpoint_byte(HANDLE process, uintptr_t address, unsigned char byte) + { + DWORD protection = 0, ignored = 0; + checked(VirtualProtectEx(process, reinterpret_cast(address), 1, PAGE_EXECUTE_READWRITE, &protection) != 0, + "Protect owned child's entrypoint"); + SIZE_T written = 0; + const bool success = WriteProcessMemory(process, reinterpret_cast(address), &byte, 1, &written) && + written == 1; + const DWORD write_error = GetLastError(); + checked(VirtualProtectEx(process, reinterpret_cast(address), 1, protection, &ignored) != 0, + "Restore owned child's entrypoint protection"); + checked(success, "Write owned child's entrypoint breakpoint", write_error); + checked(FlushInstructionCache(process, reinterpret_cast(address), 1) != 0, "Flush entrypoint"); + } + + owned_handle* redirected_file(const std::string& path, STARTUPINFOA& startup, owned_handle& input) + { + SECURITY_ATTRIBUTES attributes = {sizeof(attributes), NULL, TRUE}; + std::unique_ptr output(new owned_handle(CreateFileA(path.c_str(), GENERIC_WRITE, + FILE_SHARE_READ, &attributes, CREATE_NEW, FILE_ATTRIBUTE_NORMAL, NULL))); + checked(output->value != INVALID_HANDLE_VALUE, "Create child output"); + input.value = CreateFileA("NUL", GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, + &attributes, OPEN_EXISTING, 0, NULL); + checked(input.value != INVALID_HANDLE_VALUE, "Open null input"); + startup.cb = sizeof(startup); + startup.dwFlags = STARTF_USESTDHANDLES | STARTF_USESHOWWINDOW; + startup.wShowWindow = SW_HIDE; + startup.hStdInput = input.value; + startup.hStdOutput = startup.hStdError = output->value; + return output.release(); + } + + DWORD dump_child(HANDLE job, DWORD pid, uintptr_t base, const char* dumper, const char* folder, + bool imports, const char* report) + { + std::ostringstream address; + address << "0x" << std::hex << base; + std::string command = quote(dumper) + " -pid " + std::to_string(pid) + " -a " + address.str() + + " -o " + quote(folder) + " -db ignore -nep -nc -nt -nh -v" + (imports ? "" : " -ni"); + char repair[2] = {}; + if (GetEnvironmentVariableA("PD_REEXEC_PREPARE", repair, 2) == 1 && repair[0] == '1') + command += " -reexec"; + STARTUPINFOA startup = {}; + owned_handle input; + std::unique_ptr output(redirected_file(std::string(report) + ".dump.log", startup, input)); + PROCESS_INFORMATION child = {}; + checked(CreateProcessA(dumper, &command[0], NULL, NULL, TRUE, CREATE_SUSPENDED | CREATE_NO_WINDOW, + NULL, folder, &startup, &child) != 0, "Start dumper"); + owned_handle process(child.hProcess), thread(child.hThread); + if (!AssignProcessToJobObject(job, process.value)) + { + const DWORD error = GetLastError(); + checked(TerminateProcess(process.value, 1) != 0, "Stop unassigned dumper"); + checked(false, "Assign dumper job", error); + } + checked(ResumeThread(thread.value) != MAXDWORD, "Resume dumper"); + const DWORD wait = WaitForSingleObject(process.value, 30000); + if (wait == WAIT_TIMEOUT) + { + checked(TerminateProcess(process.value, WAIT_TIMEOUT) != 0, "Stop timed-out dumper"); + checked(WaitForSingleObject(process.value, 5000) == WAIT_OBJECT_0, "Wait for stopped dumper"); + return WAIT_TIMEOUT; + } + checked(wait == WAIT_OBJECT_0, "Wait for dumper"); + DWORD exit = 0; + checked(GetExitCodeProcess(process.value, &exit) != 0, "Read dumper status"); + return exit; + } + + observation observe(int argc, char** argv) + { + // exe, report, duration-ms, capture-phase, dumper, dump-folder, imports, target arguments... + const char *exe = argv[0], *report = argv[1], *phase = argv[3]; + char snaps_value[2] = {}; + const bool loader_snaps = GetEnvironmentVariableA("PD_REEXEC_LOADER_SNAPS", snaps_value, 2) == 1 && + snaps_value[0] == '1'; + char* end = NULL; + const unsigned long duration = strtoul(argv[2], &end, 10); + require(end != argv[2] && *end == 0 && duration >= 100 && duration <= 30000, "duration must be 100..30000 ms"); + require(strcmp(phase, "none") == 0 || strcmp(phase, "entry") == 0 || + strcmp(phase, "ready") == 0 || strcmp(phase, "idle") == 0, "unsupported capture phase"); + require(strcmp(argv[6], "0") == 0 || strcmp(argv[6], "1") == 0, "imports must be 0 or 1"); + DWORD binary_type = 0; + checked(GetBinaryTypeA(exe, &binary_type) != 0, "Read executable type"); + require(binary_type == (sizeof(void*) == 8 ? SCS_64BIT_BINARY : SCS_32BIT_BINARY), + "use the probe matching the executable architecture"); + std::string command = quote(exe); + for (int i = 7; i < argc; ++i) command += " " + quote(argv[i]); + STARTUPINFOA startup = {}; + owned_handle input; + std::unique_ptr output(redirected_file(std::string(report) + ".stdout", startup, input)); + owned_handle job(CreateJobObjectA(NULL, NULL)); + checked(job.value != NULL, "Create owned-child job"); + JOBOBJECT_EXTENDED_LIMIT_INFORMATION limits = {}; + limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; + checked(SetInformationJobObject(job.value, JobObjectExtendedLimitInformation, &limits, sizeof(limits)) != 0, + "Set owned-child job cleanup"); + PROCESS_INFORMATION child = {}; + checked(CreateProcessA(exe, &command[0], NULL, NULL, TRUE, + DEBUG_ONLY_THIS_PROCESS | CREATE_SUSPENDED | CREATE_NO_WINDOW, NULL, NULL, &startup, &child) != 0, + "Start owned debug child"); + owned_handle process(child.hProcess), thread(child.hThread); + if (!AssignProcessToJobObject(job.value, process.value)) + { + const DWORD error = GetLastError(); + checked(TerminateProcess(process.value, 1) != 0, "Stop unassigned debug child"); + checked(false, "Assign owned-child job", error); + } + checked(ResumeThread(thread.value) != MAXDWORD, "Resume owned debug child"); + observation result; + std::map modules; + uintptr_t entry = 0; + unsigned char original_byte = 0; + bool armed = false, stopped = false, exited = false; + const ULONGLONG started = GetTickCount64(); + ULONGLONG paused = 0, stopping = 0; + const auto capture = [&]() { + result.capture_attempted = true; + const ULONGLONG before = GetTickCount64(); + result.dump_exit = dump_child(job.value, child.dwProcessId, result.image_base, argv[4], argv[5], + strcmp(argv[6], "1") == 0, report); + paused += GetTickCount64() - before; + result.captured = result.dump_exit == 0; + }; + while (!exited) + { + if (!stopped && GetTickCount64() - started - paused >= duration) + { + result.timeout = true; + checked(TerminateJobObject(job.value, WAIT_TIMEOUT) != 0, "Stop owned timed-out children"); + stopped = true; + stopping = GetTickCount64(); + } + require(!stopped || GetTickCount64() - stopping < 10000, "debug child cleanup exceeded deadline"); + DEBUG_EVENT event = {}; + if (!WaitForDebugEvent(&event, 50)) + { + checked(GetLastError() == ERROR_SEM_TIMEOUT, "Wait for owned child debug event"); + if (!stopped && result.entry && WaitForInputIdle(process.value, 0) == 0) + { + result.idle = true; + if (!result.capture_attempted && strcmp(phase, "idle") == 0) capture(); + } + continue; + } + DWORD disposition = DBG_CONTINUE; + if (event.dwDebugEventCode == CREATE_PROCESS_DEBUG_EVENT || event.dwDebugEventCode == LOAD_DLL_DEBUG_EVENT) + { + const bool main = event.dwDebugEventCode == CREATE_PROCESS_DEBUG_EVENT; + owned_handle file(main ? event.u.CreateProcessInfo.hFile : event.u.LoadDll.hFile); + const uintptr_t base = reinterpret_cast(main ? + event.u.CreateProcessInfo.lpBaseOfImage : event.u.LoadDll.lpBaseOfDll); + IMAGE_NT_HEADERS header = {}; + const bool valid = image_header(process.value, base, header); + char path[MAX_PATH] = {}; + const DWORD length = GetFinalPathNameByHandleA(file.value, path, MAX_PATH, 0); + std::string name = length > 0 && length < MAX_PATH ? path : "unknown"; + const size_t slash = name.find_last_of("\\/"); + if (slash != std::string::npos) name.erase(0, slash + 1); + modules[base] = {base, valid ? header.OptionalHeader.SizeOfImage : 0, name}; + if (main) + { + if (loader_snaps) + { + PROCESS_BASIC_INFORMATION basic = {}; + require(NtQueryInformationProcess(process.value, ProcessBasicInformation, &basic, + sizeof(basic), NULL) >= 0, "query owned-child PEB failed"); + // Diagnostic-only Windows x86/x64 PEB layout; never writes a system-wide setting. + const uintptr_t flags_address = reinterpret_cast(basic.PebBaseAddress) + + (sizeof(void*) == 8 ? 0xbc : 0x68); + DWORD flags = 0; + checked(read_remote(process.value, flags_address, flags), "Read owned-child loader flags"); + flags |= 2; // FLG_SHOW_LDR_SNAPS + SIZE_T written = 0; + checked(WriteProcessMemory(process.value, reinterpret_cast(flags_address), &flags, + sizeof(flags), &written) && written == sizeof(flags), "Set owned-child loader diagnostics"); + } + require(valid && header.OptionalHeader.AddressOfEntryPoint != 0 && + header.OptionalHeader.AddressOfEntryPoint < header.OptionalHeader.SizeOfImage, "invalid native entrypoint"); + result.image_base = base; + result.entry_rva = header.OptionalHeader.AddressOfEntryPoint; + entry = base + result.entry_rva; + checked(read_remote(process.value, entry, original_byte), "Read child entrypoint byte"); + require(original_byte != 0xcc, "entrypoint already contains a breakpoint"); + breakpoint_byte(process.value, entry, 0xcc); + armed = true; + } + else ++result.dlls; + } + else if (event.dwDebugEventCode == CREATE_THREAD_DEBUG_EVENT) ++result.threads; + else if (event.dwDebugEventCode == UNLOAD_DLL_DEBUG_EVENT) + modules.erase(reinterpret_cast(event.u.UnloadDll.lpBaseOfDll)); + else if (event.dwDebugEventCode == EXCEPTION_DEBUG_EVENT) + { + const auto& exception = event.u.Exception.ExceptionRecord; + const uintptr_t address = reinterpret_cast(exception.ExceptionAddress); + std::string name = "unknown"; + uintptr_t rva = address; + auto found = modules.upper_bound(address); + if (found != modules.begin()) + { + --found; + if (address - found->first < found->second.size) + { + name = found->second.name; + rva = address - found->first; + } + } + if (exception.ExceptionCode == EXCEPTION_BREAKPOINT && armed && address == entry) + { + checked(event.dwThreadId == child.dwThreadId, "Unexpected entrypoint thread"); + breakpoint_byte(process.value, entry, original_byte); + CONTEXT context = {}; + context.ContextFlags = CONTEXT_CONTROL; + checked(GetThreadContext(thread.value, &context) != 0, "Read entrypoint context"); +#ifdef _WIN64 + context.Rip = entry; +#else + context.Eip = static_cast(entry); +#endif + checked(SetThreadContext(thread.value, &context) != 0, "Restore entrypoint instruction pointer"); + armed = false; + result.entry = true; + if (strcmp(phase, "entry") == 0) capture(); + } + else if (exception.ExceptionCode == EXCEPTION_BREAKPOINT && event.u.Exception.dwFirstChance && + !result.loader && !result.entry && _stricmp(name.c_str(), "ntdll.dll") == 0) + result.loader = true; + else + { + disposition = DBG_EXCEPTION_NOT_HANDLED; + if (result.exceptions.size() < 32) + { + std::ostringstream item; + item << "{\"code\":" << exception.ExceptionCode << ",\"first_chance\":" + << event.u.Exception.dwFirstChance << ",\"module\":" << json_string(name) + << ",\"rva\":" << rva << ",\"parameters\":["; + for (DWORD i = 0; i < exception.NumberParameters && i < EXCEPTION_MAXIMUM_PARAMETERS; ++i) + item << (i ? "," : "") << exception.ExceptionInformation[i]; + item << "]}"; + result.exceptions.push_back(item.str()); + } + if (!event.u.Exception.dwFirstChance && !stopped) + { + checked(TerminateJobObject(job.value, exception.ExceptionCode) != 0, "Stop failed owned child"); + stopped = true; + stopping = GetTickCount64(); + } + } + } + else if (event.dwDebugEventCode == OUTPUT_DEBUG_STRING_EVENT) + { + const auto& info = event.u.DebugString; + if (!info.fUnicode && info.nDebugStringLength > 0 && info.nDebugStringLength <= 4096) + { + char text[4097] = {}; + SIZE_T read = 0; + if (ReadProcessMemory(process.value, info.lpDebugStringData, text, info.nDebugStringLength, &read) && + read == info.nDebugStringLength) + { + if (loader_snaps) + { + if (result.loader_messages.size() == 256) + { + result.loader_messages.erase(result.loader_messages.begin() + 32); + ++result.dropped_messages; + } + result.loader_messages.push_back(text); + } + if (strncmp(text, "PD_REEXEC_", 10) == 0) + { + if (result.markers.size() < 32) result.markers.push_back(text); + if (!result.capture_attempted && strcmp(phase, "ready") == 0 && strcmp(text, "PD_REEXEC_READY") == 0) + capture(); + } + } + } + } + else if (event.dwDebugEventCode == EXIT_PROCESS_DEBUG_EVENT) + { + result.exit_code = event.u.ExitProcess.dwExitCode; + exited = true; + } + checked(ContinueDebugEvent(event.dwProcessId, event.dwThreadId, disposition) != 0, "Continue owned debug child"); + } + return result; + } +} + +int run_reexecution_probe(int argc, char** argv) +{ + try + { + require(argc >= 7, "usage: --reexecute exe report duration-ms none|entry|ready|idle dumper dump-folder 0|1 [args]"); + const observation result = observe(argc, argv); + write_report(argv[1], result.json()); + return 0; + } + catch (const std::exception& error) + { + fprintf(stderr, "Reexecution probe failed: %s\n", error.what()); + return 1; + } +} + +int run_reexecution_fixture() +{ + OutputDebugStringA("PD_REEXEC_READY"); + std::thread worker([] { OutputDebugStringA("PD_REEXEC_THREAD"); }); + worker.join(); + puts("PD_REEXEC_OK"); + OutputDebugStringA("PD_REEXEC_DONE"); + return 0; +} diff --git a/tests/reexecution_tests.cpp b/tests/reexecution_tests.cpp new file mode 100644 index 0000000..930de31 --- /dev/null +++ b/tests/reexecution_tests.cpp @@ -0,0 +1,330 @@ +#include "test_support.h" + +namespace +{ + template + NT& nt(std::vector& bytes) + { + return *reinterpret_cast(bytes.data() + 0x80); + } + + template + std::vector fixture(bool win64) + { + auto bytes = pe_fixture(win64); + bytes.resize(0x3200); + auto& header = nt(bytes); + header.FileHeader.NumberOfSections = 2; + header.OptionalHeader.SizeOfImage = 0x4000; + header.OptionalHeader.DllCharacteristics = IMAGE_DLLCHARACTERISTICS_GUARD_CF | + IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE | IMAGE_DLLCHARACTERISTICS_NX_COMPAT; + auto& section = reinterpret_cast(&header + 1)[1]; + memcpy(section.Name, ".data", 5); + section.VirtualAddress = 0x2000; + section.Misc.VirtualSize = section.SizeOfRawData = 0x2000; + section.PointerToRawData = 0x1200; + section.Characteristics = IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE; + return bytes; + } + + template + SIZE_T file_offset(std::vector& bytes, DWORD rva) + { + const auto& header = nt(bytes); + const auto* sections = reinterpret_cast(&header + 1); + for (WORD i = 0; i < header.FileHeader.NumberOfSections; ++i) + if (rva >= sections[i].VirtualAddress && + range_fits(sections[i].SizeOfRawData, rva - sections[i].VirtualAddress, 1)) + return sections[i].PointerToRawData + rva - sections[i].VirtualAddress; + throw std::runtime_error("test RVA not present in reconstructed file"); + } + + std::vector serialize(pe_header& image, bool expected_success = true) + { + temporary_file output; + require(image.process_pe_header() && image.process_sections() && image.process_import_directory(), + "reexecution fixture parsing failed"); + export_list exports; + char library[] = "example.dll", name[] = "Example"; + export_entry entry(library, name, 1, 0x1000, 0x12345678, image.is_64()); + exports.add_export(0x12345678, &entry); + require(image.process_disk_image(&exports, NULL) == expected_success, "unexpected preparation result"); + if (!expected_success) + return {}; + require(DeleteFileA(output.path) && image.write_image(output.path), "prepared image write failed"); + return output.read(); + } + + std::vector reconstruct(const std::vector& bytes, bool prepare, bool imports, + bool expected_success = true) + { + temporary_file input; + input.write(bytes); + test_options options; + options.Reexecution = prepare; + options.ImportRec = imports; + pe_header image(input.path, &options); + auto result = serialize(image, expected_success); + require(input.read() == bytes, "preparation changed the source file"); + return result; + } + + template + std::vector import_slots(std::vector& bytes); + + template + void cookie_state(bool win64) + { + auto bytes = fixture(win64); + const SIZE_T width = win64 ? 8 : 4; + auto& config = *reinterpret_cast(bytes.data() + 0x1400); + config.Size = sizeof(config); + config.SecurityCookie = nt(bytes).OptionalHeader.ImageBase + 0x2a00; + nt(bytes).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG] = {0x2200, sizeof(config)}; + const ULONGLONG captured = win64 ? 0x123456789abcULL : 0x12345678ULL; + memcpy(bytes.data() + 0x1c00, &captured, width); + bytes[0x1bff] = bytes[0x1c00 + width] = 0x5a; + for (bool imports : {false, true}) + for (bool prepare : {false, true}) + { + auto dump = reconstruct(bytes, prepare, imports); + const SIZE_T cookie = file_offset(dump, 0x2a00); + ULONGLONG actual = 0; + memcpy(&actual, dump.data() + cookie, width); + const ULONGLONG expected = prepare ? (win64 ? 0x2b992ddfa232ULL : 0xbb40e64eULL) : captured; + require(actual == expected && dump[cookie - 1] == 0x5a && dump[cookie + width] == 0x5a, + "cookie reset was missing, unconditional, or exceeded the pointer width"); + require(nt(dump).OptionalHeader.DllCharacteristics == nt(bytes).OptionalHeader.DllCharacteristics, + "preparation disabled a loader mitigation"); + if (prepare && imports) + { + const auto slots = import_slots(dump); + require(std::find(slots.begin(), slots.end(), 0x2a00) == slots.end(), + "GS cookie was mistaken for a function pointer"); + } + } + } + + template + void cookie_bounds(bool win64) + { + auto bytes = fixture(win64); + const SIZE_T width = win64 ? 8 : 4; + auto& config = *reinterpret_cast(bytes.data() + 0x1400); + config.Size = sizeof(config); + nt(bytes).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG] = {0x2200, sizeof(config)}; + const auto base = nt(bytes).OptionalHeader.ImageBase; + for (DWORD rva : {0x100u, 0x1100u, 0x4000u, 0x3fffu}) + { + config.SecurityCookie = base + rva; + reconstruct(bytes, true, false, false); + } + config.SecurityCookie = base - 1; + reconstruct(bytes, true, false, false); + config.SecurityCookie = 0; + reconstruct(bytes, true, false); + config.SecurityCookie = base + 0x2a00; + config.Size = offsetof(CONFIG, SecurityCookie); + bytes[0x1c00] = 0x5a; + auto old_version = reconstruct(bytes, true, false); + require(old_version[file_offset(old_version, 0x2a00)] == 0x5a, "older load-config version was overread"); + config.Size = sizeof(config); + config.SecurityCookie = base + 0x4000 - width; + bytes[bytes.size() - width] = 0x5a; + auto terminal = reconstruct(bytes, true, false); + ULONGLONG actual = 0; + memcpy(&actual, terminal.data() + file_offset(terminal, static_cast(0x4000 - width)), width); + require(actual == (win64 ? 0x2b992ddfa232ULL : 0xbb40e64eULL), "terminal cookie was rejected or truncated"); + nt(bytes).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG].VirtualAddress = 0x3ff0; + reconstruct(bytes, true, false, false); + } + + template + std::vector import_slots(std::vector& bytes) + { + const auto& directory = nt(bytes).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT]; + const SIZE_T offset = file_offset(bytes, directory.VirtualAddress); + std::vector result; + for (SIZE_T i = 0; i < directory.Size / sizeof(IMAGE_IMPORT_DESCRIPTOR); ++i) + { + IMAGE_IMPORT_DESCRIPTOR item; + memcpy(&item, bytes.data() + offset + i * sizeof(item), sizeof(item)); + if (item.FirstThunk == 0) + return result; + result.push_back(item.FirstThunk); + } + throw std::runtime_error("prepared imports lack a terminator"); + } + + template + void put_fixups(std::vector& bytes, bool win64) + { + const SIZE_T width = win64 ? 8 : 4; + const ULONGLONG address = 0x12345678; + const DWORD slots[] = {0x1100, 0x2400, static_cast(0x2400 + width), static_cast(0x4000 - width)}; + for (DWORD rva : slots) + memcpy(bytes.data() + file_offset(bytes, rva), &address, width); + } + + template + void conservative_imports(bool win64) + { + auto bytes = fixture(win64); + put_fixups(bytes, win64); + const DWORD width = win64 ? 8 : 4; + auto normal = reconstruct(bytes, false, true); + auto prepared = reconstruct(bytes, true, true); + require(import_slots(normal) == std::vector({0x1100, 0x2400, 0x2400 + width, 0x4000 - width}), + "normal aggressive reconstruction changed"); + require(import_slots(prepared) == std::vector({0x2400, 0x2400 + width, 0x4000 - width}), + "conservative imports dropped independent writable references or included code"); + require(memcmp(prepared.data() + file_offset(prepared, 0x1100), bytes.data() + 0x300, width) == 0, + "conservative scanning rewrote code bytes"); + auto& data = reinterpret_cast(&nt(bytes) + 1)[1]; + for (DWORD flags : {static_cast(IMAGE_SCN_MEM_READ), + static_cast(IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE | IMAGE_SCN_MEM_EXECUTE)}) + { + data.Characteristics = flags; + auto excluded = reconstruct(bytes, true, true); + require(import_slots(excluded).empty(), "readonly or executable import destination included"); + } + } + + template + void delay_imports(bool win64, bool malformed) + { + auto bytes = fixture(win64); + put_fixups(bytes, win64); + auto& directory = nt(bytes).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT]; + directory = {0x2800, 64}; + auto* descriptor = reinterpret_cast(bytes.data() + 0x1a00); + descriptor[0] = 1; + descriptor[3] = 0x2400; + if (malformed) + { + for (DWORD iat : {0u, 0x4000u, 0xffffffffu}) + { + descriptor[3] = iat; + reconstruct(bytes, true, true, false); + } + descriptor[3] = 0x2400; + directory.Size = 32; + reconstruct(bytes, true, true, false); + directory = {0x3ff0, 64}; + reconstruct(bytes, true, true, false); + return; + } + auto prepared = reconstruct(bytes, true, true); + require(import_slots(prepared).empty(), "protected delay-IAT section became a normal import destination"); + auto normal = reconstruct(bytes, false, true); + require(import_slots(normal).size() == 4, "delay imports changed the default analysis mode"); + if (!win64) + { + descriptor[0] = 0; + descriptor[3] = static_cast(nt(bytes).OptionalHeader.ImageBase) + 0x2400; + auto legacy = reconstruct(bytes, true, true); + require(import_slots(legacy).empty(), "legacy VA delay IAT was not excluded"); + } + } + + template + void zero_fill(bool win64) + { + for (DWORD flags : {static_cast(IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE), + static_cast(IMAGE_SCN_MEM_READ), static_cast(IMAGE_SCN_MEM_WRITE | IMAGE_SCN_MEM_EXECUTE)}) + for (DWORD raw_size : {0u, 0x1000u}) + for (bool imports : {false, true}) + for (bool occupied : {false, true}) + for (bool prepare : {false, true}) + { + auto file = fixture(win64); + std::vector memory(0x4000); + memcpy(memory.data(), file.data(), 0x200); + memcpy(memory.data() + 0x1000, file.data() + 0x200, 0x1000); + auto& header = nt(memory); + auto* sections = reinterpret_cast(&header + 1); + auto& data = sections[1]; + data.SizeOfRawData = raw_size; + data.Characteristics = flags | (raw_size ? IMAGE_SCN_CNT_INITIALIZED_DATA : IMAGE_SCN_CNT_UNINITIALIZED_DATA); + if (occupied) + reinterpret_cast(sections + 2)[0] = 0x5a; + const bool writable = (flags & IMAGE_SCN_MEM_WRITE) && !(flags & IMAGE_SCN_MEM_EXECUTE); + const SIZE_T width = win64 ? 8 : 4; + const ULONGLONG address = 0x12345678, captured_cookie = 0x1122334455667788ULL; + for (DWORD rva : {0x2400u, 0x3200u}) + memcpy(memory.data() + rva, &address, width); + memory[0x2800] = 0x5a; + if (writable) + { + auto& config = *reinterpret_cast(memory.data() + 0x1400); + memset(&config, 0, sizeof(config)); + config.Size = sizeof(config); + config.SecurityCookie = static_cast( + reinterpret_cast(memory.data()) + 0x4000 - width); + header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG] = {0x1400, sizeof(config)}; + memcpy(memory.data() + 0x4000 - width, &captured_cookie, width); + } + const auto before = memory; + test_options options; + options.Reexecution = prepare; + options.ImportRec = imports; + pe_header image(GetCurrentProcess(), memory.data(), NULL, &options); + auto dump = serialize(image); + require(memory == before, "zero-fill restoration changed captured memory"); + const auto& output_data = reinterpret_cast(&nt(dump) + 1)[1]; + const auto byte_at = [&](DWORD rva) { + return rva - output_data.VirtualAddress < output_data.SizeOfRawData ? + dump[file_offset(dump, rva)] : static_cast(0); + }; + const bool reset = prepare && writable; + require(byte_at(0x3200) == (reset ? 0 : 0x78), "zero-fill data was not restored only when requested"); + require(byte_at(0x2800) == (reset && raw_size == 0 ? 0 : 0x5a), + "file-backed data or a fully zero-fill section was handled incorrectly"); + if (imports) + { + const std::vector expected = !prepare ? std::vector({0x2400, 0x3200}) : + (writable && raw_size != 0 ? std::vector({0x2400}) : std::vector()); + require(import_slots(dump) == expected, "zero-fill data became a speculative import"); + } + if (writable) + { + ULONGLONG actual = 0; + memcpy(&actual, dump.data() + file_offset(dump, static_cast(0x4000 - width)), width); + const ULONGLONG expected = prepare ? (win64 ? 0x2b992ddfa232ULL : 0xbb40e64eULL) : + (win64 ? captured_cookie : static_cast(captured_cookie)); + require(actual == expected, "a cookie in zero-fill storage was lost or not initialized"); + } + } + auto file = fixture(win64); + std::vector memory(0x4000); + memcpy(memory.data(), file.data(), 0x200); + auto& header = nt(memory); + header.FileHeader.NumberOfSections = 3; + auto* sections = reinterpret_cast(&header + 1); + sections[1].SizeOfRawData = 0x1000; + sections[1].Misc.VirtualSize = 0x3000; + sections[2] = {}; + sections[2].VirtualAddress = 0x3000; + sections[2].Misc.VirtualSize = 0x1000; + test_options options; + options.Reexecution = true; + pe_header image(GetCurrentProcess(), memory.data(), NULL, &options); + serialize(image, false); + } +} + +void append_reexecution_tests(std::vector& tests) +{ + tests.push_back({"pe32-reexecution-cookie", [] { cookie_state(false); }}); + tests.push_back({"pe64-reexecution-cookie", [] { cookie_state(true); }}); + tests.push_back({"pe32-reexecution-cookie-bounds", [] { cookie_bounds(false); }}); + tests.push_back({"pe64-reexecution-cookie-bounds", [] { cookie_bounds(true); }}); + tests.push_back({"pe32-reexecution-imports", [] { conservative_imports(false); }}); + tests.push_back({"pe64-reexecution-imports", [] { conservative_imports(true); }}); + tests.push_back({"pe32-reexecution-delay-imports", [] { delay_imports(false, false); }}); + tests.push_back({"pe64-reexecution-delay-imports", [] { delay_imports(true, false); }}); + tests.push_back({"pe32-reexecution-delay-bounds", [] { delay_imports(false, true); }}); + tests.push_back({"pe64-reexecution-delay-bounds", [] { delay_imports(true, true); }}); + tests.push_back({"pe32-reexecution-zero-fill", [] { zero_fill(false); }}); + tests.push_back({"pe64-reexecution-zero-fill", [] { zero_fill(true); }}); +} diff --git a/tests/scheduling_tests.cpp b/tests/scheduling_tests.cpp new file mode 100644 index 0000000..d01f965 --- /dev/null +++ b/tests/scheduling_tests.cpp @@ -0,0 +1,164 @@ +#include "test_support.h" +#include "system_work.h" +#include "module_work.h" +#include +#include +#include + +namespace +{ + void scheduler_lifetimes() + { + for (int threads : {1, 4}) + { + work_pool pool(threads); + std::atomic completed(0), active(0), peak(0); + for (int process = 0; process < 20; ++process) + pool.submit([&] { + std::vector modules(513, 1); + pool.parallel_for(modules, [&](int item) { + const int running = ++active; + int observed = peak.load(); + while (running > observed && !peak.compare_exchange_weak(observed, running)) {} + completed += item; + --active; + }); + require(active >= 0, "module lifetime ended early"); + }); + pool.wait(); + require(completed == 20 * 513 && active == 0 && peak <= threads, + "nested module work lost jobs, returned early or exceeded the worker limit"); + pool.submit([&] { ++completed; }); + pool.wait(); + require(completed == 20 * 513 + 1, "workers did not survive a drained queue"); + } + } + + void scheduler_errors() + { + work_pool pool(2); + std::atomic completed(0); + pool.submit([&] { + std::vector modules = {1, 2, 3, 4}; + pool.parallel_for(modules, [&](int item) { + ++completed; + if (item == 2) throw std::runtime_error("expected module failure"); + }); + }); + bool reported = false; + try { pool.wait(); } + catch (const std::runtime_error& error) { reported = strcmp(error.what(), "expected module failure") == 0; } + require(reported && completed == 4, "module error was hidden or siblings were not drained"); + } + + void final_process_uses_idle_workers() + { + work_pool pool(4); + std::mutex mutex; + std::condition_variable ready; + int active = 0, peak = 0; + pool.submit([&] { + std::vector modules(32); + pool.parallel_for(modules, [&](int) { + std::unique_lock lock(mutex); + peak = (std::max)(peak, ++active); + ready.notify_all(); + const bool concurrent = ready.wait_for(lock, std::chrono::seconds(5), [&] { return peak >= 2; }); + --active; + require(concurrent, "only one worker processed the final process's modules"); + }); + }); + pool.wait(); + require(peak >= 2 && peak <= 4 && active == 0, "tail module concurrency or lifetime is incorrect"); + } + + void final_process_scan() + { + work_pool pool(3); + std::atomic completed(0); + int scans = 0; + system_work(pool, true, [&] { + if (++scans == 1) return std::vector{1, 2, 2}; + require(completed == 2, "final discovery ran before in-flight module work completed"); + return std::vector{2, 3}; + }, [&](DWORD pid, work_pool& workers) { + std::vector modules{pid, pid}; + std::atomic count(0); + workers.parallel_for(modules, [&](DWORD) { ++count; }); + require(count == 2, "process finished before module jobs"); + ++completed; + }); + require(scans == 2 && completed == 3, "final pass missed new processes or repeated old PIDs"); + scans = 0; + system_work(pool, false, [&] { ++scans; return std::vector{}; }, + [](DWORD, work_pool&) { throw std::runtime_error("empty inventory dispatched work"); }); + require(scans == 1, "hash generation unexpectedly rescanned processes"); + } + + void heap_boundaries() + { + std::set heaps{0x1000, 0x2000, 0x3000, 0x4000}; + exclude_module_heaps(heaps, 0x2000, 0x2000); + require(heaps == std::set({0x1000, 0x4000}), "adjacent loose heap was removed"); + } + + void parallel_database_files() + { + temporary_file location, source; + require(DeleteFileA(location.path) && CreateDirectoryA(location.path, NULL), "create database test folder"); + const std::string root(location.path), nested = root + "\\nested"; + std::vector files; + const auto cleanup = [&](void*) { + for (const auto& path : files) DeleteFileA(path.c_str()); + RemoveDirectoryA(nested.c_str()); + RemoveDirectoryA(root.c_str()); + }; + std::unique_ptr guard(reinterpret_cast(1), cleanup); + require(CreateDirectoryA(nested.c_str(), NULL) != 0, "create nested database test folder"); + auto first = pe_fixture(false), second = pe_fixture(true); + auto* second_header = reinterpret_cast(second.data() + 0x80); + memcpy(reinterpret_cast(second_header + 1)->Name, ".other", 6); + first[0x200] = second[0x200] = 0x31; + first[0x201] = second[0x201] = 0xc0; + for (int i = 0; i < 270; ++i) + { + source.write(first); + const auto path = root + "\\" + std::to_string(i) + ".exe"; + require(CopyFileA(source.path, path.c_str(), TRUE) != 0, "copy database fixture"); + files.push_back(path); + } + source.write(second); + for (const auto& path : {nested + "\\second.exe", root + "\\ignored.bin"}) + { + require(CopyFileA(source.path, path.c_str(), TRUE) != 0, "copy nested database fixture"); + files.push_back(path); + } + test_options options; + options.EntryPointHash = true; + pe_header other(source.path, &options); + require(other.process_pe_header() && other.process_sections(), "database reference parse"); + const auto second_hash = other.get_hash(); + for (bool recursive : {false, true}) + { + temporary_file clean1, full1, short1, clean4, full4, short4; + pe_hash_database serial(clean1.path, full1.path, short1.path), parallel(clean4.path, full4.path, short4.path); + require(serial.add_folder(location.path, L"*.exe", recursive, 1) && + parallel.add_folder(location.path, L"*.exe", recursive, 4), "folder traversal failed"); + require(serial.count() == parallel.count() && serial.contains(second_hash) == recursive && + parallel.contains(second_hash) == recursive && + serial.snapshot_entrypoints()->full == parallel.snapshot_entrypoints()->full && + serial.snapshot_entrypoints()->short_hashes == parallel.snapshot_entrypoints()->short_hashes, + "parallel traversal changed hashes, recursion or filename filtering"); + } + } +} + +void append_scheduling_tests(std::vector& tests) +{ + tests.push_back({"module-scheduler-lifetimes", scheduler_lifetimes}); + tests.push_back({"module-scheduler-errors", scheduler_errors}); + tests.push_back({"system-tail-module-concurrency", final_process_uses_idle_workers}); + tests.push_back({"system-final-scan", final_process_scan}); + tests.push_back({"module-heap-boundaries", heap_boundaries}); + tests.push_back({"parallel-database-folders", parallel_database_files}); +} diff --git a/tests/system_benchmark.cpp b/tests/system_benchmark.cpp new file mode 100644 index 0000000..c691152 --- /dev/null +++ b/tests/system_benchmark.cpp @@ -0,0 +1,102 @@ +#include "test_support.h" +#include "dump_process.h" +#include "system_work.h" +#include +#include +#include + +int run_system_fixture() +{ + try + { + std::vector>> images; + for (int i = 0; i < 64; ++i) + { + auto bytes = pe_fixture(sizeof(void*) == 8); + void* memory = VirtualAlloc(NULL, 0x202000, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); + require(memory != NULL, "allocate inert system fixture"); + images.emplace_back(memory, [](void* block) { VirtualFree(block, 0, MEM_RELEASE); }); + auto* base = static_cast(memory); + memcpy(base, bytes.data(), 0x200); + memcpy(base + 0x1000, bytes.data() + 0x200, 0x1000); + auto* header = reinterpret_cast(base + 0x80); + header->OptionalHeader.ImageBase = reinterpret_cast(memory); + header->OptionalHeader.SizeOfImage = 0x202000; + auto* section = reinterpret_cast(header + 1); + sprintf_s(reinterpret_cast(section->Name), 8, "mod%03d", i); + section->Misc.VirtualSize = 0x201000; + } + puts("PD_SYSTEM_READY"); + fflush(stdout); + Sleep(600000); + return 0; + } + catch (const std::exception& error) + { + fprintf(stderr, "System fixture failed: %s\n", error.what()); + return 1; + } +} + +int run_system_benchmark(int argc, char** argv) +{ + try + { + require(argc == 6, "usage: --system-work pid dump|hash threads output report flags(i,c,g,r)"); + const DWORD pid = strtoul(argv[0], NULL, 10); + const int threads = atoi(argv[2]); + require(pid != 0 && threads >= 0 && threads <= 64, "invalid process or worker count"); + test_options options; + options.EntryPointHash = true; + options.ImportRec = strchr(argv[5], 'i') != NULL; + options.DumpChunks = strchr(argv[5], 'c') != NULL; + options.ForceGenHeader = strchr(argv[5], 'g') != NULL; + options.Reexecution = strchr(argv[5], 'r') != NULL; + delete[] options.output_path; + options.output_path = new char[strlen(argv[3]) + 1]; + strcpy(options.output_path, argv[3]); + temporary_file clean, ep, prefix; + pe_hash_database database(clean.path, ep.path, prefix.path); + const bool hashing = strcmp(argv[1], "hash") == 0; + require(hashing || strcmp(argv[1], "dump") == 0, "invalid system workload"); + unordered_set hashes, full, prefixes; + const auto action = [&](work_pool* pool) { + dump_process process(pid, &database, &options, true); + if (!hashing) process.dump_all(pool); + process.get_all_hashes(&hashes, hashing ? &full : NULL, hashing ? &prefixes : NULL, pool); + database.add_hashes(hashes); + if (hashing) database.add_hashes_eps(full, prefixes); + }; + const auto started = std::chrono::steady_clock::now(); + if (threads) + { + work_pool pool(threads); + system_work(pool, !hashing, [=] { return std::vector{pid}; }, + [&](DWORD, work_pool& workers) { action(&workers); }); + } + else action(NULL); + const double ms = std::chrono::duration(std::chrono::steady_clock::now() - started).count(); + std::ofstream output(argv[4]); + require(output.good(), "cannot open system benchmark report"); + output << "{\"milliseconds\":" << ms << ",\"hashes\":["; + const auto write = [&](const unordered_set& values) { + bool first = true; + for (auto value : std::set(values.begin(), values.end())) + { + if (!first) output << ','; + first = false; + output << value; + } + }; + write(hashes); output << "],\"full\":["; + write(full); output << "],\"prefixes\":["; + write(prefixes); output << "]}\n"; + require(output.good(), "cannot write system benchmark report"); + return 0; + } + catch (const std::exception& error) + { + fprintf(stderr, "System benchmark failed: %s\n", error.what()); + return 1; + } +} diff --git a/tests/test_benchmark_dumps.py b/tests/test_benchmark_dumps.py new file mode 100644 index 0000000..528559f --- /dev/null +++ b/tests/test_benchmark_dumps.py @@ -0,0 +1,291 @@ +import struct +import tempfile +import unittest +from unittest.mock import patch +from pathlib import Path +from types import SimpleNamespace +import subprocess +import json + +from benchmark_dumps import benchmark_running, compare_code, compare_reports, file_metrics, pe_info +from pe_quality import compare_quality, PE +from benchmark_performance import METRICS, compare as compare_performance, read_result, summarize + + +def fixture(): + data = bytearray(0x1200) + data[:2] = b"MZ" + struct.pack_into("& tests); void append_name_safety_tests(std::vector& tests); void append_hook_tests(std::vector& tests); void append_output_tests(std::vector& tests); +void append_reconstruction_tests(std::vector& tests); +void append_pipeline_tests(std::vector& tests); +void append_performance_tests(std::vector& tests); +void append_reexecution_tests(std::vector& tests); +void append_scheduling_tests(std::vector& tests); +void append_entrypoint_tests(std::vector& tests); +int run_performance_benchmark(int argc, char** argv); +int run_reexecution_probe(int argc, char** argv); +int run_reexecution_fixture(); +int run_entrypoint_benchmark(int argc, char** argv); +int run_system_fixture(); +int run_system_benchmark(int argc, char** argv); inline void require(bool condition, const char* message) {