diff --git a/.github/workflows/build-test.yml b/.github/workflows/build-test.yml index 2c7ef92e..3db31f4c 100644 --- a/.github/workflows/build-test.yml +++ b/.github/workflows/build-test.yml @@ -14,7 +14,7 @@ jobs: - uses: actions/setup-go@v7 with: - go-version: 1.24.x + go-version: 1.25.x - name: Test run: go test ./... @@ -23,19 +23,7 @@ jobs: run: go build . working-directory: cmd/shuffledns/ - # Installing massdns - - uses: actions/checkout@v7 - with: - repository: blechschmidt/massdns - path: ./massdns - - - name: Install massdns - working-directory: ./massdns - run: | - make all - sudo make install - - # Tests + # Integration tests (native resolver; massdns binary not required) - name: Bruteforce Mode Test run: go run . -v -d scanme.sh -r ../../tests/resolvers.txt -w ../../tests/wordlist.txt -mode bruteforce working-directory: cmd/shuffledns/ diff --git a/.github/workflows/lint-test.yml b/.github/workflows/lint-test.yml index 4aea1696..53302da0 100644 --- a/.github/workflows/lint-test.yml +++ b/.github/workflows/lint-test.yml @@ -11,7 +11,7 @@ jobs: - name: Set up Go uses: actions/setup-go@v7 with: - go-version: 1.24.x + go-version: 1.25.x - name: Checkout code uses: actions/checkout@v7 diff --git a/.github/workflows/release-binary.yml b/.github/workflows/release-binary.yml index c857669d..f27f58e6 100644 --- a/.github/workflows/release-binary.yml +++ b/.github/workflows/release-binary.yml @@ -13,7 +13,7 @@ jobs: - uses: actions/checkout@v7 - uses: actions/setup-go@v7 with: - go-version: 1.24.x + go-version: 1.25.x - uses: goreleaser/goreleaser-action@v7 with: args: "release --clean" diff --git a/.github/workflows/release-test.yml b/.github/workflows/release-test.yml index 1a11dc71..d0ec7f40 100644 --- a/.github/workflows/release-test.yml +++ b/.github/workflows/release-test.yml @@ -19,7 +19,7 @@ jobs: - name: Set up Go uses: actions/setup-go@v7 with: - go-version: 1.24.x + go-version: 1.25.x - name: release test uses: goreleaser/goreleaser-action@v7 diff --git a/README.md b/README.md index ebb0e958..1f1af2de 100644 --- a/README.md +++ b/README.md @@ -3,8 +3,7 @@
-

massDNS wrapper to bruteforce and resolve the subdomains with wildcard handling support

- +

Native high-throughput DNS bruteforce and resolve with wildcard handling

@@ -13,50 +12,59 @@

- +

- Feature • - Install • + Features • + Install • + Usage • Run • - Wildcard • - License • - Discord + Wildcards • + massdns parity • + License

--- +`shuffledns` bruteforces and resolves subdomains with multi-level wildcard filtering. Resolution is done by a **native Go stub resolver** (massdns-compatible design and output). An optional **iterative** mode walks from the DNS roots and needs no public resolver list. -`shuffleDNS` is a wrapper around `massdns`, written in go, that allows you to enumerate valid subdomains using active bruteforce, as well as resolve subdomains with wildcard handling and easy input-output support. - -Based on the work on `massdns` project by [@blechschmidt](https://github.com/blechschmidt). +Inspired by [massdns](https://github.com/blechschmidt/massdns) by [@blechschmidt](https://github.com/blechschmidt). The massdns binary is **not required**. - # Features +# Features

shuffledns

- - Simple and modular code base making it easy to contribute. - - Fast And Simple active subdomain scanning. - - Handles wildcard subdomains in a smart manner. - - Optimized for **ease of use** - - **Stdin** and **stdout** support for integrating in workflows +- Native high-throughput resolver (no massdns binary) +- Bruteforce, resolve, and filter modes with stdin/stdout +- Multi-level wildcard handling +- Optional iterative resolution from the DNS roots (`-it`) +- Resolver health scoring, adaptive concurrency, Linux sendmmsg/recvmmsg batching +- Shard / resume for distributed or long runs + +# Installation + +`go1.24+` required: + +```bash +go install -v github.com/projectdiscovery/shuffledns/cmd/shuffledns@latest +``` + +You still need a resolver list for the default stub mode (e.g. from [dnsvalidator](https://github.com/vortexau/dnsvalidator)) unless you use `-it` / `--iterative`. # Usage ```bash shuffledns -h ``` -This will display help for the tool. Here are all the switches it supports. ```yaml -shuffleDNS is a wrapper around massdns written in go that allows you to enumerate valid subdomains using active bruteforce as well as resolve subdomains with wildcard handling and easy input-output support. +shuffleDNS is a high-throughput DNS bruteforcer and resolver with wildcard handling. It uses a native Go stub resolver (massdns-compatible) and optional iterative resolution from the DNS roots. Usage: ./shuffledns [flags] -Flags: Flags: INPUT: -d, -domain string[] Domain to find or resolve subdomains for @@ -65,15 +73,12 @@ INPUT: -w, -wordlist string File containing words to bruteforce for domain -r, -resolver string File containing list of resolvers for enumeration -tr, -trusted-resolver string File containing list of trusted resolvers - -ri, -raw-input string Validate raw full massdns output + -ri, -raw-input string Filter wildcards from an existing massdns-format output file -mode string Execution mode (bruteforce, resolve, filter) RATE-LIMIT: - -t int Number of concurrent massdns resolves (default 10000) - -UPDATE: - -up, -update update shuffledns to latest version - -duc, -disable-update-check disable automatic shuffledns update check + -t int Max concurrent in-flight DNS queries (default 10000) + -qps int Max outbound DNS queries per second (0 = unlimited) OUTPUT: -o, -output string File to write output to (optional) @@ -81,14 +86,34 @@ OUTPUT: -wo, -wildcard-output string Dump wildcard ips to output file CONFIGURATIONS: - -m, -massdns string Path to the massdns binary - -mcmd, -massdns-cmd string Optional massdns commands to run (example '-i 10') + -m, -massdns string Deprecated (ignored): massdns binary is not used + -mcmd, -massdns-cmd string Deprecated (ignored) -directory string Temporary directory for enumeration OPTIMIZATIONS: - -retries int Number of retries for dns enumeration (default 5) - -sw, -strict-wildcard Perform wildcard check on all found subdomains - -wt int Number of concurrent wildcard checks (default 250) + -retries int Number of retries for dns enumeration (default 5) + -sw, -strict-wildcard Perform wildcard check on all found subdomains + -wt int Number of concurrent wildcard checks (default 250) + -filter-internal-ips Filter out internal/private IP addresses + +RESOLVER: + -rt, -type string DNS record type (A, AAAA, CNAME, NS, PTR, MX, TXT, SOA) + -bm, -batch-mode string sendmmsg/recvmmsg: off | on | adaptive (Linux, default off) + -sc, -socket-count int UDP sockets per run (0 = scale to cores) + -udp-size int EDNS0 UDP payload size (0 = 1232) + -norecurse Send non-recursive queries (RD=0) + -sticky Do not rotate resolver on retry + -rhz, -resolver-health De-weight failing resolvers + -acy, -adaptive-concurrency Adapt in-flight concurrency to packet loss + -cc, -cross-check Re-verify positive answers on a second resolver + -ei, -extended-input Parse 'name [resolver ...]' input lines + -no-verify-ip Disable reply source-IP verification + -no-tcp-fallback Disable TCP fallback on truncated answers + -it, -iterative Recurse from root servers (no -r needed) + +DISTRIBUTED: + -shard string Process only shard m of n (e.g. 2/8) + -rs, -resume string Checkpoint file for crash-safe stop/resume DEBUG: -silent Show only subdomains in output @@ -97,83 +122,96 @@ DEBUG: -nc, -no-color Don't Use colors in output ``` - - - - -
- -## Prerequisite +# Running shuffledns -`shuffledns` requires `massdns` to be installed in order to perform its operations. You can see the installation instructions at [massdns project](https://github.com/blechschmidt/massdns#compilation). If you place the binary in `/usr/bin/massdns` or `/usr/local/bin/massdns`, the tool will auto-detect the presence of the binary and use it. On Windows, you need to supply the path to the binary for the tool to work. +### Resolve -The tool also needs a list of valid resolvers. The [dnsvalidator](https://github.com/vortexau/dnsvalidator) project can be used to generate these lists. You also need to provide wordlist, you can use a custom wordlist or use the [commonspeak2-wordlist](https://wordlists-cdn.assetnote.io/data/manual/best-dns-wordlist.txt). - -
+```bash +shuffledns -d example.com -list example-subdomains.txt -r resolvers.txt -mode resolve +``` -## Installation Instructions +```bash +subfinder -d example.com | shuffledns -d example.com -r resolvers.txt -mode resolve +``` -`shuffledns` requires `go1.24+` to install successfully. Run the following command to install the latest version: +### Bruteforce ```bash -go install -v github.com/projectdiscovery/shuffledns/cmd/shuffledns@latest +shuffledns -d hackerone.com -w wordlist.txt -r resolvers.txt -mode bruteforce ``` -## Running shuffledns +```bash +echo hackerone.com | shuffledns -w wordlist.txt -r resolvers.txt -mode bruteforce +``` -`shuffledns` supports two types of operations: +### Iterative (no public resolvers) -**Subdomain resolving** +```bash +shuffledns -d example.com -w wordlist.txt -mode bruteforce -it +``` -To resolve a list of subdomains, you can pass the list of subdomains via the `-list` option. +### Filter existing massdns-format output ```bash -shuffledns -d example.com -list example-subdomains.txt -r resolvers.txt -mode resolve +shuffledns -d example.com -ri massdns-output.txt -mode filter ``` -This will run the tool against subdomains in `example-subdomains.txt` and returns the results. The tool uses the resolvers specified with `-r` flag to do the resolving. +### Tuning -You can also pass the list of subdomains at standard input (STDIN). This allows for easy integration in automation pipelines. +- `-t` caps **in-flight** queries (massdns `-s`). More concurrency does not help once resolvers or RTT are the limit. +- `-qps` caps **send rate** when you need to stay under resolver/abuse limits. +- `-bm adaptive` helps on high-RTT / bursty paths; leave `off` on LAN/low latency (default). +- `-rhz` / `-acy` help when public resolvers drop or rate-limit. -```bash -subfinder -d example.com | shuffledns -d example.com -r resolvers.txt -mode resolve -``` +Live QPS is roughly `min(-t / RTT, -qps, resolver capacity)`. For loopback and public-resolver head-to-heads see [`bench/`](bench/). -This uses the subdomains found passively by `subfinder` and resolves them with `shuffledns` returning only the unique and valid subdomains. +# Handling Wildcards -**Subdomain Bruteforcing** +`shuffledns` tracks how many names map to each IP. Past a small threshold it walks hostname labels for that IP and filters wildcard answers with few extra DNS requests. Wildcard filtering requires `-d` / domain input. -`shuffledns` also supports bruteforce of a target with a given wordlist. You can use the `w` flag to pass a wordlist which will be used to generate permutations that will be resolved using massdns. +# massdns parity -```bash -shuffledns -d hackerone.com -w wordlist.txt -r resolvers.txt -mode bruteforce -``` +`go run ./cmd/resolve` aims to be a **massdns CLI drop-in** for common stub workloads: + +| Area | Status | +|---|---| +| Stub resolve (`-s/-c/-i/-t/-r`, sticky, norecurse, verify-ip, extended-input, socket-count) | Yes | +| Output `-o` S / F / L / J / **B** (+ modifiers) | Yes (`pkg/output`) | +| `--bindto`, `--rcvbuf`, `--sndbuf`, `--predictable`, `--flush`, `--filter/--ignore/--retry` | Yes | +| `--status-format`, `-q`, `-l` | Yes | +| PTR / validate / AXFR / NSEC(3) / iterative / shard / resume | Yes (native extras) | +| `--drop-user` / `--drop-group` / `--root` | Yes (Unix; after sockets open) | +| `--rand-src-ipv6` / `--rand-src-ipv6-file` | Yes (Linux + `CAP_NET_RAW`; IPv6 resolvers; not with `--bindto`) | +| `--processes`, `--busy-poll` | Accepted, **ignored** | -This will run the tool against `hackerone.com` with the wordlist `wordlist.txt`. The domain bruteforce can also be done with standard input as in previous example for resolving the subdomains. +**shuffledns** itself is not a massdns replacement (hostname list + wildcards). Use `cmd/resolve` when you need massdns-compatible output and flags. + +Example: ```bash -echo hackerone.com | shuffledns -w wordlist.txt -r resolvers.txt -mode bruteforce +go run ./cmd/resolve -r resolvers.txt -t AAAA -o Snl -w out.txt names.txt +# same shape as: massdns -r resolvers.txt -t AAAA -o Snl -w out.txt names.txt ``` ---- - - - - - -
+# Throughput -## Handling Wildcards +Local loopback bench (`RESOLVE_BENCH=1 go test ./pkg/resolve -run TestResolverBenchmark`, 50k names, 8 simulated resolvers, `-t 10000`): -A special feature of `shuffleDNS` is its ability to handle multi-level DNS based wildcards, and do it so with a very reduced number of DNS requests. Sometimes all the subdomains would resolve, leading to lots of garbage in the results. The way `shuffleDNS` handles this is by keeping track of how many subdomains point to an IP, and if the number of subdomains increase beyond a certain small threshold, it checks for wildcard on all the levels of the hosts for that IP iteratively. +| scenario | ~qps | +|---|---| +| lan-fast (~0.5ms) | ~90k | +| wan-typical (15±10ms, 0.5% loss) | ~40k | +| wan-lossy (25±20ms, 5% loss) | ~27k | +| rate-limited (3k qps/resolver) | ~40k | -
+Compare against the massdns binary with [`bench/`](bench/) (Docker). Numbers are workload- and resolver-bound; public resolver lists will land closer to the wan/rate-limited rows than lan-fast. -### Notes +# Notes -- Wildcard filter feature works with domain (`-d`) input only. -- Resolving or Brute-forcing only one operation can be done at a time. +- Resolving and bruteforcing are separate modes (`-mode`). +- `-m` / `-mcmd` / `-retain-stderr` / `-batch-size` are accepted for compatibility and ignored. +- README usage dump may lag slightly behind `-h` as flags evolve. ### License -`shuffledns` is distributed under [GPL v3 License](https://github.com/projectdiscovery/shuffledns/blob/main/LICENSE.md) +`shuffledns` is distributed under [GPL v3 License](https://github.com/projectdiscovery/shuffledns/blob/main/LICENSE.md). diff --git a/bench/Dockerfile b/bench/Dockerfile new file mode 100644 index 00000000..78168a9a --- /dev/null +++ b/bench/Dockerfile @@ -0,0 +1,35 @@ +# Benchmark image: builds massdns from source and the dnsbench harness, then +# benchmarks the native resolver against massdns using a battery of loopback +# resolvers. No DNS traffic leaves the container (resolvers are 127.0.0.1). +# +# Build (context must be the repository root): +# docker build -f bench/Dockerfile -t shuffledns-dnsbench . +# +# Run: +# docker run --rm shuffledns-dnsbench -names 200000 -resolvers 16 -hit 5 +# +# The container needs no network at runtime; everything is loopback. + +FROM golang:1.24-bookworm AS build + +# --- build massdns from source ------------------------------------------------- +RUN apt-get update \ + && apt-get install -y --no-install-recommends git ca-certificates gcc make libc6-dev \ + && rm -rf /var/lib/apt/lists/* +RUN git clone --depth 1 https://github.com/blechschmidt/massdns.git /opt/massdns \ + && make -C /opt/massdns + +# --- build the dnsbench harness ------------------------------------------------ +WORKDIR /src +COPY go.mod go.sum ./ +RUN go mod download +COPY . . +RUN CGO_ENABLED=0 go build -o /usr/local/bin/dnsbench ./cmd/dnsbench + +# --- runtime image ------------------------------------------------------------- +FROM debian:bookworm-slim +COPY --from=build /opt/massdns/bin/massdns /usr/local/bin/massdns +COPY --from=build /usr/local/bin/dnsbench /usr/local/bin/dnsbench + +ENTRYPOINT ["/usr/local/bin/dnsbench", "-massdns", "/usr/local/bin/massdns"] +CMD ["-names", "100000", "-resolvers", "8", "-hit", "5"] diff --git a/bench/Dockerfile.iterbench b/bench/Dockerfile.iterbench new file mode 100644 index 00000000..6b55f808 --- /dev/null +++ b/bench/Dockerfile.iterbench @@ -0,0 +1,25 @@ +# Offline iterative-vs-stub accuracy/throughput benchmark. Builds the iterbench +# harness and runs it entirely on loopback: a real-socket authoritative +# hierarchy (127.0.0.x) for the iterative engine and a simulated flaky +# recursive-resolver battery (127.0.0.1) for the stub engine. No DNS traffic +# leaves the container. +# +# Build (context must be the repository root): +# docker build -f bench/Dockerfile.iterbench -t shuffledns-iterbench . +# +# Run: +# docker run --rm shuffledns-iterbench -domains 200 -hosts 250 -hijack 0.15 +# +# Needs Linux loopback (127.0.0.x); requires no network at runtime. + +FROM golang:1.24-bookworm AS build +WORKDIR /src +COPY go.mod go.sum ./ +RUN go mod download +COPY . . +RUN CGO_ENABLED=0 go build -o /usr/local/bin/iterbench ./cmd/iterbench + +FROM debian:bookworm-slim +COPY --from=build /usr/local/bin/iterbench /usr/local/bin/iterbench +ENTRYPOINT ["/usr/local/bin/iterbench"] +CMD ["-domains", "200", "-hosts", "250", "-hijack", "0.15"] diff --git a/bench/README.md b/bench/README.md new file mode 100644 index 00000000..870cd410 --- /dev/null +++ b/bench/README.md @@ -0,0 +1,193 @@ +# DNS resolver benchmark + +Benchmarks the native Go resolver (`pkg/resolve`) against the external +`massdns` binary using a battery of **loopback** DNS servers that simulate +remote recursive resolvers (`internal/simdns`). + +No DNS traffic ever leaves the host: every resolver an engine talks to is a +`127.0.0.1:` UDP server started in-process. The simulated resolvers +model real remote conditions: + +- response latency + jitter +- packet loss (exercises the timeout/retransmit path) +- SERVFAIL rate (exercises retry-on-rotation) +- per-resolver QPS cap (models throttling public resolvers) +- a deterministic, mostly-NXDOMAIN hit ratio (realistic bruteforce workload) + +The answered/NXDOMAIN decision and synthetic A record are deterministic per +name, so both engines resolve the exact same workload under identical +conditions. + +## Scenarios + +| scenario | latency | loss | servfail | per-resolver QPS | +|---------------|--------------|------|----------|------------------| +| `lan-fast` | ~0.5 ms | - | - | unlimited | +| `wan-typical` | 15 ± 10 ms | 0.5% | - | unlimited | +| `wan-lossy` | 25 ± 20 ms | 5% | 2% | unlimited | +| `rate-limited`| 10 ± 10 ms | - | - | 3000 / resolver | + +## Running with Docker (recommended) + +`massdns` is Linux-only (epoll), so the benchmark runs in a container that +builds massdns from source. The build context must be the repository root: + +```bash +docker build -f bench/Dockerfile -t shuffledns-dnsbench . +docker run --rm shuffledns-dnsbench -names 200000 -resolvers 16 -hit 5 +``` + +Useful flags (`dnsbench -h`): + +- `-names` names to resolve per scenario (default 50000) +- `-resolvers` number of simulated loopback resolvers (default 8) +- `-hit` percent of names that resolve, rest NXDOMAIN (default 5) +- `-concurrency` in-flight concurrency / massdns hashmap size (default 10000) +- `-retries` retry budget per name, applied to both engines (default 5) +- `-sockets` native udp socket count (0 = scale to cores) +- `-batch-mode` native batching: `off` (default) | `on` | `adaptive` +- `-engines` `native,massdns` (default both) +- `-scenarios` comma-separated scenario names or `all` + +### A note on `-batch-mode` (sendmmsg/recvmmsg) + +The native resolver can batch datagrams per syscall using Linux +`sendmmsg`/`recvmmsg` (via `golang.org/x/net/ipv4`, IPv4-only resolver sets). +Batching is workload-dependent, so there are three modes: + +- `off` (default): always one datagram per syscall — the proven, lowest-latency + path. Best on loopback / low-RTT links. +- `on`: always batch. Helps on **high-latency / bursty** links, where many + in-flight responses arrive clustered in time and fill batches, amortizing the + syscall cost. **Counterproductive on loopback / low-RTT** links: batches stay + tiny (1–2 datagrams), so the per-call message-array setup costs more than a + plain `sendto`/`recvfrom`, and it adds first-response latency. +- `adaptive`: start in single mode and let a lightweight controller toggle + batching at runtime based on observed conditions (see below). + +On this loopback benchmark you can see both effects: `-batch-mode on` roughly +halves `lan-fast` throughput but modestly improves the higher-latency scenarios, +while `-batch-mode adaptive` stays near single-path performance on `lan-fast` +and ramps batching up only on the latency-/loss-bound scenarios. + +#### How `adaptive` decides + +A controller goroutine samples every 200ms and toggles batching with hysteresis: + +- **Smoothed RTT** (EWMA of observed round-trip times): batching needs latency + for packets to cluster. Engage above ~3ms, disengage below ~1ms. +- **In-flight depth** (`len(sem)`): batches can only fill if the pipeline is + deep. Engage only when depth ≥ 2×batch size; disengage when it falls below a + batch. +- **Interval packet loss** (retransmits ÷ sends per tick): rising loss often + means a resolver or kernel buffer is already saturated, and bursty `sendmmsg` + makes that worse — so loss above ~15% forces batching **off** to spread sends + out. + +Toggling only changes the send/recv strategy; results are identical either way. +When batching is disengaged the send path uses plain `sendto` (it also falls +back to single sends for batches smaller than 4 datagrams) and the read path +uses single `recvfrom`, so "adaptive-off" costs the same as the default path. + +Cross-platform behaviour of the batch API (`ReadBatch`/`WriteBatch`): + +- **Linux**: real `recvmmsg`/`sendmmsg` (batched). +- **macOS, Windows, *BSD, others**: the same calls transparently fall back to a + single `recvmsg`/`sendmsg` per call, so the code is portable and correct + everywhere; only Linux gets the kernel batch syscalls. + +### Read-path parallelism (the client-side equivalent of `SO_REUSEPORT`) + +`SO_REUSEPORT` is a *server* mechanism for sharing one well-known port across +sockets so the kernel fans incoming packets out to multiple readers. A stub +resolver is a *client*: it already opens N independent sockets on distinct +ephemeral ports, each drained by its own reader goroutine, which delivers the +same per-core receive parallelism without a shared in-flight map or port +juggling. The socket pool therefore scales with `GOMAXPROCS` by default +(`-sockets 0`); binding those sockets to a single port via `SO_REUSEPORT` would +yield identical queue/reader counts with no throughput gain, so it is +deliberately not used. + +Both engines are given matching timeout/interval and retry budgets so the +comparison is fair; the exact massdns flags used are documented in +`cmd/dnsbench/main.go`. + +## Native-only (no Docker) + +The native engine is pure Go and runs anywhere: + +```bash +go run ./cmd/dnsbench -engines native -names 200000 -resolvers 16 +``` + +The same battery is also wired into a Go test, which can toggle the adaptive +resolver features so you can see their effect per scenario: + +```bash +# baseline +RESOLVE_BENCH=1 go test ./pkg/resolve -run TestResolverBenchmark -v \ + -bench.names 100000 -bench.resolvers 16 + +# adaptive batching (helps deep/high-latency pipelines) +RESOLVE_BENCH=1 go test ./pkg/resolve -run TestResolverBenchmark -v \ + -bench.names 100000 -bench.resolvers 16 -bench.batch adaptive + +# health scoring + adaptive concurrency (helps lossy / rate-limited resolvers) +RESOLVE_BENCH=1 go test ./pkg/resolve -run TestResolverBenchmark -v \ + -bench.names 100000 -bench.resolvers 16 -bench.health -bench.adaptconc +``` + +## Resolver capabilities (parity with / beyond massdns) + +Core massdns stub features plus several massdns TODOs. Main shuffledns flags: + +| Capability | shuffledns flag | Notes | +|---|---|---| +| Record type | `-rt A\|AAAA\|CNAME\|NS\|PTR\|MX\|TXT\|SOA` | massdns `-t`; bruteforce store is A/AAAA-oriented | +| In-flight cap | `-t` | massdns `-s` | +| Send-rate cap | `-qps` | token bucket; 0 = unlimited | +| Non-recursive | `-norecurse` | massdns `--norecurse` (RD=0) | +| Sticky resolver | `-sticky` | massdns `--sticky` | +| EDNS0 UDP size | `-udp-size N` | OPT advertised payload (default 1232) | +| Source-IP verify | on by default, `-no-verify-ip` to disable | massdns `--verify-ip` (opt-out here) | +| TCP fallback | on by default, `-no-tcp-fallback` to disable | follows truncated (TC) answers | +| Batching | `-bm off\|on\|adaptive` | sendmmsg/recvmmsg (Linux) | +| Socket count | `-sc N` | massdns `--socket-count` | +| Per-resolver health | `-rhz` | de-weights failing resolvers (massdns TODO) | +| Adaptive concurrency | `-acy` | AIMD in-flight cap on loss (massdns TODO) | +| Cross-resolver check | `-cc` | poisoning/spam detection (massdns TODO) | +| Extended input | `-ei` | `name [resolver ...]` lines (massdns `--extended-input`) | +| Iterative from roots | `-it` | no public resolver list | + +### Privilege drop & raw IPv6 source (cmd/resolve) + +| Feature | Flags | Notes | +|---|---|---| +| Drop root after open | `--drop-user` / `--drop-group` / `--root` | Unix; defaults to `nobody` when euid=0 unless `--root` | +| Random IPv6 source | `--rand-src-ipv6 PREFIX` / `--rand-src-ipv6-file` | Linux `SOCK_RAW` + `IPV6_HDRINCL`; needs `CAP_NET_RAW`; IPv6 resolvers only; incompatible with `--bindto` | + +### Intentionally ignored (accepted for CLI drop-in) + +`--processes`, `--busy-poll`. + +`-o B`, `--bindto`, `--rcvbuf`/`--sndbuf`, `--predictable`, `--flush` are +implemented on `cmd/resolve`. Full `-o` formatting lives there, not in the +shuffledns hostname-list output. + +### QPS notes + +- Effective throughput ≈ `min(-t / RTT, -qps, Σ resolver capacity)`. +- Default `-bm off` is correct for low-RTT; forced `on` often hurts lan-fast. +- Public-resolver runs are usually RTT/rate-limit bound, not CPU bound. + +Example native-only numbers (50k names, 8 sim resolvers, `-t 10000`): + +| scenario | ~qps | +|---|---| +| lan-fast | ~90k | +| wan-typical | ~40k | +| wan-lossy | ~27k | +| rate-limited (3k/resolver) | ~40k | + +Retry rcode policy matches massdns (retry everything except NOERROR / NXDOMAIN). +Progress is available via `OnProgress` / `Stats` on the resolver client. diff --git a/cmd/dnsbench/main.go b/cmd/dnsbench/main.go new file mode 100644 index 00000000..1ae1343f --- /dev/null +++ b/cmd/dnsbench/main.go @@ -0,0 +1,383 @@ +// Command dnsbench benchmarks DNS resolution engines against a battery of +// loopback resolvers that simulate remote ones (see internal/simdns). It can +// drive both the native resolver and the external massdns binary against the +// exact same simulated conditions, so the comparison is apples-to-apples. +// +// No DNS traffic leaves the host: every resolver the engines talk to is a +// 127.0.0.1: UDP server started in-process. massdns is Linux-only +// (epoll), so this is intended to run inside the provided Docker image; the +// native engine runs anywhere. +// +// Example (inside the container): +// +// dnsbench -massdns /opt/massdns/bin/massdns -names 200000 -resolvers 16 -hit 5 +package main + +import ( + "bufio" + "context" + "flag" + "fmt" + "os" + "os/exec" + "strings" + "sync/atomic" + "time" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/internal/simdns" + "github.com/projectdiscovery/shuffledns/pkg/resolve" +) + +var ( + names = flag.Int("names", 50000, "number of names to resolve per scenario") + resolvers = flag.Int("resolvers", 8, "number of simulated loopback resolvers") + hit = flag.Int("hit", 5, "percentage of names that resolve (rest NXDOMAIN)") + concurrency = flag.Int("concurrency", 10000, "in-flight concurrency / massdns hashmap size") + retries = flag.Int("retries", 5, "retry budget per name (both engines)") + sockets = flag.Int("sockets", 0, "native resolver udp socket count (0 = scale to cores)") + batchMode = flag.String("batch-mode", "off", "native batching: off | on | adaptive (sendmmsg/recvmmsg, Linux)") + massdnsPath = flag.String("massdns", "", "path to the massdns binary (empty = skip massdns)") + engines = flag.String("engines", "native,massdns", "comma-separated engines to run") + scenarios = flag.String("scenarios", "all", "comma-separated scenario names or 'all'") + keepFiles = flag.Bool("keep", false, "keep generated names/resolvers files") +) + +type scenario struct { + name string + cfg simdns.Config +} + +func parseBatchMode(s string) resolve.BatchMode { + switch strings.ToLower(strings.TrimSpace(s)) { + case "on", "enabled", "true": + return resolve.BatchEnabled + case "adaptive", "auto": + return resolve.BatchAdaptive + default: + return resolve.BatchDisabled + } +} + +func allScenarios(hitPct int) []scenario { + return []scenario{ + {"lan-fast", simdns.Config{BaseLatency: 200 * time.Microsecond, Jitter: 300 * time.Microsecond, HitPercent: hitPct}}, + {"wan-typical", simdns.Config{BaseLatency: 15 * time.Millisecond, Jitter: 10 * time.Millisecond, LossRate: 0.005, HitPercent: hitPct}}, + {"wan-lossy", simdns.Config{BaseLatency: 25 * time.Millisecond, Jitter: 20 * time.Millisecond, LossRate: 0.05, ServfailRate: 0.02, HitPercent: hitPct}}, + {"rate-limited", simdns.Config{BaseLatency: 10 * time.Millisecond, Jitter: 10 * time.Millisecond, QPSPerServer: 3000, HitPercent: hitPct}}, + } +} + +type result struct { + engine string + ttfr time.Duration + wall time.Duration + qps float64 + resolved int64 +} + +func main() { + flag.Parse() + + selectedEngines := splitCSV(*engines) + wantMassdns := contains(selectedEngines, "massdns") + wantNative := contains(selectedEngines, "native") + + if wantMassdns && *massdnsPath == "" { + // try to find it on PATH + if p, err := exec.LookPath("massdns"); err == nil { + *massdnsPath = p + } else { + fmt.Fprintln(os.Stderr, "massdns requested but -massdns not set and not found on PATH; skipping massdns") + wantMassdns = false + } + } + + scns := filterScenarios(allScenarios(*hit), *scenarios) + if len(scns) == 0 { + fmt.Fprintln(os.Stderr, "no scenarios selected") + os.Exit(1) + } + + // write the shared names file once + namesFile, err := writeNamesFile(*names) + if err != nil { + fmt.Fprintf(os.Stderr, "could not write names file: %v\n", err) + os.Exit(1) + } + if !*keepFiles { + defer func() { _ = os.Remove(namesFile) }() + } + + fmt.Printf("dnsbench: names=%d resolvers=%d hit=%d%% concurrency=%d retries=%d\n", + *names, *resolvers, *hit, *concurrency, *retries) + fmt.Printf("no traffic leaves the host; resolvers are loopback (127.0.0.1)\n\n") + fmt.Printf("%-14s %-9s %10s %10s %12s %10s\n", "scenario", "engine", "ttfr", "wall", "qps", "resolved") + fmt.Printf("%s\n", strings.Repeat("-", 70)) + + for _, sc := range scns { + var rows []result + if wantNative { + r, err := runNative(sc, namesFile) + if err != nil { + fmt.Fprintf(os.Stderr, "native %s failed: %v\n", sc.name, err) + } else { + rows = append(rows, r) + } + } + if wantMassdns { + r, err := runMassdns(sc, namesFile) + if err != nil { + fmt.Fprintf(os.Stderr, "massdns %s failed: %v\n", sc.name, err) + } else { + rows = append(rows, r) + } + } + for _, r := range rows { + fmt.Printf("%-14s %-9s %10s %10s %12.0f %10d\n", + sc.name, r.engine, + r.ttfr.Round(100*time.Microsecond), + r.wall.Round(time.Millisecond), + r.qps, r.resolved) + } + fmt.Printf("%s\n", strings.Repeat("-", 70)) + } +} + +// perAttemptTimeout sizes a timeout/interval to a few RTTs so lossy scenarios +// complete via retransmission instead of stalling. Both engines use it. +func perAttemptTimeout(cfg simdns.Config) time.Duration { + rtt := cfg.BaseLatency + cfg.Jitter + t := 6 * rtt + if t < 500*time.Millisecond { + t = 500 * time.Millisecond + } + return t +} + +func runNative(sc scenario, namesFile string) (result, error) { + battery, err := simdns.Start(*resolvers, sc.cfg) + if err != nil { + return result{}, err + } + defer battery.Stop() + + timeout := perAttemptTimeout(sc.cfg) + + // pre-load names into memory (outside the timed region) so the producer is + // never the bottleneck; we want to measure the resolver, not the scanner. + nameList, err := loadNames(namesFile) + if err != nil { + return result{}, err + } + + var resolved atomic.Int64 + var ttfrNanos atomic.Int64 + start := time.Now() + + client, err := resolve.New(resolve.Options{ + Resolvers: battery.Addrs, + Concurrency: *concurrency, + SocketCount: *sockets, + Batch: parseBatchMode(*batchMode), + Timeout: timeout, + MaxRetries: *retries, + OnResult: func(r resolve.Result) { + ttfrNanos.CompareAndSwap(0, int64(time.Since(start))) + if r.Rcode == dns.RcodeSuccess && len(r.A) > 0 { + resolved.Add(1) + } + }, + }) + if err != nil { + return result{}, err + } + defer client.Close() + + input := make(chan string, 8192) + go func() { + defer close(input) + for _, n := range nameList { + input <- n + } + }() + + if err := client.Run(context.Background(), input); err != nil { + return result{}, err + } + wall := time.Since(start) + + printServerStats("native", sc.name, battery.Stats) + return result{ + engine: "native", + ttfr: time.Duration(ttfrNanos.Load()), + wall: wall, + qps: float64(*names) / wall.Seconds(), + resolved: resolved.Load(), + }, nil +} + +func runMassdns(sc scenario, namesFile string) (result, error) { + battery, err := simdns.Start(*resolvers, sc.cfg) + if err != nil { + return result{}, err + } + defer battery.Stop() + + resolversFile, err := writeResolversFile(battery.Addrs) + if err != nil { + return result{}, err + } + if !*keepFiles { + defer func() { _ = os.Remove(resolversFile) }() + } + + timeout := perAttemptTimeout(sc.cfg) + intervalMs := int(timeout / time.Millisecond) + + // -o S: simple output (answer RRs only), one line per record. + // -s : hashmap size / concurrent lookups. + // -i : retransmit interval for a name (match native per-attempt timeout). + // -c : resolve attempts before giving up (match native retry budget). + args := []string{ + "-r", resolversFile, + "-t", "A", + "-o", "S", + "-s", fmt.Sprintf("%d", *concurrency), + "-i", fmt.Sprintf("%d", intervalMs), + "-c", fmt.Sprintf("%d", *retries), + "--flush", + namesFile, + } + + cmd := exec.Command(*massdnsPath, args...) + stdout, err := cmd.StdoutPipe() + if err != nil { + return result{}, err + } + cmd.Stderr = nil // discard massdns progress/error log + + start := time.Now() + if err := cmd.Start(); err != nil { + return result{}, err + } + + var ttfr time.Duration + var resolved int64 + scanner := bufio.NewScanner(stdout) + scanner.Buffer(make([]byte, 0, 1024*1024), 4*1024*1024) + for scanner.Scan() { + line := scanner.Text() + if line == "" { + continue + } + if ttfr == 0 { + ttfr = time.Since(start) + } + // simple-format A record lines look like: "name. A 10.1.2.3" + if strings.Contains(line, " A ") { + resolved++ + } + } + if err := cmd.Wait(); err != nil { + return result{}, fmt.Errorf("massdns exited: %w", err) + } + wall := time.Since(start) + + printServerStats("massdns", sc.name, battery.Stats) + return result{ + engine: "massdns", + ttfr: ttfr, + wall: wall, + qps: float64(*names) / wall.Seconds(), + resolved: resolved, + }, nil +} + +func printServerStats(engine, scenario string, s *simdns.Stats) { + fmt.Fprintf(os.Stderr, " [%s/%s] server-side queries=%d answered=%d dropped=%d servfail=%d ratelimited=%d\n", + engine, scenario, s.Queries.Load(), s.Answered.Load(), s.Dropped.Load(), + s.Servfail.Load(), s.RateLimited.Load()) +} + +func loadNames(path string) ([]string, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer func() { _ = f.Close() }() + var names []string + scanner := bufio.NewScanner(f) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for scanner.Scan() { + if n := strings.TrimSpace(scanner.Text()); n != "" { + names = append(names, n) + } + } + return names, scanner.Err() +} + +func writeNamesFile(n int) (string, error) { + f, err := os.CreateTemp("", "dnsbench-names-*.txt") + if err != nil { + return "", err + } + defer func() { _ = f.Close() }() + w := bufio.NewWriter(f) + for i := 0; i < n; i++ { + if _, err := fmt.Fprintf(w, "host%d.bench.example.com\n", i); err != nil { + return "", err + } + } + return f.Name(), w.Flush() +} + +func writeResolversFile(addrs []string) (string, error) { + f, err := os.CreateTemp("", "dnsbench-resolvers-*.txt") + if err != nil { + return "", err + } + defer func() { _ = f.Close() }() + w := bufio.NewWriter(f) + for _, a := range addrs { + if _, err := fmt.Fprintln(w, a); err != nil { + return "", err + } + } + return f.Name(), w.Flush() +} + +func splitCSV(s string) []string { + parts := strings.Split(s, ",") + out := make([]string, 0, len(parts)) + for _, p := range parts { + p = strings.TrimSpace(p) + if p != "" { + out = append(out, p) + } + } + return out +} + +func contains(haystack []string, needle string) bool { + for _, h := range haystack { + if strings.EqualFold(h, needle) { + return true + } + } + return false +} + +func filterScenarios(all []scenario, sel string) []scenario { + if strings.TrimSpace(sel) == "" || strings.EqualFold(strings.TrimSpace(sel), "all") { + return all + } + want := splitCSV(sel) + var out []scenario + for _, s := range all { + if contains(want, s.name) { + out = append(out, s) + } + } + return out +} diff --git a/cmd/iterbench/main.go b/cmd/iterbench/main.go new file mode 100644 index 00000000..20fa6521 --- /dev/null +++ b/cmd/iterbench/main.go @@ -0,0 +1,171 @@ +// Command iterbench is a fully-local, offline benchmark that compares the +// iterative-from-root resolver against the stub resolver on the SAME workload, +// measuring both throughput and accuracy (false negatives). No packet leaves +// the host: +// +// - the iterative engine resolves against a real-socket authoritative +// hierarchy (internal/authsim) bound to 127.0.0.x loopback addresses; +// - the stub engine resolves against a battery of simulated recursive +// resolvers (internal/simdns) modelling real-world public-resolver +// conditions (latency, loss, rate-limiting and, crucially, rate-limit / +// hijack induced false NXDOMAIN — the massdns #117 failure mode). +// +// Both layers answer identical synthetic IPs, so any name the stub fails to +// resolve while the iterative engine succeeds is a measured false negative +// attributable to the public-resolver dependency that iterative removes. +// +// Must run on Linux (binding 127.0.0.x). The repo Dockerfile builds it. +package main + +import ( + "context" + "flag" + "fmt" + "os" + "sync/atomic" + "time" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/internal/authsim" + "github.com/projectdiscovery/shuffledns/internal/simdns" + "github.com/projectdiscovery/shuffledns/pkg/iterative" + "github.com/projectdiscovery/shuffledns/pkg/resolve" +) + +func main() { + domains := flag.Int("domains", 200, "registrable domains under .com") + hosts := flag.Int("hosts", 250, "hostnames per domain") + port := flag.Int("port", 5354, "shared UDP port for the authoritative hierarchy") + iterWorkers := flag.Int("iter-workers", 256, "iterative resolver worker count") + stubConc := flag.Int("stub-concurrency", 10000, "stub resolver in-flight concurrency") + recursors := flag.Int("recursors", 16, "simulated recursive resolvers for the stub") + retries := flag.Int("retries", 5, "retry budget (both engines)") + hijack := flag.Float64("hijack", 0.15, "fraction of existing names a recursor falsely answers NXDOMAIN (massdns #117)") + loss := flag.Float64("loss", 0.01, "recursor packet loss rate") + servfail := flag.Float64("servfail", 0.01, "recursor SERVFAIL rate") + flag.Parse() + + fmt.Printf("iterbench: domains=%d hosts/domain=%d names=%d\n", *domains, *hosts, *domains*(*hosts)) + fmt.Printf("no traffic leaves the host; authoritative tree on 127.0.0.x:%d, recursors on 127.0.0.1\n\n", *port) + + // ----- authoritative hierarchy (reliable, for iterative) ----- + hier, err := authsim.Build(*domains, *hosts, *port) + if err != nil { + fatal("could not start authoritative hierarchy: %v", err) + } + defer hier.Stop() + names := hier.Names + total := len(names) + + // ----- simulated recursive resolvers (flaky, for stub) ----- + battery, err := simdns.Start(*recursors, simdns.Config{ + BaseLatency: 1 * time.Millisecond, + Jitter: 2 * time.Millisecond, + LossRate: *loss, + ServfailRate: *servfail, + HijackRate: *hijack, + HitPercent: 100, // every name exists; any miss is a false negative + }) + if err != nil { + fatal("could not start recursor battery: %v", err) + } + defer battery.Stop() + + fmt.Printf("%-12s %-10s %-10s %-12s %-10s %-12s\n", "engine", "ttfr", "wall", "qps", "resolved", "missed(FN)") + fmt.Println("--------------------------------------------------------------------------") + + // ----- iterative engine ----- + itResolved, itTTFR, itWall := runIterative(hier, names, *iterWorkers, *retries) + report("iterative", itTTFR, itWall, total, itResolved) + fmt.Printf(" authoritative queries=%d queries/name=%.3f\n", hier.Queries.Load(), float64(hier.Queries.Load())/float64(total)) + + // ----- stub engine ----- + stResolved, stTTFR, stWall := runStub(battery.Addrs, names, *stubConc, *retries) + report("stub", stTTFR, stWall, total, stResolved) + fmt.Printf(" recursor hijacked(false NXDOMAIN)=%d servfail=%d dropped=%d ratelimited=%d\n", + battery.Stats.Hijacked.Load(), battery.Stats.Servfail.Load(), battery.Stats.Dropped.Load(), battery.Stats.RateLimited.Load()) + + fmt.Println("\nAccuracy delta (the point):") + fmt.Printf(" iterative missed %d/%d (%.2f%%)\n", total-int(itResolved), total, 100*float64(total-int(itResolved))/float64(total)) + fmt.Printf(" stub missed %d/%d (%.2f%%) <- false negatives from flaky/lying recursors\n", + total-int(stResolved), total, 100*float64(total-int(stResolved))/float64(total)) +} + +func runIterative(hier *authsim.Hierarchy, names []string, workers, retries int) (int64, time.Duration, time.Duration) { + r, err := iterative.New(iterative.Options{ + RootServers: []string{hier.RootAddr}, + QueryType: dns.TypeA, + Concurrency: workers, + Retries: retries, + Timeout: 3 * time.Second, + }) + if err != nil { + fatal("could not create iterative resolver: %v", err) + } + in := feed(names) + var resolved, ttfrSet atomic.Int64 + var ttfr time.Duration + start := time.Now() + _ = r.ResolveStream(context.Background(), in, iterative.StreamConfig{ + QueryType: dns.TypeA, + OnResult: func(res *resolve.Result) { + if res.Rcode == dns.RcodeSuccess && len(res.A) > 0 { + if ttfrSet.CompareAndSwap(0, 1) { + ttfr = time.Since(start) + } + resolved.Add(1) + } + }, + }) + return resolved.Load(), ttfr, time.Since(start) +} + +func runStub(resolvers, names []string, concurrency, retries int) (int64, time.Duration, time.Duration) { + var resolved, ttfrSet atomic.Int64 + var ttfr time.Duration + start := time.Now() + client, err := resolve.New(resolve.Options{ + Resolvers: resolvers, + QueryType: dns.TypeA, + Concurrency: concurrency, + MaxRetries: retries, + Timeout: 2 * time.Second, + OnResult: func(r resolve.Result) { + if r.Rcode == dns.RcodeSuccess && len(r.A) > 0 { + if ttfrSet.CompareAndSwap(0, 1) { + ttfr = time.Since(start) + } + resolved.Add(1) + } + }, + }) + if err != nil { + fatal("could not create stub resolver: %v", err) + } + defer client.Close() + _ = client.Run(context.Background(), feed(names)) + return resolved.Load(), ttfr, time.Since(start) +} + +func feed(names []string) <-chan string { + ch := make(chan string, 8192) + go func() { + defer close(ch) + for _, n := range names { + ch <- n + } + }() + return ch +} + +func report(engine string, ttfr, wall time.Duration, total int, resolved int64) { + qps := float64(total) / wall.Seconds() + missed := total - int(resolved) + fmt.Printf("%-12s %-10s %-10s %-12.0f %-10d %-12d\n", + engine, ttfr.Round(time.Microsecond), wall.Round(time.Millisecond), qps, resolved, missed) +} + +func fatal(format string, args ...interface{}) { + fmt.Fprintf(os.Stderr, format+"\n", args...) + os.Exit(1) +} diff --git a/cmd/resolve/flags.go b/cmd/resolve/flags.go new file mode 100644 index 00000000..ca4bb90f --- /dev/null +++ b/cmd/resolve/flags.go @@ -0,0 +1,147 @@ +package main + +import ( + "flag" + "fmt" + "os" + "strconv" + "strings" + "time" + + "github.com/miekg/dns" +) + +// parseFlags accepts the massdns CLI surface (plus native extensions). Flags that +// only applied to the C binary's multi-process / epoll busy-poll model are +// accepted and ignored; raw IPv6 source spoofing and privilege drop are wired. +func parseFlags() config { + var cfg config + var ptrList string + var intervalMs int + var retryCSV string + var ignoreBusyPoll bool + var processes int + + fs := flag.NewFlagSet("resolve", flag.ContinueOnError) + fs.SetOutput(os.Stderr) + fs.Usage = func() { + fmt.Fprintf(os.Stderr, `resolve — native massdns-compatible bulk DNS resolver + +Usage: resolve [options] [domainlist] + +Massdns-compatible flags are accepted. --processes and --busy-poll are parsed +and ignored. Native extensions (iterative, ptr, validate, zonewalk, ...) are +documented below. + +`) + fs.PrintDefaults() + } + + // Core massdns flags (short + long aliases where Go's flag package allows). + fs.StringVar(&cfg.resolversFile, "r", "", "file with resolver IPs (massdns -r/--resolvers)") + fs.StringVar(&cfg.resolversFile, "resolvers", "", "alias of -r") + fs.StringVar(&cfg.qtype, "t", "A", "record type (massdns -t/--type)") + fs.StringVar(&cfg.qtype, "type", "A", "alias of -t") + fs.StringVar(&cfg.format, "o", "F", "output flags (massdns -o): S/F/L/J/B + modifiers") + fs.StringVar(&cfg.format, "output", "F", "alias of -o") + fs.StringVar(&cfg.outFile, "w", "", "write output to file (massdns -w/--outfile)") + fs.StringVar(&cfg.outFile, "outfile", "", "alias of -w") + fs.IntVar(&cfg.concurrency, "s", 10000, "concurrent lookups (massdns -s/--hashmap-size)") + fs.IntVar(&cfg.concurrency, "hashmap-size", 10000, "alias of -s") + fs.IntVar(&cfg.retries, "c", 50, "resolve attempts before giving up (massdns -c/--resolve-count)") + fs.IntVar(&cfg.retries, "resolve-count", 50, "alias of -c") + fs.IntVar(&intervalMs, "i", 500, "retransmit interval in ms (massdns -i/--interval)") + fs.IntVar(&intervalMs, "interval", 500, "alias of -i") + fs.DurationVar(&cfg.timeout, "timeout", 0, "per-attempt timeout (0 = 2×interval)") + fs.BoolVar(&cfg.norecurse, "norecurse", false, "send non-recursive queries (RD=0)") + fs.BoolVar(&cfg.sticky, "sticky", false, "do not rotate resolver on retry") + fs.BoolVar(&cfg.predictable, "predictable", false, "use resolvers incrementally") + fs.BoolVar(&cfg.verifyIP, "verify-ip", false, "verify reply source IP (massdns --verify-ip)") + fs.BoolVar(&cfg.extendedInput, "extended-input", false, "input lines: name [resolver ...]") + fs.BoolVar(&cfg.flush, "flush", false, "flush output after every reply") + fs.BoolVar(&cfg.quiet, "q", false, "quiet mode (suppress status)") + fs.BoolVar(&cfg.quiet, "quiet", false, "alias of -q") + fs.StringVar(&cfg.errorLog, "l", "", "error log file path (default stderr)") + fs.StringVar(&cfg.errorLog, "error-log", "", "alias of -l") + fs.StringVar(&cfg.statusFormat, "status-format", "ansi", "status updates: ansi|json|none") + fs.StringVar(&cfg.bindAddr, "b", "", "local bind address (massdns -b/--bindto)") + fs.StringVar(&cfg.bindAddr, "bindto", "", "alias of -b") + fs.IntVar(&cfg.rcvbuf, "rcvbuf", 0, "SO_RCVBUF bytes (0 = default large buffer)") + fs.IntVar(&cfg.sndbuf, "sndbuf", 0, "SO_SNDBUF bytes (0 = OS default)") + fs.IntVar(&cfg.socketCount, "socket-count", 0, "UDP sockets (0 = scale to cores)") + fs.StringVar(&cfg.filter, "filter", "", "only output these response codes") + fs.StringVar(&cfg.ignore, "ignore", "", "drop these response codes") + fs.StringVar(&retryCSV, "retry", "", "response codes that trigger retry (default: all but NOERROR,NXDOMAIN)") + + // Rate / native engine controls. + fs.IntVar(&cfg.qps, "qps", 0, "max outbound queries per second (0 = unlimited)") + fs.StringVar(&cfg.batchMode, "batch-mode", "off", "datagram batching: off|on|adaptive (Linux)") + fs.BoolVar(&cfg.noTCPFallback, "no-tcp-fallback", false, "disable TCP fallback on truncated answers") + fs.BoolVar(&cfg.resolverHealth, "resolver-health", false, "de-weight failing resolvers") + fs.BoolVar(&cfg.adaptiveConc, "adaptive-concurrency", false, "adapt in-flight concurrency to loss") + fs.BoolVar(&cfg.crossCheck, "cross-check", false, "re-verify positives on a second resolver") + + // Native extensions (not in massdns). + fs.BoolVar(&cfg.iterative, "iterative", false, "recurse from root servers (no -r needed)") + fs.StringVar(&ptrList, "ptr", "", "reverse-PTR sweep targets: IPs/CIDRs/ranges") + fs.BoolVar(&cfg.onlyType, "only-type", false, "output only answer records matching the queried type") + fs.BoolVar(&cfg.validate, "validate", false, "validate the -r resolver list and print the good ones") + fs.StringVar(&cfg.validateDomain, "validate-domain", "", "known-good domains for --validate") + fs.StringVar(&cfg.zone, "zonewalk", "", "NSEC zone-walk the given zone") + fs.StringVar(&cfg.axfr, "axfr", "", "AXFR/IXFR the given zone") + fs.StringVar(&cfg.nsec3Dict, "nsec3-dict", "", "wordlist to crack NSEC3 from --zonewalk") + fs.StringVar(&cfg.shard, "shard", "", "process shard m/n (e.g. 2/8)") + fs.StringVar(&cfg.resume, "resume", "", "checkpoint file for crash-safe resume") + + // Privilege drop (after sockets open) and Linux raw IPv6 source spoofing. + fs.StringVar(&cfg.dropUser, "drop-user", "", "drop privileges to user after open (default nobody when root)") + fs.StringVar(&cfg.dropGroup, "drop-group", "", "drop privileges to group after open (default nobody when root)") + fs.BoolVar(&cfg.keepRoot, "root", false, "do not drop privileges when running as root") + fs.StringVar(&cfg.randSrcIPv6, "rand-src-ipv6", "", "random IPv6 source from prefix (Linux, CAP_NET_RAW; e.g. 2001:db8::/32)") + fs.StringVar(&cfg.randSrcIPv6File, "rand-src-ipv6-file", "", "file of IPv6 source addresses (Linux, CAP_NET_RAW)") + + // Accepted and ignored (massdns multi-process / epoll busy-poll). + fs.BoolVar(&ignoreBusyPoll, "busy-poll", false, "ignored (epoll busy-poll; not applicable)") + fs.IntVar(&processes, "processes", 1, "ignored (use -s / sockets instead of processes)") + _ = ignoreBusyPoll + + if err := fs.Parse(os.Args[1:]); err != nil { + os.Exit(2) + } + cfg.args = fs.Args() + cfg.ptrTargets = splitCSV(ptrList) + cfg.interval = time.Duration(intervalMs) * time.Millisecond + if cfg.timeout <= 0 { + cfg.timeout = 2 * cfg.interval + if cfg.timeout < time.Second { + cfg.timeout = time.Second + } + } + if retryCSV != "" { + cfg.retryRcodes = parseRetryRcodes(retryCSV) + } + if processes > 1 && !cfg.quiet { + fmt.Fprintf(os.Stderr, "note: --processes=%d ignored; raise -s/--socket-count instead\n", processes) + } + if (cfg.randSrcIPv6 != "" || cfg.randSrcIPv6File != "") && cfg.bindAddr != "" { + fmt.Fprintln(os.Stderr, "error: --bindto and --rand-src-ipv6 cannot be used together") + os.Exit(2) + } + // massdns --verify-ip is opt-in; without it, skip source verification. + cfg.disableVerifyIP = !cfg.verifyIP + return cfg +} + +func parseRetryRcodes(csv string) []int { + var out []int + for _, p := range splitCSV(csv) { + if v, ok := dns.StringToRcode[strings.ToUpper(p)]; ok { + out = append(out, v) + continue + } + if n, err := strconv.Atoi(p); err == nil { + out = append(out, n) + } + } + return out +} diff --git a/cmd/resolve/integration_test.go b/cmd/resolve/integration_test.go new file mode 100644 index 00000000..2742048f --- /dev/null +++ b/cmd/resolve/integration_test.go @@ -0,0 +1,244 @@ +package main + +import ( + "bytes" + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/internal/simdns" + "github.com/projectdiscovery/shuffledns/pkg/output" + "github.com/projectdiscovery/shuffledns/pkg/ptr" + "github.com/projectdiscovery/shuffledns/pkg/resolve" +) + +// TestReversePTRSweepEndToEnd drives the full reverse-PTR pipeline against the +// loopback resolver battery: ptr generator -> native resolver (PTR) -> massdns +// simple output. No traffic leaves the host. +func TestReversePTRSweepEndToEnd(t *testing.T) { + battery, err := simdns.Start(4, simdns.Config{ + BaseLatency: 200 * time.Microsecond, + HitPercent: 100, + }) + if err != nil { + t.Fatalf("simdns.Start: %v", err) + } + defer battery.Stop() + + var buf bytes.Buffer + w, err := output.NewWriter(&buf, "Snl") + if err != nil { + t.Fatal(err) + } + + client, err := resolve.New(resolve.Options{ + Resolvers: battery.Addrs, + QueryType: dns.TypePTR, + Concurrency: 1000, + MaxRetries: 3, + Timeout: time.Second, + OnResult: func(r resolve.Result) { + _ = w.Write(r) + }, + }) + if err != nil { + t.Fatalf("resolve.New: %v", err) + } + defer client.Close() + + input := make(chan string, 256) + go func() { + defer close(input) + _ = ptr.Stream(context.Background(), []string{"192.0.2.0/28"}, input) + }() + + if err := client.Run(context.Background(), input); err != nil { + t.Fatalf("Run: %v", err) + } + _ = w.Flush() + + out := buf.String() + // 192.0.2.0/28 = 16 addresses; all resolve (HitPercent 100). + ptrLines := 0 + for _, line := range strings.Split(out, "\n") { + if strings.Contains(line, " PTR ") { + ptrLines++ + if !strings.Contains(line, "in-addr.arpa.") || !strings.Contains(line, ".ptr.example.com.") { + t.Fatalf("unexpected PTR line: %q", line) + } + } + } + if ptrLines != 16 { + t.Fatalf("expected 16 PTR answers, got %d in:\n%s", ptrLines, out) + } +} + +// TestCLIAgainstBattery builds the resolve binary and runs it end-to-end against +// the loopback battery, validating flag parsing, stdin input and simple output. +func TestCLIAgainstBattery(t *testing.T) { + if testing.Short() { + t.Skip("skipping binary build in -short mode") + } + + battery, err := simdns.Start(4, simdns.Config{BaseLatency: 200 * time.Microsecond, HitPercent: 100}) + if err != nil { + t.Fatalf("simdns.Start: %v", err) + } + defer battery.Stop() + + dir := t.TempDir() + bin := filepath.Join(dir, "resolve") + if out, err := exec.Command("go", "build", "-o", bin, ".").CombinedOutput(); err != nil { + t.Fatalf("build failed: %v\n%s", err, out) + } + + resolversFile := filepath.Join(dir, "resolvers.txt") + if err := os.WriteFile(resolversFile, []byte(strings.Join(battery.Addrs, "\n")+"\n"), 0o600); err != nil { + t.Fatal(err) + } + + var names strings.Builder + for i := 0; i < 20; i++ { + names.WriteString("host") + names.WriteByte(byte('0' + i%10)) + names.WriteString(".bench.example.com\n") + } + + cmd := exec.Command(bin, "-r", resolversFile, "-t", "A", "-o", "Snl", "-timeout", "1s") + cmd.Stdin = strings.NewReader(names.String()) + out, err := cmd.Output() + if err != nil { + t.Fatalf("resolve binary failed: %v", err) + } + if !strings.Contains(string(out), " A 10.") { + t.Fatalf("expected A-record output, got:\n%s", out) + } +} + +// buildResolve compiles the resolve binary into a temp dir and returns its path. +func buildResolve(t *testing.T) string { + t.Helper() + bin := filepath.Join(t.TempDir(), "resolve") + if out, err := exec.Command("go", "build", "-o", bin, ".").CombinedOutput(); err != nil { + t.Fatalf("build failed: %v\n%s", err, out) + } + return bin +} + +func writeResolversFile(t *testing.T, addrs []string) string { + t.Helper() + f := filepath.Join(t.TempDir(), "resolvers.txt") + if err := os.WriteFile(f, []byte(strings.Join(addrs, "\n")+"\n"), 0o600); err != nil { + t.Fatal(err) + } + return f +} + +// resolvedNames extracts the queried names from "Snl" output lines. +func resolvedNames(out string) map[string]struct{} { + set := map[string]struct{}{} + for _, line := range strings.Split(out, "\n") { + f := strings.Fields(line) + if len(f) == 3 && f[1] == "A" { + set[strings.TrimSuffix(f[0], ".")] = struct{}{} + } + } + return set +} + +func benchNames(n int) string { + var b strings.Builder + for i := 0; i < n; i++ { + fmt.Fprintf(&b, "host%d.bench.example.com\n", i) + } + return b.String() +} + +// TestShardingPartitionsWork runs two shards over the same input and asserts the +// shards resolve disjoint subsets whose union is the full input. +func TestShardingPartitionsWork(t *testing.T) { + if testing.Short() { + t.Skip("skipping binary build in -short mode") + } + battery, err := simdns.Start(4, simdns.Config{BaseLatency: 200 * time.Microsecond, HitPercent: 100}) + if err != nil { + t.Fatalf("simdns.Start: %v", err) + } + defer battery.Stop() + + bin := buildResolve(t) + resolversFile := writeResolversFile(t, battery.Addrs) + const total = 50 + names := benchNames(total) + + run := func(shard string) map[string]struct{} { + cmd := exec.Command(bin, "-r", resolversFile, "-t", "A", "-o", "Snl", "-timeout", "1s", "-shard", shard) + cmd.Stdin = strings.NewReader(names) + out, err := cmd.Output() + if err != nil { + t.Fatalf("shard %s failed: %v", shard, err) + } + return resolvedNames(string(out)) + } + + a := run("1/2") + b := run("2/2") + + // disjoint + for n := range a { + if _, ok := b[n]; ok { + t.Fatalf("name %q appeared in both shards", n) + } + } + // union == full set + if len(a)+len(b) != total { + t.Fatalf("shard union = %d, want %d (a=%d b=%d)", len(a)+len(b), total, len(a), len(b)) + } + if len(a) == 0 || len(b) == 0 { + t.Fatalf("a shard got nothing (a=%d b=%d)", len(a), len(b)) + } +} + +// TestResumeSkipsCompleted runs once to populate a checkpoint, then re-runs with +// the same checkpoint and asserts everything is skipped. +func TestResumeSkipsCompleted(t *testing.T) { + if testing.Short() { + t.Skip("skipping binary build in -short mode") + } + battery, err := simdns.Start(4, simdns.Config{BaseLatency: 200 * time.Microsecond, HitPercent: 100}) + if err != nil { + t.Fatalf("simdns.Start: %v", err) + } + defer battery.Stop() + + bin := buildResolve(t) + resolversFile := writeResolversFile(t, battery.Addrs) + resumeFile := filepath.Join(t.TempDir(), "resume.log") + names := benchNames(30) + + run := func() string { + cmd := exec.Command(bin, "-r", resolversFile, "-t", "A", "-o", "Snl", "-timeout", "1s", "-resume", resumeFile) + cmd.Stdin = strings.NewReader(names) + out, err := cmd.Output() + if err != nil { + t.Fatalf("run failed: %v", err) + } + return string(out) + } + + first := run() + if len(resolvedNames(first)) != 30 { + t.Fatalf("first run resolved %d, want 30", len(resolvedNames(first))) + } + + second := run() + if got := len(resolvedNames(second)); got != 0 { + t.Fatalf("resume run should skip all completed names, but resolved %d", got) + } +} diff --git a/cmd/resolve/main.go b/cmd/resolve/main.go new file mode 100644 index 00000000..303ce651 --- /dev/null +++ b/cmd/resolve/main.go @@ -0,0 +1,485 @@ +// Command resolve is a native, dependency-free massdns-compatible bulk resolver. +// It accepts the massdns CLI surface (obsolete flags are ignored) and writes +// results in massdns output formats via pkg/output. +// +// Examples: +// +// resolve -r resolvers.txt -t A -o Snl names.txt > out.txt +// resolve -r resolvers.txt --ptr 192.0.2.0/24 -o J > ptr.ndjson +package main + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "os" + "os/signal" + "strings" + "syscall" + "time" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/pkg/axfr" + "github.com/projectdiscovery/shuffledns/pkg/checkpoint" + "github.com/projectdiscovery/shuffledns/pkg/iterative" + "github.com/projectdiscovery/shuffledns/pkg/output" + "github.com/projectdiscovery/shuffledns/pkg/ptr" + "github.com/projectdiscovery/shuffledns/pkg/resolve" + "github.com/projectdiscovery/shuffledns/pkg/shard" + "github.com/projectdiscovery/shuffledns/pkg/zonewalk" +) + +func main() { + cfg := parseFlags() + + errOut := io.Writer(os.Stderr) + if cfg.errorLog != "" { + f, err := os.OpenFile(cfg.errorLog, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644) + if err != nil { + fatal("could not open error log: %v", err) + } + defer func() { _ = f.Close() }() + errOut = f + } + + var resolvers []string + var err error + if cfg.resolversFile != "" { + resolvers, err = readLines(cfg.resolversFile) + if err != nil { + fatal("could not read resolvers: %v", err) + } + } + if !cfg.iterative && len(resolvers) == 0 { + fatal("no resolvers provided (-r); or use --iterative to recurse from root") + } + + qtype := dns.TypeA + if cfg.qtype != "" { + t, ok := dns.StringToType[strings.ToUpper(cfg.qtype)] + if !ok { + fatal("unknown record type %q", cfg.qtype) + } + qtype = t + } + if len(cfg.ptrTargets) > 0 { + qtype = dns.TypePTR + } + + out := os.Stdout + if cfg.outFile != "" { + out, err = os.Create(cfg.outFile) + if err != nil { + fatal("could not create output file: %v", err) + } + defer func() { _ = out.Close() }() + } + + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + if cfg.validate { + runValidate(ctx, cfg, resolvers, out) + return + } + if cfg.axfr != "" { + runAXFR(ctx, cfg, resolvers, out, errOut) + return + } + if cfg.zone != "" { + runZonewalk(ctx, cfg, resolvers, out, errOut) + return + } + + writer, err := output.NewWriterWithConfig(out, output.Config{ + Format: cfg.format, + FilterRcodes: splitCSV(cfg.filter), + IgnoreRcodes: splitCSV(cfg.ignore), + OnlyQueryType: cfg.onlyType, + FlushEach: cfg.flush, + }) + if err != nil { + fatal("invalid output format: %v", err) + } + defer func() { _ = writer.Flush() }() + + shardCfg, err := shard.Parse(cfg.shard) + if err != nil { + fatal("%v", err) + } + var ckpt *checkpoint.Checkpoint + if cfg.resume != "" { + ckpt, err = checkpoint.Open(cfg.resume) + if err != nil { + fatal("could not open resume checkpoint: %v", err) + } + defer func() { _ = ckpt.Close() }() + if n := ckpt.Resumed(); n > 0 && !cfg.quiet { + _, _ = fmt.Fprintf(errOut, "resuming: skipping %d already-completed names\n", n) + } + } + markDone := func(name string) { + if ckpt != nil { + _ = ckpt.Done(name) + } + } + + onResult := func(r resolve.Result) { + markDone(r.Name) + if err := writer.Write(r); err != nil { + _, _ = fmt.Fprintf(errOut, "write error: %v\n", err) + } + } + onError := func(name string, err error) { + markDone(name) + _ = writer.WriteFailure(name, qtype, "resolution failed") + if err != nil && !cfg.quiet { + _, _ = fmt.Fprintf(errOut, "%s: %v\n", name, err) + } + } + + input := make(chan string, 8192) + go func() { + defer close(input) + raw := make(chan string, 8192) + go func() { + defer close(raw) + if len(cfg.ptrTargets) > 0 { + if err := ptr.Stream(ctx, cfg.ptrTargets, raw); err != nil { + _, _ = fmt.Fprintf(errOut, "ptr generation error: %v\n", err) + } + return + } + produceNames(ctx, cfg.args, raw, errOut) + }() + for name := range raw { + if !shardCfg.Owns(name) { + continue + } + if ckpt != nil && ckpt.Has(name) { + continue + } + // Stop feeding once the consumer is gone (interrupt), otherwise this + // send blocks forever on a full buffer and leaks the producer. + select { + case input <- name: + case <-ctx.Done(): + return + } + } + }() + + if cfg.iterative { + ir, err := iterative.New(iterative.Options{ + QueryType: qtype, + Concurrency: iterativeWorkers(cfg.concurrency), + Timeout: cfg.timeout, + IPv6: qtype == dns.TypeAAAA, + }) + if err != nil { + fatal("could not create iterative resolver: %v", err) + } + err = ir.ResolveStream(ctx, input, iterative.StreamConfig{ + QueryType: qtype, + OnResult: func(r *resolve.Result) { onResult(*r) }, + OnError: onError, + }) + // A cancelled context is a user interrupt (Ctrl-C / SIGTERM), not a + // failure: fall through so buffered output and the checkpoint are flushed + // by the deferred cleanups instead of being lost to os.Exit. + if err != nil && !errors.Is(err, context.Canceled) { + fatal("resolution failed: %v", err) + } + } else { + client, err := resolve.New(resolve.Options{ + Resolvers: resolvers, + QueryType: qtype, + Concurrency: cfg.concurrency, + QPS: cfg.qps, + MaxRetries: cfg.retries, + Timeout: cfg.timeout, + Interval: cfg.interval, + NoRecurse: cfg.norecurse, + Sticky: cfg.sticky, + Predictable: cfg.predictable, + ExtendedInput: cfg.extendedInput, + DisableIPVerification: cfg.disableVerifyIP, + DisableTCPFallback: cfg.noTCPFallback, + Batch: parseBatchMode(cfg.batchMode), + SocketCount: cfg.socketCount, + BindAddr: cfg.bindAddr, + ReadBuffer: cfg.rcvbuf, + WriteBuffer: cfg.sndbuf, + RandSrcIPv6: cfg.randSrcIPv6, + RandSrcIPv6File: cfg.randSrcIPv6File, + RetryRcodes: cfg.retryRcodes, + ResolverHealth: cfg.resolverHealth, + AdaptiveConcurrency: cfg.adaptiveConc, + CrossCheck: cfg.crossCheck, + OnResult: onResult, + OnError: onError, + OnProgress: statusReporter(cfg, errOut), + }) + if err != nil { + fatal("could not create resolver: %v", err) + } + defer client.Close() + // Drop root after sockets are open (massdns --drop-user/--drop-group/--root). + if err := resolve.DropPrivileges(cfg.dropUser, cfg.dropGroup, cfg.keepRoot); err != nil { + fatal("privilege drop: %v", err) + } + if err := client.Run(ctx, input); err != nil && !errors.Is(err, context.Canceled) { + fatal("resolution failed: %v", err) + } + } + + if err := writer.Flush(); err != nil { + fatal("flush failed: %v", err) + } +} + +func statusReporter(cfg config, errOut io.Writer) func(resolve.Stats) { + if cfg.quiet || strings.EqualFold(cfg.statusFormat, "none") { + return nil + } + jsonMode := strings.EqualFold(cfg.statusFormat, "json") + return func(s resolve.Stats) { + if jsonMode { + _, _ = fmt.Fprintf(errOut, `{"queries":%d,"retransmits":%d,"answered":%d,"inflight":%d,"concurrency":%d,"loss":%.4f,"rtt_ms":%.2f}`+"\n", + s.Queries, s.Retransmits, s.Answered, s.InflightDepth, s.ConcurrencyCap, s.LossRate, float64(s.RTT.Microseconds())/1000) + return + } + _, _ = fmt.Fprintf(errOut, "\rprocessed: %d | answered: %d | inflight: %d | conc: %d | loss: %.1f%% | rtt: %s", + s.Queries, s.Answered, s.InflightDepth, s.ConcurrencyCap, s.LossRate*100, s.RTT.Round(time.Microsecond)) + } +} + +func iterativeWorkers(concurrency int) int { + const max = 1024 + if concurrency <= 0 { + return 200 + } + if concurrency > max { + return max + } + return concurrency +} + +type config struct { + resolversFile string + qtype string + format string + outFile string + concurrency int + qps int + retries int + timeout time.Duration + interval time.Duration + norecurse bool + sticky bool + predictable bool + verifyIP bool + disableVerifyIP bool + extendedInput bool + flush bool + quiet bool + errorLog string + statusFormat string + bindAddr string + rcvbuf int + sndbuf int + socketCount int + dropUser string + dropGroup string + keepRoot bool + randSrcIPv6 string + randSrcIPv6File string + retryRcodes []int + batchMode string + noTCPFallback bool + resolverHealth bool + adaptiveConc bool + crossCheck bool + iterative bool + ptrTargets []string + filter string + ignore string + onlyType bool + validate bool + validateDomain string + zone string + axfr string + nsec3Dict string + shard string + resume string + args []string +} + +func splitCSV(s string) []string { + var out []string + for _, p := range strings.Split(s, ",") { + if p = strings.TrimSpace(p); p != "" { + out = append(out, p) + } + } + return out +} + +func runValidate(ctx context.Context, cfg config, resolvers []string, out *os.File) { + good, report, err := resolve.ValidateResolvers(ctx, resolve.ValidateConfig{ + Resolvers: resolvers, + GoodDomains: splitCSV(cfg.validateDomain), + Timeout: cfg.timeout, + }) + if err != nil { + fatal("validation failed: %v", err) + } + w := bufio.NewWriter(out) + defer func() { _ = w.Flush() }() + for _, r := range good { + _, _ = fmt.Fprintln(w, r) + } + _, _ = fmt.Fprintf(os.Stderr, "validated %d resolvers: %d good, %d rejected\n", + len(report), len(good), len(report)-len(good)) +} + +func runAXFR(ctx context.Context, cfg config, resolvers []string, out *os.File, errOut io.Writer) { + w := bufio.NewWriter(out) + defer func() { _ = w.Flush() }() + res, err := axfr.Attempt(ctx, axfr.Config{ + Zone: cfg.axfr, + Resolvers: resolvers, + Timeout: cfg.timeout, + OnName: func(name string) { _, _ = fmt.Fprintln(w, name) }, + OnNameserver: func(ns string, names int, err error) { + if err != nil { + _, _ = fmt.Fprintf(errOut, "axfr %s: refused/failed (%v)\n", ns, err) + } else { + _, _ = fmt.Fprintf(errOut, "axfr %s: transferred %d names\n", ns, names) + } + }, + }) + if err != nil { + fatal("zone transfer failed: %v", err) + } + _, _ = fmt.Fprintf(errOut, "AXFR of %s via %s transferred %d names (%d records)\n", + cfg.axfr, res.Nameserver, len(res.Names), res.Records) +} + +func runZonewalk(ctx context.Context, cfg config, resolvers []string, out *os.File, errOut io.Writer) { + w := bufio.NewWriter(out) + defer func() { _ = w.Flush() }() + res, err := zonewalk.Walk(ctx, zonewalk.Config{ + Zone: cfg.zone, + Resolvers: resolvers, + Timeout: cfg.timeout, + OnName: func(name string) { _, _ = fmt.Fprintln(w, name) }, + }) + if err != nil { + fatal("zone walk failed: %v", err) + } + if res.NSEC3 { + _, _ = fmt.Fprintf(errOut, "zone %s is NSEC3-signed (salt=%s iterations=%d)\n", + cfg.zone, res.NSEC3Param.Salt, res.NSEC3Param.Iterations) + if cfg.nsec3Dict == "" { + _, _ = fmt.Fprintf(errOut, "supply --nsec3-dict to harvest and crack the NSEC3 ring\n") + return + } + candidates, rerr := readLines(cfg.nsec3Dict) + if rerr != nil { + fatal("could not read nsec3 wordlist: %v", rerr) + } + cres, cerr := zonewalk.CrackNSEC3(ctx, zonewalk.CrackConfig{ + Zone: cfg.zone, + Resolvers: resolvers, + Timeout: cfg.timeout, + Candidates: candidates, + OnName: func(name string) { _, _ = fmt.Fprintln(w, name) }, + }) + if cerr != nil { + fatal("nsec3 crack failed: %v", cerr) + } + _, _ = fmt.Fprintf(errOut, "NSEC3 crack of %s: harvested %d hashes, recovered %d/%d names (saturated=%t)\n", + cfg.zone, cres.HarvestedHashes, len(cres.Names), len(candidates), cres.Saturated) + return + } + _, _ = fmt.Fprintf(errOut, "zone walk of %s discovered %d names\n", cfg.zone, len(res.Names)) +} + +func produceNames(ctx context.Context, files []string, out chan<- string, errOut io.Writer) { + emit := func(line string) bool { + line = strings.TrimSpace(line) + if line == "" { + return true + } + select { + case <-ctx.Done(): + return false + case out <- line: + return true + } + } + + if len(files) == 0 { + scanLines(os.Stdin, emit) + return + } + for _, fname := range files { + f, err := os.Open(fname) + if err != nil { + _, _ = fmt.Fprintf(errOut, "could not open %s: %v\n", fname, err) + continue + } + cont := scanLines(f, emit) + _ = f.Close() + if !cont { + return + } + } +} + +func scanLines(f *os.File, emit func(string) bool) bool { + scanner := bufio.NewScanner(f) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for scanner.Scan() { + if !emit(scanner.Text()) { + return false + } + } + return true +} + +func readLines(path string) ([]string, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer func() { _ = f.Close() }() + var lines []string + scanner := bufio.NewScanner(f) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for scanner.Scan() { + if l := strings.TrimSpace(scanner.Text()); l != "" { + lines = append(lines, l) + } + } + return lines, scanner.Err() +} + +func parseBatchMode(s string) resolve.BatchMode { + switch strings.ToLower(strings.TrimSpace(s)) { + case "on", "enabled", "true": + return resolve.BatchEnabled + case "adaptive", "auto": + return resolve.BatchAdaptive + default: + return resolve.BatchDisabled + } +} + +func fatal(format string, args ...interface{}) { + _, _ = fmt.Fprintf(os.Stderr, format+"\n", args...) + os.Exit(1) +} diff --git a/cmd/shuffledns/main.go b/cmd/shuffledns/main.go index e87386b2..f44e2b9c 100644 --- a/cmd/shuffledns/main.go +++ b/cmd/shuffledns/main.go @@ -14,6 +14,9 @@ func main() { gologger.Fatal().Msgf("Could not create runner: %s\n", err) } - massdnsRunner.RunEnumeration() + if err := massdnsRunner.RunEnumeration(); err != nil { + massdnsRunner.Close() + gologger.Fatal().Msgf("Could not run enumeration: %s\n", err) + } massdnsRunner.Close() } diff --git a/go.mod b/go.mod index ba8f11b1..b6117f30 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/projectdiscovery/shuffledns go 1.25.0 require ( + github.com/hashicorp/golang-lru/v2 v2.0.7 github.com/miekg/dns v1.1.62 github.com/projectdiscovery/dnsx v1.3.0 github.com/projectdiscovery/goflags v0.1.74 @@ -11,7 +12,6 @@ require ( github.com/remeh/sizedwaitgroup v1.0.0 github.com/rs/xid v1.5.0 github.com/stretchr/testify v1.11.1 - github.com/syndtr/goleveldb v1.0.0 github.com/weppos/publicsuffix-go v0.50.3-0.20260104170930-90713dec78f2 ) @@ -30,8 +30,8 @@ require ( github.com/projectdiscovery/utils v0.11.1 github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d // indirect golang.org/x/mod v0.35.0 // indirect - golang.org/x/net v0.55.0 // indirect - golang.org/x/sys v0.45.0 // indirect + golang.org/x/net v0.55.0 + golang.org/x/sys v0.45.0 golang.org/x/tools v0.44.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) @@ -52,7 +52,6 @@ require ( github.com/fatih/color v1.15.0 // indirect github.com/gaissmai/bart v0.28.0 // indirect github.com/go-ole/go-ole v1.2.6 // indirect - github.com/golang/snappy v0.0.4 // indirect github.com/google/go-github/v30 v30.1.0 // indirect github.com/google/go-querystring v1.1.0 // indirect github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect @@ -90,5 +89,4 @@ require ( golang.org/x/term v0.43.0 // indirect golang.org/x/text v0.37.0 // indirect gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect - gopkg.in/yaml.v2 v2.2.2 // indirect ) diff --git a/go.sum b/go.sum index 3f25e603..ef791319 100644 --- a/go.sum +++ b/go.sum @@ -42,16 +42,11 @@ github.com/ebitengine/purego v0.10.0 h1:QIw4xfpWT6GWTzaW5XEKy3HXoqrJGx1ijYHzTF0/ github.com/ebitengine/purego v0.10.0/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ= github.com/fatih/color v1.15.0 h1:kOqh6YHBtK8aywxGerMG2Eq3H6Qgoqeo13Bk2Mv/nBs= github.com/fatih/color v1.15.0/go.mod h1:0h5ZqXfHYED7Bhv2ZJamyIOUej9KtShiJESRwBDUSsw= -github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= github.com/gaissmai/bart v0.28.0 h1:89yZLo8NmyqD0RYgJ3QO9HhqqGGw+oWhf90cZm69Lko= github.com/gaissmai/bart v0.28.0/go.mod h1:GREWQfTLRWz/c5FTOsIw+KkscuFkIV5t8Rp7Nd1Td5c= github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= -github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= -github.com/golang/snappy v0.0.0-20180518054509-2e65f85255db/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= -github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM= -github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= @@ -68,10 +63,10 @@ github.com/google/uuid v1.3.1 h1:KjJaJ9iWZ3jOFZIf1Lqf4laDRCasjl0BCmnEGxkdLb4= github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8= github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0= +github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= -github.com/hpcloud/tail v1.0.0 h1:nfCOvKYfkgYP8hkirhJocXT2+zOD8yUNjXaWfTlyFKI= -github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI= @@ -108,11 +103,6 @@ github.com/muesli/reflow v0.3.0 h1:IFsN6K9NfGtjeggFP+68I4chLZV2yIKsXJFNZ+eWh6s= github.com/muesli/reflow v0.3.0/go.mod h1:pbwTDkVPibjO2kyvBQRBxTWEEGDGq0FlB1BIKtnHY/8= github.com/muesli/termenv v0.15.3-0.20240618155329-98d742f6907a h1:2MaM6YC3mGu54x+RKAA6JiFFHlHDY1UbkxqppT7wYOg= github.com/muesli/termenv v0.15.3-0.20240618155329-98d742f6907a/go.mod h1:hxSnBBYLK21Vtq/PHd0S2FYCxBXzBua8ov5s1RobyRQ= -github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= -github.com/onsi/ginkgo v1.7.0 h1:WSHQ+IS43OoUrWtD1/bbclrwK8TTH5hzp+umCiuxHgs= -github.com/onsi/ginkgo v1.7.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= -github.com/onsi/gomega v1.4.3 h1:RE1xgDvH7imwFD45h+u2SgIfERHlS2yNG4DObb5BSKU= -github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= @@ -152,8 +142,6 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= -github.com/syndtr/goleveldb v1.0.0 h1:fBdIW9lB4Iz0n9khmH8w27SJ3QEJ7+IgjPEwGSZiFdE= -github.com/syndtr/goleveldb v1.0.0/go.mod h1:ZVVdQEZoIme9iO1Ch2Jdy24qqXrMMOU6lpPAyBWyWuQ= github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= @@ -187,7 +175,6 @@ golang.org/x/exp v0.0.0-20250106191152-7588d65b2ba8 h1:yqrTHse8TCMW1M1ZCP+VAR/l0 golang.org/x/exp v0.0.0-20250106191152-7588d65b2ba8/go.mod h1:tujkw807nyEEAamNbDrEGzRav+ilXA7PCRAd6xsmwiU= golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM= golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU= -golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= @@ -196,10 +183,8 @@ golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.27.0 h1:da9Vo7/tDv5RH/7nZDz1eMGS/q1Vv1N/7FCrBhI9I3M= golang.org/x/oauth2 v0.27.0/go.mod h1:onh5ek6nERTohokkhCD/y2cV4Do3fxFHFuAejCkRWT8= -golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= -golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -229,12 +214,5 @@ google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9Ywl gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= -gopkg.in/fsnotify.v1 v1.4.7 h1:xOHLXZwVvI9hhs+cLKq5+I5onOuwQLhQwiu63xxlHs4= -gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= -gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= -gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= -gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= -gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw= -gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/authsim/authsim.go b/internal/authsim/authsim.go new file mode 100644 index 00000000..7d0e565e --- /dev/null +++ b/internal/authsim/authsim.go @@ -0,0 +1,236 @@ +// Package authsim provides a real-socket, loopback authoritative DNS hierarchy +// (root -> TLD -> per-domain authoritative servers) for benchmarking the +// iterative-from-root resolver entirely offline. Every server binds a distinct +// 127.0.0.x address (Linux routes all of 127/8 to loopback without aliasing), +// so the iterative resolver follows real referrals and glue over real UDP +// without a single packet leaving the host. +// +// Answers use simdns.SyntheticIP, identical to the simulated recursive-resolver +// battery, so the iterative engine (talking to these authoritative servers) and +// the stub engine (talking to recursors) resolve the exact same workload and +// their results can be compared for accuracy (false negatives). +package authsim + +import ( + "fmt" + "net" + "strings" + "sync" + "sync/atomic" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/internal/simdns" +) + +type nsEntry struct { + name string + ip string // "" == glueless +} + +type zoneServer struct { + ip string + port int + conn *net.UDPConn + authZones []string + delegations map[string][]nsEntry + closeCh chan struct{} + wg sync.WaitGroup + queries *atomic.Int64 +} + +// Hierarchy is a running authoritative tree. +type Hierarchy struct { + RootAddr string // "127.0.0.1:port" to seed the iterative resolver's roots + Names []string // all resolvable leaf names (host{n}.d{k}.com) + Port int + Queries atomic.Int64 + + servers []*zoneServer +} + +// Build starts an authoritative hierarchy with `domains` registrable domains +// under .com, each served by its own authoritative server with `hosts` names. +// All servers listen on `port` across distinct loopback IPs. +func Build(domains, hosts, port int) (*Hierarchy, error) { + h := &Hierarchy{Port: port} + + const rootIP = "127.0.0.1" + const comIP = "127.0.0.2" + + root := h.newServer(rootIP, port, ".") + root.delegations["com."] = []nsEntry{{name: "a.gtld-servers.net.", ip: comIP}} + + com := h.newServer(comIP, port, "com.") + + for d := 0; d < domains; d++ { + zone := fmt.Sprintf("d%d.com.", d) + authIP := domainIP(d) + nsName := fmt.Sprintf("ns.%s", zone) + com.delegations[zone] = []nsEntry{{name: nsName, ip: authIP}} + h.newServer(authIP, port, zone) + for n := 0; n < hosts; n++ { + h.Names = append(h.Names, fmt.Sprintf("host%d.d%d.com", n, d)) + } + } + + for _, s := range h.servers { + if err := s.start(); err != nil { + h.Stop() + return nil, fmt.Errorf("bind %s:%d: %w (on macOS only 127.0.0.1 is available; run in Docker/Linux)", s.ip, port, err) + } + } + h.RootAddr = net.JoinHostPort(rootIP, fmt.Sprint(port)) + return h, nil +} + +func (h *Hierarchy) newServer(ip string, port int, zones ...string) *zoneServer { + s := &zoneServer{ + ip: ip, + port: port, + authZones: zones, + delegations: map[string][]nsEntry{}, + closeCh: make(chan struct{}), + queries: &h.Queries, + } + h.servers = append(h.servers, s) + return s +} + +// Stop shuts down all servers. +func (h *Hierarchy) Stop() { + for _, s := range h.servers { + select { + case <-s.closeCh: + default: + close(s.closeCh) + } + if s.conn != nil { + _ = s.conn.Close() + } + } + for _, s := range h.servers { + s.wg.Wait() + } +} + +func domainIP(d int) string { + return net.IPv4(127, 1, byte(d>>8), byte(d)).String() +} + +func (s *zoneServer) start() error { + addr := &net.UDPAddr{IP: net.ParseIP(s.ip), Port: s.port} + conn, err := net.ListenUDP("udp", addr) + if err != nil { + return err + } + _ = conn.SetReadBuffer(8 * 1024 * 1024) + _ = conn.SetWriteBuffer(8 * 1024 * 1024) + s.conn = conn + s.wg.Add(1) + go s.serve() + return nil +} + +func (s *zoneServer) serve() { + defer s.wg.Done() + buf := make([]byte, 4096) + for { + n, from, err := s.conn.ReadFromUDP(buf) + if err != nil { + select { + case <-s.closeCh: + return + default: + continue + } + } + pkt := make([]byte, n) + copy(pkt, buf[:n]) + s.handle(pkt, from) + } +} + +func (s *zoneServer) handle(pkt []byte, from *net.UDPAddr) { + s.queries.Add(1) + req := new(dns.Msg) + if req.Unpack(pkt) != nil || len(req.Question) == 0 { + return + } + q := req.Question[0] + qname := canonical(q.Name) + + m := new(dns.Msg) + m.SetReply(req) + m.Authoritative = true + + // most specific authoritative zone and delegation for the name + authZone := "" + for _, z := range s.authZones { + if inBailiwick(qname, z) && len(z) > len(authZone) { + authZone = z + } + } + bestChild := "" + for child := range s.delegations { + if inBailiwick(qname, child) && len(child) > len(bestChild) { + bestChild = child + } + } + + switch { + case bestChild != "" && len(bestChild) > len(authZone): + m.Authoritative = false + for _, ns := range s.delegations[bestChild] { + m.Ns = append(m.Ns, mustRR(fmt.Sprintf("%s 3600 IN NS %s", bestChild, ns.name))) + if ns.ip != "" { + m.Extra = append(m.Extra, mustRR(fmt.Sprintf("%s 3600 IN A %s", ns.name, ns.ip))) + } + } + case authZone == "": + m.Rcode = dns.RcodeServerFailure + case q.Qtype == dns.TypeA && isHostName(qname): + m.Answer = append(m.Answer, mustRR(fmt.Sprintf("%s 60 IN A %s", qname, simdns.SyntheticIP(qname)))) + case q.Qtype == dns.TypeA && isNSName(qname): + // answer the server's own NS address if asked (glue self-lookup) + m.Answer = append(m.Answer, mustRR(fmt.Sprintf("%s 60 IN A %s", qname, s.ip))) + default: + m.Rcode = dns.RcodeNameError + m.Ns = append(m.Ns, mustRR(fmt.Sprintf("%s 3600 IN SOA ns.%s hostmaster.%s 1 3600 600 86400 60", authZone, authZone, authZone))) + } + + out, err := m.Pack() + if err != nil { + return + } + _, _ = s.conn.WriteToUDP(out, from) +} + +func isHostName(name string) bool { return strings.HasPrefix(name, "host") } +func isNSName(name string) bool { return strings.HasPrefix(name, "ns.") } + +func canonical(name string) string { + if name == "" || name == "." { + return "." + } + name = strings.ToLower(name) + if !strings.HasSuffix(name, ".") { + name += "." + } + return name +} + +func inBailiwick(child, parent string) bool { + child, parent = canonical(child), canonical(parent) + if parent == "." { + return true + } + return child == parent || strings.HasSuffix(child, "."+parent) +} + +func mustRR(s string) dns.RR { + rr, err := dns.NewRR(s) + if err != nil { + panic(err) + } + return rr +} diff --git a/internal/simdns/simdns.go b/internal/simdns/simdns.go new file mode 100644 index 00000000..564145da --- /dev/null +++ b/internal/simdns/simdns.go @@ -0,0 +1,277 @@ +// Package simdns provides a battery of loopback UDP DNS servers that simulate +// remote recursive resolvers. It is used to benchmark DNS resolvers (the native +// resolver as well as the massdns binary) without sending any traffic to the +// public internet. +// +// Each simulated resolver models realistic remote conditions: response latency +// with jitter, packet loss, a SERVFAIL rate, and an optional per-resolver QPS +// cap (over-budget queries are dropped, like a throttling public resolver). The +// answered/NXDOMAIN decision is deterministic per name, so the workload is +// stable and reproducible across runs and across engines. +package simdns + +import ( + "fmt" + "math/rand" + "net" + "sync" + "sync/atomic" + "time" + + "github.com/miekg/dns" +) + +// Config describes the simulated behaviour of one remote resolver. +type Config struct { + BaseLatency time.Duration // minimum response delay + Jitter time.Duration // additional uniform-random delay on top of base + LossRate float64 // fraction of queries silently dropped (0..1) + ServfailRate float64 // fraction of queries answered with SERVFAIL (0..1) + QPSPerServer int // per-resolver rate cap; over-budget queries are dropped (0 = unlimited) + HitPercent int // fraction of names that return an A record + // HijackRate models a misbehaving recursive resolver that returns NXDOMAIN + // for a name that actually exists (e.g. rate-limit-induced false negatives + // or NXDOMAIN hijacking). This is the massdns #117 failure mode: the stub + // resolver believes the NXDOMAIN (terminal, not retried) and misses the + // name. Fraction 0..1 of otherwise-hitting queries answered NXDOMAIN. + HijackRate float64 +} + +// Stats aggregates what the battery actually did, for sanity reporting. +type Stats struct { + Queries atomic.Int64 + Answered atomic.Int64 + Dropped atomic.Int64 + Servfail atomic.Int64 + RateLimited atomic.Int64 + Hijacked atomic.Int64 // existing names falsely answered NXDOMAIN +} + +// tokenBucket is a tiny non-blocking rate limiter used to model a resolver that +// rate-limits and drops excess traffic. +type tokenBucket struct { + mu sync.Mutex + tokens float64 + max float64 + refill float64 // tokens per second + lastFill time.Time +} + +func newTokenBucket(qps int) *tokenBucket { + if qps <= 0 { + return nil + } + return &tokenBucket{ + tokens: float64(qps), + max: float64(qps), + refill: float64(qps), + lastFill: time.Now(), + } +} + +func (b *tokenBucket) allow() bool { + b.mu.Lock() + defer b.mu.Unlock() + now := time.Now() + b.tokens += now.Sub(b.lastFill).Seconds() * b.refill + if b.tokens > b.max { + b.tokens = b.max + } + b.lastFill = now + if b.tokens >= 1 { + b.tokens-- + return true + } + return false +} + +// resolver is a single loopback UDP DNS server modelling a remote resolver. +type resolver struct { + conn *net.UDPConn + cfg Config + stats *Stats + bucket *tokenBucket + closeCh chan struct{} + wg sync.WaitGroup +} + +// Battery is a running set of simulated resolvers. +type Battery struct { + Addrs []string // resolver addresses in host:port form (loopback) + Stats *Stats + resolvers []*resolver +} + +// Start launches n simulated resolvers on ephemeral 127.0.0.1 ports. +func Start(n int, cfg Config) (*Battery, error) { + b := &Battery{Stats: &Stats{}} + + for i := 0; i < n; i++ { + conn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}) + if err != nil { + b.Stop() + return nil, err + } + _ = conn.SetReadBuffer(8 * 1024 * 1024) + _ = conn.SetWriteBuffer(8 * 1024 * 1024) + r := &resolver{ + conn: conn, + cfg: cfg, + stats: b.Stats, + bucket: newTokenBucket(cfg.QPSPerServer), + closeCh: make(chan struct{}), + } + b.resolvers = append(b.resolvers, r) + b.Addrs = append(b.Addrs, conn.LocalAddr().String()) + r.wg.Add(1) + go r.serve() + } + + return b, nil +} + +// Stop shuts down all simulated resolvers and waits for their goroutines. +func (b *Battery) Stop() { + for _, r := range b.resolvers { + select { + case <-r.closeCh: + default: + close(r.closeCh) + } + _ = r.conn.Close() + } + for _, r := range b.resolvers { + r.wg.Wait() + } +} + +func (r *resolver) serve() { + defer r.wg.Done() + buf := make([]byte, 4096) + for { + n, addr, err := r.conn.ReadFromUDP(buf) + if err != nil { + select { + case <-r.closeCh: + return + default: + continue + } + } + pkt := make([]byte, n) + copy(pkt, buf[:n]) + r.wg.Add(1) + go r.handle(pkt, addr) + } +} + +func (r *resolver) handle(pkt []byte, addr *net.UDPAddr) { + defer r.wg.Done() + + r.stats.Queries.Add(1) + + // per-resolver rate cap: drop what we cannot serve, like a throttling + // public resolver. The client should rotate to another resolver and retry. + if r.bucket != nil && !r.bucket.allow() { + r.stats.RateLimited.Add(1) + return + } + + // simulate packet loss: silently drop, forcing a retransmit on timeout. + if r.cfg.LossRate > 0 && rand.Float64() < r.cfg.LossRate { + r.stats.Dropped.Add(1) + return + } + + req := new(dns.Msg) + if err := req.Unpack(pkt); err != nil || len(req.Question) == 0 { + return + } + q := req.Question[0] + + // simulate RTT + delay := r.cfg.BaseLatency + if r.cfg.Jitter > 0 { + delay += time.Duration(rand.Int63n(int64(r.cfg.Jitter) + 1)) + } + if delay > 0 { + t := time.NewTimer(delay) + select { + case <-t.C: + case <-r.closeCh: + t.Stop() + return + } + } + + m := new(dns.Msg) + m.SetReply(req) + + switch { + case r.cfg.ServfailRate > 0 && rand.Float64() < r.cfg.ServfailRate: + m.Rcode = dns.RcodeServerFailure + r.stats.Servfail.Add(1) + case q.Qtype == dns.TypeA && NameHits(q.Name, r.cfg.HitPercent) && r.cfg.HijackRate > 0 && rand.Float64() < r.cfg.HijackRate: + // existing name falsely reported as NXDOMAIN (rate-limit/hijack); the + // stub believes it and produces a false negative. + m.Rcode = dns.RcodeNameError + r.stats.Hijacked.Add(1) + case q.Qtype == dns.TypeA && NameHits(q.Name, r.cfg.HitPercent): + rr, err := dns.NewRR(fmt.Sprintf("%s 60 IN A %s", q.Name, SyntheticIP(q.Name))) + if err == nil { + m.Answer = append(m.Answer, rr) + } + case q.Qtype == dns.TypePTR && NameHits(q.Name, r.cfg.HitPercent): + rr, err := dns.NewRR(fmt.Sprintf("%s 60 IN PTR %s", q.Name, SyntheticPTR(q.Name))) + if err == nil { + m.Answer = append(m.Answer, rr) + } + default: + m.Rcode = dns.RcodeNameError + } + + out, err := m.Pack() + if err != nil { + return + } + if _, err := r.conn.WriteToUDP(out, addr); err == nil { + r.stats.Answered.Add(1) + } +} + +// NameHits deterministically decides whether a name resolves, so the hit ratio +// is stable across runs and across resolvers (every resolver agrees). +func NameHits(name string, hitPercent int) bool { + if hitPercent >= 100 { + return true + } + if hitPercent <= 0 { + return false + } + return int(fnv32(name)%100) < hitPercent +} + +// SyntheticIP returns a deterministic plausible host address for a name. +func SyntheticIP(name string) string { + h := fnv32(name) + last := byte(h%254) + 1 // avoid .0/.255 + return fmt.Sprintf("10.%d.%d.%d", byte(h>>16), byte(h>>8), last) +} + +// SyntheticPTR returns a deterministic plausible hostname for a reverse query. +func SyntheticPTR(name string) string { + return fmt.Sprintf("host-%d.ptr.example.com.", fnv32(name)%100000) +} + +func fnv32(s string) uint32 { + const ( + offset = 2166136261 + prime = 16777619 + ) + h := uint32(offset) + for i := 0; i < len(s); i++ { + h ^= uint32(s[i]) + h *= prime + } + return h +} diff --git a/pkg/axfr/axfr.go b/pkg/axfr/axfr.go new file mode 100644 index 00000000..10e89ed7 --- /dev/null +++ b/pkg/axfr/axfr.go @@ -0,0 +1,246 @@ +// Package axfr attempts DNS zone transfers (AXFR, with IXFR fallback) against a +// zone's authoritative nameservers. A misconfigured nameserver that allows +// transfers from arbitrary clients returns the entire zone — every owner name — +// in a single exchange, which is the highest-payoff subdomain enumeration +// shortcut: no wordlist, no guessing, complete and authoritative. +// +// Open AXFR is uncommon on well-run zones (low single-digit percent), but it is +// nearly free to attempt and total when it lands, so it belongs as a first pass +// before bruteforce. Secondary/forgotten nameservers are frequently laxer than +// the primary, so every nameserver of the zone is tried. +package axfr + +import ( + "context" + "fmt" + "net" + "strings" + "time" + + "github.com/miekg/dns" +) + +// Config controls a zone-transfer attempt. +type Config struct { + // Zone is the apex to transfer (e.g. "example.com"). + Zone string + // Resolvers are used to discover the zone's NS set and resolve nameserver + // addresses when Nameservers is not supplied (host or host:port). + Resolvers []string + // Nameservers optionally lists explicit nameserver addresses to try + // (host or host:port); when set, NS discovery is skipped. + Nameservers []string + // Timeout is the per-nameserver transfer timeout. Default 10s. + Timeout time.Duration + // MaxNames caps the owner names kept from a single transfer, bounding memory + // against a hostile server that streams an endless zone. Default 1,000,000. + MaxNames int + // OnName fires for each newly discovered owner name (deduplicated). + OnName func(string) + // OnNameserver fires after each nameserver attempt with its outcome. + OnNameserver func(ns string, names int, err error) +} + +// Result is the outcome of a successful zone transfer. +type Result struct { + Zone string + Nameserver string // nameserver that allowed the transfer + Names []string // unique owner names within the zone (no trailing dot) + Records int // total resource records transferred +} + +// Attempt tries to transfer the zone from each of its nameservers and returns +// the first successful transfer. It returns an error only when no nameserver +// allowed the transfer (the common case for well-configured zones). +func Attempt(ctx context.Context, cfg Config) (*Result, error) { + if strings.TrimSpace(cfg.Zone) == "" { + return nil, fmt.Errorf("zone is required") + } + if cfg.Timeout <= 0 { + cfg.Timeout = 10 * time.Second + } + if cfg.MaxNames <= 0 { + cfg.MaxNames = 1_000_000 + } + apex := dns.Fqdn(strings.ToLower(cfg.Zone)) + + servers := cfg.Nameservers + if len(servers) == 0 { + var err error + servers, err = discoverNameservers(ctx, apex, cfg.Resolvers, cfg.Timeout) + if err != nil { + return nil, err + } + } + if len(servers) == 0 { + return nil, fmt.Errorf("no nameservers found for %s", cfg.Zone) + } + + var lastErr error + for _, ns := range servers { + res, err := transferFrom(ctx, apex, withPort(ns), cfg) + if cfg.OnNameserver != nil { + n := 0 + if res != nil { + n = len(res.Names) + } + cfg.OnNameserver(ns, n, err) + } + if err == nil && res != nil && len(res.Names) > 0 { + res.Zone = cfg.Zone + return res, nil + } + if err != nil { + lastErr = err + } + } + if lastErr == nil { + lastErr = fmt.Errorf("no nameserver allowed transfer of %s", cfg.Zone) + } + return nil, lastErr +} + +// transferFrom performs the AXFR (falling back to IXFR) against one nameserver. +func transferFrom(ctx context.Context, apex, nsAddr string, cfg Config) (*Result, error) { + run := func(setup func(*dns.Msg)) (*Result, error) { + t := &dns.Transfer{DialTimeout: cfg.Timeout, ReadTimeout: cfg.Timeout, WriteTimeout: cfg.Timeout} + m := new(dns.Msg) + setup(m) + ch, err := t.In(m, nsAddr) + if err != nil { + return nil, err + } + res := &Result{Nameserver: nsAddr} + seen := map[string]struct{}{} + // abandoning `ch` mid-stream would block miekg's sender goroutine forever + // on its unbuffered channel send (leaking the goroutine and the TCP conn). + // stop() unblocks that send by draining and closes the conn so the sender + // returns promptly instead of waiting out another read timeout. + stopped := false + stop := func() { + if stopped { + return + } + stopped = true + go func() { + for range ch { + } + }() + if t.Conn != nil { + _ = t.Close() + } + } + for env := range ch { + if env.Error != nil { + if len(res.Names) > 0 { + return res, nil // partial transfer is still useful + } + return nil, env.Error + } + for _, rr := range env.RR { + res.Records++ + owner := strings.TrimSuffix(strings.ToLower(rr.Header().Name), ".") + name := dns.Fqdn(strings.ToLower(rr.Header().Name)) + if !inZone(name, apex) { + continue + } + if _, dup := seen[owner]; dup { + continue + } + seen[owner] = struct{}{} + res.Names = append(res.Names, owner) + if cfg.OnName != nil { + cfg.OnName(owner) + } + if len(res.Names) >= cfg.MaxNames { + stop() + return res, nil + } + } + if err := ctx.Err(); err != nil { + stop() + return res, err + } + } + return res, nil + } + + res, err := run(func(m *dns.Msg) { m.SetAxfr(apex) }) + if err == nil && res != nil && len(res.Names) > 0 { + return res, nil + } + // some servers refuse AXFR but permit IXFR; try it as a fallback. + if res2, err2 := run(func(m *dns.Msg) { m.SetIxfr(apex, 0, "", "") }); err2 == nil && res2 != nil && len(res2.Names) > 0 { + return res2, nil + } + return res, err +} + +// discoverNameservers resolves the zone's NS set and their addresses. +func discoverNameservers(ctx context.Context, apex string, resolvers []string, timeout time.Duration) ([]string, error) { + if len(resolvers) == 0 { + return nil, fmt.Errorf("resolvers are required to discover nameservers (or pass Nameservers)") + } + server := withPort(resolvers[0]) + client := &dns.Client{Timeout: timeout} + + m := new(dns.Msg) + m.SetQuestion(apex, dns.TypeNS) + m.RecursionDesired = true + resp, _, err := client.ExchangeContext(ctx, m, server) + if err != nil { + return nil, fmt.Errorf("NS lookup for %s failed: %w", apex, err) + } + + var nsNames []string + for _, rr := range resp.Answer { + if ns, ok := rr.(*dns.NS); ok { + nsNames = append(nsNames, ns.Ns) + } + } + + var addrs []string + seen := map[string]struct{}{} + for _, ns := range nsNames { + for _, qt := range []uint16{dns.TypeA, dns.TypeAAAA} { + am := new(dns.Msg) + am.SetQuestion(dns.Fqdn(ns), qt) + am.RecursionDesired = true + ar, _, aerr := client.ExchangeContext(ctx, am, server) + if aerr != nil { + continue + } + for _, rr := range ar.Answer { + var ip string + switch v := rr.(type) { + case *dns.A: + ip = v.A.String() + case *dns.AAAA: + ip = v.AAAA.String() + } + if ip == "" { + continue + } + if _, dup := seen[ip]; dup { + continue + } + seen[ip] = struct{}{} + addrs = append(addrs, ip) + } + } + } + return addrs, nil +} + +func inZone(name, apex string) bool { + name, apex = dns.Fqdn(strings.ToLower(name)), dns.Fqdn(strings.ToLower(apex)) + return name == apex || strings.HasSuffix(name, "."+apex) +} + +func withPort(s string) string { + s = strings.TrimSpace(s) + if _, _, err := net.SplitHostPort(s); err != nil { + return net.JoinHostPort(s, "53") + } + return s +} diff --git a/pkg/axfr/axfr_test.go b/pkg/axfr/axfr_test.go new file mode 100644 index 00000000..fec1137a --- /dev/null +++ b/pkg/axfr/axfr_test.go @@ -0,0 +1,102 @@ +package axfr + +import ( + "context" + "net" + "testing" + + "github.com/miekg/dns" +) + +// startTCPServer starts a TCP DNS server (AXFR runs over TCP) and returns its addr. +func startTCPServer(t *testing.T, handler dns.HandlerFunc) (string, func()) { + t.Helper() + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + srv := &dns.Server{Listener: l, Handler: handler} + go func() { _ = srv.ActivateAndServe() }() + return l.Addr().String(), func() { _ = srv.Shutdown() } +} + +func mustRR(t *testing.T, s string) dns.RR { + t.Helper() + rr, err := dns.NewRR(s) + if err != nil { + t.Fatalf("NewRR(%q): %v", s, err) + } + return rr +} + +func TestAttemptAXFR(t *testing.T) { + zone := "example.com." + records := func() []dns.RR { + return []dns.RR{ + mustRR(t, "example.com. 3600 IN SOA ns1.example.com. admin.example.com. 1 3600 600 86400 60"), + mustRR(t, "example.com. 3600 IN NS ns1.example.com."), + mustRR(t, "www.example.com. 3600 IN A 93.184.216.34"), + mustRR(t, "mail.example.com. 3600 IN A 93.184.216.35"), + mustRR(t, "ftp.example.com. 3600 IN CNAME www.example.com."), + mustRR(t, "example.com. 3600 IN SOA ns1.example.com. admin.example.com. 1 3600 600 86400 60"), + } + } + + addr, stop := startTCPServer(t, func(w dns.ResponseWriter, r *dns.Msg) { + if r.Question[0].Qtype == dns.TypeAXFR { + ch := make(chan *dns.Envelope) + tr := new(dns.Transfer) + go func() { + ch <- &dns.Envelope{RR: records()} + close(ch) + }() + _ = tr.Out(w, r, ch) + return + } + m := new(dns.Msg) + m.SetReply(r) + m.Rcode = dns.RcodeRefused + _ = w.WriteMsg(m) + }) + defer stop() + + var streamed []string + res, err := Attempt(context.Background(), Config{ + Zone: zone, + Nameservers: []string{addr}, + OnName: func(n string) { streamed = append(streamed, n) }, + }) + if err != nil { + t.Fatalf("Attempt: %v", err) + } + want := map[string]bool{"www.example.com": false, "mail.example.com": false, "ftp.example.com": false, "example.com": false} + for _, n := range res.Names { + if _, ok := want[n]; ok { + want[n] = true + } + } + for n, found := range want { + if !found { + t.Errorf("expected %q in transferred names %v", n, res.Names) + } + } + if len(streamed) != len(res.Names) { + t.Errorf("OnName fired %d times, result has %d names", len(streamed), len(res.Names)) + } +} + +func TestAttemptRefused(t *testing.T) { + // a server that refuses transfers must yield an error, not a panic/partial. + addr, stop := startTCPServer(t, func(w dns.ResponseWriter, r *dns.Msg) { + m := new(dns.Msg) + m.SetReply(r) + m.Rcode = dns.RcodeRefused + _ = w.WriteMsg(m) + }) + defer stop() + + _, err := Attempt(context.Background(), Config{Zone: "secure.example.", Nameservers: []string{addr}}) + if err == nil { + t.Fatal("expected error when transfer is refused") + } +} diff --git a/pkg/checkpoint/checkpoint.go b/pkg/checkpoint/checkpoint.go new file mode 100644 index 00000000..a9f54c0e --- /dev/null +++ b/pkg/checkpoint/checkpoint.go @@ -0,0 +1,151 @@ +// Package checkpoint provides crash-safe stop/resume for long resolution runs. +// +// The model is an append-only log of completed names. A name is "completed" once +// it reaches a terminal state (an answer or a definitive give-up), at which point +// the consumer calls Done. On a later run, Open replays the log so the producer +// can skip names that already finished via Has. +// +// Semantics are at-least-once: only names recorded as done are skipped, so any +// name that was in flight when the process died is simply re-resolved on resume. +// Nothing is ever lost (no false negatives); at worst a small in-flight window is +// repeated, which is harmless because downstream dedup absorbs it. +package checkpoint + +import ( + "bufio" + "fmt" + "os" + "strings" + "sync" + "time" +) + +// flushEvery bounds how many completions buffer before a forced flush, so a +// crash loses at most this many log records (which are re-resolved anyway). +const flushEvery = 1024 + +// Checkpoint is a resumable completed-name log. It is safe for concurrent use. +type Checkpoint struct { + mu sync.Mutex + done map[string]struct{} + f *os.File + w *bufio.Writer + pending int + closed bool +} + +// Open opens (creating if needed) the checkpoint at path and loads any previously +// completed names from it. Subsequent writes append to the same file. +func Open(path string) (*Checkpoint, error) { + if strings.TrimSpace(path) == "" { + return nil, fmt.Errorf("checkpoint path is required") + } + c := &Checkpoint{done: make(map[string]struct{})} + + // load existing entries (resume) + if f, err := os.Open(path); err == nil { + sc := bufio.NewScanner(f) + sc.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for sc.Scan() { + if line := strings.TrimSpace(sc.Text()); line != "" { + c.done[line] = struct{}{} + } + } + _ = f.Close() + if err := sc.Err(); err != nil { + return nil, fmt.Errorf("could not read checkpoint: %w", err) + } + } else if !os.IsNotExist(err) { + return nil, fmt.Errorf("could not open checkpoint: %w", err) + } + + f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600) + if err != nil { + return nil, fmt.Errorf("could not open checkpoint for append: %w", err) + } + c.f = f + c.w = bufio.NewWriterSize(f, 64*1024) + return c, nil +} + +// Resumed reports the number of completed names loaded from a prior run. +func (c *Checkpoint) Resumed() int { + c.mu.Lock() + defer c.mu.Unlock() + return len(c.done) +} + +// Has reports whether name already completed in this or a prior run. +func (c *Checkpoint) Has(name string) bool { + c.mu.Lock() + _, ok := c.done[name] + c.mu.Unlock() + return ok +} + +// Done records name as completed. It returns false if the name was already +// recorded (so callers can detect duplicates). Records are buffered and flushed +// periodically; Close guarantees durability of everything recorded. +func (c *Checkpoint) Done(name string) bool { + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return false + } + if _, ok := c.done[name]; ok { + return false + } + c.done[name] = struct{}{} + _, _ = c.w.WriteString(name) + _ = c.w.WriteByte('\n') + c.pending++ + if c.pending >= flushEvery { + _ = c.w.Flush() + c.pending = 0 + } + return true +} + +// Flush flushes buffered records to disk. +func (c *Checkpoint) Flush() error { + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return nil + } + c.pending = 0 + return c.w.Flush() +} + +// Close flushes and closes the underlying file. +func (c *Checkpoint) Close() error { + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return nil + } + c.closed = true + if err := c.w.Flush(); err != nil { + _ = c.f.Close() + return err + } + return c.f.Close() +} + +// FlushPeriodically flushes the log on the given interval until stop is closed. +// Useful for bounding data loss on a crash during very long, low-completion runs. +func (c *Checkpoint) FlushPeriodically(interval time.Duration, stop <-chan struct{}) { + if interval <= 0 { + return + } + t := time.NewTicker(interval) + defer t.Stop() + for { + select { + case <-stop: + return + case <-t.C: + _ = c.Flush() + } + } +} diff --git a/pkg/checkpoint/checkpoint_test.go b/pkg/checkpoint/checkpoint_test.go new file mode 100644 index 00000000..a55cbb08 --- /dev/null +++ b/pkg/checkpoint/checkpoint_test.go @@ -0,0 +1,128 @@ +package checkpoint + +import ( + "path/filepath" + "sync" + "testing" +) + +func TestDoneAndHas(t *testing.T) { + path := filepath.Join(t.TempDir(), "cp.log") + c, err := Open(path) + if err != nil { + t.Fatal(err) + } + if c.Has("a.com") { + t.Fatal("fresh checkpoint should not have anything") + } + if !c.Done("a.com") { + t.Fatal("first Done should report newly-completed") + } + if c.Done("a.com") { + t.Fatal("second Done should report already-completed") + } + if !c.Has("a.com") { + t.Fatal("Has should see completed name") + } + if err := c.Close(); err != nil { + t.Fatal(err) + } +} + +func TestResume(t *testing.T) { + path := filepath.Join(t.TempDir(), "cp.log") + c, err := Open(path) + if err != nil { + t.Fatal(err) + } + for _, n := range []string{"a.com", "b.com", "c.com"} { + c.Done(n) + } + if err := c.Close(); err != nil { + t.Fatal(err) + } + + // reopen: prior completions must be visible + c2, err := Open(path) + if err != nil { + t.Fatal(err) + } + defer func() { _ = c2.Close() }() + if c2.Resumed() != 3 { + t.Fatalf("expected 3 resumed, got %d", c2.Resumed()) + } + for _, n := range []string{"a.com", "b.com", "c.com"} { + if !c2.Has(n) { + t.Fatalf("resumed checkpoint missing %q", n) + } + } + if c2.Has("d.com") { + t.Fatal("should not have uncompleted name") + } +} + +// TestCrashSafety simulates a crash (no Close) after enough completions to force +// at least one buffer flush; those must survive on reopen. +func TestCrashSafety(t *testing.T) { + path := filepath.Join(t.TempDir(), "cp.log") + c, err := Open(path) + if err != nil { + t.Fatal(err) + } + const n = flushEvery + 10 + for i := 0; i < n; i++ { + c.Done(itoa(i)) + } + // no Close(): only flushed records (>= flushEvery) are guaranteed durable. + c2, err := Open(path) + if err != nil { + t.Fatal(err) + } + defer func() { _ = c2.Close() }() + if c2.Resumed() < flushEvery { + t.Fatalf("expected at least %d durable records, got %d", flushEvery, c2.Resumed()) + } +} + +func TestConcurrentDone(t *testing.T) { + path := filepath.Join(t.TempDir(), "cp.log") + c, err := Open(path) + if err != nil { + t.Fatal(err) + } + defer func() { _ = c.Close() }() + var wg sync.WaitGroup + for g := 0; g < 8; g++ { + wg.Add(1) + go func(g int) { + defer wg.Done() + for i := 0; i < 1000; i++ { + c.Done(itoa(g*1000 + i)) + } + }(g) + } + wg.Wait() + if got := c.Resumed(); got != 8000 { + t.Fatalf("expected 8000 unique completions, got %d", got) + } +} + +func TestOpenRequiresPath(t *testing.T) { + if _, err := Open(""); err == nil { + t.Fatal("expected error for empty path") + } +} + +func itoa(i int) string { + if i == 0 { + return "0" + } + var b [20]byte + pos := len(b) + for i > 0 { + pos-- + b[pos] = byte('0' + i%10) + i /= 10 + } + return "host-" + string(b[pos:]) +} diff --git a/pkg/iterative/bench_test.go b/pkg/iterative/bench_test.go new file mode 100644 index 00000000..c6226acf --- /dev/null +++ b/pkg/iterative/bench_test.go @@ -0,0 +1,130 @@ +package iterative + +import ( + "context" + "fmt" + "os" + "sync/atomic" + "testing" + "time" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/pkg/resolve" +) + +// buildBenchHierarchy creates a wide tree: one root, one .com TLD, and `domains` +// registrable domains each served by its own authoritative server with `hosts` +// names. This models a real bruteforce workload (many names under relatively +// few registrable domains), where the delegation cache should collapse to ~1 +// upstream query per name after warmup. +func buildBenchHierarchy(domains, hosts int) (*memHierarchy, []string) { + h := newHierarchy() + const rootIP = "127.0.0.1" + const comIP = "127.0.1.1" + + root := h.server(rootIP).authoritative(".") + root.delegate("com.", memNS{name: "a.gtld.net.", ip: comIP}) + + com := h.server(comIP).authoritative("com.") + + var names []string + for d := 0; d < domains; d++ { + zone := fmt.Sprintf("d%d.com.", d) + authIP := ip4(4096 + d) // distinct 127.0.x.y per domain + com.delegate(zone, memNS{name: fmt.Sprintf("ns.%s", zone), ip: authIP}) + auth := h.server(authIP).authoritative(zone) + auth.a(fmt.Sprintf("ns.%s", zone), authIP) + for n := 0; n < hosts; n++ { + name := fmt.Sprintf("host%d.d%d.com", n, d) + auth.a(canonical(name), fmt.Sprintf("10.%d.%d.%d", byte(d>>8), byte(d), n%254+1)) + names = append(names, name) + } + } + return h, names +} + +// TestIterativeBenchmark reports throughput, TTFR, and the cache win +// (upstream queries per name). Gated behind ITER_BENCH=1 so it doesn't run in +// normal CI. No real traffic: the hierarchy is fully in-memory. +// +// ITER_BENCH=1 go test ./pkg/iterative -run TestIterativeBenchmark -v +func TestIterativeBenchmark(t *testing.T) { + if os.Getenv("ITER_BENCH") == "" { + t.Skip("set ITER_BENCH=1 to run the iterative resolver benchmark") + } + domains, hosts, concurrency := 200, 500, 256 + h, names := buildBenchHierarchy(domains, hosts) + + r, err := New(Options{ + RootServers: []string{"127.0.0.1"}, + QueryType: dns.TypeA, + Concurrency: concurrency, + }) + if err != nil { + t.Fatal(err) + } + r.newExchanger = h.factory() + + in := make(chan string, 4096) + go func() { + defer close(in) + for _, n := range names { + in <- n + } + }() + + var resolved, ttfrOnce atomic.Int64 + start := time.Now() + var ttfr time.Duration + err = r.ResolveStream(context.Background(), in, StreamConfig{ + OnResult: func(res *resolve.Result) { + if len(res.A) > 0 { + if ttfrOnce.CompareAndSwap(0, 1) { + ttfr = time.Since(start) + } + resolved.Add(1) + } + }, + }) + if err != nil { + t.Fatal(err) + } + wall := time.Since(start) + + total := len(names) + upstream := h.queries.Load() + t.Logf("names=%d domains=%d hosts/domain=%d concurrency=%d", total, domains, hosts, concurrency) + t.Logf("resolved=%d wall=%s qps=%.0f ttfr=%s", resolved.Load(), wall.Round(time.Millisecond), float64(total)/wall.Seconds(), ttfr.Round(time.Microsecond)) + t.Logf("upstream queries=%d queries/name=%.3f (cache win: ->1.0 means root/TLD/zone walked once)", upstream, float64(upstream)/float64(total)) + + if resolved.Load() != int64(total) { + t.Fatalf("expected all %d names resolved, got %d", total, resolved.Load()) + } +} + +// BenchmarkIterativeStream is a standard go benchmark of warm-cache resolution +// throughput (algorithm + cache + concurrency overhead, no network). +func BenchmarkIterativeStream(b *testing.B) { + h, names := buildBenchHierarchy(50, 200) + r, _ := New(Options{RootServers: []string{"127.0.0.1"}, QueryType: dns.TypeA, Concurrency: 128}) + r.newExchanger = h.factory() + + // warm the cache once + warm := make(chan string, len(names)) + for _, n := range names { + warm <- n + } + close(warm) + _ = r.ResolveStream(context.Background(), warm, StreamConfig{OnResult: func(*resolve.Result) {}}) + + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + in := make(chan string, len(names)) + for _, n := range names { + in <- n + } + close(in) + _ = r.ResolveStream(context.Background(), in, StreamConfig{OnResult: func(*resolve.Result) {}}) + } +} diff --git a/pkg/iterative/cache.go b/pkg/iterative/cache.go new file mode 100644 index 00000000..4319ab39 --- /dev/null +++ b/pkg/iterative/cache.go @@ -0,0 +1,109 @@ +package iterative + +import ( + "net" + "strings" + "time" + + lru "github.com/hashicorp/golang-lru/v2" +) + +// nsEntry is a single nameserver of a delegation: its name plus any glue +// addresses learned from the parent zone (empty when the delegation is +// glueless and the address must be resolved separately). +type nsEntry struct { + name string + addrs []net.IP +} + +// delegation is the set of nameservers authoritative for (or one step closer +// to) a zone, as learned from a referral. It is the unit cached and shared +// across all in-flight resolutions, which is what amortizes the cost of +// re-walking root/TLD for every name (the key win over a per-name resolver). +type delegation struct { + zone string // canonical zone (lowercase, trailing dot) + ns []nsEntry // nameservers, some possibly with glue + expiry time.Time // zero == never expires (root hints) +} + +func (d *delegation) expired(now time.Time) bool { + return !d.expiry.IsZero() && now.After(d.expiry) +} + +// cache is a bounded, TTL-aware store of delegations keyed by canonical zone. +// It is safe for concurrent use (the underlying LRU is locked) and shared by +// every worker so that common ancestors (root, TLDs, popular zones) are walked +// once and reused. +type cache struct { + lru *lru.Cache[string, *delegation] +} + +func newCache(size int) *cache { + if size <= 0 { + size = 1 << 16 + } + l, _ := lru.New[string, *delegation](size) + return &cache{lru: l} +} + +func (c *cache) put(d *delegation) { + if d == nil || d.zone == "" { + return + } + c.lru.Add(canonical(d.zone), d) +} + +// best returns the deepest (closest to name) non-expired cached delegation, +// walking ancestor labels of name. It returns nil when nothing is cached. +func (c *cache) best(name string) *delegation { + now := time.Now() + for z := canonical(name); ; z = parentZone(z) { + if d, ok := c.lru.Get(z); ok { + if d.expired(now) { + c.lru.Remove(z) + } else { + return d + } + } + if z == "." { + return nil + } + } +} + +// canonical lowercases a name and ensures a single trailing dot. The root is ".". +func canonical(name string) string { + if name == "" || name == "." { + return "." + } + name = strings.ToLower(name) + if !strings.HasSuffix(name, ".") { + name += "." + } + return name +} + +// parentZone returns the parent of a canonical zone ("a.b.c." -> "b.c."); the +// parent of a TLD (or root) is the root ".". +func parentZone(zone string) string { + if zone == "." || zone == "" { + return "." + } + zone = strings.TrimSuffix(zone, ".") + i := strings.IndexByte(zone, '.') + if i < 0 { + return "." + } + return zone[i+1:] + "." +} + +// inBailiwick reports whether child is equal to or a subdomain of parent. It is +// used to reject out-of-bailiwick referrals and glue (a core anti-poisoning +// check): an authority can only delegate names within its own zone. +func inBailiwick(child, parent string) bool { + child, parent = canonical(child), canonical(parent) + if parent == "." { + return true + } + return child == parent || strings.HasSuffix(child, "."+parent) +} diff --git a/pkg/iterative/cache_test.go b/pkg/iterative/cache_test.go new file mode 100644 index 00000000..69abf5cb --- /dev/null +++ b/pkg/iterative/cache_test.go @@ -0,0 +1,70 @@ +package iterative + +import ( + "testing" + "time" +) + +func TestCanonicalAndParent(t *testing.T) { + cases := map[string]string{ + "WWW.Example.COM": "www.example.com.", + "example.com.": "example.com.", + "": ".", + ".": ".", + } + for in, want := range cases { + if got := canonical(in); got != want { + t.Errorf("canonical(%q)=%q want %q", in, got, want) + } + } + parents := map[string]string{ + "a.b.c.": "b.c.", + "example.com.": "com.", + "com.": ".", + ".": ".", + } + for in, want := range parents { + if got := parentZone(in); got != want { + t.Errorf("parentZone(%q)=%q want %q", in, got, want) + } + } +} + +func TestInBailiwick(t *testing.T) { + yes := [][2]string{{"a.example.com", "example.com"}, {"example.com", "example.com"}, {"x.y.com", "com"}, {"anything", "."}} + no := [][2]string{{"example.org", "example.com"}, {"com", "example.com"}, {"notexample.com", "example.com"}} + for _, c := range yes { + if !inBailiwick(c[0], c[1]) { + t.Errorf("inBailiwick(%q,%q) = false, want true", c[0], c[1]) + } + } + for _, c := range no { + if inBailiwick(c[0], c[1]) { + t.Errorf("inBailiwick(%q,%q) = true, want false", c[0], c[1]) + } + } +} + +func TestCacheDeepestAncestor(t *testing.T) { + c := newCache(1024) + c.put(&delegation{zone: "com."}) + c.put(&delegation{zone: "example.com."}) + + if d := c.best("www.example.com"); d == nil || d.zone != "example.com." { + t.Fatalf("expected deepest match example.com., got %#v", d) + } + if d := c.best("host.other.com"); d == nil || d.zone != "com." { + t.Fatalf("expected com., got %#v", d) + } + if d := c.best("nothing.org"); d != nil { + t.Fatalf("expected nil for uncached tree, got %#v", d) + } +} + +func TestCacheTTLExpiry(t *testing.T) { + c := newCache(1024) + c.put(&delegation{zone: "example.com.", expiry: time.Now().Add(-time.Second)}) + if d := c.best("a.example.com"); d != nil { + t.Fatalf("expired delegation should not be returned, got %#v", d) + } +} diff --git a/pkg/iterative/hierarchy_test.go b/pkg/iterative/hierarchy_test.go new file mode 100644 index 00000000..c1bfb3a6 --- /dev/null +++ b/pkg/iterative/hierarchy_test.go @@ -0,0 +1,183 @@ +package iterative + +import ( + "context" + "fmt" + "net" + "net/netip" + "sync/atomic" + + "github.com/miekg/dns" +) + +// memHierarchy is an in-memory authoritative DNS tree used to test and +// benchmark the iterative resolver without any real sockets or internet +// traffic. Each "server" is addressed by a synthetic loopback IP; the resolver +// reaches it through the injected exchanger, which routes by destination IP and +// role-plays the appropriate authoritative behaviour (referral / answer / +// NODATA / NXDOMAIN). +type memHierarchy struct { + servers map[string]*memServer // keyed by IP string + queries atomic.Int64 // total queries served (for cache-win assertions) +} + +type memNS struct { + name string + ip string // "" == glueless +} + +type memServer struct { + ip string + authZones []string + delegations map[string][]memNS // childZone -> nameservers + records map[string]map[uint16][]string +} + +func newHierarchy() *memHierarchy { return &memHierarchy{servers: map[string]*memServer{}} } + +func (h *memHierarchy) server(ip string) *memServer { + s, ok := h.servers[ip] + if !ok { + s = &memServer{ip: ip, delegations: map[string][]memNS{}, records: map[string]map[uint16][]string{}} + h.servers[ip] = s + } + return s +} + +func (s *memServer) authoritative(zones ...string) *memServer { + for _, z := range zones { + s.authZones = append(s.authZones, canonical(z)) + } + return s +} + +func (s *memServer) delegate(child string, ns ...memNS) *memServer { + s.delegations[canonical(child)] = append(s.delegations[canonical(child)], ns...) + return s +} + +func (s *memServer) rr(name string, qtype uint16, data ...string) *memServer { + name = canonical(name) + if s.records[name] == nil { + s.records[name] = map[uint16][]string{} + } + s.records[name][qtype] = append(s.records[name][qtype], data...) + return s +} + +func (s *memServer) a(name string, ips ...string) *memServer { return s.rr(name, dns.TypeA, ips...) } +func (s *memServer) cname(name, target string) *memServer { return s.rr(name, dns.TypeCNAME, target) } + +func (s *memServer) hasAny(name string) bool { + _, ok := s.records[canonical(name)] + return ok +} + +// exchange implements the exchanger interface for the resolver. +func (h *memHierarchy) exchange(ctx context.Context, server netip.AddrPort, msg *dns.Msg) (*dns.Msg, error) { + h.queries.Add(1) + srv := h.servers[server.Addr().Unmap().String()] + m := new(dns.Msg) + m.SetReply(msg) + m.RecursionAvailable = false + if srv == nil { + return nil, errTimeout // unreachable server (no response) + } + q := msg.Question[0] + qname := canonical(q.Name) + qtype := q.Qtype + + // most specific authoritative zone for the name + authZone := "" + for _, z := range srv.authZones { + if inBailiwick(qname, z) && len(z) > len(authZone) { + authZone = z + } + } + // most specific delegation that is an ancestor of the name + bestChild := "" + for child := range srv.delegations { + if inBailiwick(qname, child) && len(child) > len(bestChild) { + bestChild = child + } + } + + // referral: a delegation more specific than what we serve authoritatively + if bestChild != "" && len(bestChild) > len(authZone) { + for _, ns := range srv.delegations[bestChild] { + m.Ns = append(m.Ns, mustRR(fmt.Sprintf("%s 3600 IN NS %s", bestChild, ns.name))) + if ns.ip != "" { + m.Extra = append(m.Extra, mustRR(fmt.Sprintf("%s 3600 IN A %s", ns.name, ns.ip))) + } + } + return m, nil + } + + if authZone == "" { + m.Rcode = dns.RcodeServerFailure + return m, nil + } + m.Authoritative = true + + if recs := srv.records[qname][qtype]; len(recs) > 0 { + for _, d := range recs { + m.Answer = append(m.Answer, mustRR(rrString(qname, qtype, d))) + } + return m, nil + } + // CNAME indirection + if cns := srv.records[qname][dns.TypeCNAME]; len(cns) > 0 && qtype != dns.TypeCNAME { + tgt := canonical(cns[0]) + m.Answer = append(m.Answer, mustRR(fmt.Sprintf("%s 3600 IN CNAME %s", qname, tgt))) + // include in-zone target records (as a real authoritative server would) + if inBailiwick(tgt, authZone) { + for _, d := range srv.records[tgt][qtype] { + m.Answer = append(m.Answer, mustRR(rrString(tgt, qtype, d))) + } + } + return m, nil + } + + // NODATA (name exists, type doesn't) vs NXDOMAIN + if !srv.hasAny(qname) { + m.Rcode = dns.RcodeNameError + } + m.Ns = append(m.Ns, mustRR(fmt.Sprintf("%s 3600 IN SOA ns.%s hostmaster.%s 1 3600 600 86400 60", authZone, authZone, authZone))) + return m, nil +} + +func (h *memHierarchy) close() {} + +// factory returns a newExchanger function that hands every worker the shared +// in-memory hierarchy (no per-worker socket). +func (h *memHierarchy) factory() func() (exchanger, error) { + return func() (exchanger, error) { return h, nil } +} + +func rrString(name string, qtype uint16, data string) string { + switch qtype { + case dns.TypeA: + return fmt.Sprintf("%s 60 IN A %s", name, data) + case dns.TypeAAAA: + return fmt.Sprintf("%s 60 IN AAAA %s", name, data) + case dns.TypeTXT: + return fmt.Sprintf("%s 60 IN TXT \"%s\"", name, data) + case dns.TypeNS: + return fmt.Sprintf("%s 60 IN NS %s", name, data) + case dns.TypeMX: + return fmt.Sprintf("%s 60 IN MX 10 %s", name, data) + default: + return fmt.Sprintf("%s 60 IN A %s", name, data) + } +} + +func mustRR(s string) dns.RR { + rr, err := dns.NewRR(s) + if err != nil { + panic(err) + } + return rr +} + +// ip4 generates distinct synthetic loopback IPs for hierarchy servers. +func ip4(n int) string { return net.IPv4(127, 0, byte(n>>8), byte(n)).String() } diff --git a/pkg/iterative/parse.go b/pkg/iterative/parse.go new file mode 100644 index 00000000..5bd881f6 --- /dev/null +++ b/pkg/iterative/parse.go @@ -0,0 +1,165 @@ +package iterative + +import ( + "strings" + "time" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/pkg/resolve" +) + +// chaseCNAME follows a CNAME chain present in answer starting at start and +// returns the final target name plus the chain of CNAME targets traversed. If +// there is no CNAME for start, it returns (start, nil). +func chaseCNAME(answer []dns.RR, start string) (string, []string) { + cur := canonical(start) + var chain []string + seen := map[string]struct{}{cur: {}} + for { + var next string + for _, rr := range answer { + c, ok := rr.(*dns.CNAME) + if !ok { + continue + } + if canonical(c.Header().Name) == cur { + next = canonical(c.Target) + break + } + } + if next == "" { + return cur, chain + } + if _, dup := seen[next]; dup { + return cur, chain + } + seen[next] = struct{}{} + chain = append(chain, strings.TrimSuffix(next, ".")) + cur = next + } +} + +// answersType reports whether answer contains at least one record of qtype +// owned by name. +func answersType(answer []dns.RR, name string, qtype uint16) bool { + name = canonical(name) + for _, rr := range answer { + if rr.Header().Rrtype == qtype && canonical(rr.Header().Name) == name { + return true + } + } + return false +} + +// buildResult converts a final authoritative response into a resolve.Result, +// collecting every record type so the output formatter can render faithfully. +// cnameChain holds CNAME targets traversed before the final answer. +func buildResult(name string, qtype uint16, resp *dns.Msg, server string, cnameChain []string) *resolve.Result { + r := &resolve.Result{ + Name: strings.TrimSuffix(canonical(name), "."), + Type: qtype, + Rcode: resp.Rcode, + Resolver: server, + Msg: resp, + Timestamp: time.Now(), + CNAME: append([]string{}, cnameChain...), + } + // Only harvest records owned by the queried name or a name reached through + // its CNAME chain. Without this an in-bailiwick-but-hostile (or lame) + // authoritative could staple unrelated A/AAAA records for other owners into + // this name's result. + validOwners := map[string]struct{}{canonical(name): {}} + for _, c := range cnameChain { + validOwners[canonical(c)] = struct{}{} + } + for changed := true; changed; { + changed = false + for _, rr := range resp.Answer { + c, ok := rr.(*dns.CNAME) + if !ok { + continue + } + if _, in := validOwners[canonical(c.Header().Name)]; !in { + continue + } + if tgt := canonical(c.Target); !mapHas(validOwners, tgt) { + validOwners[tgt] = struct{}{} + changed = true + } + } + } + for _, rr := range resp.Answer { + if _, ok := validOwners[canonical(rr.Header().Name)]; !ok { + continue + } + switch v := rr.(type) { + case *dns.A: + r.A = append(r.A, v.A.String()) + case *dns.AAAA: + r.AAAA = append(r.AAAA, v.AAAA.String()) + case *dns.CNAME: + tgt := strings.TrimSuffix(canonical(v.Target), ".") + if !contains(r.CNAME, tgt) { + r.CNAME = append(r.CNAME, tgt) + } + case *dns.NS: + r.NS = append(r.NS, strings.TrimSuffix(canonical(v.Ns), ".")) + case *dns.PTR: + r.PTR = append(r.PTR, strings.TrimSuffix(canonical(v.Ptr), ".")) + case *dns.MX: + r.MX = append(r.MX, strings.TrimSuffix(canonical(v.Mx), ".")) + case *dns.TXT: + r.TXT = append(r.TXT, strings.Join(v.Txt, "")) + case *dns.SOA: + r.SOA = append(r.SOA, strings.TrimSuffix(canonical(v.Ns), ".")) + } + } + if len(r.CNAME) == 0 { + r.CNAME = nil + } + return r +} + +// mergeCNAME stitches a CNAME-followed sub-resolution back onto the original +// query name, preserving the original name/type but carrying the final answer's +// records and the full CNAME chain. +func mergeCNAME(name string, qtype uint16, chain []string, server string, sub *resolve.Result) *resolve.Result { + out := &resolve.Result{ + Name: strings.TrimSuffix(canonical(name), "."), + Type: qtype, + Rcode: sub.Rcode, + Resolver: server, + Msg: sub.Msg, + Timestamp: sub.Timestamp, + A: sub.A, + AAAA: sub.AAAA, + PTR: sub.PTR, + MX: sub.MX, + TXT: sub.TXT, + NS: sub.NS, + SOA: sub.SOA, + } + cn := append([]string{}, chain...) + cn = append(cn, sub.CNAME...) + if len(cn) > 0 { + out.CNAME = cn + } + if out.Resolver == "" { + out.Resolver = sub.Resolver + } + return out +} + +func mapHas(m map[string]struct{}, k string) bool { + _, ok := m[k] + return ok +} + +func contains(s []string, v string) bool { + for _, x := range s { + if x == v { + return true + } + } + return false +} diff --git a/pkg/iterative/poison_test.go b/pkg/iterative/poison_test.go new file mode 100644 index 00000000..7275dd92 --- /dev/null +++ b/pkg/iterative/poison_test.go @@ -0,0 +1,73 @@ +package iterative + +import ( + "context" + "testing" + + "github.com/miekg/dns" +) + +// TestReferralBailiwickRejection ensures a parent that tries to delegate (or +// glue) a name outside its own zone is ignored, so it cannot redirect the +// resolver to an attacker-controlled server (cache poisoning). +func TestReferralBailiwickRejection(t *testing.T) { + r, _ := New(Options{RootServers: []string{"127.0.0.1"}}) + + resp := new(dns.Msg) + // querying under com.; a hostile com server tries to delegate evil.org. + resp.Ns = []dns.RR{ + mustRR("evil.org. 3600 IN NS ns.evil.org."), + mustRR("example.com. 3600 IN NS ns.example.com."), // legitimate, in-bailiwick + } + resp.Extra = []dns.RR{ + mustRR("ns.evil.org. 3600 IN A 6.6.6.6"), // out-of-bailiwick glue + mustRR("ns.example.com. 3600 IN A 127.0.4.0"), // in-bailiwick glue + } + + s := &session{r: r} + child := s.parseReferral(resp, "com.", "www.example.com.") + if child == nil { + t.Fatal("expected a valid in-bailiwick referral") + } + if child.zone != "example.com." { + t.Fatalf("accepted wrong delegation zone %q (out-of-bailiwick leak?)", child.zone) + } + for _, ns := range child.ns { + if ns.name == "ns.evil.org." { + t.Fatal("accepted out-of-bailiwick nameserver") + } + for _, ip := range ns.addrs { + if ip.String() == "6.6.6.6" { + t.Fatal("accepted out-of-bailiwick glue") + } + } + } +} + +// TestReferralMustBeCloser ensures a referral to the same or a higher zone is +// rejected (prevents infinite referral loops). +func TestReferralMustBeCloser(t *testing.T) { + r, _ := New(Options{RootServers: []string{"127.0.0.1"}}) + resp := new(dns.Msg) + resp.Ns = []dns.RR{mustRR("com. 3600 IN NS a.gtld.net.")} + s := &session{r: r} + if child := s.parseReferral(resp, "com.", "www.example.com."); child != nil { + t.Fatalf("referral to same zone com. should be rejected, got %#v", child) + } +} + +func TestCNAMELoopDetected(t *testing.T) { + h := newHierarchy() + h.server("127.0.0.1").authoritative(".").delegate("com.", memNS{name: "ns.com.", ip: "127.0.2.0"}) + ex := h.server("127.0.2.0").authoritative("com.", "loop.com.") + ex.cname("a.loop.com.", "b.loop.com.") + ex.cname("b.loop.com.", "a.loop.com.") + + r, _ := New(Options{RootServers: []string{"127.0.0.1"}}) + r.newExchanger = h.factory() + + _, err := r.Resolve(context.Background(), "a.loop.com", dns.TypeA) + if err != ErrCNAMELoop { + t.Fatalf("expected ErrCNAMELoop, got %v", err) + } +} diff --git a/pkg/iterative/resolver.go b/pkg/iterative/resolver.go new file mode 100644 index 00000000..f53eec19 --- /dev/null +++ b/pkg/iterative/resolver.go @@ -0,0 +1,458 @@ +// Package iterative implements a high-throughput iterative DNS resolver that +// recurses from the root servers itself, caching delegations (NS + glue) in a +// shared, bounded cache. Unlike a stub resolver it does not depend on +// third-party recursive resolvers, which is the root cause of the +// false-negative / poisoning problems that plague large bruteforce runs: there +// is no resolver list to curate, no rate-limited public resolver returning a +// bogus NOERROR/NXDOMAIN, and every answer comes straight from the zone's +// authoritative servers. +// +// The cache is the performance lever. The first name in a zone walks +// root -> TLD -> authoritative; every subsequent name in that zone (or any +// already-seen ancestor) reuses the cached delegation and costs a single round +// trip to the authoritative server. Across a typical bruteforce workload (many +// names under few registrable domains) this collapses to ~1 query per name +// after warmup. +// +// Anti-poisoning: referrals and glue are accepted only when in-bailiwick (a +// server may only delegate names within its own zone), responses are matched +// by transaction id + question, and replies are accepted only from the address +// the query was sent to. +package iterative + +import ( + "context" + "errors" + "fmt" + "math/rand/v2" + "net" + "net/netip" + "strings" + "time" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/pkg/resolve" +) + +// Resolver errors. +var ( + ErrMaxReferrals = errors.New("iterative: too many referral hops") + ErrMaxDepth = errors.New("iterative: max recursion depth exceeded") + ErrNoNameserver = errors.New("iterative: no responsive nameserver for zone") + ErrCNAMELoop = errors.New("iterative: CNAME loop detected") +) + +// QueryInfo describes a single authoritative query for the hooks. +type QueryInfo struct { + Name string // name being looked up (no trailing dot) + Type uint16 + Zone string // delegation zone the query targets + Server string // nameserver address (ip:port) + Depth int // glue-resolution recursion depth +} + +// Hooks are optional observation points for SDK/observability use. All may be +// nil. They fire on the resolution path and must be cheap and concurrency-safe. +type Hooks struct { + // OnQuery fires before each authoritative query is sent. + OnQuery func(QueryInfo) + // OnReferral fires when a referral is accepted and cached (descending a level). + OnReferral func(parentZone, childZone string, ns []string) + // OnCacheHit fires when resolution starts from a cached delegation (not root). + OnCacheHit func(name, zone string) + // OnResponse fires for every accepted authoritative response (read-only). + OnResponse func(QueryInfo, *dns.Msg) +} + +// Options configures the iterative resolver. +type Options struct { + // QueryType is the record type to resolve (default dns.TypeA). + QueryType uint16 + // Timeout is the per-attempt query timeout (default 3s). + Timeout time.Duration + // Retries is the number of UDP retransmits to the same server before + // moving to the next nameserver (default 2). + Retries int + // Concurrency is the number of parallel workers for ResolveStream + // (default 100). Each worker reuses a single UDP socket. + Concurrency int + // MaxReferrals caps referral hops per name (loop guard, default 30). + MaxReferrals int + // MaxDepth caps nested glue-resolution recursion (default 8). + MaxDepth int + // UDPSize is the EDNS0 advertised payload size (default 1232). <512 disables EDNS0. + UDPSize int + // IPv6 enables using IPv6 glue and transport in addition to IPv4. + IPv6 bool + // DisableTCPFallback disables retrying truncated (TC) answers over TCP. + DisableTCPFallback bool + // CacheSize bounds the delegation cache (default 65536 zones). + CacheSize int + // Port is the destination port for nameservers (default 53). Overridable + // for testing against a loopback hierarchy. + Port int + // RootServers overrides the built-in root hints with "ip" or "ip:port" + // addresses (testing / split-horizon roots). + RootServers []string + // Hooks holds optional lifecycle callbacks. + Hooks Hooks +} + +func (o *Options) setDefaults() { + if o.QueryType == 0 { + o.QueryType = dns.TypeA + } + if o.Timeout <= 0 { + o.Timeout = 3 * time.Second + } + if o.Retries < 0 { + o.Retries = 0 + } else if o.Retries == 0 { + o.Retries = 2 + } + if o.Concurrency <= 0 { + o.Concurrency = 100 + } + if o.MaxReferrals <= 0 { + o.MaxReferrals = 30 + } + if o.MaxDepth <= 0 { + o.MaxDepth = 8 + } + if o.UDPSize == 0 { + o.UDPSize = 1232 + } + if o.CacheSize <= 0 { + o.CacheSize = 1 << 16 + } + if o.Port <= 0 { + o.Port = 53 + } +} + +// exchanger performs a single query/response with one nameserver. The default +// implementation uses UDP with TCP fallback; tests inject an in-memory hierarchy. +type exchanger interface { + exchange(ctx context.Context, server netip.AddrPort, msg *dns.Msg) (*dns.Msg, error) + close() +} + +// Resolver is a shared, concurrency-safe iterative resolver. Create one with +// New and reuse it across many names; the delegation cache is shared by all +// resolutions and is what makes repeated lookups cheap. +type Resolver struct { + opts Options + cache *cache + root *delegation + + // newExchanger builds a per-worker exchanger (own socket). Overridable in tests. + newExchanger func() (exchanger, error) +} + +// New creates a Resolver with the given options. +func New(opts Options) (*Resolver, error) { + opts.setDefaults() + r := &Resolver{ + opts: opts, + cache: newCache(opts.CacheSize), + } + r.root = r.buildRoot() + r.newExchanger = func() (exchanger, error) { return newUDPExchanger(&r.opts) } + return r, nil +} + +func (r *Resolver) buildRoot() *delegation { + if len(r.opts.RootServers) == 0 { + return defaultRootDelegation(r.opts.IPv6) + } + // Custom roots (testing / split-horizon). Accept "ip" or "ip:port"; an + // explicit port overrides the resolver's default port (last one wins, and + // custom root sets are normally uniform). + d := &delegation{zone: "."} + for i, s := range r.opts.RootServers { + host := s + if ap, err := netip.ParseAddrPort(s); err == nil { + host = ap.Addr().String() + r.opts.Port = int(ap.Port()) + } + ip := net.ParseIP(host) + ns := nsEntry{name: fmt.Sprintf("root-%d.", i)} + if ip != nil { + ns.addrs = append(ns.addrs, ip) + } + d.ns = append(d.ns, ns) + } + return d +} + +// addrPort converts a glue IP to a netip.AddrPort using the resolver's port. +func (r *Resolver) addrPort(ip net.IP) (netip.AddrPort, bool) { + a, ok := netip.AddrFromSlice(ip) + if !ok { + return netip.AddrPort{}, false + } + a = a.Unmap() + if a.Is6() && !r.opts.IPv6 { + return netip.AddrPort{}, false + } + return netip.AddrPortFrom(a, uint16(r.opts.Port)), true +} + +// Resolve performs a one-off iterative resolution of name/qtype. For bulk work +// prefer ResolveStream, which reuses sockets and shares cache warmth across a +// pool of workers. +func (r *Resolver) Resolve(ctx context.Context, name string, qtype uint16) (*resolve.Result, error) { + if qtype == 0 { + qtype = r.opts.QueryType + } + ex, err := r.newExchanger() + if err != nil { + return nil, err + } + defer ex.close() + s := &session{r: r, ex: ex} + return s.resolve(ctx, name, qtype) +} + +// session is a single worker's resolution context: it owns one exchanger +// (socket) reused across the sequential queries a resolution requires. +type session struct { + r *Resolver + ex exchanger +} + +func (s *session) resolve(ctx context.Context, name string, qtype uint16) (*resolve.Result, error) { + return s.resolveDepth(ctx, name, qtype, 0, map[string]struct{}{}) +} + +func (s *session) resolveDepth(ctx context.Context, name string, qtype uint16, depth int, cnameSeen map[string]struct{}) (*resolve.Result, error) { + if depth > s.r.opts.MaxDepth { + return nil, ErrMaxDepth + } + sname := canonical(name) + + del := s.r.cache.best(sname) + if del == nil { + del = s.r.root + } else if s.r.opts.Hooks.OnCacheHit != nil { + s.r.opts.Hooks.OnCacheHit(strings.TrimSuffix(sname, "."), del.zone) + } + zone := del.zone + + var cnameChain []string + + for hops := 0; hops < s.r.opts.MaxReferrals; hops++ { + if err := ctx.Err(); err != nil { + return nil, err + } + + resp, server, err := s.queryZone(ctx, del, zone, sname, qtype, depth) + if err != nil { + return nil, err + } + if s.r.opts.Hooks.OnResponse != nil { + s.r.opts.Hooks.OnResponse(QueryInfo{Name: strings.TrimSuffix(sname, "."), Type: qtype, Zone: zone, Server: server, Depth: depth}, resp) + } + + // definitive negative + if resp.Rcode == dns.RcodeNameError { + return buildResult(name, qtype, resp, server, cnameChain), nil + } + + // follow any CNAME chain present in this answer + final, cnames := chaseCNAME(resp.Answer, sname) + if len(cnames) > 0 { + cnameChain = append(cnameChain, cnames...) + } + + // direct answer of the requested type for the (possibly chased) name? + if answersType(resp.Answer, final, qtype) { + return buildResult(name, qtype, resp, server, cnameChain), nil + } + + // CNAME points outside what this answer resolves: restart for the target. + if final != sname && qtype != dns.TypeCNAME && qtype != dns.TypeANY { + if _, seen := cnameSeen[final]; seen { + return nil, ErrCNAMELoop + } + cnameSeen[final] = struct{}{} + sub, err := s.resolveDepth(ctx, final, qtype, depth+1, cnameSeen) + if err != nil { + return nil, err + } + return mergeCNAME(name, qtype, cnameChain, server, sub), nil + } + + // referral to a closer zone? + if child := s.parseReferral(resp, zone, sname); child != nil { + s.r.cache.put(child) + if s.r.opts.Hooks.OnReferral != nil { + s.r.opts.Hooks.OnReferral(zone, child.zone, nsNames(child)) + } + del = child + zone = child.zone + continue + } + + // NOERROR with no answer and no usable referral == NODATA (name exists, + // type doesn't) or an empty/lame response we can't progress past. + return buildResult(name, qtype, resp, server, cnameChain), nil + } + return nil, ErrMaxReferrals +} + +// queryZone tries the nameservers of a delegation until one returns a usable +// response. Glueless nameservers have their addresses resolved on demand. +func (s *session) queryZone(ctx context.Context, del *delegation, zone, sname string, qtype uint16, depth int) (*dns.Msg, string, error) { + order := rand.Perm(len(del.ns)) + + // first pass: nameservers that already have glue (no extra round trips) + for _, gluedOnly := range []bool{true, false} { + for _, idx := range order { + ns := del.ns[idx] + addrs := ns.addrs + if len(addrs) == 0 { + if gluedOnly { + continue + } + // glueless: resolve the nameserver's address from the root. + addrs = s.resolveNSAddrs(ctx, ns.name, depth) + if len(addrs) == 0 { + continue + } + } else if !gluedOnly { + // already tried in the glued pass + continue + } + for _, ip := range addrs { + ap, ok := s.r.addrPort(ip) + if !ok { + continue + } + if s.r.opts.Hooks.OnQuery != nil { + s.r.opts.Hooks.OnQuery(QueryInfo{Name: strings.TrimSuffix(sname, "."), Type: qtype, Zone: zone, Server: ap.String(), Depth: depth}) + } + resp, err := s.ex.exchange(ctx, ap, s.newQuery(sname, qtype)) + if err != nil { + continue + } + return resp, ap.String(), nil + } + } + } + return nil, "", ErrNoNameserver +} + +// resolveNSAddrs resolves the A (and AAAA when enabled) addresses of a glueless +// nameserver, bounded by depth to prevent runaway recursion. +func (s *session) resolveNSAddrs(ctx context.Context, nsName string, depth int) []net.IP { + if depth+1 > s.r.opts.MaxDepth { + return nil + } + var out []net.IP + if res, err := s.resolveDepth(ctx, nsName, dns.TypeA, depth+1, map[string]struct{}{}); err == nil { + for _, a := range res.A { + if ip := net.ParseIP(a); ip != nil { + out = append(out, ip) + } + } + } + if s.r.opts.IPv6 { + if res, err := s.resolveDepth(ctx, nsName, dns.TypeAAAA, depth+1, map[string]struct{}{}); err == nil { + for _, a := range res.AAAA { + if ip := net.ParseIP(a); ip != nil { + out = append(out, ip) + } + } + } + } + return out +} + +// newQuery builds an iterative (RD=0) query message with EDNS0. +func (s *session) newQuery(sname string, qtype uint16) *dns.Msg { + m := new(dns.Msg) + m.Id = dns.Id() + m.RecursionDesired = false + m.Question = []dns.Question{{Name: sname, Qtype: qtype, Qclass: dns.ClassINET}} + if s.r.opts.UDPSize >= 512 { + m.SetEdns0(uint16(s.r.opts.UDPSize), false) + } + return m +} + +// parseReferral extracts a closer, in-bailiwick delegation from a response's +// authority (NS) and additional (glue) sections. Returns nil when the response +// is not a usable referral (no NS, out-of-bailiwick, or not closer than zone). +func (s *session) parseReferral(resp *dns.Msg, parentZone, sname string) *delegation { + var child string + nsByOwner := map[string][]string{} + var minTTL uint32 = 0xffffffff + for _, rr := range resp.Ns { + ns, ok := rr.(*dns.NS) + if !ok { + continue + } + owner := canonical(ns.Header().Name) + // the referral zone must be within the parent and a strict descendant + // (progress), and an ancestor of the queried name. + if !inBailiwick(owner, parentZone) || owner == canonical(parentZone) { + continue + } + if !inBailiwick(sname, owner) { + continue + } + child = owner + nsByOwner[owner] = append(nsByOwner[owner], canonical(ns.Ns)) + if ns.Header().Ttl < minTTL { + minTTL = ns.Header().Ttl + } + } + if child == "" { + return nil + } + + // collect in-bailiwick glue for the chosen child's nameservers + glue := map[string][]net.IP{} + for _, rr := range resp.Extra { + var name string + var ip net.IP + switch a := rr.(type) { + case *dns.A: + name, ip = canonical(a.Header().Name), a.A + case *dns.AAAA: + if !s.r.opts.IPv6 { + continue + } + name, ip = canonical(a.Header().Name), a.AAAA + default: + continue + } + // accept glue only within the parent zone's bailiwick (anti-poisoning). + if !inBailiwick(name, parentZone) { + continue + } + glue[name] = append(glue[name], ip) + } + + d := &delegation{zone: child} + if minTTL == 0xffffffff || minTTL < 1 { + minTTL = 60 + } + d.expiry = time.Now().Add(time.Duration(minTTL) * time.Second) + for _, nsname := range nsByOwner[child] { + d.ns = append(d.ns, nsEntry{name: nsname, addrs: glue[nsname]}) + } + if len(d.ns) == 0 { + return nil + } + return d +} + +func nsNames(d *delegation) []string { + out := make([]string, 0, len(d.ns)) + for _, ns := range d.ns { + out = append(out, ns.name) + } + return out +} diff --git a/pkg/iterative/resolver_test.go b/pkg/iterative/resolver_test.go new file mode 100644 index 00000000..c1f6789d --- /dev/null +++ b/pkg/iterative/resolver_test.go @@ -0,0 +1,194 @@ +package iterative + +import ( + "context" + "fmt" + "sync" + "testing" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/pkg/resolve" +) + +// buildTestHierarchy wires a small but representative tree: +// +// root(127.0.0.1) -> com.(127.0.2.0), net.(127.0.3.0) +// com. -> example.com.(127.0.4.0, glued) +// -> glueless.com. (NS ns.ext.net., NO glue) +// net. -> ext.net.(127.0.5.0, glued) +// example.com. : www A, many hostN A, alias CNAME->www, (NODATA/NXDOMAIN otherwise) +// ext.net. : ns.ext.net. A (so glueless delegation resolves), svc A +func buildTestHierarchy() *memHierarchy { + h := newHierarchy() + const ( + rootIP = "127.0.0.1" + comIP = "127.0.2.0" + netIP = "127.0.3.0" + exIP = "127.0.4.0" + extIP = "127.0.5.0" + ) + h.server(rootIP).authoritative("."). + delegate("com.", memNS{name: "a.gtld.net.", ip: comIP}). + delegate("net.", memNS{name: "a.gtld-net.net.", ip: netIP}) + + h.server(comIP).authoritative("com."). + delegate("example.com.", memNS{name: "ns.example.com.", ip: exIP}). + delegate("glueless.com.", memNS{name: "ns.ext.net."}) // glueless: no IP + + h.server(netIP).authoritative("net."). + delegate("ext.net.", memNS{name: "ns.ext.net.", ip: extIP}) + + ex := h.server(exIP).authoritative("example.com.") + ex.a("www.example.com.", "93.184.216.34") + ex.cname("alias.example.com.", "www.example.com.") + for i := 0; i < 50; i++ { + ex.a(fmt.Sprintf("host%d.example.com.", i), fmt.Sprintf("198.51.100.%d", i+1)) + } + + extn := h.server(extIP).authoritative("ext.net.", "glueless.com.") + extn.a("ns.ext.net.", extIP) + extn.a("svc.glueless.com.", "203.0.113.7") + return h +} + +func newTestResolver(t *testing.T, h *memHierarchy, hooks Hooks) *Resolver { + t.Helper() + r, err := New(Options{ + RootServers: []string{"127.0.0.1"}, + QueryType: dns.TypeA, + Concurrency: 8, + Hooks: hooks, + }) + if err != nil { + t.Fatal(err) + } + r.newExchanger = h.factory() + return r +} + +func TestIterativeBasicResolution(t *testing.T) { + h := buildTestHierarchy() + r := newTestResolver(t, h, Hooks{}) + + res, err := r.Resolve(context.Background(), "www.example.com", dns.TypeA) + if err != nil { + t.Fatal(err) + } + if len(res.A) != 1 || res.A[0] != "93.184.216.34" { + t.Fatalf("unexpected A records: %#v", res.A) + } + if res.Rcode != dns.RcodeSuccess { + t.Fatalf("rcode = %d, want NOERROR", res.Rcode) + } +} + +func TestIterativeCNAME(t *testing.T) { + h := buildTestHierarchy() + r := newTestResolver(t, h, Hooks{}) + + res, err := r.Resolve(context.Background(), "alias.example.com", dns.TypeA) + if err != nil { + t.Fatal(err) + } + if len(res.A) != 1 || res.A[0] != "93.184.216.34" { + t.Fatalf("CNAME target A not resolved: %#v", res.A) + } + if len(res.CNAME) == 0 || res.CNAME[0] != "www.example.com" { + t.Fatalf("expected CNAME chain to www.example.com, got %#v", res.CNAME) + } +} + +func TestIterativeNXDOMAIN(t *testing.T) { + h := buildTestHierarchy() + r := newTestResolver(t, h, Hooks{}) + + res, err := r.Resolve(context.Background(), "nope.example.com", dns.TypeA) + if err != nil { + t.Fatal(err) + } + if res.Rcode != dns.RcodeNameError { + t.Fatalf("rcode = %d, want NXDOMAIN", res.Rcode) + } + if len(res.A) != 0 { + t.Fatalf("NXDOMAIN should have no A records: %#v", res.A) + } +} + +func TestIterativeGluelessDelegation(t *testing.T) { + h := buildTestHierarchy() + r := newTestResolver(t, h, Hooks{}) + + // svc.glueless.com is served by ns.ext.net., which has NO glue at com.; + // the resolver must resolve the nameserver's address first. + res, err := r.Resolve(context.Background(), "svc.glueless.com", dns.TypeA) + if err != nil { + t.Fatalf("glueless resolution failed: %v", err) + } + if len(res.A) != 1 || res.A[0] != "203.0.113.7" { + t.Fatalf("unexpected A for glueless: %#v", res.A) + } +} + +func TestIterativeCacheReuseReducesQueries(t *testing.T) { + h := buildTestHierarchy() + r := newTestResolver(t, h, Hooks{}) + ctx := context.Background() + + // cold: root -> com -> example.com -> answer + if _, err := r.Resolve(ctx, "host0.example.com", dns.TypeA); err != nil { + t.Fatal(err) + } + cold := h.queries.Load() + + // warm: example.com delegation cached, expect a single authoritative query + before := h.queries.Load() + if _, err := r.Resolve(ctx, "host1.example.com", dns.TypeA); err != nil { + t.Fatal(err) + } + warm := h.queries.Load() - before + + if cold < 3 { + t.Fatalf("expected cold path to take >=3 queries (root,tld,auth), got %d", cold) + } + if warm != 1 { + t.Fatalf("expected warm path to take exactly 1 query (cached delegation), got %d", warm) + } +} + +func TestIterativeStreamConcurrent(t *testing.T) { + h := buildTestHierarchy() + r := newTestResolver(t, h, Hooks{}) + + names := make(chan string) + go func() { + defer close(names) + for i := 0; i < 50; i++ { + names <- fmt.Sprintf("host%d.example.com", i) + } + }() + + var mu sync.Mutex + got := map[string]string{} + err := r.ResolveStream(context.Background(), names, StreamConfig{ + OnResult: func(res *resolve.Result) { + mu.Lock() + if len(res.A) > 0 { + got[res.Name] = res.A[0] + } + mu.Unlock() + }, + }) + if err != nil { + t.Fatal(err) + } + if len(got) != 50 { + t.Fatalf("expected 50 resolved names, got %d", len(got)) + } + for i := 0; i < 50; i++ { + name := fmt.Sprintf("host%d.example.com", i) + want := fmt.Sprintf("198.51.100.%d", i+1) + if got[name] != want { + t.Errorf("%s = %q, want %q", name, got[name], want) + } + } +} diff --git a/pkg/iterative/roots.go b/pkg/iterative/roots.go new file mode 100644 index 00000000..c2406243 --- /dev/null +++ b/pkg/iterative/roots.go @@ -0,0 +1,49 @@ +package iterative + +import "net" + +// rootServer is a single root nameserver hint (name + addresses). +type rootServer struct { + name string + v4 string + v6 string +} + +// rootHints is the static list of the 13 DNS root servers (IANA root hints). +// They bootstrap iterative resolution: the resolver starts here when its +// delegation cache has no closer ancestor for a name. Addresses change very +// rarely; if one is stale the resolver simply rotates to another root. +var rootHints = []rootServer{ + {"a.root-servers.net.", "198.41.0.4", "2001:503:ba3e::2:30"}, + {"b.root-servers.net.", "199.9.14.201", "2001:500:200::b"}, + {"c.root-servers.net.", "192.33.4.12", "2001:500:2::c"}, + {"d.root-servers.net.", "199.7.91.13", "2001:500:2d::d"}, + {"e.root-servers.net.", "192.203.230.10", "2001:500:a8::e"}, + {"f.root-servers.net.", "192.5.5.241", "2001:500:2f::f"}, + {"g.root-servers.net.", "192.112.36.4", "2001:500:12::d0d"}, + {"h.root-servers.net.", "198.97.190.53", "2001:500:1::53"}, + {"i.root-servers.net.", "192.36.148.17", "2001:7fe::53"}, + {"j.root-servers.net.", "192.58.128.30", "2001:503:c27::2:30"}, + {"k.root-servers.net.", "193.0.14.129", "2001:7fd::1"}, + {"l.root-servers.net.", "199.7.83.42", "2001:500:9f::42"}, + {"m.root-servers.net.", "202.12.27.33", "2001:dc3::35"}, +} + +// defaultRootDelegation builds the "." delegation from the static hints. wantV6 +// controls whether IPv6 glue is included. +func defaultRootDelegation(wantV6 bool) *delegation { + d := &delegation{zone: "."} + for _, r := range rootHints { + ns := nsEntry{name: r.name} + if ip := net.ParseIP(r.v4); ip != nil { + ns.addrs = append(ns.addrs, ip) + } + if wantV6 { + if ip := net.ParseIP(r.v6); ip != nil { + ns.addrs = append(ns.addrs, ip) + } + } + d.ns = append(d.ns, ns) + } + return d +} diff --git a/pkg/iterative/stream.go b/pkg/iterative/stream.go new file mode 100644 index 00000000..7b94dba3 --- /dev/null +++ b/pkg/iterative/stream.go @@ -0,0 +1,95 @@ +package iterative + +import ( + "context" + "strings" + "sync" + + "github.com/projectdiscovery/shuffledns/pkg/resolve" +) + +// StreamConfig wires a bulk iterative resolution. +type StreamConfig struct { + // OnResult is called for every name that produced a final response + // (success, NODATA, or NXDOMAIN). Must be concurrency-safe. + OnResult func(*resolve.Result) + // OnError is called when a name could not be resolved (no responsive + // nameserver, loop, context cancelled). Must be concurrency-safe. + OnError func(name string, err error) + // QueryType overrides the resolver's default query type for this stream. + QueryType uint16 +} + +// ResolveStream consumes names from the channel and resolves them iteratively +// using a pool of Concurrency workers. Every worker reuses a single UDP socket +// and shares the resolver's delegation cache, so cache warmth (root, TLDs, +// popular zones) is amortized across the whole workload. It blocks until the +// input channel is closed and all in-flight work drains, or ctx is cancelled. +func (r *Resolver) ResolveStream(ctx context.Context, names <-chan string, cfg StreamConfig) error { + qtype := cfg.QueryType + if qtype == 0 { + qtype = r.opts.QueryType + } + + var wg sync.WaitGroup + wg.Add(r.opts.Concurrency) + for i := 0; i < r.opts.Concurrency; i++ { + go func() { + defer wg.Done() + ex, err := r.newExchanger() + if err != nil { + // a worker that cannot open a socket simply drains its share; + // other workers continue. Report once per failed name. Stop on + // cancellation so an idle producer can't wedge the drain. + for { + select { + case <-ctx.Done(): + return + case name, ok := <-names: + if !ok { + return + } + if cfg.OnError != nil { + cfg.OnError(strings.TrimSpace(name), err) + } + } + } + } + defer ex.close() + s := &session{r: r, ex: ex} + for { + // Select on ctx.Done() as well as the channel: a plain range + // blocks on an idle-but-open producer, so a cancelled ctx would + // never unblock the worker and ResolveStream would hang. + var ( + name string + ok bool + ) + select { + case <-ctx.Done(): + return + case name, ok = <-names: + if !ok { + return + } + } + name = strings.TrimSpace(name) + if name == "" { + continue + } + res, rerr := s.resolve(ctx, name, qtype) + if rerr != nil { + if cfg.OnError != nil { + cfg.OnError(name, rerr) + } + continue + } + if cfg.OnResult != nil { + cfg.OnResult(res) + } + } + }() + } + wg.Wait() + return ctx.Err() +} diff --git a/pkg/iterative/udp.go b/pkg/iterative/udp.go new file mode 100644 index 00000000..a87fc2e8 --- /dev/null +++ b/pkg/iterative/udp.go @@ -0,0 +1,127 @@ +package iterative + +import ( + "context" + "errors" + "net" + "net/netip" + "time" + + "github.com/miekg/dns" +) + +// udpExchanger performs authoritative queries over a single reused, unconnected +// UDP socket (the client-side equivalent of zdns's socket reuse: one socket per +// worker, reused for every destination, avoiding per-query socket setup). It +// verifies the reply's source address and transaction id, and falls back to TCP +// on truncation. +type udpExchanger struct { + opts *Options + conn *net.UDPConn + buf []byte +} + +func newUDPExchanger(opts *Options) (exchanger, error) { + network := "udp4" + if opts.IPv6 { + network = "udp" + } + conn, err := net.ListenUDP(network, nil) + if err != nil { + return nil, err + } + _ = conn.SetReadBuffer(4 * 1024 * 1024) + _ = conn.SetWriteBuffer(4 * 1024 * 1024) + return &udpExchanger{opts: opts, conn: conn, buf: make([]byte, 64*1024)}, nil +} + +func (u *udpExchanger) close() { + if u.conn != nil { + _ = u.conn.Close() + } +} + +func (u *udpExchanger) exchange(ctx context.Context, server netip.AddrPort, msg *dns.Msg) (*dns.Msg, error) { + packed, err := msg.Pack() + if err != nil { + return nil, err + } + dst := net.UDPAddrFromAddrPort(server) + + for attempt := 0; attempt <= u.opts.Retries; attempt++ { + deadline := time.Now().Add(u.opts.Timeout) + if d, ok := ctx.Deadline(); ok && d.Before(deadline) { + deadline = d + } + if _, err := u.conn.WriteToUDP(packed, dst); err != nil { + return nil, err + } + _ = u.conn.SetReadDeadline(deadline) + + // read until a matching reply arrives or the deadline expires; stray + // or spoofed packets (wrong source / id / question) are ignored. + for { + n, from, rerr := u.conn.ReadFromUDP(u.buf) + if rerr != nil { + if ne, ok := rerr.(net.Error); ok && ne.Timeout() { + break // retransmit + } + return nil, rerr + } + // On a dual-stack ("udp") socket an IPv4 server's reply arrives as a + // 4-in-6 address, so compare against the unmapped form; otherwise every + // IPv4 response is dropped when IPv6 mode is enabled. + fromAP := from.AddrPort() + fromNorm := netip.AddrPortFrom(fromAP.Addr().Unmap(), fromAP.Port()) + if !fromNorm.Addr().IsValid() || fromNorm != server { + continue // source-address verification (anti off-path spoofing) + } + resp := new(dns.Msg) + if resp.Unpack(u.buf[:n]) != nil { + continue + } + if resp.Id != msg.Id || !sameQuestion(resp, msg) { + continue + } + if resp.Truncated && !u.opts.DisableTCPFallback { + if tcp, terr := u.exchangeTCP(ctx, server, msg); terr == nil { + return tcp, nil + } + } + return resp, nil + } + if err := ctx.Err(); err != nil { + return nil, err + } + } + return nil, errTimeout +} + +func (u *udpExchanger) exchangeTCP(ctx context.Context, server netip.AddrPort, msg *dns.Msg) (*dns.Msg, error) { + d := net.Dialer{Timeout: u.opts.Timeout} + conn, err := d.DialContext(ctx, "tcp", server.String()) + if err != nil { + return nil, err + } + defer func() { _ = conn.Close() }() + co := &dns.Conn{Conn: conn} + if dl, ok := ctx.Deadline(); ok { + _ = conn.SetDeadline(dl) + } else { + _ = conn.SetDeadline(time.Now().Add(u.opts.Timeout)) + } + if err := co.WriteMsg(msg); err != nil { + return nil, err + } + return co.ReadMsg() +} + +var errTimeout = errors.New("iterative: query timed out") + +func sameQuestion(a, b *dns.Msg) bool { + if len(a.Question) != 1 || len(b.Question) != 1 { + return false + } + qa, qb := a.Question[0], b.Question[0] + return qa.Qtype == qb.Qtype && qa.Qclass == qb.Qclass && canonical(qa.Name) == canonical(qb.Name) +} diff --git a/pkg/massdns/massdns.go b/pkg/massdns/massdns.go index 016f8fe7..701d78ad 100644 --- a/pkg/massdns/massdns.go +++ b/pkg/massdns/massdns.go @@ -20,39 +20,51 @@ type Options struct { Domains []string // Retries is the number of retries for dns Retries int - // MassdnsPath is the path to the binary - MassdnsPath string - // Threads is the hashmap size for massdns + // Threads is the number of concurrent in-flight dns queries Threads int - // InputFile is the file to use for massdns input - InputFile string // ResolversFile is the file with the resolvers ResolversFile string // TrustedResolvers is the file with the trusted resolvers TrustedResolvers string - // TempDir is a temporary directory for storing massdns misc files + // TempDir is a temporary directory for storing misc files TempDir string - // OutputFile is the file to use for massdns output + // OutputFile is the file to write results to OutputFile string // Json is format ouput to ndjson format Json bool // WildcardsThreads is the number of wildcards concurrent threads WildcardsThreads int - // MassdnsRaw perform wildcards filtering from an existing massdns output file + // MassdnsRaw performs wildcards filtering from an existing massdns output file MassdnsRaw string // StrictWildcard controls whether the wildcard check should be performed on each result StrictWildcard bool // WildcardOutputFile is the file where the list of wildcards is dumped WildcardOutputFile string - // MassDnsCmd supports massdns flags - MassDnsCmd string - // KeepStderr controls whether to capture and store massdns stderr output - KeepStderr bool - // BatchSize controls the number of lines per chunk for incremental processing - BatchSize int // FilterInternalIPs controls whether to filter out internal/private IP addresses FilterInternalIPs bool + // Native resolver tuning (forwarded to pkg/resolve). + QueryType string // DNS record type to resolve (A, AAAA, ...). Default A. + BatchMode string // sendmmsg/recvmmsg batching: off | on | adaptive + SocketCount int // UDP sockets per run (0 = scale to cores) + UDPSize int // EDNS0 advertised UDP payload size (0 = default; <512 disables) + QPS int // outbound query rate limit (0 = unlimited) + NoRecurse bool // send non-recursive queries (RD=0) + Sticky bool // do not rotate resolver on retry + ResolverHealth bool // per-resolver health scoring / de-weighting + AdaptiveConcurrency bool // shrink/grow in-flight cap based on packet loss + CrossCheck bool // re-verify positive answers on a second resolver + ExtendedInput bool // parse "name [resolver ...]" input lines + NoVerifyIP bool // disable reply source-IP verification + NoTCPFallback bool // disable TCP fallback on truncated answers + // Iterative resolves from the root servers directly (no recursive resolver + // list needed), caching delegations. Removes the public-resolver dependency. + Iterative bool + + // Distributed resolution and resume. + Shard string // "m/n": process only shard m of n + ResumeFile string // checkpoint file for crash-safe stop/resume + OnResult func(*retryabledns.DNSData) } diff --git a/pkg/massdns/process.go b/pkg/massdns/process.go index d68119d6..875dd930 100644 --- a/pkg/massdns/process.go +++ b/pkg/massdns/process.go @@ -8,15 +8,18 @@ import ( "fmt" "io" "os" - "os/exec" - "strconv" "strings" "sync/atomic" "time" + "github.com/miekg/dns" "github.com/projectdiscovery/dnsx/libs/dnsx" "github.com/projectdiscovery/gologger" + "github.com/projectdiscovery/shuffledns/pkg/checkpoint" + "github.com/projectdiscovery/shuffledns/pkg/iterative" "github.com/projectdiscovery/shuffledns/pkg/parser" + "github.com/projectdiscovery/shuffledns/pkg/resolve" + "github.com/projectdiscovery/shuffledns/pkg/shard" "github.com/projectdiscovery/shuffledns/pkg/store" "github.com/projectdiscovery/shuffledns/pkg/wildcards" "github.com/projectdiscovery/utils/batcher" @@ -27,68 +30,29 @@ import ( "github.com/weppos/publicsuffix-go/publicsuffix" ) -// runs massdns binary with the specified options -func (instance *Instance) RunWithContext(ctx context.Context) (stdout, stderr string, took time.Duration, err error) { - start := time.Now() - - // Create temporary file for massdns output - stdoutFile, err := os.CreateTemp(instance.options.TempDir, "massdns-stdout-") - if err != nil { - return "", "", 0, fmt.Errorf("could not create temp file for massdns output: %w", err) - } - defer func() { - _ = stdoutFile.Close() - }() - - // Handle stderr based on KeepStderr option - var stderrFile *os.File - if instance.options.KeepStderr { - stderrFile, err = os.CreateTemp(instance.options.TempDir, "massdns-stderr-") - if err != nil { - return "", "", 0, fmt.Errorf("could not create temp file for massdns stderr: %w", err) - } - defer func() { - _ = stderrFile.Close() - }() - } - - // Run the command on a temp file and wait for the output - args := []string{"-r", instance.options.ResolversFile, "-o", "Snl", "--retry", "REFUSED", "--retry", "SERVFAIL", "-t", "A", instance.options.InputFile, "-s", strconv.Itoa(instance.options.Threads)} - if instance.options.MassDnsCmd != "" { - args = append(args, strings.Fields(instance.options.MassDnsCmd)...) - } - - cmd := exec.CommandContext(ctx, instance.options.MassdnsPath, args...) - cmd.Stdout = stdoutFile - - // Set stderr based on KeepStderr option - if instance.options.KeepStderr { - cmd.Stderr = stderrFile - } else { - // Discard stderr by sending it to /dev/null - cmd.Stderr = nil - } - - err = cmd.Run() - - // Return stderr filename only if it was captured - stderrFilename := "" - if instance.options.KeepStderr { - stderrFilename = stderrFile.Name() +// parseBatchMode maps a CLI string to the resolver batching mode. +func parseBatchMode(s string) resolve.BatchMode { + switch strings.ToLower(strings.TrimSpace(s)) { + case "on", "enabled", "true": + return resolve.BatchEnabled + case "adaptive", "auto": + return resolve.BatchAdaptive + default: + return resolve.BatchDisabled } - - return stdoutFile.Name(), stderrFilename, time.Since(start), err } +// Run processes an existing raw massdns output file for wildcard filtering. +// This is the backward-compatible path for the -ri/--raw-input flag; live +// resolution no longer shells out to massdns and is handled by the streaming +// methods below. func (instance *Instance) Run(ctx context.Context) error { - // Process a created list or the massdns input - inputFile := instance.options.InputFile - if instance.options.MassdnsRaw != "" { - inputFile = instance.options.MassdnsRaw + if instance.options.MassdnsRaw == "" { + return errors.New("streaming processing should be used for new resolution runs") } - // Check for blank input file or non-existent input file - blank, err := fileutil.IsEmpty(inputFile) + // Check for blank or non-existent input file + blank, err := fileutil.IsEmpty(instance.options.MassdnsRaw) if err != nil { return err } @@ -96,113 +60,257 @@ func (instance *Instance) Run(ctx context.Context) error { return errors.New("blank input file specified") } - // Check if we need to run massdns - if instance.options.MassdnsRaw == "" { - // This case is now handled by the streaming methods in the runner - // The Run method is only called for raw massdns output processing - return errors.New("streaming processing should be used for new massdns runs") - } - - // Create a store for storing ip metadata shstore, err := store.New(instance.options.TempDir) if err != nil { return fmt.Errorf("could not create store: %w", err) } defer shstore.Close() - // parse the input file gologger.Info().Msgf("Started parsing massdns input\n") now := time.Now() - err = instance.parseMassDNSOutputFile(instance.options.MassdnsRaw, shstore) - if err != nil { + if err := instance.parseMassDNSOutputFile(instance.options.MassdnsRaw, shstore); err != nil { return fmt.Errorf("could not parse massdns input: %w", err) } gologger.Info().Msgf("Massdns input parsing completed in %s\n", time.Since(now)) - if instance.options.AutoExtractRootDomains { - gologger.Info().Msgf("Started extracting root domains\n") - now := time.Now() - err = instance.autoExtractRootDomains(shstore) + return instance.postProcess(shstore) +} + +// resolveNames resolves a stream of hostnames produced by the produce function +// using the native asynchronous resolver, storing successful answers in the +// store. No temporary files are created and no external binary is invoked. +// +// When configured, the input is partitioned by shard (distributed coordination) +// and filtered against a resume checkpoint (crash-safe stop/resume) before being +// dispatched; both filters run in the producer goroutine, off the resolver's hot +// path. +func (instance *Instance) resolveNames(ctx context.Context, shstore *store.Store, produce func(out chan<- string) error) error { + // Iterative mode recurses from the root servers and needs no resolver list; + // the stub path requires one. + var resolvers []string + var err error + if !instance.options.Iterative { + resolvers, err = wildcards.LoadResolversFromFile(instance.options.ResolversFile) if err != nil { - return fmt.Errorf("could not extract root domains: %w", err) + return fmt.Errorf("could not load resolvers: %w", err) } - gologger.Info().Msgf("Root domain extraction completed in %s\n", time.Since(now)) } - // Perform wildcard filtering only if domain name has been specified - if len(instance.options.Domains) > 0 { - gologger.Info().Msgf("Started removing wildcards records\n") - now := time.Now() - err = instance.filterWildcards(shstore) + shardCfg, err := shard.Parse(instance.options.Shard) + if err != nil { + return err + } + + var ckpt *checkpoint.Checkpoint + if instance.options.ResumeFile != "" { + ckpt, err = checkpoint.Open(instance.options.ResumeFile) if err != nil { - return fmt.Errorf("could not filter wildcards: %w", err) + return fmt.Errorf("could not open resume checkpoint: %w", err) + } + defer func() { _ = ckpt.Close() }() + if n := ckpt.Resumed(); n > 0 { + gologger.Info().Msgf("Resuming: skipping %d already-completed names\n", n) } - gologger.Info().Msgf("Wildcard removal completed in %s\n", time.Since(now)) } - gologger.Info().Msgf("Finished enumeration, started writing output\n") + var resolved atomic.Int64 - // Write the final elaborated list out - now = time.Now() - err = instance.writeOutput(shstore) - if err != nil { - return fmt.Errorf("could not write output: %w", err) + // The shuffledns runner keys its store by IP address, so only address + // record types are meaningful here; other types (NS, MX, TXT, ...) would + // resolve successfully but be silently dropped at store time. Reject them + // with a clear pointer to cmd/resolve, which renders arbitrary types. + qtype := dns.TypeA + if instance.options.QueryType != "" { + t, ok := dns.StringToType[strings.ToUpper(instance.options.QueryType)] + if !ok { + return fmt.Errorf("unknown query type %q", instance.options.QueryType) + } + if t != dns.TypeA && t != dns.TypeAAAA { + return fmt.Errorf("query type %q is not supported by the bruteforce runner (only A/AAAA); use the resolve command for arbitrary record types", instance.options.QueryType) + } + qtype = t } - gologger.Info().Msgf("Output written in %s\n", time.Since(now)) - return nil -} -// runChunk runs massdns on a specific chunk file -func (instance *Instance) runChunk(ctx context.Context, chunkFile string) (stdout, stderr string, took time.Duration, err error) { - start := time.Now() + markDone := func(name string) { + if ckpt != nil { + _ = ckpt.Done(name) + } + } - // Create temporary file for massdns output - stdoutFile, err := os.CreateTemp(instance.options.TempDir, "massdns-chunk-stdout-") - if err != nil { - return "", "", 0, fmt.Errorf("could not create temp file for massdns output: %w", err) + // storeResult is the engine-agnostic result handler: it marks the name done + // (for resume), then stores any address answers keyed by IP. Shared by both + // the stub resolver and the iterative-from-root resolver. + storeResult := func(name string, rcode int, a, aaaa []string) { + markDone(name) + if rcode != dns.RcodeSuccess { + return + } + ips := a + if len(aaaa) > 0 { + ips = append(append([]string{}, a...), aaaa...) + } + if len(ips) == 0 { + return + } + for _, ip := range ips { + if instance.shouldFilterIP(ip) { + continue + } + _ = shstore.Append(ip, name) + } + resolved.Add(1) } - defer func() { - _ = stdoutFile.Close() - }() - // Handle stderr based on KeepStderr option - var stderrFile *os.File - if instance.options.KeepStderr { - stderrFile, err = os.CreateTemp(instance.options.TempDir, "massdns-chunk-stderr-") - if err != nil { - return "", "", 0, fmt.Errorf("could not create temp file for massdns stderr: %w", err) + // admit applies shard ownership and resume skipping to the name. For + // extended-input lines ("name resolver ..."), only the leading name is keyed. + admit := func(line string) bool { + name := line + if instance.options.ExtendedInput { + if i := strings.IndexAny(name, " \t"); i >= 0 { + name = name[:i] + } } - defer func() { - _ = stderrFile.Close() + if !shardCfg.Owns(name) { + return false + } + if ckpt != nil && ckpt.Has(name) { + return false + } + return true + } + + // produceErr is written by the producer goroutine before it closes its + // output channel; the channel-close -> Run-drains -> Run-returns chain + // establishes happens-before, so the read below (after Run) is safe. + var produceErr error + + input := make(chan string, 4096) + if !shardCfg.Enabled() && ckpt == nil { + // fast path: no filtering, producer writes straight to the resolver. + go func() { + defer close(input) + produceErr = produce(input) + }() + } else { + raw := make(chan string, 4096) + go func() { + defer close(raw) + produceErr = produce(raw) + }() + go func() { + defer close(input) + for name := range raw { + if admit(name) { + input <- name + } + } }() } - // Run the command on the chunk file - args := []string{"-r", instance.options.ResolversFile, "-o", "Snl", "--retry", "REFUSED", "--retry", "SERVFAIL", "-t", "A", chunkFile, "-s", strconv.Itoa(instance.options.Threads)} - if instance.options.MassDnsCmd != "" { - args = append(args, strings.Fields(instance.options.MassDnsCmd)...) + if err := instance.runEngine(ctx, resolvers, qtype, input, storeResult, markDone); err != nil { + return err + } + if produceErr != nil { + return fmt.Errorf("could not read input: %w", produceErr) } - cmd := exec.CommandContext(ctx, instance.options.MassdnsPath, args...) - cmd.Stdout = stdoutFile + gologger.Info().Msgf("Resolved %d hosts\n", resolved.Load()) + return nil +} - // Set stderr based on KeepStderr option - if instance.options.KeepStderr { - cmd.Stderr = stderrFile - } else { - // Discard stderr by sending it to /dev/null - cmd.Stderr = nil +// runEngine consumes names from input using the configured resolution engine: +// the iterative-from-root resolver (no resolver list required) when +// Options.Iterative is set, otherwise the asynchronous stub resolver against +// the supplied recursive resolvers. Both feed answers to storeResult. +func (instance *Instance) runEngine(ctx context.Context, resolvers []string, qtype uint16, input <-chan string, storeResult func(name string, rcode int, a, aaaa []string), markDone func(string)) error { + if instance.options.Iterative { + // Each iterative worker holds one reused UDP socket; clamp the stub + // thread count (which can be very large) to avoid fd exhaustion. + workers := instance.options.Threads + if workers <= 0 { + workers = 200 + } else if workers > 1024 { + workers = 1024 + } + ir, err := iterative.New(iterative.Options{ + QueryType: qtype, + Concurrency: workers, + IPv6: qtype == dns.TypeAAAA, + }) + if err != nil { + return fmt.Errorf("could not create iterative resolver: %w", err) + } + gologger.Info().Msgf("Using iterative-from-root resolver (no recursive resolvers needed)\n") + return ir.ResolveStream(ctx, input, iterative.StreamConfig{ + QueryType: qtype, + OnResult: func(r *resolve.Result) { + storeResult(r.Name, r.Rcode, r.A, r.AAAA) + }, + OnError: func(name string, _ error) { + markDone(name) + }, + }) + } + + client, err := resolve.New(resolve.Options{ + Resolvers: resolvers, + QueryType: qtype, + MaxRetries: instance.options.Retries, + Concurrency: instance.options.Threads, + QPS: instance.options.QPS, + Batch: parseBatchMode(instance.options.BatchMode), + SocketCount: instance.options.SocketCount, + UDPSize: instance.options.UDPSize, + NoRecurse: instance.options.NoRecurse, + Sticky: instance.options.Sticky, + ResolverHealth: instance.options.ResolverHealth, + AdaptiveConcurrency: instance.options.AdaptiveConcurrency, + CrossCheck: instance.options.CrossCheck, + ExtendedInput: instance.options.ExtendedInput, + DisableIPVerification: instance.options.NoVerifyIP, + DisableTCPFallback: instance.options.NoTCPFallback, + OnResult: func(r resolve.Result) { + storeResult(r.Name, r.Rcode, r.A, r.AAAA) + }, + OnError: func(name string, _ error) { + markDone(name) + }, + }) + if err != nil { + return fmt.Errorf("could not create resolver: %w", err) } + defer client.Close() + return client.Run(ctx, input) +} - err = cmd.Run() +// postProcess performs the common steps after the store has been populated: +// optional root-domain extraction, wildcard removal, and output writing. +func (instance *Instance) postProcess(shstore *store.Store) error { + if instance.options.AutoExtractRootDomains { + gologger.Info().Msgf("Started extracting root domains\n") + now := time.Now() + if err := instance.autoExtractRootDomains(shstore); err != nil { + return fmt.Errorf("could not extract root domains: %w", err) + } + gologger.Info().Msgf("Root domain extraction completed in %s\n", time.Since(now)) + } - // Return stderr filename only if it was captured - stderrFilename := "" - if instance.options.KeepStderr { - stderrFilename = stderrFile.Name() + if len(instance.options.Domains) > 0 { + gologger.Info().Msgf("Started removing wildcards records\n") + now := time.Now() + if err := instance.filterWildcards(shstore); err != nil { + return fmt.Errorf("could not filter wildcards: %w", err) + } + gologger.Info().Msgf("Wildcard removal completed in %s\n", time.Since(now)) } - return stdoutFile.Name(), stderrFilename, time.Since(start), err + gologger.Info().Msgf("Finished enumeration, started writing output\n") + now := time.Now() + if err := instance.writeOutput(shstore); err != nil { + return fmt.Errorf("could not write output: %w", err) + } + gologger.Info().Msgf("Output written in %s\n", time.Since(now)) + return nil } type item struct { @@ -210,6 +318,8 @@ type item struct { domain string } +// parseMassDNSOutputFile parses a raw massdns output file (used by the +// -ri/--raw-input compatibility path) into the store. func (instance *Instance) parseMassDNSOutputFile(tmpFile string, store *store.Store) error { flushToDisk := func(ip string, domains []string) error { if err := store.Append(ip, domains...); err != nil { @@ -238,7 +348,6 @@ func (instance *Instance) parseMassDNSOutputFile(tmpFile string, store *store.St err := parser.ParseFile(tmpFile, func(domain string, ips []string) error { for _, ip := range ips { - // Filter out 0.0.0.0 always, and internal IPs if flag is set if instance.shouldFilterIP(ip) { continue } @@ -248,7 +357,6 @@ func (instance *Instance) parseMassDNSOutputFile(tmpFile string, store *store.St }) bulkWriter.Stop() - bulkWriter.WaitDone() if err != nil { @@ -412,12 +520,13 @@ func (instance *Instance) writeOutput(store *store.Store) error { defer swg.Done() if dnsResolver != nil { - if resp, err := dnsResolver.QueryOne(hostname); err != nil || len(resp.A) == 0 { + if resp, err := dnsResolver.QueryOne(hostname); err != nil || (len(resp.A) == 0 && len(resp.AAAA) == 0) { gologger.Info().Msgf("not resolved with trusted resolver - skipping: %s", hostname) return } else { // perform a last check on wildcards ip in case some hosts sneaked due to bad resolvers - for _, ip := range resp.A { + addrs := append(append([]string{}, resp.A...), resp.AAAA...) + for _, ip := range addrs { if instance.wildcardStore.Has(ip) { gologger.Info().Msgf("resolved with trusted resolver but is a wildcard - skipping: %s", hostname) return @@ -470,292 +579,86 @@ func (instance *Instance) writeOutput(store *store.Store) error { return nil } -// ProcessDomainStreaming processes domain bruteforce using streaming with batcher +// ProcessDomainStreaming resolves a domain bruteforce by streaming generated +// permutations directly into the native resolver (no chunk files). func (instance *Instance) ProcessDomainStreaming(ctx context.Context, wordlistFile *os.File) error { - // Create a store for storing ip metadata shstore, err := store.New(instance.options.TempDir) if err != nil { return fmt.Errorf("could not create store: %w", err) } defer shstore.Close() - // Create batcher for streaming permutations - chunkNum := 0 - permutationCount := 0 - - bulkWriter := batcher.New[string]( - batcher.WithMaxCapacity[string](instance.options.BatchSize), - batcher.WithFlushInterval[string](10*time.Second), - batcher.WithFlushCallback[string](func(permutations []string) { - chunkNum++ - if len(permutations) == 0 { - return - } - - gologger.Info().Msgf("Processing chunk %d (%d permutations, total: %d)\n", - chunkNum, len(permutations), permutationCount) - - // Create temporary chunk file - chunkFile, err := os.CreateTemp(instance.options.TempDir, fmt.Sprintf("chunk-%d-", chunkNum)) - if err != nil { - gologger.Error().Msgf("Could not create chunk file: %s\n", err) - return - } - - // Write permutations to chunk file - writer := bufio.NewWriter(chunkFile) - for _, permutation := range permutations { - _, err := writer.WriteString(permutation + "\n") - if err != nil { - gologger.Error().Msgf("Could not write to chunk file: %s\n", err) - _ = chunkFile.Close() - _ = os.Remove(chunkFile.Name()) - return - } - } - _ = writer.Flush() - _ = chunkFile.Close() + gologger.Info().Msgf("Started bruteforce resolution with native resolver\n") + now := time.Now() - // Run massdns on this chunk - chunkStart := time.Now() - stdoutFile, stderrFile, took, err := instance.runChunk(ctx, chunkFile.Name()) - if err != nil { - gologger.Error().Msgf("Could not execute massdns on chunk %d: %s\n", chunkNum, err) - _ = os.Remove(chunkFile.Name()) - return + var permutationCount atomic.Int64 + err = instance.resolveNames(ctx, shstore, func(out chan<- string) error { + scanner := bufio.NewScanner(wordlistFile) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for scanner.Scan() { + // RFC4343 - case insensitive domain + text := strings.ToLower(strings.TrimSpace(scanner.Text())) + if text == "" { + continue } - - gologger.Info().Msgf("Chunk %d massdns execution took %s\n", chunkNum, took) - - // Parse the chunk output immediately - parseStart := time.Now() - err = instance.parseMassDNSOutputFile(stdoutFile, shstore) - if err != nil { - gologger.Error().Msgf("Could not parse massdns output for chunk %d: %s\n", chunkNum, err) - _ = os.Remove(chunkFile.Name()) - _ = os.Remove(stdoutFile) - if stderrFile != "" { - _ = os.Remove(stderrFile) + for _, domain := range instance.options.Domains { + // Stop on cancellation so the producer can't wedge on a full + // channel once the resolver has stopped draining (interrupt). + select { + case out <- text + "." + domain: + permutationCount.Add(1) + case <-ctx.Done(): + return ctx.Err() } - return } - - gologger.Info().Msgf("Chunk %d parsing completed in %s\n", chunkNum, time.Since(parseStart)) - - // Clean up chunk files immediately - _ = os.Remove(chunkFile.Name()) - _ = os.Remove(stdoutFile) - if stderrFile != "" { - _ = os.Remove(stderrFile) - } - - gologger.Info().Msgf("Chunk %d completed in %s\n", chunkNum, time.Since(chunkStart)) - }), - ) - - bulkWriter.Run() - - // Read wordlist and generate permutations on-the-fly - scanner := bufio.NewScanner(wordlistFile) - for scanner.Scan() { - // RFC4343 - case insensitive domain - text := strings.ToLower(scanner.Text()) - if text == "" { - continue - } - - // Generate permutations for each domain - for _, domain := range instance.options.Domains { - permutation := text + "." + domain - bulkWriter.Append(permutation) - permutationCount++ - } - } - - // Stop the batcher and wait for completion - bulkWriter.Stop() - bulkWriter.WaitDone() - - if err := scanner.Err(); err != nil { - return fmt.Errorf("error reading wordlist: %w", err) - } - - gologger.Info().Msgf("Total permutations generated: %d\n", permutationCount) - - // Perform post-processing steps - if instance.options.AutoExtractRootDomains { - gologger.Info().Msgf("Started extracting root domains\n") - now := time.Now() - err = instance.autoExtractRootDomains(shstore) - if err != nil { - return fmt.Errorf("could not extract root domains: %w", err) - } - gologger.Info().Msgf("Root domain extraction completed in %s\n", time.Since(now)) - } - - // Perform wildcard filtering only if domain name has been specified - if len(instance.options.Domains) > 0 { - gologger.Info().Msgf("Started removing wildcards records\n") - now := time.Now() - err = instance.filterWildcards(shstore) - if err != nil { - return fmt.Errorf("could not filter wildcards: %w", err) } - gologger.Info().Msgf("Wildcard removal completed in %s\n", time.Since(now)) - } - - gologger.Info().Msgf("Finished enumeration, started writing output\n") - - // Write the final elaborated list out - now := time.Now() - err = instance.writeOutput(shstore) + return scanner.Err() + }) if err != nil { - return fmt.Errorf("could not write output: %w", err) + return fmt.Errorf("could not resolve permutations: %w", err) } - gologger.Info().Msgf("Output written in %s\n", time.Since(now)) - return nil + gologger.Info().Msgf("Resolved %d permutations in %s\n", permutationCount.Load(), time.Since(now)) + + return instance.postProcess(shstore) } -// ProcessSubdomainsStreaming processes subdomain list using streaming with batcher +// ProcessSubdomainsStreaming resolves a list of subdomains by streaming them +// directly into the native resolver (no chunk files). func (instance *Instance) ProcessSubdomainsStreaming(ctx context.Context, subdomainReader io.Reader) error { - // Create a store for storing ip metadata shstore, err := store.New(instance.options.TempDir) if err != nil { return fmt.Errorf("could not create store: %w", err) } defer shstore.Close() - // Create batcher for streaming subdomains - chunkNum := 0 - subdomainCount := 0 - - bulkWriter := batcher.New[string]( - batcher.WithMaxCapacity[string](instance.options.BatchSize), - batcher.WithFlushInterval[string](10*time.Second), - batcher.WithFlushCallback[string](func(subdomains []string) { - chunkNum++ - if len(subdomains) == 0 { - return - } - - gologger.Info().Msgf("Processing chunk %d (%d subdomains, total: %d)\n", - chunkNum, len(subdomains), subdomainCount) - - // Create temporary chunk file - chunkFile, err := os.CreateTemp(instance.options.TempDir, fmt.Sprintf("chunk-%d-", chunkNum)) - if err != nil { - gologger.Error().Msgf("Could not create chunk file: %s\n", err) - return - } - - // Write subdomains to chunk file - writer := bufio.NewWriter(chunkFile) - for _, subdomain := range subdomains { - _, err := writer.WriteString(subdomain + "\n") - if err != nil { - gologger.Error().Msgf("Could not write to chunk file: %s\n", err) - _ = chunkFile.Close() - _ = os.Remove(chunkFile.Name()) - return - } - } - _ = writer.Flush() - _ = chunkFile.Close() - - // Run massdns on this chunk - chunkStart := time.Now() - stdoutFile, stderrFile, took, err := instance.runChunk(ctx, chunkFile.Name()) - if err != nil { - gologger.Error().Msgf("Could not execute massdns on chunk %d: %s\n", chunkNum, err) - _ = os.Remove(chunkFile.Name()) - return - } - - gologger.Info().Msgf("Chunk %d massdns execution took %s\n", chunkNum, took) + gologger.Info().Msgf("Started resolving subdomains with native resolver\n") + now := time.Now() - // Parse the chunk output immediately - parseStart := time.Now() - err = instance.parseMassDNSOutputFile(stdoutFile, shstore) - if err != nil { - gologger.Error().Msgf("Could not parse massdns output for chunk %d: %s\n", chunkNum, err) - _ = os.Remove(chunkFile.Name()) - _ = os.Remove(stdoutFile) - if stderrFile != "" { - _ = os.Remove(stderrFile) - } - return + var subdomainCount atomic.Int64 + err = instance.resolveNames(ctx, shstore, func(out chan<- string) error { + scanner := bufio.NewScanner(subdomainReader) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for scanner.Scan() { + // RFC4343 - case insensitive domain + subdomain := strings.ToLower(strings.TrimSpace(scanner.Text())) + if subdomain == "" { + continue } - - gologger.Info().Msgf("Chunk %d parsing completed in %s\n", chunkNum, time.Since(parseStart)) - - // Clean up chunk files immediately - _ = os.Remove(chunkFile.Name()) - _ = os.Remove(stdoutFile) - if stderrFile != "" { - _ = os.Remove(stderrFile) + select { + case out <- subdomain: + subdomainCount.Add(1) + case <-ctx.Done(): + return ctx.Err() } - - gologger.Info().Msgf("Chunk %d completed in %s\n", chunkNum, time.Since(chunkStart)) - }), - ) - - bulkWriter.Run() - - // Read subdomains and stream them to batcher - scanner := bufio.NewScanner(subdomainReader) - for scanner.Scan() { - // RFC4343 - case insensitive domain - subdomain := strings.ToLower(strings.TrimSpace(scanner.Text())) - if subdomain == "" { - continue } - - bulkWriter.Append(subdomain) - subdomainCount++ - } - - // Stop the batcher and wait for completion - bulkWriter.Stop() - bulkWriter.WaitDone() - - if err := scanner.Err(); err != nil { - return fmt.Errorf("error reading subdomains: %w", err) - } - - gologger.Info().Msgf("Total subdomains processed: %d\n", subdomainCount) - - // Perform post-processing steps - if instance.options.AutoExtractRootDomains { - gologger.Info().Msgf("Started extracting root domains\n") - now := time.Now() - err = instance.autoExtractRootDomains(shstore) - if err != nil { - return fmt.Errorf("could not extract root domains: %w", err) - } - gologger.Info().Msgf("Root domain extraction completed in %s\n", time.Since(now)) - } - - // Perform wildcard filtering only if domain name has been specified - if len(instance.options.Domains) > 0 { - gologger.Info().Msgf("Started removing wildcards records\n") - now := time.Now() - err = instance.filterWildcards(shstore) - if err != nil { - return fmt.Errorf("could not filter wildcards: %w", err) - } - gologger.Info().Msgf("Wildcard removal completed in %s\n", time.Since(now)) - } - - gologger.Info().Msgf("Finished enumeration, started writing output\n") - - // Write the final elaborated list out - now := time.Now() - err = instance.writeOutput(shstore) + return scanner.Err() + }) if err != nil { - return fmt.Errorf("could not write output: %w", err) + return fmt.Errorf("could not resolve subdomains: %w", err) } - gologger.Info().Msgf("Output written in %s\n", time.Since(now)) - return nil + gologger.Info().Msgf("Resolved input of %d subdomains in %s\n", subdomainCount.Load(), time.Since(now)) + + return instance.postProcess(shstore) } diff --git a/pkg/output/binary.go b/pkg/output/binary.go new file mode 100644 index 00000000..30af0b36 --- /dev/null +++ b/pkg/output/binary.go @@ -0,0 +1,135 @@ +package output + +import ( + "encoding/binary" + "fmt" + "io" + "net" + "strconv" + "time" + + "github.com/projectdiscovery/shuffledns/pkg/resolve" +) + +// massdns binary format (see massdns binfile_write_head / OUTPUT_BINARY). +// +// The file is platform-descriptive: a header records native sizes/offsets so +// readers such as scripts/dnsparse.py can parse records written by this writer. +// We emit a stable, documented layout that matches common Linux sockaddr field +// offsets so files are interchangeable with Linux massdns when written there. + +const ( + binaryVersion = uint32(0) + binarySockaddrSize = 128 + binaryFamilyOffset = 0 + binaryFamilySize = 2 + binaryPortSize = 2 + binaryTimeSize = 8 + binarySizeTSize = 8 + binaryAFInet = uint16(2) // AF_INET on Linux + binaryAFInet6 = uint16(10) // AF_INET6 on Linux + binarySinAddrOff = 4 + binarySinPortOff = 2 + binarySin6AddrOff = 8 + binarySin6PortOff = 2 +) + +func writeBinaryHeader(w io.Writer) error { + ne := binary.NativeEndian + if _, err := w.Write([]byte("massdns\x00")); err != nil { + return err + } + if err := binary.Write(w, ne, uint32(0x12345678)); err != nil { + return err + } + if err := binary.Write(w, ne, binaryVersion); err != nil { + return err + } + if _, err := w.Write([]byte{byte(binarySizeTSize)}); err != nil { + return err + } + fields := []uint64{ + binaryTimeSize, + binarySockaddrSize, + binaryFamilyOffset, + binaryFamilySize, + binaryPortSize, + } + for _, v := range fields { + if err := binary.Write(w, ne, v); err != nil { + return err + } + } + if err := binary.Write(w, ne, binaryAFInet); err != nil { + return err + } + if err := binary.Write(w, ne, uint64(binarySinAddrOff)); err != nil { + return err + } + if err := binary.Write(w, ne, uint64(binarySinPortOff)); err != nil { + return err + } + if err := binary.Write(w, ne, binaryAFInet6); err != nil { + return err + } + if err := binary.Write(w, ne, uint64(binarySin6AddrOff)); err != nil { + return err + } + if err := binary.Write(w, ne, uint64(binarySin6PortOff)); err != nil { + return err + } + return nil +} + +func writeBinaryRecord(w io.Writer, r resolve.Result) error { + if r.Msg == nil { + return nil + } + raw, err := r.Msg.Pack() + if err != nil { + return err + } + if len(raw) > 0xffff { + return fmt.Errorf("dns message too large for binary format: %d", len(raw)) + } + + ts := r.Timestamp + if ts.IsZero() { + ts = time.Now() + } + ne := binary.NativeEndian + if err := binary.Write(w, ne, uint64(ts.Unix())); err != nil { + return err + } + + var ss [binarySockaddrSize]byte + host, portStr, splitErr := net.SplitHostPort(r.Resolver) + if splitErr != nil { + host = r.Resolver + } + port := uint16(53) + if portStr != "" { + if p, convErr := strconv.Atoi(portStr); convErr == nil && p > 0 && p < 65536 { + port = uint16(p) + } + } + if ip := net.ParseIP(host); ip != nil { + if v4 := ip.To4(); v4 != nil { + ne.PutUint16(ss[binaryFamilyOffset:], binaryAFInet) + binary.BigEndian.PutUint16(ss[binarySinPortOff:], port) + copy(ss[binarySinAddrOff:], v4) + } else { + ne.PutUint16(ss[binaryFamilyOffset:], binaryAFInet6) + binary.BigEndian.PutUint16(ss[binarySin6PortOff:], port) + copy(ss[binarySin6AddrOff:], ip.To16()) + } + } + if _, err := w.Write(ss[:]); err != nil { + return err + } + if err := binary.Write(w, ne, uint16(len(raw))); err != nil { + return err + } + _, err = w.Write(raw) + return err +} diff --git a/pkg/output/binary_test.go b/pkg/output/binary_test.go new file mode 100644 index 00000000..0bb3dc8e --- /dev/null +++ b/pkg/output/binary_test.go @@ -0,0 +1,69 @@ +package output + +import ( + "bytes" + "encoding/binary" + "testing" + "time" + + "github.com/miekg/dns" + "github.com/stretchr/testify/require" +) + +func TestBinaryOutputRoundTripHeaderAndRecord(t *testing.T) { + r := sampleResult() + r.Timestamp = time.Unix(1513458347, 0) + r.Resolver = "192.0.2.1:53" + + var buf bytes.Buffer + w, err := NewWriterWithConfig(&buf, Config{Format: "B"}) + require.NoError(t, err) + require.NoError(t, w.Write(r)) + require.NoError(t, w.Flush()) + + data := buf.Bytes() + require.True(t, bytes.HasPrefix(data, []byte("massdns\x00"))) + + // magic(8) + endian(4) + version(4) + size_len(1) = 17 + off := 8 + require.Equal(t, uint32(0x12345678), binary.NativeEndian.Uint32(data[off:off+4])) + off += 4 + require.Equal(t, uint32(0), binary.NativeEndian.Uint32(data[off:off+4])) + off += 4 + require.Equal(t, byte(8), data[off]) + off++ + + // skip descriptive sizes/offsets (5*8 + family/offsets…) + // After size_len: time, sockaddr, family_off, family_size, port_size (5 uint64) + off += 5 * 8 + // family_inet (2) + sin_addr_off (8) + sin_port_off (8) + off += 2 + 8 + 8 + // family_inet6 (2) + sin6_addr_off (8) + sin6_port_off (8) + off += 2 + 8 + 8 + + // record: time(8) + sockaddr(128) + len(2) + payload + require.Greater(t, len(data), off+8+128+2) + ts := binary.NativeEndian.Uint64(data[off : off+8]) + require.Equal(t, uint64(1513458347), ts) + off += 8 + ss := data[off : off+128] + off += 128 + require.Equal(t, binaryAFInet, binary.NativeEndian.Uint16(ss[0:2])) + require.Equal(t, uint16(53), binary.BigEndian.Uint16(ss[binarySinPortOff:binarySinPortOff+2])) + require.Equal(t, netIPv4(192, 0, 2, 1), ss[binarySinAddrOff:binarySinAddrOff+4]) + + msgLen := binary.NativeEndian.Uint16(data[off : off+2]) + off += 2 + raw := data[off : off+int(msgLen)] + msg := new(dns.Msg) + require.NoError(t, msg.Unpack(raw)) + require.Equal(t, "example.com.", msg.Question[0].Name) +} + +func netIPv4(a, b, c, d byte) []byte { return []byte{a, b, c, d} } + +func TestParseFormatBinary(t *testing.T) { + f, err := ParseFormat("B") + require.NoError(t, err) + require.Equal(t, modeBinary, f.mode) +} diff --git a/pkg/output/output.go b/pkg/output/output.go new file mode 100644 index 00000000..ed8c33e6 --- /dev/null +++ b/pkg/output/output.go @@ -0,0 +1,598 @@ +// Package output renders resolver results in massdns-compatible output formats. +// +// It mirrors the massdns -o flag: a format string whose first recognized +// major letter selects the mode (S simple, F full, L list, J ndjson) and whose +// remaining letters toggle mode-specific options. The goal is drop-in parity so +// existing massdns post-processing pipelines keep working against the native +// resolver. +// +// Supported major modes: +// +// S - simple text (default; "name. TYPE data" per record, the form shuffledns parses) +// F - full text (dig-like packet dump with ;; Server/Size/Unix time header) +// L - domain list (tab-separated query/type/data, one record per line) +// J - ndjson (one JSON object per reply) +// B - binary (massdns binary packet stream; readable by dnsparse.py) +package output + +import ( + "bufio" + "encoding/json" + "fmt" + "io" + "strconv" + "strings" + "sync" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/pkg/resolve" +) + +type mode int + +const ( + modeSimple mode = iota + modeFull + modeList + modeNDJSON + modeBinary +) + +// Format is a parsed massdns -o specification. +type Format struct { + mode mode + + // simple-mode section selectors + answer bool // n + authority bool // u + additional bool // d + + // simple-mode modifiers + ttlClass bool // t: include TTL and class + question bool // q: print the question line + meta bool // r: prepend resolver, unix ts and rcode to the question + sepReplies bool // l: blank line between replies + matchQ bool // m: only records whose owner matches the question name + indent bool // i: indent reply records by a tab + sepSections bool // s: blank line between sections + + listIncludeEmpty bool // 0 (list mode): include NOERROR replies without answers + ndjsonFailures bool // e (ndjson mode): emit a record for terminal failures +} + +// ParseFormat parses a massdns-style format string (e.g. "Snl", "F", "J"). +// An empty string defaults to simple answer-record output ("Sn"). +func ParseFormat(spec string) (*Format, error) { + f := &Format{mode: modeSimple} + modeSet := false + for _, r := range spec { + switch r { + case 'S': + f.mode, modeSet = modeSimple, true + case 'F': + f.mode, modeSet = modeFull, true + case 'L': + f.mode, modeSet = modeList, true + case 'J': + f.mode, modeSet = modeNDJSON, true + case 'B': + f.mode, modeSet = modeBinary, true + case 'n': + f.answer = true + case 'u': + f.authority = true + case 'd': + f.additional = true + case 't': + f.ttlClass = true + case 'q': + f.question = true + case 'r': + f.meta = true + case 'l': + f.sepReplies = true + case 'm': + f.matchQ = true + case 'i': + f.indent = true + case 's': + f.sepSections = true + case '0': + f.listIncludeEmpty = true + case 'e': + f.ndjsonFailures = true + case ' ', '\t': + // ignore separators + default: + return nil, fmt.Errorf("unknown output flag %q", string(r)) + } + } + _ = modeSet + // Simple mode with no explicit section defaults to the answer section, which + // matches plain `-o S`. + if f.mode == modeSimple && !f.answer && !f.authority && !f.additional { + f.answer = true + } + return f, nil +} + +// EmitsFailures reports whether terminal (post-retry) failures should be handed +// to WriteFailure. Only ndjson with the 'e' flag records them. +func (f *Format) EmitsFailures() bool { return f.mode == modeNDJSON && f.ndjsonFailures } + +// Config bundles a format spec with response-filtering options. It is the +// SDK-friendly way to build a Writer with massdns --filter/--ignore semantics +// and the "only the queried record type" behaviour (massdns issue #1). +type Config struct { + // Format is the massdns-style -o spec (see ParseFormat). Empty = "Sn". + Format string + // FilterRcodes, when non-empty, restricts output to replies whose response + // code is in the set (massdns --filter). Names are response-code strings or + // numbers, e.g. "NOERROR", "NXDOMAIN", "3". + FilterRcodes []string + // IgnoreRcodes drops replies whose response code is in the set (massdns + // --ignore). Applied after FilterRcodes. + IgnoreRcodes []string + // OnlyQueryType emits only answer records whose type matches the queried + // type, dropping CNAME/glue chains from the output (massdns issue #1). + OnlyQueryType bool + // FlushEach flushes after every written reply (massdns --flush). + FlushEach bool +} + +// rcodeValue resolves a response-code name or number to its numeric value. +func rcodeValue(s string) (int, bool) { + s = strings.TrimSpace(s) + if s == "" { + return 0, false + } + if v, ok := dns.StringToRcode[strings.ToUpper(s)]; ok { + return v, true + } + if n, err := strconv.Atoi(s); err == nil { + return n, true + } + return 0, false +} + +func rcodeSet(names []string) (map[int]struct{}, error) { + if len(names) == 0 { + return nil, nil + } + set := make(map[int]struct{}, len(names)) + for _, n := range names { + v, ok := rcodeValue(n) + if !ok { + return nil, fmt.Errorf("invalid response code %q", n) + } + set[v] = struct{}{} + } + return set, nil +} + +// Writer renders results to an underlying writer. It is safe for concurrent use +// (the resolver delivers results from multiple goroutines). +type Writer struct { + mu sync.Mutex + bw *bufio.Writer + f *Format + + filter map[int]struct{} // only these rcodes (nil = all) + ignore map[int]struct{} // drop these rcodes + onlyQType bool // emit only answer records matching the question type + flushEach bool + binHeader bool // binary mode header already written +} + +// NewWriter builds a Writer for the given format spec. +func NewWriter(w io.Writer, spec string) (*Writer, error) { + return NewWriterWithConfig(w, Config{Format: spec}) +} + +// NewWriterWithConfig builds a Writer from a full Config (format + filters). +func NewWriterWithConfig(w io.Writer, cfg Config) (*Writer, error) { + f, err := ParseFormat(cfg.Format) + if err != nil { + return nil, err + } + filter, err := rcodeSet(cfg.FilterRcodes) + if err != nil { + return nil, err + } + ignore, err := rcodeSet(cfg.IgnoreRcodes) + if err != nil { + return nil, err + } + return &Writer{ + bw: bufio.NewWriterSize(w, 64*1024), + f: f, + filter: filter, + ignore: ignore, + onlyQType: cfg.OnlyQueryType, + flushEach: cfg.FlushEach, + }, nil +} + +// suppressed reports whether a reply with the given rcode should be dropped by +// the configured --filter/--ignore sets. +func (w *Writer) suppressed(rcode int) bool { + if w.filter != nil { + if _, ok := w.filter[rcode]; !ok { + return true + } + } + if w.ignore != nil { + if _, ok := w.ignore[rcode]; ok { + return true + } + } + return false +} + +// Flush flushes buffered output. +func (w *Writer) Flush() error { + w.mu.Lock() + defer w.mu.Unlock() + return w.bw.Flush() +} + +// Write renders a single result. +func (w *Writer) Write(r resolve.Result) error { + if w.suppressed(r.Rcode) { + return nil + } + w.mu.Lock() + defer w.mu.Unlock() + var err error + switch w.f.mode { + case modeSimple: + err = w.writeSimple(r) + case modeFull: + err = w.writeFull(r) + case modeList: + err = w.writeList(r) + case modeNDJSON: + err = w.writeNDJSON(r) + case modeBinary: + err = w.writeBinary(r) + } + if err == nil && w.flushEach { + err = w.bw.Flush() + } + return err +} + +func (w *Writer) writeBinary(r resolve.Result) error { + if !w.binHeader { + if err := writeBinaryHeader(w.bw); err != nil { + return err + } + w.binHeader = true + } + return writeBinaryRecord(w.bw, r) +} + +// answerRecords returns the answer section, optionally filtered to records whose +// type matches the queried type (the OnlyQueryType behaviour). +func (w *Writer) answerRecords(r resolve.Result) []dns.RR { + if r.Msg == nil { + return nil + } + if !w.onlyQType { + return r.Msg.Answer + } + out := make([]dns.RR, 0, len(r.Msg.Answer)) + for _, rr := range r.Msg.Answer { + if rr.Header().Rrtype == r.Type { + out = append(out, rr) + } + } + return out +} + +// WriteFailure records a terminal query failure (ndjson 'e' flag only). +func (w *Writer) WriteFailure(name string, qtype uint16, reason string) error { + if !w.f.EmitsFailures() { + return nil + } + w.mu.Lock() + defer w.mu.Unlock() + rec := jsonReply{ + Name: dns.Fqdn(name), + Type: typeString(qtype), + Class: "IN", + Status: "ERROR", + Error: reason, + } + return w.encodeJSON(rec) +} + +// ---- simple ---- + +func (w *Writer) writeSimple(r resolve.Result) error { + qname := questionName(r) + + if w.f.meta { + // resolver, unix timestamp and rcode prepended to the question + if _, err := fmt.Fprintf(w.bw, "%s %d %s %s %s %s\n", + emptyDash(r.Resolver), r.Timestamp.Unix(), rcodeString(r.Rcode), + qname, "IN", typeString(r.Type)); err != nil { + return err + } + } else if w.f.question { + if _, err := fmt.Fprintf(w.bw, "%s %s %s\n", qname, "IN", typeString(r.Type)); err != nil { + return err + } + } + + wrote := false + emit := func(rrs []dns.RR) error { + if w.f.sepSections && wrote && len(rrs) > 0 { + if err := w.bw.WriteByte('\n'); err != nil { + return err + } + } + for _, rr := range rrs { + if w.f.matchQ && !strings.EqualFold(rr.Header().Name, qname) { + continue + } + if w.f.indent { + if err := w.bw.WriteByte('\t'); err != nil { + return err + } + } + if _, err := w.bw.WriteString(w.simpleLine(rr)); err != nil { + return err + } + wrote = true + } + return nil + } + + if r.Msg != nil { + if w.f.answer { + if err := emit(w.answerRecords(r)); err != nil { + return err + } + } + if w.f.authority { + if err := emit(r.Msg.Ns); err != nil { + return err + } + } + if w.f.additional { + if err := emit(filterOPT(r.Msg.Extra)); err != nil { + return err + } + } + } + + // 'l': separate replies with a blank line (the classic -o Snl form). + if w.f.sepReplies { + if err := w.bw.WriteByte('\n'); err != nil { + return err + } + } + return nil +} + +// simpleLine renders one RR as "name TYPE data" (or "name TTL CLASS TYPE data" +// with the 't' flag), terminated by a newline. +func (w *Writer) simpleLine(rr dns.RR) string { + h := rr.Header() + data := rdata(rr) + if w.f.ttlClass { + return fmt.Sprintf("%s %d %s %s %s\n", h.Name, h.Ttl, classString(h.Class), typeString(h.Rrtype), data) + } + return fmt.Sprintf("%s %s %s\n", h.Name, typeString(h.Rrtype), data) +} + +// ---- full ---- + +func (w *Writer) writeFull(r resolve.Result) error { + size := 0 + body := "" + if r.Msg != nil { + size = r.Msg.Len() + body = r.Msg.String() + } + if _, err := fmt.Fprintf(w.bw, ";; Server: %s\n;; Size: %d\n;; Unix time: %d\n%s\n\n", + emptyDash(r.Resolver), size, r.Timestamp.Unix(), body); err != nil { + return err + } + return nil +} + +// ---- list ---- + +func (w *Writer) writeList(r resolve.Result) error { + if r.Msg == nil { + return nil + } + qname := questionName(r) + answers := w.answerRecords(r) + if len(answers) == 0 { + if w.f.listIncludeEmpty && r.Rcode == dns.RcodeSuccess { + if _, err := fmt.Fprintf(w.bw, "%s\n", strings.TrimSuffix(qname, ".")); err != nil { + return err + } + } + return nil + } + for _, rr := range answers { + h := rr.Header() + if _, err := fmt.Fprintf(w.bw, "%s\t%s\t%s\n", + strings.TrimSuffix(h.Name, "."), typeString(h.Rrtype), rdata(rr)); err != nil { + return err + } + } + return nil +} + +// ---- ndjson ---- + +type jsonRecord struct { + Name string `json:"name"` + Type string `json:"type"` + Class string `json:"class"` + TTL uint32 `json:"ttl"` + Data string `json:"data"` +} + +type jsonData struct { + Answers []jsonRecord `json:"answers"` + Authorities []jsonRecord `json:"authorities,omitempty"` + Additionals []jsonRecord `json:"additionals,omitempty"` +} + +type jsonReply struct { + Name string `json:"name"` + Type string `json:"type"` + Class string `json:"class"` + Status string `json:"status"` + RxTs int64 `json:"rx_ts,omitempty"` + Resolver string `json:"resolver,omitempty"` + Flags []string `json:"flags,omitempty"` + Data jsonData `json:"data"` + Error string `json:"error,omitempty"` +} + +func (w *Writer) writeNDJSON(r resolve.Result) error { + rec := jsonReply{ + Name: questionName(r), + Type: typeString(r.Type), + Class: "IN", + Status: rcodeString(r.Rcode), + Resolver: r.Resolver, + } + if !r.Timestamp.IsZero() { + rec.RxTs = r.Timestamp.Unix() + } + if r.Msg != nil { + rec.Flags = msgFlags(r.Msg) + rec.Data.Answers = jsonRecords(w.answerRecords(r)) + rec.Data.Authorities = jsonRecords(r.Msg.Ns) + rec.Data.Additionals = jsonRecords(filterOPT(r.Msg.Extra)) + } + if rec.Data.Answers == nil { + rec.Data.Answers = []jsonRecord{} + } + return w.encodeJSON(rec) +} + +func (w *Writer) encodeJSON(rec jsonReply) error { + b, err := json.Marshal(rec) + if err != nil { + return err + } + if _, err := w.bw.Write(b); err != nil { + return err + } + return w.bw.WriteByte('\n') +} + +func jsonRecords(rrs []dns.RR) []jsonRecord { + if len(rrs) == 0 { + return nil + } + out := make([]jsonRecord, 0, len(rrs)) + for _, rr := range rrs { + h := rr.Header() + out = append(out, jsonRecord{ + Name: h.Name, + Type: typeString(h.Rrtype), + Class: classString(h.Class), + TTL: h.Ttl, + Data: rdata(rr), + }) + } + return out +} + +// ---- helpers ---- + +// rdata returns the record-specific data portion of an RR (everything after the +// "name ttl class type" header that miekg/dns prints). +func rdata(rr dns.RR) string { + full := rr.String() + header := rr.Header().String() + return strings.TrimPrefix(full, header) +} + +// filterOPT drops EDNS0 OPT pseudo-records from the additional section; they are +// transport metadata, not answer data, and massdns does not print them. +func filterOPT(rrs []dns.RR) []dns.RR { + if len(rrs) == 0 { + return nil + } + out := rrs[:0:0] + for _, rr := range rrs { + if _, ok := rr.(*dns.OPT); ok { + continue + } + out = append(out, rr) + } + return out +} + +func msgFlags(m *dns.Msg) []string { + var f []string + if m.Response { + f = append(f, "qr") + } + if m.Authoritative { + f = append(f, "aa") + } + if m.Truncated { + f = append(f, "tc") + } + if m.RecursionDesired { + f = append(f, "rd") + } + if m.RecursionAvailable { + f = append(f, "ra") + } + if m.AuthenticatedData { + f = append(f, "ad") + } + if m.CheckingDisabled { + f = append(f, "cd") + } + return f +} + +// questionName returns the query name with a trailing dot (massdns convention). +func questionName(r resolve.Result) string { + if r.Msg != nil && len(r.Msg.Question) > 0 { + return r.Msg.Question[0].Name + } + return dns.Fqdn(r.Name) +} + +func typeString(t uint16) string { + if s, ok := dns.TypeToString[t]; ok { + return s + } + return fmt.Sprintf("TYPE%d", t) +} + +func classString(c uint16) string { + if s, ok := dns.ClassToString[c]; ok { + return s + } + return fmt.Sprintf("CLASS%d", c) +} + +func rcodeString(rc int) string { + if s, ok := dns.RcodeToString[rc]; ok { + return s + } + return fmt.Sprintf("RCODE%d", rc) +} + +func emptyDash(s string) string { + if s == "" { + return "-" + } + return s +} diff --git a/pkg/output/output_test.go b/pkg/output/output_test.go new file mode 100644 index 00000000..d8fe4135 --- /dev/null +++ b/pkg/output/output_test.go @@ -0,0 +1,261 @@ +package output + +import ( + "bytes" + "encoding/json" + "strings" + "testing" + "time" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/pkg/parser" + "github.com/projectdiscovery/shuffledns/pkg/resolve" +) + +// sampleResult builds a resolve.Result with a fully populated *dns.Msg for +// example.com -> 93.184.216.34 (A) with an authority NS record. +func sampleResult() resolve.Result { + m := new(dns.Msg) + m.SetQuestion("example.com.", dns.TypeA) + m.Response = true + m.RecursionDesired = true + m.RecursionAvailable = true + a, _ := dns.NewRR("example.com. 45929 IN A 93.184.216.34") + m.Answer = []dns.RR{a} + ns, _ := dns.NewRR("example.com. 24852 IN NS a.iana-servers.net.") + m.Ns = []dns.RR{ns} + + return resolve.Result{ + Name: "example.com", + Type: dns.TypeA, + Rcode: dns.RcodeSuccess, + A: []string{"93.184.216.34"}, + Resolver: "192.0.2.1:53", + Msg: m, + Timestamp: time.Unix(1513458347, 0), + } +} + +func render(t *testing.T, spec string, r resolve.Result) string { + t.Helper() + var buf bytes.Buffer + w, err := NewWriter(&buf, spec) + if err != nil { + t.Fatalf("NewWriter(%q): %v", spec, err) + } + if err := w.Write(r); err != nil { + t.Fatalf("Write: %v", err) + } + if err := w.Flush(); err != nil { + t.Fatalf("Flush: %v", err) + } + return buf.String() +} + +func TestSimpleDefault(t *testing.T) { + got := render(t, "S", sampleResult()) + want := "example.com. A 93.184.216.34\n" + if got != want { + t.Fatalf("got %q want %q", got, want) + } +} + +func TestSimpleSnlRoundTripsThroughParser(t *testing.T) { + // The classic massdns -o Snl form must be consumable by pkg/parser. + got := render(t, "Snl", sampleResult()) + if !strings.HasSuffix(got, "\n\n") { + t.Fatalf("Snl should end replies with a blank line: %q", got) + } + + var domain string + var ips []string + err := parser.ParseReader(strings.NewReader(got), func(d string, ip []string) error { + domain = d + ips = ip + return nil + }) + if err != nil { + t.Fatal(err) + } + if domain != "example.com" { + t.Fatalf("parser domain: got %q", domain) + } + if len(ips) != 1 || ips[0] != "93.184.216.34" { + t.Fatalf("parser ips: got %v", ips) + } +} + +func TestSimpleTTLClass(t *testing.T) { + got := render(t, "St", sampleResult()) + want := "example.com. 45929 IN A 93.184.216.34\n" + if got != want { + t.Fatalf("got %q want %q", got, want) + } +} + +func TestSimpleQuestionAndAuthority(t *testing.T) { + got := render(t, "Snuq", sampleResult()) + if !strings.Contains(got, "example.com. IN A\n") { + t.Fatalf("missing question line: %q", got) + } + if !strings.Contains(got, "example.com. A 93.184.216.34\n") { + t.Fatalf("missing answer: %q", got) + } + if !strings.Contains(got, "example.com. NS a.iana-servers.net.\n") { + t.Fatalf("missing authority: %q", got) + } +} + +func TestSimpleMeta(t *testing.T) { + got := render(t, "Snr", sampleResult()) + if !strings.HasPrefix(got, "192.0.2.1:53 1513458347 NOERROR example.com. IN A\n") { + t.Fatalf("meta prefix wrong: %q", got) + } +} + +func TestFull(t *testing.T) { + got := render(t, "F", sampleResult()) + for _, want := range []string{ + ";; Server: 192.0.2.1:53", + ";; Unix time: 1513458347", + "ANSWER SECTION", + "93.184.216.34", + } { + if !strings.Contains(got, want) { + t.Fatalf("full output missing %q in:\n%s", want, got) + } + } +} + +func TestList(t *testing.T) { + got := render(t, "L", sampleResult()) + want := "example.com\tA\t93.184.216.34\n" + if got != want { + t.Fatalf("got %q want %q", got, want) + } +} + +func TestNDJSON(t *testing.T) { + got := render(t, "J", sampleResult()) + var rec struct { + Name string `json:"name"` + Type string `json:"type"` + Class string `json:"class"` + Status string `json:"status"` + Resolver string `json:"resolver"` + RxTs int64 `json:"rx_ts"` + Data struct { + Answers []struct { + Name string `json:"name"` + Type string `json:"type"` + TTL uint32 `json:"ttl"` + Data string `json:"data"` + } `json:"answers"` + } `json:"data"` + Flags []string `json:"flags"` + } + if err := json.Unmarshal([]byte(strings.TrimSpace(got)), &rec); err != nil { + t.Fatalf("invalid ndjson %q: %v", got, err) + } + if rec.Name != "example.com." || rec.Type != "A" || rec.Status != "NOERROR" { + t.Fatalf("bad header fields: %+v", rec) + } + if rec.Resolver != "192.0.2.1:53" || rec.RxTs != 1513458347 { + t.Fatalf("bad meta: %+v", rec) + } + if len(rec.Data.Answers) != 1 || rec.Data.Answers[0].Data != "93.184.216.34" || rec.Data.Answers[0].TTL != 45929 { + t.Fatalf("bad answers: %+v", rec.Data.Answers) + } +} + +func TestNDJSONFailure(t *testing.T) { + var buf bytes.Buffer + w, err := NewWriter(&buf, "Je") + if err != nil { + t.Fatal(err) + } + if err := w.WriteFailure("dead.example.com", dns.TypeA, "resolution failed"); err != nil { + t.Fatal(err) + } + _ = w.Flush() + if !strings.Contains(buf.String(), `"status":"ERROR"`) || !strings.Contains(buf.String(), "dead.example.com.") { + t.Fatalf("bad failure record: %q", buf.String()) + } + + // without 'e', failures are suppressed + var buf2 bytes.Buffer + w2, _ := NewWriter(&buf2, "J") + _ = w2.WriteFailure("dead.example.com", dns.TypeA, "x") + _ = w2.Flush() + if buf2.Len() != 0 { + t.Fatalf("expected no failure output without 'e', got %q", buf2.String()) + } +} + +func renderCfg(t *testing.T, cfg Config, r resolve.Result) string { + t.Helper() + var buf bytes.Buffer + w, err := NewWriterWithConfig(&buf, cfg) + if err != nil { + t.Fatalf("NewWriterWithConfig: %v", err) + } + if err := w.Write(r); err != nil { + t.Fatalf("Write: %v", err) + } + _ = w.Flush() + return buf.String() +} + +func TestFilterRcode(t *testing.T) { + r := sampleResult() // NOERROR + // only NXDOMAIN allowed -> NOERROR reply suppressed + if got := renderCfg(t, Config{Format: "S", FilterRcodes: []string{"NXDOMAIN"}}, r); got != "" { + t.Fatalf("expected suppression, got %q", got) + } + // NOERROR allowed -> emitted + if got := renderCfg(t, Config{Format: "S", FilterRcodes: []string{"NOERROR"}}, r); got == "" { + t.Fatal("expected output for allowed rcode") + } +} + +func TestIgnoreRcode(t *testing.T) { + r := sampleResult() // NOERROR + if got := renderCfg(t, Config{Format: "S", IgnoreRcodes: []string{"NOERROR"}}, r); got != "" { + t.Fatalf("expected NOERROR to be ignored, got %q", got) + } +} + +func TestOnlyQueryType(t *testing.T) { + // build a result whose answer has a CNAME plus the queried A record + m := new(dns.Msg) + m.SetQuestion("www.example.com.", dns.TypeA) + cname, _ := dns.NewRR("www.example.com. 300 IN CNAME example.com.") + a, _ := dns.NewRR("example.com. 300 IN A 93.184.216.34") + m.Answer = []dns.RR{cname, a} + r := resolve.Result{Name: "www.example.com", Type: dns.TypeA, Rcode: dns.RcodeSuccess, Msg: m} + + full := renderCfg(t, Config{Format: "S"}, r) + if !strings.Contains(full, "CNAME") { + t.Fatalf("expected CNAME in unfiltered output: %q", full) + } + only := renderCfg(t, Config{Format: "S", OnlyQueryType: true}, r) + if strings.Contains(only, "CNAME") { + t.Fatalf("only-type output should drop CNAME: %q", only) + } + if !strings.Contains(only, "A 93.184.216.34") { + t.Fatalf("only-type output should keep the A record: %q", only) + } +} + +func TestInvalidRcodeConfig(t *testing.T) { + var buf bytes.Buffer + if _, err := NewWriterWithConfig(&buf, Config{Format: "S", FilterRcodes: []string{"BOGUS"}}); err == nil { + t.Fatal("expected error for invalid rcode name") + } +} + +func TestUnknownFlag(t *testing.T) { + if _, err := ParseFormat("SZ"); err == nil { + t.Fatal("expected error for unknown flag") + } +} diff --git a/pkg/ptr/ptr.go b/pkg/ptr/ptr.go new file mode 100644 index 00000000..25ec25e4 --- /dev/null +++ b/pkg/ptr/ptr.go @@ -0,0 +1,204 @@ +// Package ptr generates reverse-DNS (PTR) query names from IP targets, the +// native equivalent of massdns's scripts/ptr.py. Targets may be single IPs, +// CIDR blocks, or inclusive "start-end" ranges, in IPv4 or IPv6. +// +// Each emitted name is the reversed in-addr.arpa (IPv4) or ip6.arpa (IPv6) +// label that a resolver expects for a PTR lookup; e.g. 1.2.3.4 becomes +// 4.3.2.1.in-addr.arpa. These names are fed straight into the resolver with +// query type PTR. +package ptr + +import ( + "context" + "fmt" + "net" + "strings" + + "github.com/miekg/dns" +) + +// minIPv6Prefix bounds IPv6 CIDR expansion so a careless /32 doesn't try to +// enumerate an astronomically large space. /112 is 65536 addresses. +const minIPv6Prefix = 112 + +// ReverseName returns the reverse-DNS name (in-addr.arpa / ip6.arpa) for ip. +func ReverseName(ip net.IP) (string, error) { + if ip == nil { + return "", fmt.Errorf("nil ip") + } + return dns.ReverseAddr(ip.String()) +} + +// Expand walks every IP described by targets and calls emit with its reverse +// name. emit may return false to stop early. Targets are parsed as, in order: +// an inclusive range "a-b", a CIDR "ip/bits", or a single IP. +func Expand(targets []string, emit func(name string) bool) error { + for _, t := range targets { + t = strings.TrimSpace(t) + if t == "" { + continue + } + stop, err := expandOne(t, emit) + if err != nil { + return err + } + if stop { + return nil + } + } + return nil +} + +// Stream feeds reverse names produced from targets into out, honouring ctx +// cancellation. It does not close out (the caller owns it). +func Stream(ctx context.Context, targets []string, out chan<- string) error { + return Expand(targets, func(name string) bool { + select { + case <-ctx.Done(): + return false + case out <- name: + return true + } + }) +} + +// expandOne handles a single target token. It returns stop=true when emit asked +// to halt. +func expandOne(t string, emit func(string) bool) (bool, error) { + switch { + case strings.Contains(t, "-"): + return expandRange(t, emit) + case strings.Contains(t, "/"): + return expandCIDR(t, emit) + default: + ip := net.ParseIP(t) + if ip == nil { + return false, fmt.Errorf("invalid IP %q", t) + } + return emitIP(ip, emit), nil + } +} + +func expandRange(t string, emit func(string) bool) (bool, error) { + parts := strings.SplitN(t, "-", 2) + startIP := net.ParseIP(strings.TrimSpace(parts[0])) + endIP := net.ParseIP(strings.TrimSpace(parts[1])) + if startIP == nil || endIP == nil { + return false, fmt.Errorf("invalid range %q", t) + } + start, sv4 := normalize(startIP) + end, ev4 := normalize(endIP) + if sv4 != ev4 { + return false, fmt.Errorf("range %q mixes IPv4 and IPv6", t) + } + if compareBytes(start, end) > 0 { + return false, fmt.Errorf("range %q start is after end", t) + } + cur := dupIP(start) + for { + if !emit(reverse(cur, sv4)) { + return true, nil + } + if compareBytes(cur, end) == 0 { + return false, nil + } + inc(cur) + } +} + +func expandCIDR(t string, emit func(string) bool) (bool, error) { + _, ipnet, err := net.ParseCIDR(t) + if err != nil { + return false, err + } + ones, bits := ipnet.Mask.Size() + isV4 := bits == 32 + if !isV4 && ones < minIPv6Prefix { + return false, fmt.Errorf("IPv6 CIDR %q too large; use /%d or longer", t, minIPv6Prefix) + } + // network address (already masked by ParseCIDR) and the last address in the + // block: last = network | ^mask. + first := dupIP(maskedBase(ipnet.IP, isV4)) + mask := ipnet.Mask + last := dupIP(first) + for i := range last { + last[i] |= ^mask[i] + } + cur := dupIP(first) + for { + if !emit(reverse(cur, isV4)) { + return true, nil + } + if compareBytes(cur, last) == 0 { + return false, nil + } + inc(cur) + } +} + +// maskedBase normalizes a CIDR network IP to its 4- or 16-byte form. +func maskedBase(ip net.IP, v4 bool) []byte { + if v4 { + return dupIP(ip.To4()) + } + return dupIP(ip.To16()) +} + +func emitIP(ip net.IP, emit func(string) bool) bool { + b, v4 := normalize(ip) + return !emit(reverse(b, v4)) +} + +// reverse builds the reverse-DNS name for a normalized address. v4 selects the +// 4-byte vs 16-byte interpretation. +func reverse(b []byte, v4 bool) string { + var ip net.IP + if v4 { + ip = net.IPv4(b[0], b[1], b[2], b[3]) + } else { + ip = make(net.IP, len(b)) + copy(ip, b) + } + name, err := dns.ReverseAddr(ip.String()) + if err != nil { + return "" + } + return name +} + +// normalize returns the address in its minimal byte form (4 bytes for IPv4, +// 16 for IPv6) and whether it is IPv4. +func normalize(ip net.IP) ([]byte, bool) { + if v4 := ip.To4(); v4 != nil { + return dupIP(v4), true + } + return dupIP(ip.To16()), false +} + +func dupIP(ip net.IP) []byte { + b := make([]byte, len(ip)) + copy(b, ip) + return b +} + +// inc increments a big-endian byte-encoded address in place. +func inc(b []byte) { + for i := len(b) - 1; i >= 0; i-- { + b[i]++ + if b[i] != 0 { + break + } + } +} + +func compareBytes(a, b []byte) int { + for i := range a { + switch { + case a[i] < b[i]: + return -1 + case a[i] > b[i]: + return 1 + } + } + return 0 +} diff --git a/pkg/ptr/ptr_test.go b/pkg/ptr/ptr_test.go new file mode 100644 index 00000000..9a752e62 --- /dev/null +++ b/pkg/ptr/ptr_test.go @@ -0,0 +1,121 @@ +package ptr + +import ( + "context" + "net" + "testing" +) + +func collect(t *testing.T, targets ...string) []string { + t.Helper() + var got []string + if err := Expand(targets, func(name string) bool { + got = append(got, name) + return true + }); err != nil { + t.Fatalf("Expand(%v): %v", targets, err) + } + return got +} + +func TestReverseName(t *testing.T) { + name, err := ReverseName(net.ParseIP("1.2.3.4")) + if err != nil { + t.Fatal(err) + } + if name != "4.3.2.1.in-addr.arpa." { + t.Fatalf("got %q", name) + } +} + +func TestSingleIP(t *testing.T) { + got := collect(t, "192.0.2.5") + if len(got) != 1 || got[0] != "5.2.0.192.in-addr.arpa." { + t.Fatalf("got %v", got) + } +} + +func TestExpandCIDRv4(t *testing.T) { + got := collect(t, "192.0.2.0/30") + want := []string{ + "0.2.0.192.in-addr.arpa.", + "1.2.0.192.in-addr.arpa.", + "2.2.0.192.in-addr.arpa.", + "3.2.0.192.in-addr.arpa.", + } + if len(got) != len(want) { + t.Fatalf("count: got %d want %d (%v)", len(got), len(want), got) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("idx %d: got %q want %q", i, got[i], want[i]) + } + } +} + +func TestExpandRange(t *testing.T) { + got := collect(t, "192.0.2.1-192.0.2.3") + if len(got) != 3 { + t.Fatalf("count: got %d (%v)", len(got), got) + } + if got[0] != "1.2.0.192.in-addr.arpa." || got[2] != "3.2.0.192.in-addr.arpa." { + t.Fatalf("got %v", got) + } +} + +func TestExpandCIDRv6Small(t *testing.T) { + got := collect(t, "2001:db8::/126") + if len(got) != 4 { + t.Fatalf("count: got %d (%v)", len(got), got) + } + for _, n := range got { + if len(n) < len("ip6.arpa.") || n[len(n)-len("ip6.arpa."):] != "ip6.arpa." { + t.Fatalf("not an ip6.arpa name: %q", n) + } + } +} + +func TestExpandCIDRv6TooLarge(t *testing.T) { + err := Expand([]string{"2001:db8::/32"}, func(string) bool { return true }) + if err == nil { + t.Fatal("expected error for oversized IPv6 CIDR") + } +} + +func TestEarlyStop(t *testing.T) { + count := 0 + err := Expand([]string{"10.0.0.0/8"}, func(string) bool { + count++ + return count < 5 // stop after 5 + }) + if err != nil { + t.Fatal(err) + } + if count != 5 { + t.Fatalf("expected early stop at 5, got %d", count) + } +} + +func TestInvalidTarget(t *testing.T) { + if err := Expand([]string{"not-an-ip"}, func(string) bool { return true }); err == nil { + t.Fatal("expected error for invalid target") + } +} + +func TestStreamCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + out := make(chan string, 2) + done := make(chan error, 1) + go func() { done <- Stream(ctx, []string{"10.0.0.0/8"}, out) }() + + <-out // consume one + cancel() + // drain so Stream isn't blocked on send + go func() { + for range out { + } + }() + if err := <-done; err != nil { + t.Fatalf("Stream returned error: %v", err) + } +} diff --git a/pkg/resolve/bench_test.go b/pkg/resolve/bench_test.go new file mode 100644 index 00000000..69bf843b --- /dev/null +++ b/pkg/resolve/bench_test.go @@ -0,0 +1,179 @@ +package resolve + +import ( + "context" + "flag" + "fmt" + "os" + "sync/atomic" + "testing" + "time" + + "github.com/miekg/dns" + "github.com/projectdiscovery/shuffledns/internal/simdns" +) + +// This file benchmarks the native resolver against a battery of loopback UDP +// DNS servers that simulate remote resolvers (see internal/simdns). No traffic +// leaves the host. It is gated behind RESOLVE_BENCH so `go test ./...` stays +// fast. Example: +// +// RESOLVE_BENCH=1 go test ./pkg/resolve -run TestResolverBenchmark -v \ +// -bench.names 200000 -bench.resolvers 16 -bench.hit 5 + +var ( + benchNames = flag.Int("bench.names", 50000, "number of names to resolve per scenario") + benchResolvers = flag.Int("bench.resolvers", 8, "number of simulated loopback resolvers") + benchHit = flag.Int("bench.hit", 5, "percentage of names that resolve (rest are NXDOMAIN)") + benchConc = flag.Int("bench.concurrency", 10000, "client in-flight concurrency") + benchSockets = flag.Int("bench.sockets", 8, "client udp socket count") + benchBatch = flag.String("bench.batch", "off", "batching mode: off | on | adaptive") + benchHealth = flag.Bool("bench.health", false, "enable per-resolver health scoring") + benchAdaptConc = flag.Bool("bench.adaptconc", false, "enable adaptive concurrency") + benchHooks = flag.Bool("bench.hooks", false, "attach no-op lifecycle hooks (measure hook overhead)") +) + +// noopHooks returns a full set of no-op lifecycle hooks for overhead testing. +func noopHooks() Hooks { + return Hooks{ + OnQuery: func(QueryInfo) {}, + OnRetry: func(QueryInfo) {}, + OnResponse: func(QueryInfo, *dns.Msg) {}, + OnTimeout: func(QueryInfo) {}, + OnTruncated: func(QueryInfo) {}, + OnCrossCheckFailed: func(string, []string, []string) {}, + OnResolverState: func(string, bool) {}, + } +} + +func benchBatchMode(s string) BatchMode { + switch s { + case "on": + return BatchEnabled + case "adaptive": + return BatchAdaptive + default: + return BatchDisabled + } +} + +// scenario pairs a human-readable name with simulated network conditions. +type scenario struct { + name string + cfg simdns.Config +} + +func benchScenarios(hit int) []scenario { + return []scenario{ + {"lan-fast", simdns.Config{BaseLatency: 200 * time.Microsecond, Jitter: 300 * time.Microsecond, HitPercent: hit}}, + {"wan-typical", simdns.Config{BaseLatency: 15 * time.Millisecond, Jitter: 10 * time.Millisecond, LossRate: 0.005, HitPercent: hit}}, + {"wan-lossy", simdns.Config{BaseLatency: 25 * time.Millisecond, Jitter: 20 * time.Millisecond, LossRate: 0.05, ServfailRate: 0.02, HitPercent: hit}}, + {"rate-limited", simdns.Config{BaseLatency: 10 * time.Millisecond, Jitter: 10 * time.Millisecond, QPSPerServer: 3000, HitPercent: hit}}, + } +} + +func TestResolverBenchmark(t *testing.T) { + if testing.Short() { + t.Skip("skipping resolver benchmark in -short mode") + } + if os.Getenv("RESOLVE_BENCH") == "" { + t.Skip("set RESOLVE_BENCH=1 to run the loopback resolver benchmark") + } + + hit := *benchHit + + t.Logf("names=%d resolvers=%d hit=%d%% concurrency=%d sockets=%d", + *benchNames, *benchResolvers, hit, *benchConc, *benchSockets) + t.Logf("%-14s %10s %10s %10s %10s %10s %10s", + "scenario", "ttfr", "wall", "qps", "resolved", "nxdomain", "failed") + + for _, sc := range benchScenarios(hit) { + runBenchmarkScenario(t, sc) + } +} + +func runBenchmarkScenario(t *testing.T, sc scenario) { + t.Helper() + + battery, err := simdns.Start(*benchResolvers, sc.cfg) + if err != nil { + t.Fatalf("could not start sim resolvers: %v", err) + } + defer battery.Stop() + + // size the per-attempt timeout to a few RTTs so lossy scenarios still + // complete via retransmission instead of stalling. + rtt := sc.cfg.BaseLatency + sc.cfg.Jitter + timeout := 6 * rtt + if timeout < 500*time.Millisecond { + timeout = 500 * time.Millisecond + } + + var resolved, nxdomain, failed atomic.Int64 + var ttfrNanos atomic.Int64 // 0 until first result + start := time.Now() + + var hooks Hooks + if *benchHooks { + hooks = noopHooks() + } + + client, err := New(Options{ + Resolvers: battery.Addrs, + Concurrency: *benchConc, + SocketCount: *benchSockets, + Timeout: timeout, + MaxRetries: 5, + Batch: benchBatchMode(*benchBatch), + ResolverHealth: *benchHealth, + AdaptiveConcurrency: *benchAdaptConc, + Hooks: hooks, + OnResult: func(r Result) { + ttfrNanos.CompareAndSwap(0, int64(time.Since(start))) + switch { + case r.Rcode == dns.RcodeSuccess && len(r.A) > 0: + resolved.Add(1) + case r.Rcode == dns.RcodeNameError: + nxdomain.Add(1) + default: + failed.Add(1) + } + }, + OnError: func(string, error) { failed.Add(1) }, + }) + if err != nil { + t.Fatalf("could not create client: %v", err) + } + defer client.Close() + + total := *benchNames + input := make(chan string, 4096) + go func() { + defer close(input) + for i := 0; i < total; i++ { + input <- fmt.Sprintf("host%d.bench.example.com", i) + } + }() + + if err := client.Run(context.Background(), input); err != nil { + t.Fatalf("run failed: %v", err) + } + wall := time.Since(start) + + qps := float64(total) / wall.Seconds() + ttfr := time.Duration(ttfrNanos.Load()) + + t.Logf("%-14s %10s %10s %10.0f %10d %10d %10d", + sc.name, + ttfr.Round(100*time.Microsecond), + wall.Round(time.Millisecond), + qps, + resolved.Load(), + nxdomain.Load(), + failed.Load(), + ) + s := battery.Stats + t.Logf(" server-side: queries=%d answered=%d dropped=%d servfail=%d ratelimited=%d", + s.Queries.Load(), s.Answered.Load(), s.Dropped.Load(), + s.Servfail.Load(), s.RateLimited.Load()) +} diff --git a/pkg/resolve/health.go b/pkg/resolve/health.go new file mode 100644 index 00000000..29e7e695 --- /dev/null +++ b/pkg/resolve/health.go @@ -0,0 +1,72 @@ +package resolve + +import ( + "math" + "math/rand/v2" + "sync/atomic" +) + +// resolverHealth tracks a smoothed success rate per resolver so that failing or +// rate-limiting resolvers are de-weighted in selection and recover over time. +// This addresses massdns's own TODO of avoiding resolvers that refuse or stop +// answering after a while. +type resolverHealth struct { + n int + weight []atomic.Uint64 // float64 bits, EWMA success in [0,1]; starts at 1.0 +} + +const ( + healthAlpha = 0.1 // EWMA smoothing for per-resolver success rate + healthThreshold = 0.5 // success-rate boundary between healthy and unhealthy +) + +func newResolverHealth(n int) *resolverHealth { + h := &resolverHealth{n: n, weight: make([]atomic.Uint64, n)} + for i := range h.weight { + h.weight[i].Store(math.Float64bits(1.0)) + } + return h +} + +func (h *resolverHealth) score(i int) float64 { + return math.Float64frombits(h.weight[i].Load()) +} + +// record nudges a resolver's success EWMA toward 1 (answered) or 0 (failed). +// It returns whether the resolver crossed the healthy threshold and the new +// healthy state, so callers can emit state-change events. +func (h *resolverHealth) record(i int, ok bool) (transitioned, healthy bool) { + if i < 0 || i >= h.n { + return false, true + } + var target float64 + if ok { + target = 1.0 + } + for { + oldBits := h.weight[i].Load() + old := math.Float64frombits(oldBits) + nw := old + healthAlpha*(target-old) + if h.weight[i].CompareAndSwap(oldBits, math.Float64bits(nw)) { + wasHealthy := old >= healthThreshold + nowHealthy := nw >= healthThreshold + return wasHealthy != nowHealthy, nowHealthy + } + } +} + +// pick selects a resolver using the power-of-two-choices: sample two at random +// and keep the healthier one. This is lock-free, needs no running totals, and +// naturally steers load away from unhealthy resolvers while still probing them +// occasionally (so they can recover). +func (h *resolverHealth) pick() int { + if h.n == 1 { + return 0 + } + a := rand.IntN(h.n) + b := rand.IntN(h.n) + if h.score(b) > h.score(a) { + return b + } + return a +} diff --git a/pkg/resolve/limiter.go b/pkg/resolve/limiter.go new file mode 100644 index 00000000..6497e5f4 --- /dev/null +++ b/pkg/resolve/limiter.go @@ -0,0 +1,78 @@ +package resolve + +import ( + "context" + "time" +) + +// limiter is a simple token-bucket rate limiter that refills at a fixed QPS. +// take() blocks until a token is available. It is intentionally lightweight to +// avoid a hard dependency while keeping outbound query rate bounded. +type limiter struct { + tokens chan struct{} + qps int + stopCh chan struct{} +} + +func newLimiter(qps int) *limiter { + if qps <= 0 { + return nil + } + // burst capacity of one second worth of tokens, capped for memory sanity + burst := qps + if burst > 100000 { + burst = 100000 + } + return &limiter{ + tokens: make(chan struct{}, burst), + qps: qps, + stopCh: make(chan struct{}), + } +} + +func (l *limiter) start(ctx context.Context) { + // refill in small slices to smooth out bursts + slices := 100 + perTick := l.qps / slices + if perTick < 1 { + perTick = 1 + slices = l.qps + } + interval := time.Second / time.Duration(slices) + if interval <= 0 { + interval = time.Millisecond + } + + go func() { + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-l.stopCh: + return + case <-ticker.C: + for i := 0; i < perTick; i++ { + select { + case l.tokens <- struct{}{}: + default: + // bucket full + } + } + } + } + }() +} + +func (l *limiter) take() { + <-l.tokens +} + +func (l *limiter) stop() { + select { + case <-l.stopCh: + default: + close(l.stopCh) + } +} diff --git a/pkg/resolve/privdrop_test.go b/pkg/resolve/privdrop_test.go new file mode 100644 index 00000000..93711016 --- /dev/null +++ b/pkg/resolve/privdrop_test.go @@ -0,0 +1,28 @@ +package resolve + +import ( + "os" + "runtime" + "testing" +) + +func TestDropPrivilegesNonRootNoop(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("windows privilege model differs") + } + if os.Geteuid() == 0 { + t.Skip("running as root; noop path not exercised") + } + if err := DropPrivileges("", "", false); err != nil { + t.Fatal(err) + } + if err := DropPrivileges("nobody", "nobody", false); err != nil { + t.Fatal(err) + } +} + +func TestDropPrivilegesKeepRoot(t *testing.T) { + if err := DropPrivileges("does-not-exist", "does-not-exist", true); err != nil { + t.Fatal(err) + } +} diff --git a/pkg/resolve/privdrop_unix.go b/pkg/resolve/privdrop_unix.go new file mode 100644 index 00000000..2a7e2afe --- /dev/null +++ b/pkg/resolve/privdrop_unix.go @@ -0,0 +1,83 @@ +//go:build unix + +package resolve + +import ( + "fmt" + "os/user" + "strconv" + + "golang.org/x/sys/unix" +) + +// DropPrivileges drops root after sockets are opened (massdns --drop-user / +// --drop-group). No-op when not running as root. When keepRoot is true (massdns +// --root), privileges are left unchanged. +func DropPrivileges(username, groupname string, keepRoot bool) error { + if unix.Geteuid() != 0 { + return nil + } + if keepRoot { + return nil + } + if username == "" { + username = "nobody" + } + if groupname == "" { + groupname = "nobody" + } + + uid, err := lookupUID(username) + if err != nil { + return err + } + gid, err := lookupGID(groupname) + if err != nil { + return err + } + // Drop root's supplementary groups before setgid/setuid; otherwise the + // process keeps root's group memberships after the drop (a classic + // incomplete-privilege-drop gap). Must happen while still privileged. + if err := unix.Setgroups([]int{gid}); err != nil { + return fmt.Errorf("setgroups(%s): %w", groupname, err) + } + if err := unix.Setgid(gid); err != nil { + return fmt.Errorf("setgid(%s): %w", groupname, err) + } + if err := unix.Setuid(uid); err != nil { + return fmt.Errorf("setuid(%s): %w", username, err) + } + return nil +} + +func lookupUID(name string) (int, error) { + u, err := user.Lookup(name) + if err != nil { + // massdns falls back to 65534 when the default nobody user is missing, + // but errors when an explicit --drop-user is unknown. + if name == "nobody" { + return 65534, nil + } + return 0, fmt.Errorf("user %q does not exist", name) + } + id, err := strconv.Atoi(u.Uid) + if err != nil { + return 0, err + } + return id, nil +} + +func lookupGID(name string) (int, error) { + g, err := user.LookupGroup(name) + if err != nil { + if name == "nobody" || name == "nogroup" { + return 65534, nil + } + return 0, fmt.Errorf("group %q does not exist", name) + } + id, err := strconv.Atoi(g.Gid) + if err != nil { + return 0, err + } + return id, nil +} diff --git a/pkg/resolve/privdrop_windows.go b/pkg/resolve/privdrop_windows.go new file mode 100644 index 00000000..de3890ac --- /dev/null +++ b/pkg/resolve/privdrop_windows.go @@ -0,0 +1,16 @@ +//go:build windows + +package resolve + +import "fmt" + +// DropPrivileges is not supported on Windows. +func DropPrivileges(username, groupname string, keepRoot bool) error { + if keepRoot { + return nil + } + if username != "" || groupname != "" { + return fmt.Errorf("privilege drop is not supported on Windows") + } + return nil +} diff --git a/pkg/resolve/raw6_linux.go b/pkg/resolve/raw6_linux.go new file mode 100644 index 00000000..41122420 --- /dev/null +++ b/pkg/resolve/raw6_linux.go @@ -0,0 +1,126 @@ +//go:build linux + +package resolve + +import ( + "encoding/binary" + "errors" + "fmt" + "net" + "syscall" + + "golang.org/x/sys/unix" +) + +const rawSrcPort = 666 // massdns hard-coded UDP source port for HDRINCL sends + +// openRawUDPv6 creates a SOCK_RAW IPPROTO_UDP socket with IPV6_HDRINCL so we can +// forge the IPv6 source address (massdns --rand-src-ipv6). Requires CAP_NET_RAW. +func openRawUDPv6() (int, error) { + fd, err := unix.Socket(unix.AF_INET6, unix.SOCK_RAW, unix.IPPROTO_UDP) + if err != nil { + return -1, fmt.Errorf("raw ipv6 udp socket: %w (need CAP_NET_RAW?)", err) + } + if err := unix.SetsockoptInt(fd, unix.IPPROTO_IPV6, unix.IPV6_HDRINCL, 1); err != nil { + _ = unix.Close(fd) + return -1, fmt.Errorf("IPV6_HDRINCL: %w", err) + } + return fd, nil +} + +func closeRawFD(fd int) { + if fd >= 0 { + _ = unix.Close(fd) + } +} + +func setRawRecvBuffer(fd, n int) error { + if fd < 0 || n <= 0 { + return nil + } + return unix.SetsockoptInt(fd, unix.SOL_SOCKET, unix.SO_RCVBUF, n) +} + +func setRawSendBuffer(fd, n int) error { + if fd < 0 || n <= 0 { + return nil + } + return unix.SetsockoptInt(fd, unix.SOL_SOCKET, unix.SO_SNDBUF, n) +} + +// writeRawUDPv6 writes an IPv6+UDP+payload datagram with an explicit source +// address via a HDRINCL raw socket. sport is typically a fixed ephemeral-ish +// port (massdns uses 666). +func writeRawUDPv6(fd int, src, dst net.IP, sport, dport uint16, payload []byte) error { + src16 := src.To16() + dst16 := dst.To16() + if src16 == nil || dst16 == nil { + return errors.New("raw send requires IPv6 addresses") + } + udpLen := 8 + len(payload) + total := 40 + udpLen + buf := make([]byte, total) + + // Build a buffer laid out so the IPv6 UDP pseudo-header checksum can be + // computed in place (same trick as massdns write_raw_header), then fix the + // real IPv6 header fields. + binary.BigEndian.PutUint16(buf[4:6], uint16(udpLen)) + buf[6] = 0 + buf[7] = unix.IPPROTO_UDP // temp: next-header slot holds protocol for checksum + copy(buf[8:24], src16) + copy(buf[24:40], dst16) + binary.BigEndian.PutUint16(buf[40:42], sport) + binary.BigEndian.PutUint16(buf[42:44], dport) + binary.BigEndian.PutUint16(buf[44:46], uint16(udpLen)) + binary.BigEndian.PutUint16(buf[46:48], 0) + copy(buf[48:], payload) + sum := ipChecksum(buf) + binary.BigEndian.PutUint16(buf[46:48], sum) + + buf[0] = 0x60 // version 6 + buf[1], buf[2], buf[3] = 0, 0, 0 + buf[6] = unix.IPPROTO_UDP + buf[7] = 255 + + sa := &unix.SockaddrInet6{} + copy(sa.Addr[:], dst16) + // Raw IPv6 sendto rejects non-zero ports on some kernels (same as massdns). + return unix.Sendto(fd, buf, 0, sa) +} + +// readRawUDPv6 reads from a raw IPv6 UDP socket. On Linux the datagram starts +// with the 8-byte UDP header; we strip it and recover the source port. +func readRawUDPv6(fd int, buf []byte) (payload []byte, addr *net.UDPAddr, err error) { + n, from, err := unix.Recvfrom(fd, buf, 0) + if err != nil { + if errors.Is(err, syscall.EAGAIN) || errors.Is(err, syscall.EWOULDBLOCK) { + return nil, nil, err + } + return nil, nil, err + } + if n < 8 { + return nil, nil, errors.New("short raw udp read") + } + sport := binary.BigEndian.Uint16(buf[0:2]) + ip := net.IPv6zero + if sa6, ok := from.(*unix.SockaddrInet6); ok { + ip = make(net.IP, 16) + copy(ip, sa6.Addr[:]) + } + return buf[8:n], &net.UDPAddr{IP: ip, Port: int(sport)}, nil +} + +// ipChecksum computes the Internet checksum over buf (as used for IPv6 UDP). +func ipChecksum(buf []byte) uint16 { + var sum uint32 + for i := 0; i+1 < len(buf); i += 2 { + sum += uint32(buf[i])<<8 | uint32(buf[i+1]) + } + if len(buf)%2 == 1 { + sum += uint32(buf[len(buf)-1]) << 8 + } + for sum > 0xffff { + sum = (sum & 0xffff) + (sum >> 16) + } + return ^uint16(sum) +} \ No newline at end of file diff --git a/pkg/resolve/raw6_stub.go b/pkg/resolve/raw6_stub.go new file mode 100644 index 00000000..b843faea --- /dev/null +++ b/pkg/resolve/raw6_stub.go @@ -0,0 +1,28 @@ +//go:build !linux + +package resolve + +import ( + "errors" + "net" +) + +const rawSrcPort = 666 + +func openRawUDPv6() (int, error) { + return -1, errors.New("--rand-src-ipv6 is only supported on Linux") +} + +func closeRawFD(fd int) {} + +func setRawRecvBuffer(fd, n int) error { return nil } + +func setRawSendBuffer(fd, n int) error { return nil } + +func writeRawUDPv6(fd int, src, dst net.IP, sport, dport uint16, payload []byte) error { + return errors.New("--rand-src-ipv6 is only supported on Linux") +} + +func readRawUDPv6(fd int, buf []byte) ([]byte, *net.UDPAddr, error) { + return nil, nil, errors.New("--rand-src-ipv6 is only supported on Linux") +} diff --git a/pkg/resolve/resolve.go b/pkg/resolve/resolve.go new file mode 100644 index 00000000..399d37ed --- /dev/null +++ b/pkg/resolve/resolve.go @@ -0,0 +1,1716 @@ +// Package resolve implements a high-throughput asynchronous DNS stub +// resolver in pure Go. It is the native replacement for the massdns binary +// previously shelled out to by shuffledns. +// +// Design (mirrors massdns rather than the naive goroutine-per-query model): +// - a small fixed pool of UDP sockets, each drained by a single reader +// goroutine, so concurrency is bounded by an in-flight map and NOT by +// the number of goroutines; +// - queries are tracked in a per-socket in-flight table keyed by DNS +// transaction id; responses are correlated back by id + question name; +// - a timeout wheel retransmits unanswered queries to a different resolver +// up to a retry budget, and retries REFUSED/SERVFAIL like massdns does; +// - results stream out through a callback as soon as they arrive (no temp +// files, no text parsing), which is what gives a low time-to-first-result. +package resolve + +import ( + "context" + "errors" + "fmt" + "math" + "math/rand/v2" + "net" + "runtime" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/miekg/dns" + "golang.org/x/net/ipv4" + "golang.org/x/net/ipv6" +) + +// Result is a single resolved answer streamed to the caller. +type Result struct { + Name string // queried hostname (no trailing dot) + Type uint16 // dns query type (A, AAAA, ...) + Rcode int // dns response code + A []string // A records + AAAA []string // AAAA records + CNAME []string // CNAME targets + NS []string // NS targets + PTR []string // PTR targets + MX []string // MX exchange hosts + TXT []string // TXT strings + SOA []string // SOA primary nameservers + Resolver string // resolver that answered + + // Msg is the full parsed response message (all sections, TTLs, classes). + // It is retained so callers (e.g. the output formatter) can render the + // complete answer faithfully; it may be nil for synthetic results. + Msg *dns.Msg + // Timestamp is when the response was accepted as final (wall clock). + Timestamp time.Time +} + +// Options configures the resolver. +type Options struct { + // Resolvers is the list of recursive resolver addresses (host or host:port). + Resolvers []string + // QueryType is the dns record type to request (default dns.TypeA). + QueryType uint16 + // MaxRetries is the number of times a query is retransmitted (to a + // rotated resolver) before giving up. Default 3. + MaxRetries int + // Timeout is the per-attempt timeout before retransmission. Default 4s. + Timeout time.Duration + // Concurrency caps the number of in-flight queries. Default 10000. + Concurrency int + // QPS optionally rate-limits outbound queries (0 = unlimited). + QPS int + // SocketCount is the number of UDP sockets to spread load across. + // Default max(8, GOMAXPROCS) so the read path scales across cores. + SocketCount int + // Batch selects the datagram batching strategy (sendmmsg/recvmmsg, Linux + // only, IPv4-only resolver set). See BatchMode. Default BatchDisabled. + Batch BatchMode + // BatchSize is the maximum number of datagrams sent/received per syscall + // when batching is active. Default 64. + BatchSize int + + // Interval is the timeout-loop scan period (how promptly a lost packet is + // retransmitted). 0 = derive from Timeout. Mirrors massdns -i/--interval. + Interval time.Duration + + // NoRecurse sends non-recursive queries (RD=0), useful for cache snooping + // and probing authoritative servers directly. Mirrors massdns --norecurse. + NoRecurse bool + // UDPSize is the EDNS0 advertised UDP payload size. 0 = default (1232), + // which lets resolvers return larger answers without truncation. A value + // below 512 disables EDNS0 entirely. + UDPSize int + // DisableIPVerification turns off matching a reply's source address against + // the resolver the query was sent to. Verification is ON by default and + // guards against off-path answer spoofing (massdns --verify-ip is opt-in; + // here it is opt-out because we always send to a known address). + DisableIPVerification bool + // DisableTCPFallback disables re-querying over TCP when a UDP response has + // the truncation (TC) bit set. Fallback is ON by default. + DisableTCPFallback bool + // Sticky retransmits to the same resolver instead of rotating. Mirrors + // massdns --sticky. + Sticky bool + // RetryRcodes lists the response codes that trigger a retry. When nil the + // massdns default is used: retry everything except NOERROR and NXDOMAIN. + RetryRcodes []int + + // ResolverHealth enables per-resolver health scoring: resolvers that time + // out or error are de-weighted in selection and recover over time. + ResolverHealth bool + // AdaptiveConcurrency lets the controller shrink/grow the in-flight cap in + // response to packet loss (AIMD), preventing resolver/buffer flooding. + // Requires Batch == BatchAdaptive (the controller drives both). + AdaptiveConcurrency bool + // CrossCheck re-queries each positive answer on a second, randomly chosen + // resolver and drops results the two disagree on (basic poisoning/spam + // detection). Doubles query volume for names that resolve. + CrossCheck bool + // ExtendedInput parses each input line as "name [resolver ...]"; the + // per-name resolvers are tried (in order) before falling back to the global + // resolver pool. Mirrors massdns --extended-input. + ExtendedInput bool + // Predictable picks resolvers sequentially instead of randomly. Mirrors + // massdns --predictable (useful for resolver tests). + Predictable bool + // BindAddr is an optional local UDP bind address (host, host:port, or :port). + // Empty means the kernel picks an ephemeral port on all interfaces. + // Mirrors massdns -b/--bindto. + BindAddr string + // ReadBuffer, when > 0, sets SO_RCVBUF on each UDP socket (massdns --rcvbuf). + // When 0 a large default (8 MiB) is used. + ReadBuffer int + // WriteBuffer, when > 0, sets SO_SNDBUF on each UDP socket (massdns --sndbuf). + WriteBuffer int + // RandSrcIPv6 is a CIDR prefix used to forge a random IPv6 source address + // per query (massdns --rand-src-ipv6). Linux-only; requires CAP_NET_RAW. + // Incompatible with BindAddr. + RandSrcIPv6 string + // RandSrcIPv6File loads discrete IPv6 source addresses (one per line) for + // the same purpose (massdns --rand-src-ipv6-file). Mutually exclusive with + // RandSrcIPv6. + RandSrcIPv6File string + + // OnResult is invoked for every final answer (success or definitive + // failure such as NXDOMAIN). It must be safe for concurrent use. + OnResult func(Result) + // OnError is invoked when a query is abandoned after exhausting retries. + OnError func(name string, err error) + // OnProgress is invoked periodically (every Interval-ish tick) with a + // snapshot of resolver statistics, for status reporting. Optional. + OnProgress func(Stats) + + // Hooks holds optional fine-grained lifecycle callbacks for SDK/observability + // use. They are independent of OnResult/OnError/OnProgress and may all be + // left nil. See Hooks. + Hooks Hooks +} + +// Stats is a point-in-time snapshot of resolver activity. +type Stats struct { + Queries int64 // distinct names dispatched + Retransmits int64 // total retransmissions + Answered int64 // final answers delivered (success or definitive) + InflightDepth int // queries currently outstanding + ConcurrencyCap int // current adaptive in-flight cap + LossRate float64 // most recent interval loss estimate + RTT time.Duration + BatchActive bool +} + +// QueryInfo describes a single query attempt passed to lifecycle hooks. +type QueryInfo struct { + Name string // queried hostname (no trailing dot) + Type uint16 // dns query type + Attempt int // zero-based attempt number (0 = initial send) + Resolver string // resolver address this attempt targets +} + +// Hooks is a set of optional callbacks for observing the resolver's internals, +// intended for embedding the resolver as an SDK (progress UIs, metrics, +// tracing, custom retry/poisoning telemetry, etc.). +// +// IMPORTANT: hooks fire on hot paths (some per query/response). They must be +// cheap, non-blocking, and safe for concurrent use; offload heavy work to a +// channel or worker. Any hook may be nil. Hooks never alter resolver behaviour +// — they are observation points only. +type Hooks struct { + // OnQuery fires when a query is first put on the wire (initial send). + OnQuery func(QueryInfo) + // OnRetry fires before each retransmission (timeout or bad-rcode driven). + OnRetry func(QueryInfo) + // OnResponse fires for every response matched to an in-flight query, before + // the retry-or-finalize decision, exposing the raw message (read-only). + OnResponse func(QueryInfo, *dns.Msg) + // OnTimeout fires when an attempt's deadline expires (before retry/abandon). + OnTimeout func(QueryInfo) + // OnTruncated fires when a TC (truncated) response triggers TCP fallback. + OnTruncated func(QueryInfo) + // OnCrossCheckFailed fires when cross-resolver verification rejects a result; + // primary/secondary are the disagreeing A-record sets. + OnCrossCheckFailed func(name string, primary, secondary []string) + // OnResolverState fires when a resolver's health crosses the healthy + // threshold (requires ResolverHealth). healthy is the new state. + OnResolverState func(resolver string, healthy bool) +} + +// fire helpers keep the hot paths branch-cheap when hooks are unset. +func (c *Client) fireQuery(q *query) { + if c.hooks.OnQuery != nil { + c.hooks.OnQuery(c.queryInfo(q)) + } +} + +func (c *Client) fireRetry(q *query) { + if c.hooks.OnRetry != nil { + c.hooks.OnRetry(c.queryInfo(q)) + } +} + +func (c *Client) queryInfo(q *query) QueryInfo { + res := "" + if q.addr != nil { + res = q.addr.String() + } + return QueryInfo{Name: q.name, Type: c.opts.QueryType, Attempt: q.attempts, Resolver: res} +} + +// BatchMode selects how datagram batching (Linux sendmmsg/recvmmsg) is used. +// +// Batching amortizes the per-datagram syscall cost and helps on high-latency / +// bursty links where many packets cluster in time, but it is counterproductive +// on loopback / low-RTT links (tiny batches pay the message-array setup cost +// and add first-response latency). It also makes sends burstier, which can +// worsen loss when a resolver or kernel buffer is already saturated. +type BatchMode int + +const ( + // BatchDisabled always uses one datagram per syscall (default). Identical + // to the resolver's behaviour without any batching support. + BatchDisabled BatchMode = iota + // BatchEnabled forces batching on (Linux + IPv4 resolvers only). + BatchEnabled + // BatchAdaptive turns batching on/off at runtime based on observed RTT, + // in-flight depth and packet loss: it engages when the pipeline is deep and + // latency is high enough to fill batches, and backs off on low-RTT links or + // when loss climbs (to avoid making sends burstier). + BatchAdaptive +) + +const ( + defaultMaxRetries = 3 + defaultTimeout = 4 * time.Second + defaultConcurrency = 10000 + defaultSocketCount = 8 + defaultBatchSize = 64 + defaultUDPSize = 1232 // conservative EDNS0 payload (avoids v4/v6 fragmentation) + maxTxIDAttempts = 64 + + // adaptive batching controller thresholds (with hysteresis to avoid flapping) + adaptTick = 200 * time.Millisecond + adaptRTTOn = 3 * time.Millisecond // engage batching above this smoothed RTT + adaptRTTOff = 1 * time.Millisecond // disengage below this smoothed RTT + adaptLossOff = 0.15 // disengage when interval loss exceeds this + minBatchToMMSG = 4 // only use sendmmsg for batches at least this big + adaptDepthRatio = 2 // need depth >= ratio*batchSize to engage + + // adaptive concurrency (AIMD) thresholds + concLossHigh = 0.10 // multiplicative decrease above this interval loss + concLossLow = 0.02 // additive increase below this interval loss + concDecrease = 0.75 // cap *= concDecrease on high loss + concMinRatio = 0.05 // never shrink below this fraction of Concurrency +) + +// query holds the in-flight state for a single outstanding name. +type query struct { + name string + fqdn string + txid uint16 + resolverIdx int + addr *net.UDPAddr // resolver this attempt was sent to (for verify + retransmit) + extra []*net.UDPAddr // per-name resolvers (extended-input); tried before the pool + sentAt time.Time + deadline time.Time + attempts int + sock *socket +} + +// batchConn is the subset of ipv4.PacketConn / ipv6.PacketConn used for batched +// I/O. ipv4.Message and ipv6.Message are both aliases of the same underlying +// socket.Message type, so a single interface works for either family. +type batchConn interface { + ReadBatch(ms []ipv4.Message, flags int) (int, error) + WriteBatch(ms []ipv4.Message, flags int) (int, error) +} + +// socket is a UDP socket plus its in-flight table. +type socket struct { + conn *net.UDPConn + rawFD int // >=0: Linux SOCK_RAW IPv6 for --rand-src-ipv6; conn is nil + family int // 4, 6, or 0 (dual-stack) + pc batchConn // batch (sendmmsg/recvmmsg) wrapper; nil if batching unavailable + mu sync.Mutex + inflight map[uint16]*query +} + +// Client is an asynchronous DNS stub resolver. +type Client struct { + opts Options + resolvers []string + resolverAddrs []*net.UDPAddr // pre-resolved, shared read-only across goroutines + sockets []*socket + sem chan struct{} // static in-flight cap (fast path) + dynSem *adaptiveSem // adjustable in-flight cap (AdaptiveConcurrency); nil otherwise + limiter *limiter + pending sync.WaitGroup + srcRand *srcRand + sockets4 []*socket // IPv4-capable sockets (subset of sockets) + sockets6 []*socket // IPv6-capable sockets (subset of sockets) + + rd bool // recursion desired flag for outgoing queries + udpSize uint16 + verify bool // verify reply source address + tcp bool // TCP fallback on truncation + sticky bool + predictable bool + predSeq atomic.Uint64 // monotonic index for Predictable resolver selection + extInput bool // parse per-name resolvers from input lines + retryRcode [16]bool // retryRcode[rcode] => retransmit on this response code + sockIdx atomic.Uint64 + + addrCache sync.Map // string -> *net.UDPAddr (extended-input resolver cache) + health *resolverHealth // nil unless ResolverHealth enabled + hooks Hooks // optional lifecycle callbacks + + // precomputed hot-path gates so the common (no-hook) path costs a single + // bool test instead of per-query nil checks / function calls. + hasSendHooks bool // OnQuery or OnRetry set + sampleRTT bool // RTT sampling needed (adaptive batch/concurrency or OnProgress) + + // batchCapable is true when the batch I/O path is usable (mode != disabled + // and a single-family resolver set). When true, the dispatch/read loops use + // the batch-aware path and consult batchActive to decide per operation. + batchCapable bool + batchSize int + // batchActive is flipped by the adaptive controller (or pinned on/off for + // the non-adaptive modes); send/read paths read it to choose mmsg vs single. + batchActive atomic.Bool + + // adaptive metrics (cheap atomics sampled by the controller) + statQueries atomic.Int64 // distinct names dispatched (initial sends) + statRetransmit atomic.Int64 // retransmissions (timeout + servfail/refused retries) + statAnswered atomic.Int64 // final answers delivered + lastRTTNanos atomic.Int64 // most recent observed round-trip time + lastLossBits atomic.Uint64 // most recent interval loss (float64 bits) for stats +} + +// New creates a resolver client from the given options. +func New(opts Options) (*Client, error) { + if len(opts.Resolvers) == 0 { + return nil, errors.New("no resolvers provided") + } + if opts.QueryType == 0 { + opts.QueryType = dns.TypeA + } + if opts.MaxRetries <= 0 { + opts.MaxRetries = defaultMaxRetries + } + if opts.Timeout <= 0 { + opts.Timeout = defaultTimeout + } + if opts.Concurrency <= 0 { + opts.Concurrency = defaultConcurrency + } + if opts.SocketCount <= 0 { + // scale the socket pool (and thus the number of independent reader + // goroutines / kernel receive queues) with the available cores. + opts.SocketCount = defaultSocketCount + if n := runtime.GOMAXPROCS(0); n > opts.SocketCount { + opts.SocketCount = n + } + } + if opts.BatchSize <= 0 { + opts.BatchSize = defaultBatchSize + } + if opts.UDPSize == 0 { + opts.UDPSize = defaultUDPSize + } + + resolvers := make([]string, 0, len(opts.Resolvers)) + resolverAddrs := make([]*net.UDPAddr, 0, len(opts.Resolvers)) + allIPv4, allIPv6 := true, true + for _, r := range opts.Resolvers { + r = strings.TrimSpace(r) + if r == "" { + continue + } + normalized := normalizeResolver(r) + // pre-resolve once at startup so the hot send() path never parses an + // address string or allocates a *net.UDPAddr per query/retransmit. + addr, err := net.ResolveUDPAddr("udp", normalized) + if err != nil { + return nil, err + } + if addr.IP.To4() == nil { + allIPv4 = false + } else { + allIPv6 = false + } + resolvers = append(resolvers, normalized) + resolverAddrs = append(resolverAddrs, addr) + } + if len(resolvers) == 0 { + return nil, errors.New("no valid resolvers provided") + } + + c := &Client{ + opts: opts, + resolvers: resolvers, + resolverAddrs: resolverAddrs, + batchSize: opts.BatchSize, + rd: !opts.NoRecurse, + verify: !opts.DisableIPVerification, + tcp: !opts.DisableTCPFallback, + sticky: opts.Sticky, + predictable: opts.Predictable, + extInput: opts.ExtendedInput, + hooks: opts.Hooks, + } + if opts.UDPSize >= 512 { + c.udpSize = uint16(opts.UDPSize) + } + if opts.RandSrcIPv6 != "" && opts.RandSrcIPv6File != "" { + return nil, errors.New("--rand-src-ipv6 cannot be used with --rand-src-ipv6-file") + } + if (opts.RandSrcIPv6 != "" || opts.RandSrcIPv6File != "") && opts.BindAddr != "" { + return nil, errors.New("--bindto and --rand-src-ipv6 cannot be used together") + } + switch { + case opts.RandSrcIPv6 != "": + sr, err := newSrcRandFromPrefix(opts.RandSrcIPv6) + if err != nil { + return nil, err + } + c.srcRand = sr + case opts.RandSrcIPv6File != "": + sr, err := newSrcRandFromFile(opts.RandSrcIPv6File) + if err != nil { + return nil, err + } + c.srcRand = sr + } + // Batching uses an ipv4/ipv6 PacketConn, so it requires a single-family + // resolver set (and a matching udp4/udp6 socket). Mixed sets always use the + // portable single-datagram path. Raw IPv6 source randomization also disables + // batching (HDRINCL path is single-datagram only). + c.batchCapable = opts.Batch != BatchDisabled && (allIPv4 || allIPv6) && c.srcRand == nil + + // precompute hot-path hook/sampling gates (see field docs). + c.hasSendHooks = opts.Hooks.OnQuery != nil || opts.Hooks.OnRetry != nil + c.sampleRTT = (c.batchCapable && opts.Batch == BatchAdaptive) || + opts.AdaptiveConcurrency || opts.OnProgress != nil + + // retry policy: explicit list, or massdns default (retry all but NOERROR/NXDOMAIN). + if len(opts.RetryRcodes) > 0 { + for _, rc := range opts.RetryRcodes { + if rc >= 0 && rc < len(c.retryRcode) { + c.retryRcode[rc] = true + } + } + } else { + for rc := range c.retryRcode { + c.retryRcode[rc] = true + } + c.retryRcode[dns.RcodeSuccess] = false + c.retryRcode[dns.RcodeNameError] = false + } + + if opts.ResolverHealth { + c.health = newResolverHealth(len(resolvers)) + } + + if opts.AdaptiveConcurrency { + c.dynSem = newAdaptiveSem(opts.Concurrency) + } else { + c.sem = make(chan struct{}, opts.Concurrency) + } + + // BatchEnabled pins batching on; BatchAdaptive starts off and lets the + // controller engage it when conditions warrant. + c.batchActive.Store(opts.Batch == BatchEnabled && c.batchCapable) + if opts.QPS > 0 { + c.limiter = newLimiter(opts.QPS) + } + + readBuf := 8 * 1024 * 1024 + if opts.ReadBuffer > 0 { + readBuf = opts.ReadBuffer + } + + // --rand-src-ipv6 uses Linux SOCK_RAW + IPV6_HDRINCL (massdns). Requires + // IPv6 resolvers and CAP_NET_RAW; batching stays off. + if c.srcRand != nil { + if !allIPv6 { + return nil, errors.New("--rand-src-ipv6 requires IPv6 resolvers") + } + for i := 0; i < opts.SocketCount; i++ { + fd, err := openRawUDPv6() + if err != nil { + c.closeSockets() + return nil, err + } + _ = setRawRecvBuffer(fd, readBuf) + if opts.WriteBuffer > 0 { + _ = setRawSendBuffer(fd, opts.WriteBuffer) + } + s := &socket{ + rawFD: fd, + family: 6, + inflight: make(map[uint16]*query), + } + c.sockets = append(c.sockets, s) + c.sockets6 = append(c.sockets6, s) + } + return c, nil + } + + network := "udp" + family := 0 + switch { + case c.batchCapable && allIPv4: + network = "udp4" + family = 4 + case c.batchCapable && allIPv6: + network = "udp6" + family = 6 + } + var bindAddr *net.UDPAddr + if opts.BindAddr != "" { + var err error + bindAddr, err = net.ResolveUDPAddr(network, normalizeBindAddr(opts.BindAddr)) + if err != nil { + return nil, fmt.Errorf("bind address: %w", err) + } + } + for i := 0; i < opts.SocketCount; i++ { + conn, err := net.ListenUDP(network, bindAddr) + if err != nil { + c.closeSockets() + return nil, err + } + _ = conn.SetReadBuffer(readBuf) + if opts.WriteBuffer > 0 { + _ = conn.SetWriteBuffer(opts.WriteBuffer) + } + s := &socket{ + conn: conn, + rawFD: -1, + family: family, + inflight: make(map[uint16]*query), + } + if c.batchCapable { + if allIPv6 { + s.pc = ipv6.NewPacketConn(conn) + } else { + s.pc = ipv4.NewPacketConn(conn) + } + } + c.sockets = append(c.sockets, s) + switch family { + case 4: + c.sockets4 = append(c.sockets4, s) + case 6: + c.sockets6 = append(c.sockets6, s) + default: + c.sockets4 = append(c.sockets4, s) + c.sockets6 = append(c.sockets6, s) + } + } + + return c, nil +} + +// normalizeBindAddr accepts host, host:port, or :port forms used by massdns -b. +func normalizeBindAddr(addr string) string { + addr = strings.TrimSpace(addr) + if addr == "" { + return ":0" + } + if _, _, err := net.SplitHostPort(addr); err == nil { + return addr + } + // bare IPv6 needs brackets before appending port + if ip := net.ParseIP(addr); ip != nil { + if ip.To4() == nil { + return "[" + ip.String() + "]:0" + } + return ip.String() + ":0" + } + // hostname without port + if !strings.Contains(addr, ":") { + return addr + ":0" + } + return addr +} + +// Run consumes hostnames from input and resolves them, blocking until input +// is closed and all in-flight queries have completed or the context is +// cancelled. It is safe to call Run only once per Client. +func (c *Client) Run(ctx context.Context, input <-chan string) error { + if c.limiter != nil { + c.limiter.start(ctx) + defer c.limiter.stop() + } + + readerCtx, cancelReaders := context.WithCancel(ctx) + defer cancelReaders() + + var readerWg sync.WaitGroup + for _, s := range c.sockets { + readerWg.Add(1) + go func(s *socket) { + defer readerWg.Done() + c.readLoop(readerCtx, s) + }(s) + } + + timeoutDone := make(chan struct{}) + go func() { + defer close(timeoutDone) + c.timeoutLoop(readerCtx) + }() + + if c.dynSem != nil { + c.dynSem.watch(readerCtx) + } + + // the controller drives adaptive batching, adaptive concurrency, and + // periodic progress reporting; start it if any of those is requested. + if (c.batchCapable && c.opts.Batch == BatchAdaptive) || c.dynSem != nil || c.opts.OnProgress != nil { + go c.controllerLoop(readerCtx) + } + + if c.batchCapable { + c.dispatchBatched(ctx, input) + } else { + c.dispatchSingle(ctx, input) + } + + // Wait for outstanding queries to drain (success, NXDOMAIN, or retry + // exhaustion all release the pending counter). + waitDone := make(chan struct{}) + go func() { + c.pending.Wait() + close(waitDone) + }() + + select { + case <-waitDone: + case <-ctx.Done(): + } + + // stop the timeout loop first so it cannot issue further retransmits, then + // close the sockets to immediately unblock the reader goroutines. + cancelReaders() + <-timeoutDone + c.closeSockets() + readerWg.Wait() + return ctx.Err() +} + +// Close releases the underlying sockets. It must be called when the client is +// no longer needed (Run does not close sockets so the client can be reused for +// diagnostics, but typical callers Close after Run). +func (c *Client) Close() { + c.closeSockets() +} + +func (c *Client) closeSockets() { + for _, s := range c.sockets { + if s == nil { + continue + } + if s.conn != nil { + _ = s.conn.Close() + s.conn = nil + } + if s.rawFD >= 0 { + closeRawFD(s.rawFD) + s.rawFD = -1 + } + } +} + +// pickSocket selects a UDP/raw socket that can reach addr's address family. +func (c *Client) pickSocket(addr *net.UDPAddr) *socket { + pool := c.sockets + if addr != nil { + if addr.IP.To4() == nil { + if len(c.sockets6) > 0 { + pool = c.sockets6 + } + } else if len(c.sockets4) > 0 { + pool = c.sockets4 + } + } + if len(pool) == 0 { + return nil + } + return pool[int(c.sockIdx.Add(1)-1)%len(pool)] +} + +func (c *Client) acquire(ctx context.Context) error { + if c.dynSem != nil { + return c.dynSem.acquire(ctx) + } + select { + case c.sem <- struct{}{}: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (c *Client) release() { + if c.dynSem != nil { + c.dynSem.release() + return + } + <-c.sem +} + +// pickResolver chooses the resolver index for an attempt. prevIdx is the +// resolver used by the previous attempt (-1 for the initial send). +func (c *Client) pickResolver(attempt, prevIdx int) int { + // sticky: keep hitting the same resolver on retries (massdns --sticky). + if c.sticky && attempt > 0 && prevIdx >= 0 { + return prevIdx + } + // predictable: sequential assignment across the pool (massdns --predictable). + if c.predictable { + return int(c.predSeq.Add(1)-1) % len(c.resolvers) + } + // health scoring de-weights failing resolvers via power-of-two-choices. + if c.health != nil { + return c.health.pick() + } + // default: rotate per attempt so retries hit a different server. + // math/rand/v2 top-level funcs are safe for concurrent use and use a + // per-P source, so there is no shared-lock contention here. + return (rand.IntN(len(c.resolvers)) + attempt) % len(c.resolvers) +} + +// queryBufPool recycles the small byte buffers used to encode outgoing queries +// so the hot send path does not allocate per query/retransmit. +var queryBufPool = sync.Pool{New: func() any { b := make([]byte, 0, 256); return &b }} + +// chooseTarget selects the resolver for an attempt. Per-name (extended-input) +// resolvers are tried first, in order, before falling back to the global pool. +// A returned resolverIdx of -1 means a per-name resolver (not health-tracked). +func (c *Client) chooseTarget(attempt, prevIdx int, extra []*net.UDPAddr) (int, *net.UDPAddr) { + if attempt < len(extra) { + return -1, extra[attempt] + } + idx := c.pickResolver(attempt, prevIdx) + return idx, c.resolverAddrs[idx] +} + +// register allocates a transaction id and inserts the in-flight query into the +// socket table. It returns the registered query, or nil if the id space is +// saturated (in which case the name is failed). +func (c *Client) register(s *socket, name, fqdn string, attempt, resolverIdx int, addr *net.UDPAddr, extra []*net.UDPAddr) *query { + s.mu.Lock() + txid, ok := c.allocTxID(s) + if !ok { + s.mu.Unlock() + c.fail(name, errors.New("transaction id space exhausted")) + return nil + } + now := time.Now() + q := &query{ + name: name, + fqdn: fqdn, + txid: txid, + resolverIdx: resolverIdx, + addr: addr, + extra: extra, + sentAt: now, + deadline: now.Add(c.opts.Timeout), + attempts: attempt, + sock: s, + } + s.inflight[txid] = q + s.mu.Unlock() + + if c.hasSendHooks { + if attempt == 0 { + c.fireQuery(q) + } else { + c.fireRetry(q) + } + } + return q +} + +// cachedAddr resolves a resolver spec to a UDP address, caching the result so +// repeated extended-input lines don't re-parse the same string. +func (c *Client) cachedAddr(spec string) *net.UDPAddr { + if v, ok := c.addrCache.Load(spec); ok { + return v.(*net.UDPAddr) + } + addr, err := net.ResolveUDPAddr("udp", normalizeResolver(spec)) + if err != nil { + c.addrCache.Store(spec, (*net.UDPAddr)(nil)) + return nil + } + c.addrCache.Store(spec, addr) + return addr +} + +// parseLine splits an input line into a name and (in extended-input mode) its +// per-name resolver addresses. +func (c *Client) parseLine(line string) (string, []*net.UDPAddr) { + if !c.extInput { + return strings.TrimSpace(line), nil + } + fields := strings.Fields(line) + if len(fields) == 0 { + return "", nil + } + var extra []*net.UDPAddr + for _, r := range fields[1:] { + if a := c.cachedAddr(r); a != nil { + extra = append(extra, a) + } + } + return fields[0], extra +} + +// encodeQuery writes the wire-format query for q into buf, using the hand-rolled +// fast path and falling back to miekg for names it cannot encode. Returns the +// packed bytes, or nil on a fatal encoding error. +func (c *Client) encodeQuery(buf []byte, q *query) []byte { + packed, ok := packQuery(buf[:0], q.txid, q.name, c.opts.QueryType, c.rd, c.udpSize) + if ok { + return packed + } + msg := new(dns.Msg) + msg.SetQuestion(q.fqdn, c.opts.QueryType) + msg.RecursionDesired = c.rd + msg.Id = q.txid + if c.udpSize >= 512 { + msg.SetEdns0(c.udpSize, false) + } + p, err := msg.Pack() + if err != nil { + return nil + } + return p +} + +// send builds and transmits a single query, registering it in the socket +// in-flight table. attempt is the zero-based retransmission count. It is used +// for retransmissions and on the non-batched (IPv6/mixed) path. +func (c *Client) send(s *socket, name string, attempt, prevIdx int, extra []*net.UDPAddr) { + if c.limiter != nil { + c.limiter.take() + } + + idx, addr := c.chooseTarget(attempt, prevIdx, extra) + if s == nil { + s = c.pickSocket(addr) + } + if s == nil { + c.fail(name, errors.New("no socket for resolver address family")) + return + } + q := c.register(s, name, dns.Fqdn(name), attempt, idx, addr, extra) + if q == nil { + return + } + + bufp := queryBufPool.Get().(*[]byte) + packed := c.encodeQuery(*bufp, q) + if packed == nil { + queryBufPool.Put(bufp) + c.remove(s, q.txid) + c.fail(name, errors.New("could not encode query")) + return + } + *bufp = packed // keep any grown backing array for reuse + + var werr error + if s.rawFD >= 0 { + src := c.srcRand.pick() + dport := uint16(q.addr.Port) + if dport == 0 { + dport = 53 + } + werr = writeRawUDPv6(s.rawFD, src, q.addr.IP, rawSrcPort, dport, packed) + } else { + _, werr = s.conn.WriteToUDP(packed, q.addr) + } + queryBufPool.Put(bufp) + if werr != nil { + // transient write error: let the timeout loop retry it + return + } +} + +// dispatchSingle sends one datagram per syscall. Used when batching is +// unavailable (IPv6/mixed resolver set). +func (c *Client) dispatchSingle(ctx context.Context, input <-chan string) { + for { + select { + case <-ctx.Done(): + return + case name, ok := <-input: + if !ok { + return + } + pname, extra := c.parseLine(name) + if pname == "" { + continue + } + if err := c.acquire(ctx); err != nil { + return + } + c.pending.Add(1) + c.statQueries.Add(1) + // socket is chosen inside send() once the target resolver family is known + c.send(nil, pname, 0, -1, extra) + } + } +} + +// dispatchBatched coalesces outbound queries and flushes them with a single +// sendmmsg (on Linux) per socket, amortizing the per-datagram syscall cost that +// otherwise caps single-threaded send throughput. Each batch targets one socket +// (rotated per batch); individual messages may target different resolvers. +func (c *Client) dispatchBatched(ctx context.Context, input <-chan string) { + bufs := make([][]byte, c.batchSize) + msgs := make([]ipv4.Message, c.batchSize) + for i := range bufs { + bufs[i] = make([]byte, 0, 512) + msgs[i].Buffers = [][]byte{nil} + } + + var sockIdx uint64 + for { + // block for the first name of a batch + name, ok := c.nextName(ctx, input) + if !ok { + return + } + + s := c.sockets[int(atomic.AddUint64(&sockIdx, 1))%len(c.sockets)] + n := 0 + for { + if c.limiter != nil { + c.limiter.take() + } + if err := c.acquire(ctx); err != nil { + // flush what we have before bailing out + c.flushBatch(s, msgs[:n]) + return + } + c.pending.Add(1) + c.statQueries.Add(1) + + pname, extra := c.parseLine(name) + idx, addr := c.chooseTarget(0, -1, extra) + q := c.register(s, pname, dns.Fqdn(pname), 0, idx, addr, extra) + if q != nil { + packed := c.encodeQuery(bufs[n][:0], q) + if packed == nil { + c.remove(s, q.txid) + c.fail(pname, errors.New("could not encode query")) + } else { + bufs[n] = packed + msgs[n].Buffers[0] = packed + msgs[n].Addr = q.addr + n++ + } + } + + if n == c.batchSize { + break + } + // opportunistically pull more already-queued names without blocking + var more bool + name, ok, more = c.tryNextName(input) + if !ok { + // input closed: flush and finish + c.flushBatch(s, msgs[:n]) + return + } + if !more { + break // nothing immediately available; flush the partial batch + } + } + c.flushBatch(s, msgs[:n]) + } +} + +// flushBatch transmits a prepared batch. It uses sendmmsg only when batching is +// active and the batch is large enough to be worth the header-array setup; +// otherwise it sends each datagram with a plain sendto. Unsent messages +// (partial write / error) stay in the in-flight table and are recovered by the +// timeout loop. +func (c *Client) flushBatch(s *socket, msgs []ipv4.Message) { + if len(msgs) == 0 { + return + } + if !c.batchActive.Load() || len(msgs) < minBatchToMMSG { + for i := range msgs { + if addr, ok := msgs[i].Addr.(*net.UDPAddr); ok { + _, _ = s.conn.WriteToUDP(msgs[i].Buffers[0], addr) + } + } + return + } + for off := 0; off < len(msgs); { + n, err := s.pc.WriteBatch(msgs[off:], 0) + if err != nil || n <= 0 { + return + } + off += n + } +} + +// nextName blocks for the next non-empty name, returning ok=false when the +// input is closed or the context is cancelled. +func (c *Client) nextName(ctx context.Context, input <-chan string) (string, bool) { + for { + select { + case <-ctx.Done(): + return "", false + case name, ok := <-input: + if !ok { + return "", false + } + if name = strings.TrimSpace(name); name != "" { + return name, true + } + } + } +} + +// tryNextName does a non-blocking read of the next name. more=false means no +// name is immediately available (the caller should flush its partial batch); +// ok=false means the input channel is closed. +func (c *Client) tryNextName(input <-chan string) (name string, ok, more bool) { + for { + select { + case n, chOpen := <-input: + if !chOpen { + return "", false, false + } + if n = strings.TrimSpace(n); n != "" { + return n, true, true + } + // skip empty, keep trying without blocking + default: + return "", true, false + } + } +} + +// packQuery encodes a minimal DNS query (single question) into buf, returning +// the filled slice and true. rd sets the recursion-desired flag; when udpSize +// >= 512 an EDNS0 OPT record advertising that payload size is appended. It +// returns false for names it cannot encode (empty/oversized labels), so the +// caller can fall back to the general encoder. +func packQuery(buf []byte, id uint16, name string, qtype uint16, rd bool, udpSize uint16) ([]byte, bool) { + var flagHi byte + if rd { + flagHi = 0x01 // recursion desired + } + var arcount byte + if udpSize >= 512 { + arcount = 0x01 // one additional record (OPT) + } + buf = append(buf, + byte(id>>8), byte(id), + flagHi, 0x00, // flags + 0x00, 0x01, // QDCOUNT = 1 + 0x00, 0x00, // ANCOUNT + 0x00, 0x00, // NSCOUNT + 0x00, arcount, // ARCOUNT + ) + + name = strings.TrimSuffix(name, ".") + start := 0 + for i := 0; i <= len(name); i++ { + if i < len(name) && name[i] != '.' { + continue + } + l := i - start + if l == 0 || l > 63 { + // empty label (leading/double dot) or oversized label + return buf, false + } + buf = append(buf, byte(l)) + buf = append(buf, name[start:i]...) + start = i + 1 + } + buf = append(buf, 0x00) // root label terminator + + buf = append(buf, byte(qtype>>8), byte(qtype), 0x00, 0x01) // qtype, qclass IN + + if udpSize >= 512 { + // EDNS0 OPT pseudo-record: root name, type OPT(41), class=UDP payload + // size, extended-rcode/flags/version=0, rdlen=0. + buf = append(buf, + 0x00, // root name + 0x00, 0x29, // type OPT (41) + byte(udpSize>>8), byte(udpSize), // requestor UDP payload size + 0x00, 0x00, 0x00, 0x00, // ext-rcode, version, flags + 0x00, 0x00, // rdlen = 0 + ) + } + + if len(buf) > 512 { + return buf, false + } + return buf, true +} + +// allocTxID finds a free transaction id in the socket table. Caller holds s.mu. +func (c *Client) allocTxID(s *socket) (uint16, bool) { + for i := 0; i < maxTxIDAttempts; i++ { + id := uint16(rand.Uint32()) + if _, exists := s.inflight[id]; !exists { + return id, true + } + } + return 0, false +} + +func (c *Client) remove(s *socket, txid uint16) *query { + s.mu.Lock() + defer s.mu.Unlock() + q, ok := s.inflight[txid] + if !ok { + return nil + } + delete(s.inflight, txid) + return q +} + +// fail abandons a name after exhausting retries (or a fatal send error). +func (c *Client) fail(name string, err error) { + if c.opts.OnError != nil { + c.opts.OnError(name, err) + } + c.release() + c.pending.Done() +} + +// deliver emits a final result for a name. +func (c *Client) deliver(res Result) { + c.statAnswered.Add(1) + if c.opts.OnResult != nil { + c.opts.OnResult(res) + } + c.release() + c.pending.Done() +} + +// readLoop drains responses from a single socket. When batching is unavailable +// it always reads one datagram per syscall; when capable it reads a batch +// (recvmmsg on Linux) while batching is active and falls back to single reads +// otherwise, so the read strategy tracks the adaptive controller live. +func (c *Client) readLoop(ctx context.Context, s *socket) { + single := make([]byte, 4096) + + var msgs []ipv4.Message + if c.batchCapable { + msgs = make([]ipv4.Message, c.batchSize) + for i := range msgs { + msgs[i].Buffers = [][]byte{make([]byte, 4096)} + } + } + + for { + if c.batchCapable && c.batchActive.Load() { + n, err := s.pc.ReadBatch(msgs, 0) + if err != nil { + if ctx.Err() != nil { + return + } + continue + } + for i := 0; i < n; i++ { + m := &msgs[i] + if m.N == 0 { + continue + } + src, _ := m.Addr.(*net.UDPAddr) + c.handlePacket(s, m.Buffers[0][:m.N], src) + } + continue + } + + // single-datagram read (portable path / batching disengaged / raw IPv6) + if s.rawFD >= 0 { + payload, addr, err := readRawUDPv6(s.rawFD, single) + if err != nil { + if ctx.Err() != nil { + return + } + continue + } + c.handlePacket(s, payload, addr) + continue + } + n, addr, err := s.conn.ReadFromUDP(single) + if err != nil { + if ctx.Err() != nil { + return + } + continue + } + c.handlePacket(s, single[:n], addr) + } +} + +// sameUDPAddr reports whether two UDP addresses share IP and port. +func sameUDPAddr(a, b *net.UDPAddr) bool { + if a == nil || b == nil { + return false + } + return a.Port == b.Port && a.IP.Equal(b.IP) +} + +// handlePacket parses a single response datagram and dispatches it. src is the +// datagram's source address, used for anti-spoofing verification. +func (c *Client) handlePacket(s *socket, packet []byte, src *net.UDPAddr) { + resp := new(dns.Msg) + if err := resp.Unpack(packet); err != nil { + return + } + if len(resp.Question) == 0 { + return + } + q := c.matchAndRemove(s, resp, src) + if q == nil { + return + } + // sample round-trip time for the controller, only when something consumes + // it (adaptive batch/concurrency or progress reporting). + if c.sampleRTT { + c.lastRTTNanos.Store(int64(time.Since(q.sentAt))) + } + + if c.hooks.OnResponse != nil { + c.hooks.OnResponse(c.queryInfo(q), resp) + } + + from := "" + if src != nil { + from = src.String() + } + c.handleResponse(q, resp, from) +} + +// matchAndRemove correlates a response to an in-flight query and removes it. +// When source-IP verification is enabled, a reply whose source does not match +// the resolver the query was sent to is ignored and the query is left in-flight +// (so a genuine reply, or a retransmission, can still resolve it). +func (c *Client) matchAndRemove(s *socket, resp *dns.Msg, src *net.UDPAddr) *query { + s.mu.Lock() + defer s.mu.Unlock() + q, ok := s.inflight[resp.Id] + if !ok { + return nil + } + // validate the question name to guard against id collisions + if !strings.EqualFold(resp.Question[0].Name, q.fqdn) { + return nil + } + // anti-spoofing: the reply must come from the resolver we queried. + if c.verify && !sameUDPAddr(src, q.addr) { + return nil + } + delete(s.inflight, resp.Id) + return q +} + +func (c *Client) shouldRetry(rcode int) bool { + if rcode >= 0 && rcode < len(c.retryRcode) { + return c.retryRcode[rcode] + } + return false +} + +func (c *Client) recordHealth(idx int, ok bool) { + if c.health == nil { + return + } + transitioned, healthy := c.health.record(idx, ok) + if transitioned && c.hooks.OnResolverState != nil && idx >= 0 && idx < len(c.resolvers) { + c.hooks.OnResolverState(c.resolvers[idx], healthy) + } +} + +// handleResponse decides whether a response is final, must be retried, or needs +// a TCP follow-up (truncation). +func (c *Client) handleResponse(q *query, resp *dns.Msg, from string) { + // truncated UDP answer: the resolver answered but the payload didn't fit, + // so fetch the full record set over TCP (off the reader goroutine). + if resp.Truncated && c.tcp { + c.recordHealth(q.resolverIdx, true) + if c.hooks.OnTruncated != nil { + c.hooks.OnTruncated(c.queryInfo(q)) + } + go c.tcpFollowup(q, from) + return + } + + if c.shouldRetry(resp.Rcode) && q.attempts+1 < c.opts.MaxRetries { + c.recordHealth(q.resolverIdx, false) + c.retry(q) + return + } + + c.recordHealth(q.resolverIdx, true) + c.finalize(q, c.buildResult(q, resp, from), from) +} + +// finalize delivers a result, optionally cross-checking positive answers on a +// second resolver first (off the reader goroutine). +func (c *Client) finalize(q *query, res Result, from string) { + if c.opts.CrossCheck && res.Rcode == dns.RcodeSuccess && (len(res.A) > 0 || len(res.AAAA) > 0) { + go c.crossVerify(q, res) + return + } + c.deliver(res) +} + +func (c *Client) buildResult(q *query, resp *dns.Msg, from string) Result { + res := Result{ + Name: q.name, + Type: c.opts.QueryType, + Rcode: resp.Rcode, + Resolver: from, + Msg: resp, + Timestamp: time.Now(), + } + for _, rr := range resp.Answer { + switch v := rr.(type) { + case *dns.A: + res.A = append(res.A, v.A.String()) + case *dns.AAAA: + res.AAAA = append(res.AAAA, v.AAAA.String()) + case *dns.CNAME: + res.CNAME = append(res.CNAME, strings.TrimSuffix(v.Target, ".")) + case *dns.NS: + res.NS = append(res.NS, strings.TrimSuffix(v.Ns, ".")) + case *dns.PTR: + res.PTR = append(res.PTR, strings.TrimSuffix(v.Ptr, ".")) + case *dns.MX: + res.MX = append(res.MX, strings.TrimSuffix(v.Mx, ".")) + case *dns.TXT: + res.TXT = append(res.TXT, v.Txt...) + case *dns.SOA: + res.SOA = append(res.SOA, strings.TrimSuffix(v.Ns, ".")) + } + } + return res +} + +// tcpFollowup re-issues q over TCP to retrieve a full (untruncated) answer. +func (c *Client) tcpFollowup(q *query, from string) { + client := &dns.Client{Net: "tcp", Timeout: c.opts.Timeout} + m := new(dns.Msg) + m.SetQuestion(q.fqdn, c.opts.QueryType) + m.RecursionDesired = c.rd + m.Id = q.txid + + resp, _, err := client.Exchange(m, q.addr.String()) + if err != nil || resp == nil { + // TCP failed: fall back to a normal UDP retry, or give up. + if q.attempts+1 < c.opts.MaxRetries { + c.recordHealth(q.resolverIdx, false) + c.retry(q) + return + } + c.fail(q.name, fmt.Errorf("tcp fallback failed: %w", err)) + return + } + c.finalize(q, c.buildResult(q, resp, from), from) +} + +// crossVerify re-resolves a positive answer on a different resolver and only +// delivers it if the two agree on the address set, providing basic protection +// against DNS poisoning/spam from a single bad resolver. +func (c *Client) crossVerify(q *query, res Result) { + idx := c.otherResolver(q.resolverIdx) + client := &dns.Client{Timeout: c.opts.Timeout} + m := new(dns.Msg) + m.SetQuestion(q.fqdn, c.opts.QueryType) + m.RecursionDesired = c.rd + if c.udpSize >= 512 { + m.SetEdns0(c.udpSize, false) + } + + resp, _, err := client.Exchange(m, c.resolverAddrs[idx].String()) + if err != nil || resp == nil || resp.Rcode != dns.RcodeSuccess { + // can't confirm: drop conservatively + if c.hooks.OnCrossCheckFailed != nil { + c.hooks.OnCrossCheckFailed(q.name, res.A, nil) + } + c.fail(q.name, errPoisonSuspected) + return + } + confirm := c.buildResult(q, resp, c.resolverAddrs[idx].String()) + if sameStringSet(res.A, confirm.A) && sameStringSet(res.AAAA, confirm.AAAA) { + c.deliver(res) + return + } + if c.hooks.OnCrossCheckFailed != nil { + c.hooks.OnCrossCheckFailed(q.name, res.A, confirm.A) + } + c.fail(q.name, errPoisonSuspected) +} + +// otherResolver returns a resolver index different from idx (best effort). +func (c *Client) otherResolver(idx int) int { + if len(c.resolvers) == 1 { + return idx + } + for { + if j := rand.IntN(len(c.resolvers)); j != idx { + return j + } + } +} + +// retry retransmits a query. Every retransmit (timeout or bad-rcode) is counted +// as a loss signal for the adaptive controller. The resolver is rotated unless +// sticky mode is set. +func (c *Client) retry(q *query) { + c.statRetransmit.Add(1) + attempt := q.attempts + 1 + c.send(q.sock, q.name, attempt, q.resolverIdx, q.extra) +} + +// timeoutLoop scans in-flight tables and retransmits or abandons expired queries. +func (c *Client) timeoutLoop(ctx context.Context) { + // scan at a fraction of the per-attempt timeout so a lost packet is + // retransmitted close to its deadline instead of up to a full extra tick + // late. Bounded below to avoid burning CPU on tiny timeouts. An explicit + // Interval overrides the derived value (massdns -i/--interval). + interval := c.opts.Interval + if interval <= 0 { + interval = c.opts.Timeout / 8 + } + if interval < 25*time.Millisecond { + interval = 25 * time.Millisecond + } + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + c.scanExpired() + } + } +} + +func (c *Client) scanExpired() { + now := time.Now() + for _, s := range c.sockets { + var expired []*query + s.mu.Lock() + for id, q := range s.inflight { + if now.After(q.deadline) { + delete(s.inflight, id) + expired = append(expired, q) + } + } + s.mu.Unlock() + + for _, q := range expired { + c.recordHealth(q.resolverIdx, false) // timeout = resolver miss + if c.hooks.OnTimeout != nil { + c.hooks.OnTimeout(c.queryInfo(q)) + } + if q.attempts+1 < c.opts.MaxRetries { + c.retry(q) + } else { + c.fail(q.name, errExhausted) + } + } + } +} + +// inflightDepth reports the number of currently outstanding queries, reading +// whichever semaphore implementation is in use. +func (c *Client) inflightDepth() int { + if c.dynSem != nil { + return c.dynSem.inflight() + } + return len(c.sem) +} + +// controllerLoop is the runtime governor. Every adaptTick it samples smoothed +// RTT, in-flight depth and interval packet loss, then drives three independent +// (all optional) feedback mechanisms: +// +// - adaptive batching: engage sendmmsg/recvmmsg only when RTT is high enough +// for packets to cluster AND the pipeline is deep enough to fill batches AND +// loss is low; back off on low-RTT links, a shallow pipeline, or rising loss +// (bursty sendmmsg can worsen drops at a saturated buffer); +// - adaptive concurrency (AIMD): multiplicatively shrink the in-flight cap on +// high loss and additively grow it back as loss subsides, so we stop +// flooding resolvers that are dropping/refusing; +// - progress reporting: emit a Stats snapshot via OnProgress. +func (c *Client) controllerLoop(ctx context.Context) { + ticker := time.NewTicker(adaptTick) + defer ticker.Stop() + + const alpha = 0.3 // RTT EWMA smoothing factor + var rttEWMA time.Duration + lastQ, lastR := c.statQueries.Load(), c.statRetransmit.Load() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + qNow, rNow := c.statQueries.Load(), c.statRetransmit.Load() + dq, dr := qNow-lastQ, rNow-lastR + lastQ, lastR = qNow, rNow + + var loss float64 + if total := dq + dr; total > 0 { + loss = float64(dr) / float64(total) + } + c.lastLossBits.Store(math.Float64bits(loss)) + + if sample := time.Duration(c.lastRTTNanos.Load()); sample > 0 { + if rttEWMA == 0 { + rttEWMA = sample + } else { + rttEWMA = time.Duration(alpha*float64(sample) + (1-alpha)*float64(rttEWMA)) + } + } + + depth := c.inflightDepth() + + // --- adaptive batching --- + if c.batchCapable && c.opts.Batch == BatchAdaptive { + active := c.batchActive.Load() + switch { + case active: + if rttEWMA < adaptRTTOff || depth < c.batchSize || loss > adaptLossOff { + active = false + } + default: + if rttEWMA >= adaptRTTOn && depth >= adaptDepthRatio*c.batchSize && loss <= adaptLossOff { + active = true + } + } + c.batchActive.Store(active) + } + + // --- adaptive concurrency (AIMD) --- + if c.dynSem != nil { + cap := c.dynSem.capacity() + minCap := int(float64(c.opts.Concurrency) * concMinRatio) + if minCap < 1 { + minCap = 1 + } + switch { + case loss > concLossHigh: + nc := int(float64(cap) * concDecrease) + if nc < minCap { + nc = minCap + } + c.dynSem.setCap(nc) + case loss < concLossLow: + // additive increase, ~5% of max per tick + step := c.opts.Concurrency / 20 + if step < 1 { + step = 1 + } + c.dynSem.setCap(cap + step) + } + } + + // --- progress reporting --- + if c.opts.OnProgress != nil { + capNow := c.opts.Concurrency + if c.dynSem != nil { + capNow = c.dynSem.capacity() + } + c.opts.OnProgress(Stats{ + Queries: qNow, + Retransmits: rNow, + Answered: c.statAnswered.Load(), + InflightDepth: depth, + ConcurrencyCap: capNow, + LossRate: loss, + RTT: rttEWMA, + BatchActive: c.batchActive.Load(), + }) + } + } + } +} + +// Stats returns a current snapshot of resolver activity (safe to call anytime). +func (c *Client) Stats() Stats { + capNow := c.opts.Concurrency + if c.dynSem != nil { + capNow = c.dynSem.capacity() + } + return Stats{ + Queries: c.statQueries.Load(), + Retransmits: c.statRetransmit.Load(), + Answered: c.statAnswered.Load(), + InflightDepth: c.inflightDepth(), + ConcurrencyCap: capNow, + LossRate: math.Float64frombits(c.lastLossBits.Load()), + RTT: time.Duration(c.lastRTTNanos.Load()), + BatchActive: c.batchActive.Load(), + } +} + +var ( + errExhausted = errors.New("max retries exhausted") + errPoisonSuspected = errors.New("cross-resolver disagreement (possible poisoning)") +) + +// Resolve is a one-shot convenience entry point for SDK use: it builds a client +// from opts, resolves every name in names, and returns the collected results. +// Any OnResult hook set in opts is preserved (invoked in addition to +// collection). For streaming/long-running use, construct a Client with New and +// call Run directly. +func Resolve(ctx context.Context, opts Options, names []string) ([]Result, error) { + var mu sync.Mutex + out := make([]Result, 0, len(names)) + userOnResult := opts.OnResult + opts.OnResult = func(r Result) { + mu.Lock() + out = append(out, r) + mu.Unlock() + if userOnResult != nil { + userOnResult(r) + } + } + + client, err := New(opts) + if err != nil { + return nil, err + } + defer client.Close() + + input := make(chan string, 1024) + go func() { + defer close(input) + for _, n := range names { + select { + case input <- n: + case <-ctx.Done(): + return + } + } + }() + + if err := client.Run(ctx, input); err != nil { + return out, err + } + return out, nil +} + +// sameStringSet reports whether a and b contain the same elements (set +// equality, ignoring order and duplicates). +func sameStringSet(a, b []string) bool { + if len(a) == 0 && len(b) == 0 { + return true + } + set := make(map[string]struct{}, len(a)) + for _, v := range a { + set[v] = struct{}{} + } + for _, v := range b { + if _, ok := set[v]; !ok { + return false + } + } + seen := make(map[string]struct{}, len(b)) + for _, v := range b { + seen[v] = struct{}{} + } + for v := range set { + if _, ok := seen[v]; !ok { + return false + } + } + return true +} + +func normalizeResolver(r string) string { + r = strings.TrimSpace(r) + if _, _, err := net.SplitHostPort(r); err != nil { + return net.JoinHostPort(r, "53") + } + return r +} diff --git a/pkg/resolve/resolve_test.go b/pkg/resolve/resolve_test.go new file mode 100644 index 00000000..7260684f --- /dev/null +++ b/pkg/resolve/resolve_test.go @@ -0,0 +1,371 @@ +package resolve + +import ( + "context" + "fmt" + "net" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/miekg/dns" + "github.com/stretchr/testify/require" +) + +// startTestServer spins up an in-process UDP DNS server that answers A queries +// for names present in the zone map and NXDOMAIN otherwise. +func startTestServer(t *testing.T, zone map[string]string) (string, func()) { + t.Helper() + + pc, err := net.ListenPacket("udp", "127.0.0.1:0") + require.NoError(t, err) + + handler := dns.HandlerFunc(func(w dns.ResponseWriter, r *dns.Msg) { + m := new(dns.Msg) + m.SetReply(r) + if len(r.Question) > 0 { + q := r.Question[0] + if ip, ok := zone[q.Name]; ok && q.Qtype == dns.TypeA { + rr, _ := dns.NewRR(fmt.Sprintf("%s 60 IN A %s", q.Name, ip)) + m.Answer = append(m.Answer, rr) + } else { + m.Rcode = dns.RcodeNameError + } + } + _ = w.WriteMsg(m) + }) + + srv := &dns.Server{PacketConn: pc, Handler: handler} + go func() { _ = srv.ActivateAndServe() }() + + // give the server a moment to start + time.Sleep(50 * time.Millisecond) + + return pc.LocalAddr().String(), func() { _ = srv.Shutdown() } +} + +func TestResolveBasic(t *testing.T) { + zone := map[string]string{ + "a.example.com.": "1.1.1.1", + "b.example.com.": "2.2.2.2", + "c.example.com.": "3.3.3.3", + } + addr, stop := startTestServer(t, zone) + defer stop() + + var mu sync.Mutex + got := make(map[string][]string) + nx := make(map[string]struct{}) + + client, err := New(Options{ + Resolvers: []string{addr}, + Concurrency: 100, + SocketCount: 2, + Timeout: time.Second, + MaxRetries: 2, + OnResult: func(r Result) { + mu.Lock() + defer mu.Unlock() + switch r.Rcode { + case dns.RcodeSuccess: + got[r.Name] = r.A + case dns.RcodeNameError: + nx[r.Name] = struct{}{} + } + }, + }) + require.NoError(t, err) + defer client.Close() + + input := make(chan string) + go func() { + defer close(input) + for _, n := range []string{"a.example.com", "b.example.com", "c.example.com", "missing.example.com"} { + input <- n + } + }() + + require.NoError(t, client.Run(context.Background(), input)) + + mu.Lock() + defer mu.Unlock() + require.Equal(t, []string{"1.1.1.1"}, got["a.example.com"]) + require.Equal(t, []string{"2.2.2.2"}, got["b.example.com"]) + require.Equal(t, []string{"3.3.3.3"}, got["c.example.com"]) + require.Contains(t, nx, "missing.example.com") +} + +// TestResolveBatchModes verifies that every batching strategy resolves the same +// set of names correctly. Batching only changes the send/recv syscall path, so +// results must be identical regardless of mode (on platforms without +// sendmmsg/recvmmsg the batch path transparently degrades to single I/O). +func TestResolveBatchModes(t *testing.T) { + zone := make(map[string]string) + const total = 1500 + for i := 0; i < total; i++ { + zone[fmt.Sprintf("host%d.example.com.", i)] = fmt.Sprintf("10.1.%d.%d", i/256, i%256) + } + addr, stop := startTestServer(t, zone) + defer stop() + + for _, tc := range []struct { + name string + mode BatchMode + }{ + {"disabled", BatchDisabled}, + {"enabled", BatchEnabled}, + {"adaptive", BatchAdaptive}, + } { + t.Run(tc.name, func(t *testing.T) { + var mu sync.Mutex + seen := make(map[string]struct{}) + + client, err := New(Options{ + Resolvers: []string{addr}, + Concurrency: 500, + SocketCount: 4, + Timeout: time.Second, + MaxRetries: 3, + Batch: tc.mode, + BatchSize: 16, + OnResult: func(r Result) { + if r.Rcode == dns.RcodeSuccess && len(r.A) > 0 { + mu.Lock() + seen[r.Name] = struct{}{} + mu.Unlock() + } + }, + }) + require.NoError(t, err) + defer client.Close() + + input := make(chan string) + go func() { + defer close(input) + for i := 0; i < total; i++ { + input <- fmt.Sprintf("host%d.example.com", i) + } + }() + + require.NoError(t, client.Run(context.Background(), input)) + + mu.Lock() + defer mu.Unlock() + require.Equal(t, total, len(seen), "all names should resolve in %s mode", tc.name) + }) + } +} + +// TestRecordTypesAndFlags spins up a server that records the recursion-desired +// bit and EDNS0 presence, answers TXT, and verifies the resolver encodes +// queries correctly (EDNS0 on by default, RD configurable) and parses TXT. +func TestRecordTypesAndFlags(t *testing.T) { + pc, err := net.ListenPacket("udp", "127.0.0.1:0") + require.NoError(t, err) + + var sawEDNS, sawRD atomic.Bool + handler := dns.HandlerFunc(func(w dns.ResponseWriter, r *dns.Msg) { + if r.IsEdns0() != nil { + sawEDNS.Store(true) + } + if r.RecursionDesired { + sawRD.Store(true) + } + m := new(dns.Msg) + m.SetReply(r) + if len(r.Question) > 0 && r.Question[0].Qtype == dns.TypeTXT { + rr, _ := dns.NewRR(fmt.Sprintf("%s 60 IN TXT \"hello-world\"", r.Question[0].Name)) + m.Answer = append(m.Answer, rr) + } + _ = w.WriteMsg(m) + }) + srv := &dns.Server{PacketConn: pc, Handler: handler} + go func() { _ = srv.ActivateAndServe() }() + defer func() { _ = srv.Shutdown() }() + time.Sleep(50 * time.Millisecond) + addr := pc.LocalAddr().String() + + var mu sync.Mutex + var txt []string + client, err := New(Options{ + Resolvers: []string{addr}, + QueryType: dns.TypeTXT, + Concurrency: 16, + SocketCount: 1, + Timeout: time.Second, + MaxRetries: 2, + NoRecurse: true, // RD must NOT be set + OnResult: func(r Result) { + mu.Lock() + txt = append(txt, r.TXT...) + mu.Unlock() + }, + }) + require.NoError(t, err) + defer client.Close() + + input := make(chan string, 1) + input <- "example.com" + close(input) + require.NoError(t, client.Run(context.Background(), input)) + + mu.Lock() + defer mu.Unlock() + require.Contains(t, txt, "hello-world", "TXT record should be parsed") + require.True(t, sawEDNS.Load(), "queries should advertise EDNS0 by default") + require.False(t, sawRD.Load(), "NoRecurse should clear the RD bit") +} + +// TestAllFeaturesEnabled verifies that turning on the full feature set (health +// scoring, adaptive concurrency, adaptive batching, cross-check, source +// verification, TCP fallback) still resolves every name correctly. +func TestAllFeaturesEnabled(t *testing.T) { + zone := make(map[string]string) + const total = 1000 + for i := 0; i < total; i++ { + zone[fmt.Sprintf("h%d.example.com.", i)] = fmt.Sprintf("10.2.%d.%d", i/256, i%256) + } + addr1, stop1 := startTestServer(t, zone) + defer stop1() + addr2, stop2 := startTestServer(t, zone) + defer stop2() + + var mu sync.Mutex + seen := make(map[string]struct{}) + client, err := New(Options{ + Resolvers: []string{addr1, addr2}, + Concurrency: 300, + SocketCount: 4, + Timeout: time.Second, + MaxRetries: 3, + Batch: BatchAdaptive, + ResolverHealth: true, + AdaptiveConcurrency: true, + CrossCheck: true, + OnResult: func(r Result) { + if r.Rcode == dns.RcodeSuccess && len(r.A) > 0 { + mu.Lock() + seen[r.Name] = struct{}{} + mu.Unlock() + } + }, + }) + require.NoError(t, err) + defer client.Close() + + input := make(chan string) + go func() { + defer close(input) + for i := 0; i < total; i++ { + input <- fmt.Sprintf("h%d.example.com", i) + } + }() + require.NoError(t, client.Run(context.Background(), input)) + + mu.Lock() + defer mu.Unlock() + require.Equal(t, total, len(seen), "all names should resolve with all features enabled") +} + +// TestHooks verifies the lifecycle hooks fire and that the package-level +// Resolve convenience entry point collects results. +func TestHooks(t *testing.T) { + zone := map[string]string{ + "a.example.com.": "1.1.1.1", + "b.example.com.": "2.2.2.2", + } + addr, stop := startTestServer(t, zone) + defer stop() + + var queries, responses atomic.Int64 + + results, err := Resolve(context.Background(), Options{ + Resolvers: []string{addr}, + Concurrency: 16, + SocketCount: 1, + Timeout: time.Second, + MaxRetries: 2, + Hooks: Hooks{ + OnQuery: func(QueryInfo) { queries.Add(1) }, + OnResponse: func(QueryInfo, *dns.Msg) { responses.Add(1) }, + }, + }, []string{"a.example.com", "b.example.com", "missing.example.com"}) + require.NoError(t, err) + + require.Len(t, results, 3, "Resolve should collect a result per name") + require.GreaterOrEqual(t, queries.Load(), int64(3), "OnQuery should fire per name") + require.GreaterOrEqual(t, responses.Load(), int64(3), "OnResponse should fire per answered query") +} + +// TestHookTimeout verifies OnTimeout fires when a resolver never answers. +func TestHookTimeout(t *testing.T) { + // a black-hole UDP socket that accepts but never replies + pc, err := net.ListenPacket("udp", "127.0.0.1:0") + require.NoError(t, err) + defer func() { _ = pc.Close() }() + + var timeouts atomic.Int64 + client, err := New(Options{ + Resolvers: []string{pc.LocalAddr().String()}, + Concurrency: 4, + SocketCount: 1, + Timeout: 100 * time.Millisecond, + MaxRetries: 2, + Hooks: Hooks{OnTimeout: func(QueryInfo) { timeouts.Add(1) }}, + }) + require.NoError(t, err) + defer client.Close() + + input := make(chan string, 1) + input <- "dead.example.com" + close(input) + require.NoError(t, client.Run(context.Background(), input)) + require.GreaterOrEqual(t, timeouts.Load(), int64(1), "OnTimeout should fire for an unanswered query") +} + +func TestResolveManyNames(t *testing.T) { + zone := make(map[string]string) + const total = 2000 + for i := 0; i < total; i++ { + zone[fmt.Sprintf("host%d.example.com.", i)] = fmt.Sprintf("10.0.%d.%d", i/256, i%256) + } + addr, stop := startTestServer(t, zone) + defer stop() + + var count int64 + var mu sync.Mutex + seen := make(map[string]struct{}) + + client, err := New(Options{ + Resolvers: []string{addr}, + Concurrency: 500, + SocketCount: 4, + Timeout: time.Second, + MaxRetries: 3, + OnResult: func(r Result) { + if r.Rcode == dns.RcodeSuccess && len(r.A) > 0 { + mu.Lock() + seen[r.Name] = struct{}{} + count++ + mu.Unlock() + } + }, + }) + require.NoError(t, err) + defer client.Close() + + input := make(chan string) + go func() { + defer close(input) + for i := 0; i < total; i++ { + input <- fmt.Sprintf("host%d.example.com", i) + } + }() + + require.NoError(t, client.Run(context.Background(), input)) + + mu.Lock() + defer mu.Unlock() + require.Equal(t, total, len(seen), "all names should resolve") +} diff --git a/pkg/resolve/sem.go b/pkg/resolve/sem.go new file mode 100644 index 00000000..3968a5b8 --- /dev/null +++ b/pkg/resolve/sem.go @@ -0,0 +1,98 @@ +package resolve + +import ( + "context" + "sync" +) + +// adaptiveSem is a counting semaphore whose capacity can be changed at runtime. +// It is used for the adaptive-concurrency in-flight cap: the controller shrinks +// the cap under packet loss (to stop flooding a saturated resolver) and grows +// it back when loss subsides. Acquire is cancellable via context. +// +// Only one goroutine acquires (the dispatcher) while many release (readers and +// timeout/fallback paths), so the single mutex sees low contention. +type adaptiveSem struct { + mu sync.Mutex + cond *sync.Cond + cur int + cap int + maxCap int + closed bool +} + +func newAdaptiveSem(capacity int) *adaptiveSem { + if capacity < 1 { + capacity = 1 + } + s := &adaptiveSem{cur: 0, cap: capacity, maxCap: capacity} + s.cond = sync.NewCond(&s.mu) + return s +} + +// watch wires context cancellation to wake any blocked acquirers. +func (s *adaptiveSem) watch(ctx context.Context) { + go func() { + <-ctx.Done() + s.mu.Lock() + s.closed = true + s.mu.Unlock() + s.cond.Broadcast() + }() +} + +func (s *adaptiveSem) acquire(ctx context.Context) error { + s.mu.Lock() + for s.cur >= s.cap && !s.closed { + s.cond.Wait() + } + if s.closed { + s.mu.Unlock() + if err := ctx.Err(); err != nil { + return err + } + return context.Canceled + } + s.cur++ + s.mu.Unlock() + return nil +} + +func (s *adaptiveSem) release() { + s.mu.Lock() + if s.cur > 0 { + s.cur-- + } + s.mu.Unlock() + s.cond.Signal() +} + +// setCap adjusts the capacity, clamped to [1, maxCap]. Increasing it wakes +// waiters so they can proceed immediately. +func (s *adaptiveSem) setCap(n int) { + if n < 1 { + n = 1 + } + s.mu.Lock() + if n > s.maxCap { + n = s.maxCap + } + grew := n > s.cap + s.cap = n + s.mu.Unlock() + if grew { + s.cond.Broadcast() + } +} + +func (s *adaptiveSem) capacity() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.cap +} + +func (s *adaptiveSem) inflight() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.cur +} diff --git a/pkg/resolve/srcrand.go b/pkg/resolve/srcrand.go new file mode 100644 index 00000000..25710e7f --- /dev/null +++ b/pkg/resolve/srcrand.go @@ -0,0 +1,102 @@ +package resolve + +import ( + "bufio" + "fmt" + "math/rand/v2" + "net" + "os" + "strings" +) + +// srcRand picks a random IPv6 source address for each query, either from a +// prefix (massdns --rand-src-ipv6) or from a file of addresses +// (--rand-src-ipv6-file). Used with Linux raw IPv6 UDP sockets + IPV6_HDRINCL. +type srcRand struct { + fromFile bool + base net.IP // 16-byte network base (prefix bits set, host zeroed) + bits int + addrs []net.IP +} + +func newSrcRandFromPrefix(spec string) (*srcRand, error) { + spec = strings.TrimSpace(spec) + if spec == "" { + return nil, fmt.Errorf("empty IPv6 prefix") + } + if !strings.Contains(spec, "/") { + spec += "/128" + } + ip, ipnet, err := net.ParseCIDR(spec) + if err != nil { + return nil, fmt.Errorf("invalid --rand-src-ipv6 %q: %w", spec, err) + } + ip = ip.To16() + if ip == nil || ip.To4() != nil { + return nil, fmt.Errorf("--rand-src-ipv6 requires an IPv6 prefix") + } + ones, bits := ipnet.Mask.Size() + if bits != 128 { + return nil, fmt.Errorf("--rand-src-ipv6 requires an IPv6 prefix") + } + base := make(net.IP, 16) + copy(base, ip.Mask(ipnet.Mask).To16()) + return &srcRand{base: base, bits: ones}, nil +} + +func newSrcRandFromFile(path string) (*srcRand, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer func() { _ = f.Close() }() + var addrs []net.IP + sc := bufio.NewScanner(f) + for sc.Scan() { + line := strings.TrimSpace(sc.Text()) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + // allow optional /prefix; take address only + if i := strings.IndexByte(line, '/'); i >= 0 { + line = line[:i] + } + ip := net.ParseIP(line) + if ip == nil || ip.To4() != nil { + continue + } + addrs = append(addrs, ip.To16()) + } + if err := sc.Err(); err != nil { + return nil, err + } + if len(addrs) == 0 { + return nil, fmt.Errorf("no valid IPv6 addresses in %s", path) + } + return &srcRand{fromFile: true, addrs: addrs}, nil +} + +func (s *srcRand) pick() net.IP { + if s.fromFile { + return s.addrs[rand.IntN(len(s.addrs))] + } + out := make(net.IP, 16) + copy(out, s.base) + hostBits := 128 - s.bits + if hostBits <= 0 { + return out + } + // Randomize the host portion (same approach as massdns: full trailing + // bytes randomly, plus a partial byte for leftover bits). + fullBytes := hostBits / 8 + remBits := hostBits % 8 + if remBits > 0 { + idx := 16 - fullBytes - 1 + mask := byte((1 << remBits) - 1) + out[idx] = (out[idx] & ^mask) | (byte(rand.Uint32()) & mask) + } + for i := 0; i < fullBytes; i++ { + out[16-fullBytes+i] = byte(rand.Uint32()) + } + return out +} diff --git a/pkg/resolve/srcrand_test.go b/pkg/resolve/srcrand_test.go new file mode 100644 index 00000000..d5a9948f --- /dev/null +++ b/pkg/resolve/srcrand_test.go @@ -0,0 +1,92 @@ +package resolve + +import ( + "net" + "os" + "path/filepath" + "testing" +) + +func TestNewSrcRandFromPrefix(t *testing.T) { + sr, err := newSrcRandFromPrefix("2001:db8:abcd::/48") + if err != nil { + t.Fatal(err) + } + for i := 0; i < 32; i++ { + ip := sr.pick() + if ip.To16() == nil || ip.To4() != nil { + t.Fatalf("expected IPv6, got %v", ip) + } + // first 48 bits must match 2001:db8:abcd + if ip[0] != 0x20 || ip[1] != 0x01 || ip[2] != 0x0d || ip[3] != 0xb8 || ip[4] != 0xab || ip[5] != 0xcd { + t.Fatalf("prefix mismatch: %s", ip) + } + } +} + +func TestNewSrcRandFromPrefixRejectsIPv4(t *testing.T) { + if _, err := newSrcRandFromPrefix("192.0.2.0/24"); err == nil { + t.Fatal("expected error for IPv4 prefix") + } +} + +func TestNewSrcRandFromFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "srcs.txt") + content := "# comment\n2001:db8::1\n2001:db8::2/128\n\n192.0.2.1\n" + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + t.Fatal(err) + } + sr, err := newSrcRandFromFile(path) + if err != nil { + t.Fatal(err) + } + seen := map[string]bool{} + for i := 0; i < 20; i++ { + seen[sr.pick().String()] = true + } + if !seen[net.ParseIP("2001:db8::1").String()] || !seen[net.ParseIP("2001:db8::2").String()] { + t.Fatalf("expected both file addresses, got %v", seen) + } + if len(seen) != 2 { + t.Fatalf("expected only 2 addresses, got %v", seen) + } +} + +func TestNewSrcRandMutualExclusivityInOptions(t *testing.T) { + _, err := New(Options{ + Resolvers: []string{"[2001:db8::53]:53"}, + RandSrcIPv6: "2001:db8::/32", + RandSrcIPv6File: "x", + SocketCount: 1, + Concurrency: 1, + }) + if err == nil { + t.Fatal("expected mutual exclusion error") + } +} + +func TestNewSrcRandIncompatibleWithBind(t *testing.T) { + _, err := New(Options{ + Resolvers: []string{"[2001:db8::53]:53"}, + RandSrcIPv6: "2001:db8::/32", + BindAddr: "::", + SocketCount: 1, + Concurrency: 1, + }) + if err == nil { + t.Fatal("expected bind+rand-src error") + } +} + +func TestNewSrcRandRequiresIPv6Resolvers(t *testing.T) { + _, err := New(Options{ + Resolvers: []string{"1.1.1.1:53"}, + RandSrcIPv6: "2001:db8::/32", + SocketCount: 1, + Concurrency: 1, + }) + if err == nil { + t.Fatal("expected IPv6 resolvers required error") + } +} diff --git a/pkg/resolve/validate.go b/pkg/resolve/validate.go new file mode 100644 index 00000000..5218b810 --- /dev/null +++ b/pkg/resolve/validate.go @@ -0,0 +1,167 @@ +package resolve + +import ( + "context" + "fmt" + "math/rand/v2" + "sync" + "time" + + "github.com/miekg/dns" +) + +// ValidateConfig configures resolver-list curation. It is the native, embeddable +// equivalent of dnsvalidator: it weeds out resolvers that are dead, that lie by +// hijacking NXDOMAIN (returning an address for names that do not exist), or that +// cannot resolve a known-good name. Output quality for any mass-resolution run is +// bounded by resolver-list quality, so this is the most important accuracy knob. +type ValidateConfig struct { + // Resolvers is the candidate list (host or host:port; :53 assumed). + Resolvers []string + // GoodDomains are names expected to resolve to at least one A record. A + // resolver must successfully resolve one of them to pass. If empty, a small + // built-in set of stable, non-wildcard domains is used. + GoodDomains []string + // Timeout is the per-query timeout. Default 3s. + Timeout time.Duration + // Concurrency caps simultaneous resolver checks. Default 50. + Concurrency int + // MaxRTT optionally rejects resolvers slower than this on the positive + // probe. 0 disables the latency filter. + MaxRTT time.Duration + // OnResolver is an optional progress callback fired once per checked + // resolver (safe for concurrent use). + OnResolver func(ResolverCheck) +} + +// ResolverCheck is the verdict for a single candidate resolver. +type ResolverCheck struct { + Resolver string // normalized host:port + OK bool // passed all checks + Reason string // failure reason when !OK + RTT time.Duration // latency of the positive probe (when measured) +} + +// defaultGoodDomains are widely-deployed names that resolve to A records and are +// not wildcard zones, suitable as positive/negative resolver probes. +var defaultGoodDomains = []string{"google.com", "cloudflare.com", "wikipedia.org"} + +// ValidateResolvers checks each candidate resolver and returns the subset that +// passed plus the full per-resolver report. It performs a handful of low-volume +// queries per resolver (not a mass run), so it uses a simple bounded worker pool. +func ValidateResolvers(ctx context.Context, cfg ValidateConfig) (good []string, report []ResolverCheck, err error) { + if len(cfg.Resolvers) == 0 { + return nil, nil, fmt.Errorf("no resolvers to validate") + } + if cfg.Timeout <= 0 { + cfg.Timeout = 3 * time.Second + } + if cfg.Concurrency <= 0 { + cfg.Concurrency = 50 + } + goodDomains := cfg.GoodDomains + if len(goodDomains) == 0 { + goodDomains = defaultGoodDomains + } + + report = make([]ResolverCheck, len(cfg.Resolvers)) + sem := make(chan struct{}, cfg.Concurrency) + var wg sync.WaitGroup + + for i, r := range cfg.Resolvers { + select { + case <-ctx.Done(): + return nil, nil, ctx.Err() + case sem <- struct{}{}: + } + wg.Add(1) + go func(i int, r string) { + defer wg.Done() + defer func() { <-sem }() + chk := checkResolver(ctx, normalizeResolver(r), goodDomains, cfg.Timeout, cfg.MaxRTT) + report[i] = chk + if cfg.OnResolver != nil { + cfg.OnResolver(chk) + } + }(i, r) + } + wg.Wait() + + for _, chk := range report { + if chk.OK { + good = append(good, chk.Resolver) + } + } + return good, report, nil +} + +// checkResolver runs the liveness, NXDOMAIN-hijack and positive-resolution +// probes against a single resolver. +func checkResolver(ctx context.Context, server string, goodDomains []string, timeout, maxRTT time.Duration) ResolverCheck { + chk := ResolverCheck{Resolver: server} + client := &dns.Client{Timeout: timeout} + + // 1) NXDOMAIN-hijack probe: a random label under a good domain must NOT + // resolve to an address. A resolver that answers with an A is lying. + probe := randomLabel() + "." + goodDomains[0] + if resp, _, err := exchangeA(ctx, client, server, probe); err == nil && resp != nil { + if hasAddress(resp) { + chk.Reason = "hijacks NXDOMAIN (answers for nonexistent name)" + return chk + } + } + // (a transport error here is tolerated; the positive probe is the liveness + // gate, since some resolvers drop obviously bogus queries.) + + // 2) positive probe: at least one good domain must resolve to an A. + resolvedAny := false + for _, d := range goodDomains { + start := time.Now() + resp, _, err := exchangeA(ctx, client, server, d) + if err != nil || resp == nil { + continue + } + if resp.Rcode == dns.RcodeSuccess && hasAddress(resp) { + chk.RTT = time.Since(start) + resolvedAny = true + break + } + } + if !resolvedAny { + chk.Reason = "could not resolve any known-good domain" + return chk + } + if maxRTT > 0 && chk.RTT > maxRTT { + chk.Reason = fmt.Sprintf("too slow (%s > %s)", chk.RTT.Round(time.Millisecond), maxRTT) + return chk + } + + chk.OK = true + return chk +} + +func exchangeA(ctx context.Context, client *dns.Client, server, name string) (*dns.Msg, time.Duration, error) { + m := new(dns.Msg) + m.SetQuestion(dns.Fqdn(name), dns.TypeA) + m.RecursionDesired = true + return client.ExchangeContext(ctx, m, server) +} + +func hasAddress(m *dns.Msg) bool { + for _, rr := range m.Answer { + switch rr.(type) { + case *dns.A, *dns.AAAA: + return true + } + } + return false +} + +func randomLabel() string { + const alphabet = "abcdefghijklmnopqrstuvwxyz0123456789" + b := make([]byte, 12) + for i := range b { + b[i] = alphabet[rand.IntN(len(alphabet))] + } + return string(b) +} diff --git a/pkg/resolve/validate_test.go b/pkg/resolve/validate_test.go new file mode 100644 index 00000000..d930d4db --- /dev/null +++ b/pkg/resolve/validate_test.go @@ -0,0 +1,87 @@ +package resolve + +import ( + "context" + "net" + "testing" + "time" + + "github.com/miekg/dns" +) + +// startDNSServer starts a UDP DNS server on loopback with the given handler and +// returns its address and a stop function. +func startDNSServer(t *testing.T, handler dns.HandlerFunc) (string, func()) { + t.Helper() + pc, err := net.ListenPacket("udp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + srv := &dns.Server{PacketConn: pc, Handler: handler} + go func() { _ = srv.ActivateAndServe() }() + return pc.LocalAddr().String(), func() { _ = srv.Shutdown() } +} + +func TestValidateResolvers(t *testing.T) { + // good resolver: resolves good.test, NXDOMAIN for anything else. + goodAddr, stopGood := startDNSServer(t, func(w dns.ResponseWriter, r *dns.Msg) { + m := new(dns.Msg) + m.SetReply(r) + q := r.Question[0] + if q.Name == "good.test." && q.Qtype == dns.TypeA { + rr, _ := dns.NewRR("good.test. 60 IN A 1.2.3.4") + m.Answer = append(m.Answer, rr) + } else { + m.Rcode = dns.RcodeNameError + } + _ = w.WriteMsg(m) + }) + defer stopGood() + + // hijacking resolver: answers an A for every name (lies about NXDOMAIN). + hijackAddr, stopHijack := startDNSServer(t, func(w dns.ResponseWriter, r *dns.Msg) { + m := new(dns.Msg) + m.SetReply(r) + rr, _ := dns.NewRR(r.Question[0].Name + " 60 IN A 6.6.6.6") + m.Answer = append(m.Answer, rr) + _ = w.WriteMsg(m) + }) + defer stopHijack() + + // dead resolver: nothing listening on this port. + deadAddr := "127.0.0.1:1" + + good, report, err := ValidateResolvers(context.Background(), ValidateConfig{ + Resolvers: []string{goodAddr, hijackAddr, deadAddr}, + GoodDomains: []string{"good.test"}, + Timeout: 2 * time.Second, + Concurrency: 5, + }) + if err != nil { + t.Fatalf("ValidateResolvers: %v", err) + } + if len(report) != 3 { + t.Fatalf("expected 3 reports, got %d", len(report)) + } + if len(good) != 1 || good[0] != goodAddr { + t.Fatalf("expected only the good resolver to pass, got %v", good) + } + + // verify reasons for the rejected ones + byAddr := map[string]ResolverCheck{} + for _, c := range report { + byAddr[c.Resolver] = c + } + if c := byAddr[hijackAddr]; c.OK || c.Reason == "" { + t.Fatalf("hijacking resolver should be rejected with a reason, got %+v", c) + } + if c := byAddr[deadAddr]; c.OK { + t.Fatalf("dead resolver should be rejected") + } +} + +func TestValidateResolversEmpty(t *testing.T) { + if _, _, err := ValidateResolvers(context.Background(), ValidateConfig{}); err == nil { + t.Fatal("expected error for empty resolver list") + } +} diff --git a/pkg/runner/options.go b/pkg/runner/options.go index c2f0b1dd..6d4e8e6e 100644 --- a/pkg/runner/options.go +++ b/pkg/runner/options.go @@ -44,6 +44,28 @@ type Options struct { BatchSize int // BatchSize controls the number of lines per chunk for incremental processing FilterInternalIPs bool // FilterInternalIPs controls whether to filter out internal/private IP addresses + // Native resolver tuning (see pkg/resolve). + QueryType string // DNS record type to resolve (A, AAAA, CNAME, ...). Default A. + BatchMode string // sendmmsg/recvmmsg batching: off | on | adaptive + SocketCount int // UDP sockets per run (0 = scale to cores) + UDPSize int // EDNS0 advertised UDP payload size (0 = default 1232; <512 disables) + QPS int // outbound query rate limit (0 = unlimited) + NoRecurse bool // send non-recursive queries (RD=0) + Sticky bool // do not rotate resolver on retry + ResolverHealth bool // per-resolver health scoring / de-weighting + AdaptiveConcurrency bool // shrink/grow in-flight cap based on packet loss + CrossCheck bool // re-verify positive answers on a second resolver + ExtendedInput bool // parse "name [resolver ...]" input lines + NoVerifyIP bool // disable reply source-IP verification + NoTCPFallback bool // disable TCP fallback on truncated answers + // Iterative recurses from the root servers (no resolver list needed), + // caching delegations; removes the public-resolver dependency. + Iterative bool + + // Distributed resolution and resume. + Shard string // "m/n": process only shard m of n (distributed coordination) + ResumeFile string // checkpoint file for crash-safe stop/resume + OnResult func(*retryabledns.DNSData) } @@ -59,7 +81,7 @@ func ParseOptions() *Options { options := &Options{} flagSet := goflags.NewFlagSet() - flagSet.SetDescription(`shuffleDNS is a wrapper around massdns written in go that allows you to enumerate valid subdomains using active bruteforce as well as resolve subdomains with wildcard handling and easy input-output support.`) + flagSet.SetDescription(`shuffleDNS is a high-throughput DNS bruteforcer and resolver with wildcard handling. It uses a native Go stub resolver (massdns-compatible) and optional iterative resolution from the DNS roots.`) flagSet.CreateGroup("input", "Input", flagSet.StringSliceVarP(&options.Domains, "domain", "d", nil, "Domain to find or resolve subdomains for", goflags.FileCommaSeparatedStringSliceOptions), @@ -68,12 +90,13 @@ func ParseOptions() *Options { flagSet.StringVarP(&options.Wordlist, "wordlist", "w", "", "File containing words to bruteforce for domain"), flagSet.StringVarP(&options.ResolversFile, "resolver", "r", "", "File containing list of resolvers for enumeration"), flagSet.StringVarP(&options.TrustedResolvers, "trusted-resolver", "tr", "", "File containing list of trusted resolvers"), - flagSet.StringVarP(&options.MassdnsRaw, "raw-input", "ri", "", "Validate raw full massdns output"), + flagSet.StringVarP(&options.MassdnsRaw, "raw-input", "ri", "", "Filter wildcards from an existing massdns-format output file"), flagSet.StringVar(&options.Mode, "mode", "", "Execution mode (bruteforce, resolve, filter)"), ) flagSet.CreateGroup("rate-limit", "Rate-Limit", - flagSet.IntVar(&options.Threads, "t", 10000, "Number of concurrent massdns resolves"), + flagSet.IntVar(&options.Threads, "t", 10000, "Max concurrent in-flight DNS queries"), + flagSet.IntVar(&options.QPS, "qps", 0, "Max outbound DNS queries per second (0 = unlimited)"), ) flagSet.CreateGroup("update", "Update", @@ -88,8 +111,8 @@ func ParseOptions() *Options { ) flagSet.CreateGroup("configs", "Configurations", - flagSet.StringVarP(&options.MassdnsPath, "massdns", "m", "", "Path to the massdns binary"), - flagSet.StringVarP(&options.MassDnsCmd, "massdns-cmd", "mcmd", "", "Optional massdns commands to run (example '-i 10')"), + flagSet.StringVarP(&options.MassdnsPath, "massdns", "m", "", "Deprecated: massdns is no longer used (native resolver), flag ignored"), + flagSet.StringVarP(&options.MassDnsCmd, "massdns-cmd", "mcmd", "", "Deprecated: massdns is no longer used (native resolver), flag ignored"), flagSet.StringVar(&options.Directory, "directory", "", "Temporary directory for enumeration"), ) @@ -97,11 +120,32 @@ func ParseOptions() *Options { flagSet.IntVar(&options.Retries, "retries", 5, "Number of retries for dns enumeration"), flagSet.BoolVarP(&options.StrictWildcard, "strict-wildcard", "sw", false, "Perform wildcard check on all found subdomains"), flagSet.IntVar(&options.WildcardThreads, "wt", 250, "Number of concurrent wildcard checks"), - flagSet.BoolVar(&options.KeepStderr, "retain-stderr", false, "Capture and store massdns stderr output (default: discard)"), - flagSet.IntVar(&options.BatchSize, "batch-size", DefaultBatchSize, "Number of lines per chunk for incremental processing"), + flagSet.BoolVar(&options.KeepStderr, "retain-stderr", false, "Deprecated: massdns is no longer used (native resolver), flag ignored"), + flagSet.IntVar(&options.BatchSize, "batch-size", DefaultBatchSize, "Deprecated: native resolver streams without chunking, flag ignored"), flagSet.BoolVar(&options.FilterInternalIPs, "filter-internal-ips", false, "Filter out internal/private IP addresses (0.0.0.0 is always filtered)"), ) + flagSet.CreateGroup("resolver", "Resolver", + flagSet.StringVarP(&options.QueryType, "type", "rt", "A", "DNS record type to resolve (A, AAAA, CNAME, NS, PTR, MX, TXT, SOA)"), + flagSet.StringVarP(&options.BatchMode, "batch-mode", "bm", "off", "sendmmsg/recvmmsg batching: off | on | adaptive (Linux)"), + flagSet.IntVarP(&options.SocketCount, "socket-count", "sc", 0, "UDP sockets per run (0 = scale to cores)"), + flagSet.IntVar(&options.UDPSize, "udp-size", 0, "EDNS0 UDP payload size (0 = 1232; <512 disables EDNS0)"), + flagSet.BoolVar(&options.NoRecurse, "norecurse", false, "Send non-recursive queries (RD=0)"), + flagSet.BoolVar(&options.Sticky, "sticky", false, "Do not rotate resolver on retry"), + flagSet.BoolVarP(&options.ResolverHealth, "resolver-health", "rhz", false, "De-weight failing resolvers via health scoring"), + flagSet.BoolVarP(&options.AdaptiveConcurrency, "adaptive-concurrency", "acy", false, "Adapt in-flight concurrency to packet loss"), + flagSet.BoolVarP(&options.CrossCheck, "cross-check", "cc", false, "Re-verify positive answers on a second resolver"), + flagSet.BoolVarP(&options.ExtendedInput, "extended-input", "ei", false, "Parse 'name [resolver ...]' input lines"), + flagSet.BoolVar(&options.NoVerifyIP, "no-verify-ip", false, "Disable reply source-IP verification"), + flagSet.BoolVar(&options.NoTCPFallback, "no-tcp-fallback", false, "Disable TCP fallback on truncated answers"), + flagSet.BoolVarP(&options.Iterative, "iterative", "it", false, "Recurse from root servers (no resolver list needed); caches delegations"), + ) + + flagSet.CreateGroup("distributed", "Distributed", + flagSet.StringVar(&options.Shard, "shard", "", "Process only shard m of n for distributed runs (e.g. 2/8)"), + flagSet.StringVarP(&options.ResumeFile, "resume", "rs", "", "Checkpoint file for crash-safe stop/resume"), + ) + flagSet.CreateGroup("debug", "Debug", flagSet.BoolVar(&options.Silent, "silent", false, "Show only subdomains in output"), flagSet.BoolVar(&options.Version, "version", false, "Show version of shuffledns"), diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go index cd95a580..faaa9ae8 100644 --- a/pkg/runner/runner.go +++ b/pkg/runner/runner.go @@ -3,8 +3,8 @@ package runner import ( "context" "errors" + "fmt" "os" - "os/exec" "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/shuffledns/pkg/massdns" @@ -23,16 +23,6 @@ func New(options *Options) (*Runner, error) { options: options, } - // Setup the massdns binary path if none was give. - // If no valid path found, return an error - if options.MassdnsPath == "" { - options.MassdnsPath = runner.findBinary() - if options.MassdnsPath == "" { - return nil, errors.New("could not find massdns binary") - } - gologger.Debug().Msgf("Discovered massdns binary at %s\n", options.MassdnsPath) - } - // Create a temporary directory that will be removed at the end // of enumeration process. dir, err := os.MkdirTemp(options.Directory, "shuffledns-*") @@ -49,71 +39,14 @@ func (r *Runner) Close() { _ = os.RemoveAll(r.tempDir) } -// findBinary searches for massdns binary in various pre-defined paths -// only linux and macos paths are supported rn -func (r *Runner) findBinary() string { - otherCommonLocations := []string{ - "/usr/bin/massdns", - "/usr/local/bin/massdns", - "/data/data/com.termux/files/usr/bin/massdns", - } - - for _, file := range otherCommonLocations { - if fileutil.FileExists(file) { - return file - } - } - - file, err := exec.LookPath("massdns") - if err != nil { - return "" - } - - return file -} - -// RunEnumeration sets up the input layer for giving input to massdns -// binary and runs the actual enumeration -func (r *Runner) RunEnumeration() { - // Handle only wildcard filtering on existing massdns output - if r.options.MassdnsRaw != "" { - r.processExistingOutput() - return - } - - // Handle a domain to bruteforce with wordlist - if r.options.Wordlist != "" { - r.processDomain() - return - } - - // Handle a list of subdomains to resolve - if r.options.SubdomainsList != "" || fileutil.HasStdin() { - r.processSubdomains() - return - } -} - -// processDomain processes the bruteforce for a domain using a wordlist -func (r *Runner) processDomain() { - // Read the input wordlist for bruteforce generation - inputFile, err := os.Open(r.options.Wordlist) - if err != nil { - gologger.Error().Msgf("Could not read bruteforce wordlist (%s): %s\n", r.options.Wordlist, err) - return - } - defer func() { - _ = inputFile.Close() - }() - - gologger.Info().Msgf("Started generating bruteforce permutation with streaming processing\n") - - // Create massdns instance for processing chunks - massdns, err := massdns.New(massdns.Options{ +// massdnsOptions builds the massdns.Options from the runner options. Keeping it +// in one place avoids drift between the bruteforce, resolve and raw-input paths +// and makes the embeddable surface easy to reason about. +func (r *Runner) massdnsOptions() massdns.Options { + return massdns.Options{ Domains: r.options.Domains, AutoExtractRootDomains: r.options.AutoExtractRootDomains, Retries: r.options.Retries, - MassdnsPath: r.options.MassdnsPath, Threads: r.options.Threads, WildcardsThreads: r.options.WildcardThreads, ResolversFile: r.options.ResolversFile, @@ -124,132 +57,120 @@ func (r *Runner) processDomain() { MassdnsRaw: r.options.MassdnsRaw, StrictWildcard: r.options.StrictWildcard, WildcardOutputFile: r.options.WildcardOutputFile, - MassDnsCmd: r.options.MassDnsCmd, - KeepStderr: r.options.KeepStderr, - BatchSize: r.options.BatchSize, FilterInternalIPs: r.options.FilterInternalIPs, + QueryType: r.options.QueryType, + BatchMode: r.options.BatchMode, + SocketCount: r.options.SocketCount, + UDPSize: r.options.UDPSize, + QPS: r.options.QPS, + NoRecurse: r.options.NoRecurse, + Sticky: r.options.Sticky, + ResolverHealth: r.options.ResolverHealth, + AdaptiveConcurrency: r.options.AdaptiveConcurrency, + CrossCheck: r.options.CrossCheck, + ExtendedInput: r.options.ExtendedInput, + NoVerifyIP: r.options.NoVerifyIP, + NoTCPFallback: r.options.NoTCPFallback, + Iterative: r.options.Iterative, + Shard: r.options.Shard, + ResumeFile: r.options.ResumeFile, OnResult: r.options.OnResult, - }) + } +} + +// RunEnumeration sets up the input layer for giving input to the native +// resolver and runs the actual enumeration. It returns an error so the process +// is fully embeddable (the CLI is responsible for logging/exit codes). +func (r *Runner) RunEnumeration() error { + switch { + case r.options.MassdnsRaw != "": + return r.processExistingOutput() + case r.options.Wordlist != "": + return r.processDomain() + case r.options.SubdomainsList != "" || fileutil.HasStdin(): + return r.processSubdomains() + default: + return errors.New("no input provided: set a wordlist, a subdomains list, stdin, or raw massdns input") + } +} + +// dumpWildcards writes the discovered wildcard IPs when requested. +func (r *Runner) dumpWildcards(instance *massdns.Instance) { + if r.options.WildcardOutputFile != "" { + _ = instance.DumpWildcardsToFile(r.options.WildcardOutputFile) + } +} + +// processDomain processes the bruteforce for a domain using a wordlist +func (r *Runner) processDomain() error { + inputFile, err := os.Open(r.options.Wordlist) if err != nil { - gologger.Error().Msgf("Could not create massdns client: %s\n", err) - return + return fmt.Errorf("could not read bruteforce wordlist (%s): %w", r.options.Wordlist, err) } + defer func() { + _ = inputFile.Close() + }() + + gologger.Info().Msgf("Started generating bruteforce permutation with streaming processing\n") - // Use streaming processing with batcher - err = massdns.ProcessDomainStreaming(context.Background(), inputFile) + instance, err := massdns.New(r.massdnsOptions()) if err != nil { - gologger.Error().Msgf("Could not process domain with streaming: %s\n", err) - return + return fmt.Errorf("could not create massdns client: %w", err) } - if r.options.WildcardOutputFile != "" { - _ = massdns.DumpWildcardsToFile(r.options.WildcardOutputFile) + if err := instance.ProcessDomainStreaming(context.Background(), inputFile); err != nil { + return fmt.Errorf("could not process domain with streaming: %w", err) } + r.dumpWildcards(instance) gologger.Info().Msgf("Finished resolving.\n") + return nil } -// processSubdomain processes the resolving for a list of subdomains -func (r *Runner) processSubdomains() { - // Create massdns instance for processing chunks - massdns, err := massdns.New(massdns.Options{ - Domains: r.options.Domains, - AutoExtractRootDomains: r.options.AutoExtractRootDomains, - Retries: r.options.Retries, - MassdnsPath: r.options.MassdnsPath, - Threads: r.options.Threads, - WildcardsThreads: r.options.WildcardThreads, - ResolversFile: r.options.ResolversFile, - TrustedResolvers: r.options.TrustedResolvers, - TempDir: r.tempDir, - OutputFile: r.options.Output, - Json: r.options.Json, - MassdnsRaw: r.options.MassdnsRaw, - StrictWildcard: r.options.StrictWildcard, - WildcardOutputFile: r.options.WildcardOutputFile, - MassDnsCmd: r.options.MassDnsCmd, - KeepStderr: r.options.KeepStderr, - BatchSize: r.options.BatchSize, - FilterInternalIPs: r.options.FilterInternalIPs, - OnResult: r.options.OnResult, - }) +// processSubdomains processes the resolving for a list of subdomains +func (r *Runner) processSubdomains() error { + instance, err := massdns.New(r.massdnsOptions()) if err != nil { - gologger.Error().Msgf("Could not create massdns client: %s\n", err) - return + return fmt.Errorf("could not create massdns client: %w", err) } - // Handle stdin or file input if fileutil.HasStdin() && r.options.SubdomainsList == "" { - // Use streaming processing for stdin gologger.Info().Msgf("Processing subdomains from stdin with streaming\n") - err = massdns.ProcessSubdomainsStreaming(context.Background(), os.Stdin) - if err != nil { - gologger.Error().Msgf("Could not process subdomains with streaming: %s\n", err) - return + if err := instance.ProcessSubdomainsStreaming(context.Background(), os.Stdin); err != nil { + return fmt.Errorf("could not process subdomains from stdin: %w", err) } } else { - // Use streaming processing for file subdomainFile, err := os.Open(r.options.SubdomainsList) if err != nil { - gologger.Error().Msgf("Could not open subdomain list (%s): %s\n", r.options.SubdomainsList, err) - return + return fmt.Errorf("could not open subdomain list (%s): %w", r.options.SubdomainsList, err) } defer func() { _ = subdomainFile.Close() }() gologger.Info().Msgf("Processing subdomains from file with streaming\n") - err = massdns.ProcessSubdomainsStreaming(context.Background(), subdomainFile) - if err != nil { - gologger.Error().Msgf("Could not process subdomains with streaming from file: %s\n", err) - return + if err := instance.ProcessSubdomainsStreaming(context.Background(), subdomainFile); err != nil { + return fmt.Errorf("could not process subdomains from file: %w", err) } } - if r.options.WildcardOutputFile != "" { - _ = massdns.DumpWildcardsToFile(r.options.WildcardOutputFile) - } - + r.dumpWildcards(instance) gologger.Info().Msgf("Finished resolving.\n") + return nil } // processExistingOutput processes existing massdns output for wildcard filtering -func (r *Runner) processExistingOutput() { - massdns, err := massdns.New(massdns.Options{ - Domains: r.options.Domains, - AutoExtractRootDomains: r.options.AutoExtractRootDomains, - Retries: r.options.Retries, - MassdnsPath: r.options.MassdnsPath, - Threads: r.options.Threads, - WildcardsThreads: r.options.WildcardThreads, - ResolversFile: r.options.ResolversFile, - TrustedResolvers: r.options.TrustedResolvers, - TempDir: r.tempDir, - OutputFile: r.options.Output, - Json: r.options.Json, - MassdnsRaw: r.options.MassdnsRaw, - StrictWildcard: r.options.StrictWildcard, - WildcardOutputFile: r.options.WildcardOutputFile, - MassDnsCmd: r.options.MassDnsCmd, - KeepStderr: r.options.KeepStderr, - BatchSize: r.options.BatchSize, - FilterInternalIPs: r.options.FilterInternalIPs, - OnResult: r.options.OnResult, - }) +func (r *Runner) processExistingOutput() error { + instance, err := massdns.New(r.massdnsOptions()) if err != nil { - gologger.Error().Msgf("Could not create massdns client: %s\n", err) - return + return fmt.Errorf("could not create massdns client: %w", err) } - err = massdns.Run(context.Background()) - if err != nil { - gologger.Error().Msgf("Could not process existing massdns output: %s\n", err) - return - } - - if r.options.WildcardOutputFile != "" { - _ = massdns.DumpWildcardsToFile(r.options.WildcardOutputFile) + if err := instance.Run(context.Background()); err != nil { + return fmt.Errorf("could not process existing massdns output: %w", err) } + r.dumpWildcards(instance) gologger.Info().Msgf("Finished processing existing output.\n") + return nil } diff --git a/pkg/runner/validate.go b/pkg/runner/validate.go index 27813d9a..0fb603c2 100644 --- a/pkg/runner/validate.go +++ b/pkg/runner/validate.go @@ -17,18 +17,19 @@ func (options *Options) validateOptions() error { return errors.New("both verbose and silent mode specified") } - // Check if a list of resolvers was provided and it exists - if !fileutil.FileExists(options.ResolversFile) { - return errors.New("resolver file doesn't exists") - } - - // Check if resolvers are blank - if blank, err := fileutil.IsEmpty(options.ResolversFile); err == nil { - if blank { - return errors.New("empty resolver list specified") + // Iterative mode recurses from the root servers and needs no resolver list. + // For the stub-resolver path, a non-empty resolver file is required. + if !options.Iterative { + if !fileutil.FileExists(options.ResolversFile) { + return errors.New("resolver file doesn't exists") + } + if blank, err := fileutil.IsEmpty(options.ResolversFile); err == nil { + if blank { + return errors.New("empty resolver list specified") + } + } else { + return fmt.Errorf("could not read resolvers: %w", err) } - } else { - return fmt.Errorf("could not read resolvers: %w", err) } switch options.Mode { diff --git a/pkg/shard/shard.go b/pkg/shard/shard.go new file mode 100644 index 00000000..90a97182 --- /dev/null +++ b/pkg/shard/shard.go @@ -0,0 +1,94 @@ +// Package shard implements deterministic workload sharding for distributed, +// coordinated resolution. Running N instances each with a distinct shard +// (1/N .. N/N) over the same input partitions the work into disjoint subsets +// with no central coordinator: every instance independently decides which names +// it owns via a stable hash, so the union covers the input exactly once. +// +// This mirrors masscan-style sharding but for DNS names: because names cannot be +// split arithmetically like an address range, each instance reads the full input +// and cheaply hashes every candidate to decide ownership. +package shard + +import ( + "fmt" + "strconv" + "strings" +) + +// Shard identifies one partition (Index) out of Total. Index is 1-based and in +// the range [1, Total], matching the "m/n" convention. +type Shard struct { + Index int + Total int +} + +// Parse parses an "m/n" shard specification (e.g. "2/8"). An empty string or +// "1/1" yields a disabled shard that owns everything. +func Parse(s string) (Shard, error) { + s = strings.TrimSpace(s) + if s == "" { + return Shard{Index: 1, Total: 1}, nil + } + parts := strings.SplitN(s, "/", 2) + if len(parts) != 2 { + return Shard{}, fmt.Errorf("invalid shard %q: expected m/n", s) + } + idx, err1 := strconv.Atoi(strings.TrimSpace(parts[0])) + total, err2 := strconv.Atoi(strings.TrimSpace(parts[1])) + if err1 != nil || err2 != nil { + return Shard{}, fmt.Errorf("invalid shard %q: m and n must be integers", s) + } + sh := Shard{Index: idx, Total: total} + if err := sh.Validate(); err != nil { + return Shard{}, err + } + return sh, nil +} + +// Validate checks the shard is internally consistent. +func (s Shard) Validate() error { + if s.Total < 1 { + return fmt.Errorf("shard total must be >= 1, got %d", s.Total) + } + if s.Index < 1 || s.Index > s.Total { + return fmt.Errorf("shard index %d out of range [1,%d]", s.Index, s.Total) + } + return nil +} + +// Enabled reports whether sharding actually partitions the input (Total > 1). +func (s Shard) Enabled() bool { return s.Total > 1 } + +// Owns reports whether this shard is responsible for key. When sharding is +// disabled it always returns true. The hash is a stable FNV-1a over the +// lower-cased key, so the partitioning is identical across machines, runs and +// architectures. +func (s Shard) Owns(key string) bool { + if !s.Enabled() { + return true + } + return Bucket(key, s.Total) == s.Index-1 +} + +// Bucket returns the stable bucket index in [0, total) for key, using FNV-1a +// over the lower-cased key. Implemented inline (no hash.Hash allocation) so it +// stays cheap when called once per candidate name. +func Bucket(key string, total int) int { + if total <= 1 { + return 0 + } + const ( + offset64 = 14695981039346656037 // FNV-1a 64-bit offset basis + prime64 = 1099511628211 // FNV-1a 64-bit prime + ) + var h uint64 = offset64 + for i := 0; i < len(key); i++ { + c := key[i] + if c >= 'A' && c <= 'Z' { + c += 'a' - 'A' + } + h ^= uint64(c) + h *= prime64 + } + return int(h % uint64(total)) +} diff --git a/pkg/shard/shard_test.go b/pkg/shard/shard_test.go new file mode 100644 index 00000000..db23ee8e --- /dev/null +++ b/pkg/shard/shard_test.go @@ -0,0 +1,95 @@ +package shard + +import ( + "fmt" + "testing" +) + +func TestParse(t *testing.T) { + cases := map[string]struct { + idx, total int + err bool + }{ + "": {1, 1, false}, + "1/1": {1, 1, false}, + "2/8": {2, 8, false}, + "8/8": {8, 8, false}, + "0/8": {0, 0, true}, + "9/8": {0, 0, true}, + "2/0": {0, 0, true}, + "x/8": {0, 0, true}, + "2": {0, 0, true}, + "2/3/4": {0, 0, true}, + } + for in, want := range cases { + got, err := Parse(in) + if want.err { + if err == nil { + t.Errorf("Parse(%q): expected error", in) + } + continue + } + if err != nil { + t.Errorf("Parse(%q): unexpected error %v", in, err) + continue + } + if got.Index != want.idx || got.Total != want.total { + t.Errorf("Parse(%q): got %d/%d want %d/%d", in, got.Index, got.Total, want.idx, want.total) + } + } +} + +func TestEnabled(t *testing.T) { + if (Shard{1, 1}).Enabled() { + t.Fatal("1/1 should be disabled") + } + if !(Shard{1, 4}).Enabled() { + t.Fatal("1/4 should be enabled") + } +} + +func TestDisabledOwnsEverything(t *testing.T) { + s := Shard{1, 1} + for i := 0; i < 100; i++ { + if !s.Owns(fmt.Sprintf("host%d.example.com", i)) { + t.Fatal("disabled shard must own every key") + } + } +} + +// TestPartitionExactlyOnce is the core guarantee: across N shards every name is +// owned by exactly one shard. +func TestPartitionExactlyOnce(t *testing.T) { + const total = 8 + shards := make([]Shard, total) + for i := range shards { + shards[i] = Shard{Index: i + 1, Total: total} + } + counts := make([]int, total) + for i := 0; i < 10000; i++ { + name := fmt.Sprintf("sub-%d.target.com", i) + owners := 0 + for j, s := range shards { + if s.Owns(name) { + owners++ + counts[j]++ + } + } + if owners != 1 { + t.Fatalf("name %q owned by %d shards, want 1", name, owners) + } + } + // sanity: distribution should be roughly even (no empty shard) + for j, c := range counts { + if c == 0 { + t.Fatalf("shard %d got no names (bad distribution)", j+1) + } + } +} + +func TestStableAcrossCase(t *testing.T) { + s := Shard{Index: 3, Total: 8} + if s.Owns("Host.EXAMPLE.com") != s.Owns("host.example.com") { + t.Fatal("ownership must be case-insensitive/stable") + } +} diff --git a/pkg/store/store.go b/pkg/store/store.go index c6c1aed4..607b5999 100644 --- a/pkg/store/store.go +++ b/pkg/store/store.go @@ -1,129 +1,135 @@ +// Package store provides an in-memory ip -> hostnames index used for +// deduplication and wildcard removal. +// +// It previously persisted to LevelDB on disk (with a JSON marshal and +// read-modify-write on every append plus background compaction). The native +// resolver streams structured results in-process, so an in-memory map is both +// simpler and considerably faster; the public API is kept stable. package store import ( - "encoding/json" - "os" - "strings" - - mapsutil "github.com/projectdiscovery/utils/maps" - "github.com/syndtr/goleveldb/leveldb" - "github.com/syndtr/goleveldb/leveldb/opt" + "sort" + "sync" ) -const Megabyte = 1 << 20 - -// Store is a storage for ip based wildcard removal +// Store is an in-memory storage for ip based deduplication and wildcard removal. type Store struct { - DB *leveldb.DB + mu sync.RWMutex + data map[string]map[string]struct{} } -// New creates a new storage for ip based wildcard removal -func New(dbPath string) (*Store, error) { - storeDb, err := os.MkdirTemp(dbPath, "shuffledns-db-") - if err != nil { - return nil, err - } - db, err := leveldb.OpenFile(storeDb, &opt.Options{ - // Optimize for disk space reduction - CompactionTableSize: 64 * Megabyte, // Reduced from 256MB for more frequent compaction - WriteBuffer: 2 * Megabyte, // Reduced from 4MB for more frequent flushing - WriteL0SlowdownTrigger: 4, // Trigger slowdown earlier - WriteL0PauseTrigger: 8, // Trigger pause earlier - BlockSize: 2 * 1024, // Reduced from 4KB for better compression of small records - BlockCacheCapacity: 4 * Megabyte, // Reduced from 8MB to lower memory usage - }) - if err != nil { - return nil, err - } - return &Store{DB: db}, nil +// New creates a new in-memory store. The path argument is accepted for API +// compatibility and ignored. +func New(_ string) (*Store, error) { + return &Store{data: make(map[string]map[string]struct{})}, nil } -// New creates a new ip-hostname pair in the map +// New creates a new ip-hostname pair in the map. func (s *Store) New(ip, hostname string) error { - hostnameMap := map[string]struct{}{hostname: {}} - jsonData, err := json.Marshal(hostnameMap) - if err != nil { - return err - } - return s.DB.Put([]byte(ip), jsonData, nil) + return s.Append(ip, hostname) } -// Exists indicates if an IP exists in the map +// Exists indicates if an IP exists in the map. func (s *Store) Exists(ip string) bool { - ok, err := s.DB.Has([]byte(ip), nil) - return err == nil && ok + s.mu.RLock() + defer s.mu.RUnlock() + _, ok := s.data[ip] + return ok } -// Get gets the meta-information for an IP address from the map. +// GetHostnames returns the comma separated hostnames stored for an IP. func (s *Store) GetHostnames(ip string) string { - data, err := s.DB.Get([]byte(ip), nil) - if err != nil { - return "" - } + s.mu.RLock() + defer s.mu.RUnlock() - var hostnameMap map[string]struct{} - if err := json.Unmarshal(data, &hostnameMap); err != nil { + hostnameMap, ok := s.data[ip] + if !ok { return "" } - - return strings.Join(mapsutil.GetKeys(hostnameMap), ",") + hostnames := make([]string, 0, len(hostnameMap)) + for hostname := range hostnameMap { + hostnames = append(hostnames, hostname) + } + sort.Strings(hostnames) + return joinComma(hostnames) } +// Append adds one or more hostnames to an IP, deduplicating automatically. func (s *Store) Append(ip string, hostnames ...string) error { - // Get existing hostnames - var hostnameMap map[string]struct{} - existingData, err := s.DB.Get([]byte(ip), nil) - if err == nil && len(existingData) > 0 { - if err := json.Unmarshal(existingData, &hostnameMap); err != nil { - // If unmarshaling fails, start with empty map - hostnameMap = make(map[string]struct{}) - } - } else { - hostnameMap = make(map[string]struct{}) - } + s.mu.Lock() + defer s.mu.Unlock() - // Add new hostnames to map (automatic deduplication) + hostnameMap, ok := s.data[ip] + if !ok { + hostnameMap = make(map[string]struct{}, len(hostnames)) + s.data[ip] = hostnameMap + } for _, hostname := range hostnames { hostnameMap[hostname] = struct{}{} } - - // Marshal and store - jsonData, err := json.Marshal(hostnameMap) - if err != nil { - return err - } - - return s.DB.Put([]byte(ip), jsonData, nil) + return nil } -// Delete deletes the records for an IP from store. +// Delete removes the records for an IP from the store. func (s *Store) Delete(ip string) error { - return s.DB.Delete([]byte(ip), nil) + s.mu.Lock() + defer s.mu.Unlock() + delete(s.data, ip) + return nil } +// Close releases all resources held by the store. func (s *Store) Close() { - _ = s.DB.Close() + s.mu.Lock() + defer s.mu.Unlock() + s.data = nil } +// Iterate walks every ip and its hostnames. counter is the number of distinct +// hostnames pointing at the ip (used by the wildcard heuristic). func (s *Store) Iterate(f func(ip string, hostnames []string, counter int)) { - iter := s.DB.NewIterator(nil, nil) - defer iter.Release() - - for iter.Next() { - ip := string(iter.Key()) - - var hostnameMap map[string]struct{} - if err := json.Unmarshal(iter.Value(), &hostnameMap); err != nil { - continue - } - - // Convert map keys to slice + // snapshot under lock to avoid holding it during the callback (which may + // perform network I/O during wildcard filtering) + s.mu.RLock() + ips := make([]string, 0, len(s.data)) + snapshot := make(map[string][]string, len(s.data)) + for ip, hostnameMap := range s.data { hostnames := make([]string, 0, len(hostnameMap)) for hostname := range hostnameMap { hostnames = append(hostnames, hostname) } + sort.Strings(hostnames) + snapshot[ip] = hostnames + ips = append(ips, ip) + } + s.mu.RUnlock() + + // Iterate in sorted order so output is deterministic across runs (the map + // backing replaced a LevelDB store that iterated in sorted key order). + sort.Strings(ips) + for _, ip := range ips { + hostnames := snapshot[ip] + f(ip, hostnames, len(hostnames)) + } +} - counter := len(hostnames) - f(ip, hostnames, counter) +func joinComma(values []string) string { + switch len(values) { + case 0: + return "" + case 1: + return values[0] + } + n := len(values) - 1 + for _, v := range values { + n += len(v) + } + out := make([]byte, 0, n) + for i, v := range values { + if i > 0 { + out = append(out, ',') + } + out = append(out, v...) } + return string(out) } diff --git a/pkg/zonewalk/nsec3.go b/pkg/zonewalk/nsec3.go new file mode 100644 index 00000000..5f855efb --- /dev/null +++ b/pkg/zonewalk/nsec3.go @@ -0,0 +1,178 @@ +package zonewalk + +import ( + "context" + "fmt" + "math/rand/v2" + "strings" + "time" + + "github.com/miekg/dns" +) + +// maxNSEC3Iterations caps the SHA-1 iteration count we are willing to process. +// dns.HashName runs iter extra SHA-1 rounds per hash and the target controls +// iter (up to 65535); crack does len(candidates)*len(harvested) hashes, so a +// hostile zone advertising a huge count could burn billions of SHA-1 ops. +// RFC 9276 treats anything above a small number as unreasonable. +const maxNSEC3Iterations = 500 + +// CrackConfig controls an NSEC3 harvest-and-crack. +type CrackConfig struct { + // Zone is the NSEC3-signed apex to enumerate (e.g. "example.com"). + Zone string + // Resolvers are DNSSEC-aware resolvers (host or host:port); the first is used. + Resolvers []string + // Timeout is the per-query timeout. Default 5s. + Timeout time.Duration + // MaxQueries bounds the harvest phase (number of probe queries). Default 2000. + MaxQueries int + // Candidates is the wordlist of labels to crack against the harvested hash + // ring (e.g. "www", "mail"); each is hashed and matched offline. + Candidates []string + // OnName fires for each recovered (cracked) name. + OnName func(string) +} + +// CrackResult is the outcome of an NSEC3 harvest-and-crack. +type CrackResult struct { + Zone string + Salt string + Iterations uint16 + HashAlg uint8 + HarvestedHashes int // distinct NSEC3 records collected (existing-name hashes) + Names []string // recovered existing names (no trailing dot) + Saturated bool // harvest reached saturation (ring likely complete) +} + +// CrackNSEC3 enumerates an NSEC3-signed zone offline-style: it first harvests +// the zone's NSEC3 records by probing for non-existent names (each NXDOMAIN +// proof leaks NSEC3 records that cover existing-name hashes), walking until the +// hash ring saturates, then cracks the supplied wordlist by matching each +// candidate's hash against the harvested records. Only true positives are +// reported. Completeness depends on harvest coverage and the wordlist; this is +// the standard nsec3walker approach. +func CrackNSEC3(ctx context.Context, cfg CrackConfig) (*CrackResult, error) { + if strings.TrimSpace(cfg.Zone) == "" { + return nil, fmt.Errorf("zone is required") + } + if len(cfg.Resolvers) == 0 { + return nil, fmt.Errorf("at least one resolver is required") + } + if cfg.Timeout <= 0 { + cfg.Timeout = 5 * time.Second + } + if cfg.MaxQueries <= 0 { + cfg.MaxQueries = 2000 + } + server := normalize(cfg.Resolvers[0]) + apex := dns.Fqdn(strings.ToLower(cfg.Zone)) + client := &dns.Client{Timeout: cfg.Timeout} + + res := &CrackResult{Zone: cfg.Zone} + + // harvested NSEC3 records, deduplicated by owner name. + harvested := map[string]*dns.NSEC3{} + collect := func(resp *dns.Msg) int { + added := 0 + for _, rr := range append(append([]dns.RR{}, resp.Answer...), resp.Ns...) { + n3, ok := rr.(*dns.NSEC3) + if !ok { + continue + } + key := strings.ToLower(n3.Hdr.Name) + if _, dup := harvested[key]; !dup { + harvested[key] = n3 + added++ + if res.Salt == "" && res.Iterations == 0 { + res.Salt, res.Iterations, res.HashAlg = n3.Salt, n3.Iterations, n3.Hash + } + } + } + return added + } + + covered := func(name string) bool { + for _, n3 := range harvested { + if n3.Cover(name) || n3.Match(name) { + return true + } + } + return false + } + + // harvest: probe random non-existent names; each proof extends ring coverage. + const saturation = 64 // consecutive already-covered probes => ring ~complete + miss := 0 + for q := 0; q < cfg.MaxQueries; q++ { + if err := ctx.Err(); err != nil { + break + } + probe := randomLabel() + "." + apex + if len(harvested) > 0 && covered(probe) { + miss++ + if miss >= saturation { + res.Saturated = true + break + } + continue + } + miss = 0 + resp, err := queryDO(ctx, client, server, probe, dns.TypeA) + if err != nil { + continue + } + collect(resp) + // Bail before any expensive Cover/Match on the next iteration if the zone + // advertises an abusive iteration count. + if res.Iterations > maxNSEC3Iterations { + res.HarvestedHashes = len(harvested) + return res, fmt.Errorf("nsec3 iterations %d exceed cap %d; refusing to crack (DoS risk)", res.Iterations, maxNSEC3Iterations) + } + } + res.HarvestedHashes = len(harvested) + + // crack: match each candidate's hash against the harvested ring. + seen := map[string]struct{}{} + for _, label := range cfg.Candidates { + label = strings.TrimSpace(strings.ToLower(label)) + if label == "" { + continue + } + name := label + "." + apex + for _, n3 := range harvested { + if n3.Match(name) { + owner := strings.TrimSuffix(name, ".") + if _, dup := seen[owner]; !dup { + seen[owner] = struct{}{} + res.Names = append(res.Names, owner) + if cfg.OnName != nil { + cfg.OnName(owner) + } + } + break + } + } + } + return res, nil +} + +func queryDO(ctx context.Context, client *dns.Client, server, name string, qtype uint16) (*dns.Msg, error) { + m := new(dns.Msg) + m.SetQuestion(dns.Fqdn(name), qtype) + m.RecursionDesired = true + m.SetEdns0(4096, true) // DO bit + resp, _, err := client.ExchangeContext(ctx, m, server) + return resp, err +} + +const labelAlphabet = "abcdefghijklmnopqrstuvwxyz0123456789" + +func randomLabel() string { + n := 10 + rand.IntN(6) + b := make([]byte, n) + for i := range b { + b[i] = labelAlphabet[rand.IntN(len(labelAlphabet))] + } + return string(b) +} diff --git a/pkg/zonewalk/nsec3_test.go b/pkg/zonewalk/nsec3_test.go new file mode 100644 index 00000000..e8de70a8 --- /dev/null +++ b/pkg/zonewalk/nsec3_test.go @@ -0,0 +1,94 @@ +package zonewalk + +import ( + "context" + "sort" + "strings" + "testing" + + "github.com/miekg/dns" +) + +// buildNSEC3Ring builds a synthetic NSEC3 ring for the given existing names +// under apex, using the supplied salt/iterations. Returns the ring records. +func buildNSEC3Ring(t *testing.T, apex, salt string, iter uint16, names []string) []*dns.NSEC3 { + t.Helper() + type he struct { + name, hash string + } + var hs []he + for _, n := range names { + fq := dns.Fqdn(strings.ToLower(n)) + hs = append(hs, he{fq, dns.HashName(fq, dns.SHA1, iter, salt)}) + } + sort.Slice(hs, func(i, j int) bool { return hs[i].hash < hs[j].hash }) + + var ring []*dns.NSEC3 + for i, h := range hs { + next := hs[(i+1)%len(hs)].hash + ring = append(ring, &dns.NSEC3{ + Hdr: dns.RR_Header{Name: h.hash + "." + apex, Rrtype: dns.TypeNSEC3, Class: dns.ClassINET, Ttl: 60}, + Hash: dns.SHA1, + Flags: 0, + Iterations: iter, + SaltLength: uint8(len(salt) / 2), + Salt: salt, + HashLength: 20, + NextDomain: next, + TypeBitMap: []uint16{dns.TypeA, dns.TypeRRSIG}, + }) + } + return ring +} + +func TestCrackNSEC3(t *testing.T) { + apex := "example.com." + salt := "deadbeef" + var iter uint16 = 5 + // existing names in the zone (the apex is part of the ring too). + existing := []string{"example.com.", "www.example.com.", "mail.example.com.", "ftp.example.com."} + ring := buildNSEC3Ring(t, apex, salt, iter, existing) + + addr, stop := startServer(t, func(w dns.ResponseWriter, r *dns.Msg) { + m := new(dns.Msg) + m.SetReply(r) + m.Rcode = dns.RcodeNameError + // leak the full NSEC3 ring on every NXDOMAIN proof (a permissive but + // valid model: real servers leak the covering subset per query). + for _, n3 := range ring { + m.Ns = append(m.Ns, n3) + } + _ = w.WriteMsg(m) + }) + defer stop() + + var recovered []string + res, err := CrackNSEC3(context.Background(), CrackConfig{ + Zone: "example.com", + Resolvers: []string{addr}, + MaxQueries: 50, + Candidates: []string{"www", "mail", "ftp", "doesnotexist", "admin"}, + OnName: func(n string) { recovered = append(recovered, n) }, + }) + if err != nil { + t.Fatalf("CrackNSEC3: %v", err) + } + if res.HarvestedHashes != len(ring) { + t.Fatalf("expected %d harvested NSEC3 records, got %d", len(ring), res.HarvestedHashes) + } + got := map[string]bool{} + for _, n := range res.Names { + got[n] = true + } + for _, want := range []string{"www.example.com", "mail.example.com", "ftp.example.com"} { + if !got[want] { + t.Errorf("expected to crack %q, recovered %v", want, res.Names) + } + } + if got["doesnotexist.example.com"] || got["admin.example.com"] { + t.Errorf("cracked a non-existent name (false positive): %v", res.Names) + } + if len(recovered) != len(res.Names) { + t.Errorf("OnName fired %d times, result has %d", len(recovered), len(res.Names)) + } +} diff --git a/pkg/zonewalk/zonewalk.go b/pkg/zonewalk/zonewalk.go new file mode 100644 index 00000000..a4a73350 --- /dev/null +++ b/pkg/zonewalk/zonewalk.go @@ -0,0 +1,171 @@ +// Package zonewalk enumerates the contents of a DNS zone by following its NSEC +// chain (RFC 4034 §4). Zones signed with NSEC leak their full list of names: each +// NSEC record points to the next name in canonical order, so walking the chain +// from the apex recovers every owner name for free, without a wordlist. This is +// directly useful for subdomain enumeration. +// +// Zones signed with NSEC3 hash their names and cannot be walked online; Walk +// detects this and reports it (along with the NSEC3 parameters) so callers can +// decide whether to attempt offline hash cracking. +package zonewalk + +import ( + "context" + "fmt" + "net" + "strings" + "time" + + "github.com/miekg/dns" +) + +// Config controls a zone walk. +type Config struct { + // Zone is the apex to enumerate (e.g. "example.com"). + Zone string + // Resolvers are DNSSEC-aware resolvers to query (host or host:port). The + // first reachable one is used. + Resolvers []string + // Timeout is the per-query timeout. Default 5s. + Timeout time.Duration + // MaxNames caps the number of discovered names. A safety valve against + // pathological chains that emit endless distinct names. Default 1,000,000 + // when unset (<= 0); the visited-set and apex-wrap checks already stop + // well-formed chains, this bounds hostile ones. + MaxNames int + // OnName is an optional callback fired for each newly discovered name. + OnName func(string) +} + +// NSEC3Info holds the parameters of an NSEC3-signed zone. +type NSEC3Info struct { + Salt string + Iterations uint16 + HashAlg uint8 +} + +// Result is the outcome of a zone walk. +type Result struct { + Zone string + Names []string // discovered owner names (no trailing dot), excluding the apex + NSEC3 bool // zone uses NSEC3 (not walkable online) + NSEC3Param *NSEC3Info // populated when NSEC3 is true and parameters were seen +} + +// Walk follows the NSEC chain of cfg.Zone and returns the discovered names. +func Walk(ctx context.Context, cfg Config) (*Result, error) { + if strings.TrimSpace(cfg.Zone) == "" { + return nil, fmt.Errorf("zone is required") + } + if len(cfg.Resolvers) == 0 { + return nil, fmt.Errorf("at least one resolver is required") + } + if cfg.Timeout <= 0 { + cfg.Timeout = 5 * time.Second + } + if cfg.MaxNames <= 0 { + cfg.MaxNames = 1_000_000 + } + server := normalize(cfg.Resolvers[0]) + apex := dns.Fqdn(strings.ToLower(cfg.Zone)) + client := &dns.Client{Timeout: cfg.Timeout} + + res := &Result{Zone: cfg.Zone} + visited := map[string]struct{}{} + current := apex + + for { + resp, err := queryNSEC(ctx, client, server, current) + if err != nil { + if len(res.Names) > 0 { + // partial walk: return what we have rather than failing hard + return res, nil + } + return nil, fmt.Errorf("nsec query for %s failed: %w", current, err) + } + + if info, ok := nsec3Param(resp); ok { + res.NSEC3 = true + res.NSEC3Param = info + return res, nil + } + + next, ok := nextName(resp, current) + if !ok { + // no NSEC for this owner: chain ended or zone is not NSEC-signed + return res, nil + } + + // termination: the last NSEC wraps back to the apex (or we loop). + if next == apex || equalName(next, apex) { + return res, nil + } + if _, seen := visited[next]; seen { + return res, nil + } + visited[next] = struct{}{} + + name := strings.TrimSuffix(next, ".") + res.Names = append(res.Names, name) + if cfg.OnName != nil { + cfg.OnName(name) + } + if cfg.MaxNames > 0 && len(res.Names) >= cfg.MaxNames { + return res, nil + } + current = next + } +} + +func queryNSEC(ctx context.Context, client *dns.Client, server, name string) (*dns.Msg, error) { + m := new(dns.Msg) + m.SetQuestion(dns.Fqdn(name), dns.TypeNSEC) + m.RecursionDesired = true + m.SetEdns0(4096, true) // DO bit: request DNSSEC records + resp, _, err := client.ExchangeContext(ctx, m, server) + return resp, err +} + +// nextName finds the NSEC record covering `owner` and returns its NextDomain. +func nextName(resp *dns.Msg, owner string) (string, bool) { + for _, rr := range append(append([]dns.RR{}, resp.Answer...), resp.Ns...) { + if nsec, ok := rr.(*dns.NSEC); ok { + // prefer the record whose owner matches, but accept any NSEC if the + // owner-specific one is absent (covering record on NXDOMAIN proofs). + if equalName(nsec.Hdr.Name, owner) { + return dns.Fqdn(strings.ToLower(nsec.NextDomain)), true + } + } + } + // fall back to the first NSEC seen + for _, rr := range append(append([]dns.RR{}, resp.Answer...), resp.Ns...) { + if nsec, ok := rr.(*dns.NSEC); ok { + return dns.Fqdn(strings.ToLower(nsec.NextDomain)), true + } + } + return "", false +} + +func nsec3Param(resp *dns.Msg) (*NSEC3Info, bool) { + for _, rr := range append(append([]dns.RR{}, resp.Answer...), resp.Ns...) { + switch v := rr.(type) { + case *dns.NSEC3: + return &NSEC3Info{Salt: v.Salt, Iterations: v.Iterations, HashAlg: v.Hash}, true + case *dns.NSEC3PARAM: + return &NSEC3Info{Salt: v.Salt, Iterations: v.Iterations, HashAlg: v.Hash}, true + } + } + return nil, false +} + +func equalName(a, b string) bool { + return strings.EqualFold(dns.Fqdn(a), dns.Fqdn(b)) +} + +func normalize(r string) string { + r = strings.TrimSpace(r) + if _, _, err := net.SplitHostPort(r); err != nil { + return net.JoinHostPort(r, "53") + } + return r +} diff --git a/pkg/zonewalk/zonewalk_test.go b/pkg/zonewalk/zonewalk_test.go new file mode 100644 index 00000000..f5d3bd80 --- /dev/null +++ b/pkg/zonewalk/zonewalk_test.go @@ -0,0 +1,114 @@ +package zonewalk + +import ( + "context" + "net" + "strings" + "testing" + + "github.com/miekg/dns" +) + +func startServer(t *testing.T, handler dns.HandlerFunc) (string, func()) { + t.Helper() + pc, err := net.ListenPacket("udp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + srv := &dns.Server{PacketConn: pc, Handler: handler} + go func() { _ = srv.ActivateAndServe() }() + return pc.LocalAddr().String(), func() { _ = srv.Shutdown() } +} + +func TestWalkNSEC(t *testing.T) { + // a tiny NSEC-signed zone: apex -> a -> b -> c -> apex (wrap). + chain := map[string]string{ + "example.com.": "a.example.com.", + "a.example.com.": "b.example.com.", + "b.example.com.": "c.example.com.", + "c.example.com.": "example.com.", + } + addr, stop := startServer(t, func(w dns.ResponseWriter, r *dns.Msg) { + m := new(dns.Msg) + m.SetReply(r) + q := r.Question[0] + if q.Qtype == dns.TypeNSEC { + if next, ok := chain[strings.ToLower(q.Name)]; ok { + nsec := &dns.NSEC{ + Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeNSEC, Class: dns.ClassINET, Ttl: 60}, + NextDomain: next, + TypeBitMap: []uint16{dns.TypeA, dns.TypeNSEC}, + } + m.Answer = append(m.Answer, nsec) + } else { + m.Rcode = dns.RcodeNameError + } + } + _ = w.WriteMsg(m) + }) + defer stop() + + res, err := Walk(context.Background(), Config{ + Zone: "example.com", + Resolvers: []string{addr}, + }) + if err != nil { + t.Fatalf("Walk: %v", err) + } + if res.NSEC3 { + t.Fatal("did not expect NSEC3") + } + want := []string{"a.example.com", "b.example.com", "c.example.com"} + if len(res.Names) != len(want) { + t.Fatalf("got %v want %v", res.Names, want) + } + for i := range want { + if res.Names[i] != want[i] { + t.Fatalf("idx %d: got %q want %q", i, res.Names[i], want[i]) + } + } +} + +func TestNSEC3Detection(t *testing.T) { + // NSEC3PARAM in the authority section must flag the zone as NSEC3-signed. + m := new(dns.Msg) + m.Ns = append(m.Ns, &dns.NSEC3PARAM{ + Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeNSEC3PARAM, Class: dns.ClassINET, Ttl: 60}, + Hash: 1, + Iterations: 10, + Salt: "abcd", + }) + info, ok := nsec3Param(m) + if !ok || info == nil { + t.Fatalf("expected NSEC3 detection") + } + if info.Iterations != 10 || info.Salt != "abcd" { + t.Fatalf("bad NSEC3 params: %+v", info) + } +} + +func TestWalkMaxNames(t *testing.T) { + chain := map[string]string{ + "z.com.": "a.z.com.", "a.z.com.": "b.z.com.", "b.z.com.": "c.z.com.", "c.z.com.": "z.com.", + } + addr, stop := startServer(t, func(w dns.ResponseWriter, r *dns.Msg) { + m := new(dns.Msg) + m.SetReply(r) + if next, ok := chain[strings.ToLower(r.Question[0].Name)]; ok { + m.Answer = append(m.Answer, &dns.NSEC{ + Hdr: dns.RR_Header{Name: r.Question[0].Name, Rrtype: dns.TypeNSEC, Class: dns.ClassINET, Ttl: 60}, + NextDomain: next, TypeBitMap: []uint16{dns.TypeA}, + }) + } + _ = w.WriteMsg(m) + }) + defer stop() + + res, err := Walk(context.Background(), Config{Zone: "z.com", Resolvers: []string{addr}, MaxNames: 2}) + if err != nil { + t.Fatalf("Walk: %v", err) + } + if len(res.Names) != 2 { + t.Fatalf("expected MaxNames cap of 2, got %d (%v)", len(res.Names), res.Names) + } +}