diff --git a/CHANGELOG.md b/CHANGELOG.md index 41f8a11..1979aed 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - Resolved records are now captured during scans. `internal/dns` exposes `Resolve` and a `Record{Type, Value}` type; `scan.Event` carries `Records` for each resolved subdomain (A/AAAA today, extensible to CNAME and more). - `-format text|json|csv` flag (default `text`, byte-for-byte identical to prior output). JSON emits a buffered array of `{"subdomain", "records"}` objects; CSV streams `subdomain,type,value` rows with a header. The `-o` output file honors the selected format. Output formats are CLI-only for now (TUI-pending). +- `-rate ` flag (default 0 = unlimited) caps total DNS queries per second across the worker pool via a shared stdlib ticker gate inside `scan.Run`. The limiter respects context cancellation so `Ctrl+C` stays responsive. ## [0.5.1] - 2026-06-03 diff --git a/README.md b/README.md index a39091f..f1fa7cf 100644 --- a/README.md +++ b/README.md @@ -78,6 +78,7 @@ Or launch the interactive terminal UI with no flags: | Simulation Mode | Generate synthetic DNS results at a configurable hit rate, with zero network I/O | | Output Pipeline | Resolved domains to stdout (pipe-clean); progress and diagnostics to stderr | | Output Formats | Emit results as `text`, `json` (array of subdomain plus typed records), or `csv` via `-format` | +| Rate Limiting | Cap total DNS queries per second across the worker pool with `-rate` (context-aware) | | Interactive TUI | Form-based config and live-scrolling results via `-tui`; session values persisted |
@@ -200,6 +201,7 @@ make help # list all targets | `-force` | `false` | Continue scanning even if wildcard DNS is detected | | `-o ` | n/a | Write results to file in addition to stdout | | `-format ` | `text` | Output format: `text`, `json`, or `csv` | +| `-rate ` | `0` | Max DNS queries per second across all workers (0 = unlimited) | | `-v` | `false` | Verbose output: IPs, timings, per-query detail (stderr) | | `-progress` | `true` | Live progress line on stderr | | `-simulate` | `false` | Simulation mode: no real DNS queries | diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index af034e0..4dd0fcd 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -91,6 +91,7 @@ internal/tui/config.go — Session persistence (load/save ~/.config/sube * **`var wg sync.WaitGroup`**: A `sync.WaitGroup` waits for all worker goroutines to finish. * **Worker Goroutines Loop**: `cfg.Concurrency` goroutines are launched. Each reads prefixes from the channel, constructs the full domain, and calls `dns.ResolveDomainWithRetry()` (or `dns.SimulateResolution()` in simulate mode). * **Progress ticker**: A separate goroutine fires every second and emits `EventProgress` events so callers can update their display. + * **Rate limiter** (optional): when `cfg.Rate > 0`, a shared `time.Ticker` gate paces total DNS queries per second across the whole pool. Each worker waits on the gate before issuing a query, selecting on `ctx.Done()` so cancellation stays responsive. `0` means unlimited. * **Closing the Channel**: After all entries are sent, the channel is closed, signalling workers to exit. `wg.Wait()` blocks until all workers are done, then `EventDone` is emitted. * **Interactions**: `scan.Run` is the single entry point for scanning used by both the CLI output pipeline and the Bubble Tea TUI. It decouples the scan engine from any specific display layer. diff --git a/examples/advanced_usage.md b/examples/advanced_usage.md index 94a185f..88c293f 100644 --- a/examples/advanced_usage.md +++ b/examples/advanced_usage.md @@ -151,6 +151,16 @@ Simulation mode with verbose output shows fake IPs and timings: ./subenum -simulate -hit-rate 25 -v -w examples/sample_wordlist.txt example.com ``` +## Rate Limiting + +Use `-rate` to cap the total number of DNS queries per second across the whole worker pool. This is useful against rate-limited resolvers or to stay under a target query budget. `0` (the default) means unlimited: + +```bash +./subenum -w wordlist.txt -rate 50 example.com +``` + +The limiter is context-aware, so `Ctrl+C` stays responsive while workers are waiting on it. + ## Output Formats By default `subenum` prints human-readable `Found:` lines. Use `-format` to emit structured output instead. The `-o` file honors the same format. diff --git a/internal/scan/runner.go b/internal/scan/runner.go index ec7d4cc..d7b4015 100644 --- a/internal/scan/runner.go +++ b/internal/scan/runner.go @@ -21,6 +21,7 @@ type Config struct { Attempts int Force bool Verbose bool + Rate int } // EventKind categorises a scan event. @@ -74,6 +75,19 @@ func Run(ctx context.Context, cfg Config, events chan<- Event) { subdomains := make(chan string) var wg sync.WaitGroup + // Optional rate limiter: a shared ticker gate paces total queries per second + // across the whole worker pool. nil means unlimited. + var limiter <-chan time.Time + if cfg.Rate > 0 { + interval := time.Second / time.Duration(cfg.Rate) + if interval <= 0 { + interval = time.Nanosecond + } + rl := time.NewTicker(interval) + defer rl.Stop() + limiter = rl.C + } + // Progress ticker - fires every second. // tickerDone signals the goroutine to stop; tickerStopped confirms it has // fully exited so we never close events while a send is pending. @@ -113,6 +127,14 @@ func Run(ctx context.Context, cfg Config, events chan<- Event) { atomic.AddInt64(&processed, 1) continue } + if limiter != nil { + select { + case <-limiter: + case <-ctx.Done(): + atomic.AddInt64(&processed, 1) + continue + } + } fullDomain := prefix + "." + cfg.Domain var resolved bool var records []dns.Record diff --git a/internal/scan/runner_test.go b/internal/scan/runner_test.go index d6080b9..14ef353 100644 --- a/internal/scan/runner_test.go +++ b/internal/scan/runner_test.go @@ -66,6 +66,41 @@ func TestRunSimulateConcurrent(t *testing.T) { } } +// TestRunRateLimit asserts that -rate paces queries: N queries at R qps should +// take at least (N-1)/R seconds. Uses simulate mode so it is network-free. +func TestRunRateLimit(t *testing.T) { + if testing.Short() { + t.Skip("skipping timing-sensitive rate limit test in short mode") + } + + const queries = 20 + const rate = 10 // qps + cfg := Config{ + Domain: "example.com", + Entries: makeEntries(queries), + Concurrency: 8, + Timeout: time.Second, + Simulate: true, + HitRate: 50, + Attempts: 1, + Rate: rate, + } + + events := make(chan Event, 64) + start := time.Now() + go Run(context.Background(), cfg, events) + for range events { //nolint:revive // draining + } + elapsed := time.Since(start) + + // Floor: the first tick fires after one interval, so expect at least + // (queries-1)/rate seconds, with a margin for scheduling jitter. + minExpected := time.Duration(float64(queries-1) / float64(rate) * 0.8 * float64(time.Second)) + if elapsed < minExpected { + t.Errorf("rate-limited scan finished too fast: %s < %s", elapsed, minExpected) + } +} + // TestRunContextCancel cancels mid-scan and asserts Run returns and closes the // events channel promptly. func TestRunContextCancel(t *testing.T) { diff --git a/main.go b/main.go index 07399e1..bf415cc 100644 --- a/main.go +++ b/main.go @@ -101,6 +101,7 @@ type cliFlags struct { retries int force bool format string + rate int } func parseFlags() cliFlags { @@ -120,6 +121,7 @@ func parseFlags() cliFlags { flag.IntVar(&f.retries, "retries", 0, "Deprecated: use -attempts instead") flag.BoolVar(&f.force, "force", false, "Continue scanning even if wildcard DNS is detected") flag.StringVar(&f.format, "format", "text", "Output format: text, json, or csv") + flag.IntVar(&f.rate, "rate", 0, "Max DNS queries per second across all workers (0 = unlimited)") flag.Parse() return f } @@ -146,6 +148,10 @@ func validateFlags(f cliFlags, out *output.Writer, maxAttempts int) (string, boo out.Error("Attempts (-attempts) must be at least 1") return "", false } + if f.rate < 0 { + out.Error("Rate (-rate) must be 0 (unlimited) or a positive integer") + return "", false + } if !f.testMode { if err := validateDNSServer(f.dnsServer); err != nil { out.Error("DNS server %s: %v", f.dnsServer, err) @@ -311,6 +317,7 @@ func run() int { Attempts: maxAttempts, Force: f.force, Verbose: f.verbose, + Rate: f.rate, } events := make(chan scan.Event, 64)