From 6d7f08f40a09f08ef507ec0d52b1fc9c2855e262 Mon Sep 17 00:00:00 2001 From: "Chris (ChrisJr404)" <11917633+ChrisJr404@users.noreply.github.com> Date: Sun, 3 May 2026 11:52:56 -0400 Subject: [PATCH 1/4] feat(dank): add cancellable / max-results variants of GenerateAtFixedLength (#285) The pure-DFS walk in GenerateAtFixedLength has no exit conditions, no cancellation, and no result cap, so subdomain inputs that produce a large language make `alterx -enrich` (and any library consumer) hang indefinitely on the generation phase. Add two new public methods alongside the existing one: GenerateAtFixedLengthWithLimit(fixedLen, maxResults) bails when len(results) >= maxResults, returning ErrResultLimitReached with the partial sorted slice. GenerateAtFixedLengthWithContext(ctx, fixedLen, maxResults) checks ctx.Err() at every state expansion so timeouts apply across the whole walk; honours maxResults too (set 0 to disable the cap). Both share a single internal `generateAtFixedLength` so the historical `GenerateAtFixedLength` keeps working unchanged. New tests cover the backwards-compat path, the cap, the no-cap case, and a context deadline that fires partway through an exploding regex. --- internal/dank/dank.go | 82 ++++++++++++++++++++++++++++---- internal/dank/dank_test.go | 97 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 170 insertions(+), 9 deletions(-) create mode 100644 internal/dank/dank_test.go diff --git a/internal/dank/dank.go b/internal/dank/dank.go index 0c4b776e..bf4296a3 100644 --- a/internal/dank/dank.go +++ b/internal/dank/dank.go @@ -1,6 +1,8 @@ package dank import ( + "context" + "errors" "fmt" "math/big" "regexp" @@ -9,6 +11,12 @@ import ( "strings" ) +// ErrResultLimitReached is returned by GenerateAtFixedLengthWithLimit and +// GenerateAtFixedLengthWithContext when their respective max-results cap is +// hit before the DFS completes. The partial result slice is still returned +// alongside the error so callers can use it. +var ErrResultLimitReached = errors.New("dank: result limit reached") + // DankEncoder implementation matching Python's C++ backend exactly // Uses Brzozowski's algorithm for DFA minimization @@ -464,31 +472,84 @@ func (d *DankEncoder) NumWords(minLen, maxLen int) int64 { return total.Int64() } -// GenerateAtFixedLength returns all strings of exactly fixedLen +// GenerateAtFixedLength returns all strings of exactly fixedLen. +// +// Note: this method has no exit condition — it walks the entire DFA. For +// inputs that produce a very large language, prefer +// GenerateAtFixedLengthWithLimit or GenerateAtFixedLengthWithContext. +// See https://github.com/projectdiscovery/alterx/issues/285. func (d *DankEncoder) GenerateAtFixedLength(fixedLen int) []string { - var results []string - d.dfsGenerateFixed(0, "", fixedLen, &results) - sort.Strings(results) + results, _ := d.generateAtFixedLength(context.Background(), fixedLen, 0) return results } -// dfsGenerateFixed generates only strings of exact length -func (d *DankEncoder) dfsGenerateFixed(state int, curr string, remaining int, results *[]string) { +// GenerateAtFixedLengthWithLimit returns up to maxResults strings of exactly +// fixedLen. If the DFS produces more than maxResults results, generation +// stops early and ErrResultLimitReached is returned alongside the partial +// (still sorted) result slice. A maxResults <= 0 disables the cap. +func (d *DankEncoder) GenerateAtFixedLengthWithLimit(fixedLen, maxResults int) ([]string, error) { + return d.generateAtFixedLength(context.Background(), fixedLen, maxResults) +} + +// GenerateAtFixedLengthWithContext returns up to maxResults strings of +// exactly fixedLen, aborting early when ctx is cancelled or its deadline +// passes. ctx.Err() is checked at every state expansion in the DFS, so +// timeouts apply across the entire generation walk rather than only between +// top-level calls. Use maxResults <= 0 to disable the cap and rely solely on +// ctx for cancellation. The partial result slice (sorted) is returned even +// when ctx.Err() / ErrResultLimitReached fires. +func (d *DankEncoder) GenerateAtFixedLengthWithContext(ctx context.Context, fixedLen, maxResults int) ([]string, error) { + return d.generateAtFixedLength(ctx, fixedLen, maxResults) +} + +// generateAtFixedLength is the shared implementation behind the three public +// generators. ctx and maxResults can each be supplied independently +// (background / 0) to recover the historical no-limit behaviour. +func (d *DankEncoder) generateAtFixedLength(ctx context.Context, fixedLen, maxResults int) ([]string, error) { + var ( + results []string + err error + ) + d.dfsGenerateFixed(ctx, 0, "", fixedLen, maxResults, &results, &err) + sort.Strings(results) + return results, err +} + +// dfsGenerateFixed generates strings of exact length, aborting on ctx +// cancellation or once len(*results) reaches maxResults (when > 0). The +// outErr slot lets the recursion bubble up the cancellation cause so the +// public caller can distinguish "completed naturally" from "stopped early". +func (d *DankEncoder) dfsGenerateFixed(ctx context.Context, state int, curr string, remaining, maxResults int, results *[]string, outErr *error) { // Skip dead state (last state in DFA) deadState := len(d.dfa) - 1 if state == deadState { return } + // Bail if a previous branch already triggered cancellation / limit-reached. + if *outErr != nil { + return + } + + // Cooperatively respect context cancellation. ctx.Err() is cheap and + // short-circuits before the recursion fans out further. + if err := ctx.Err(); err != nil { + *outErr = err + return + } + if remaining == 0 { if d.dfa[state].IsFinal { *results = append(*results, curr) + if maxResults > 0 && len(*results) >= maxResults { + *outErr = ErrResultLimitReached + } } return } - // Iterate over actual transitions (sorted for deterministic output) - // Can't just use alphabet because pattern may have characters outside alphabet (like *) + // Iterate over actual transitions (sorted for deterministic output). + // Can't just use alphabet because pattern may have characters outside alphabet (like *). chars := []byte{} for ch := range d.dfa[state].Trans { chars = append(chars, ch) @@ -499,7 +560,10 @@ func (d *DankEncoder) dfsGenerateFixed(state int, curr string, remaining int, re next := d.dfa[state].Trans[ch] // Don't transition to dead state during generation if next != deadState { - d.dfsGenerateFixed(next, curr+string(ch), remaining-1, results) + d.dfsGenerateFixed(ctx, next, curr+string(ch), remaining-1, maxResults, results, outErr) + if *outErr != nil { + return + } } } } diff --git a/internal/dank/dank_test.go b/internal/dank/dank_test.go new file mode 100644 index 00000000..c6589785 --- /dev/null +++ b/internal/dank/dank_test.go @@ -0,0 +1,97 @@ +package dank + +import ( + "context" + "errors" + "testing" + "time" +) + +// regex that matches strings of the form "a[0-2]" (3 strings: a0, a1, a2). +const smallRegex = "a[0-2]" + +// regex that explodes: 5 alphas anywhere in a 5-char window. The fixed-length +// generation walk is bounded by the alphabet size (~39) raised to fixedLen, so +// fixedLen=6 here yields tens of millions of strings — enough that any of +// the bounded variants should bail before the unbounded one would. +const explodingRegex = "[a-z][a-z][a-z][a-z][a-z][0-9]" + +func TestGenerateAtFixedLength_BackwardsCompat(t *testing.T) { + d := NewDankEncoder(smallRegex, 16) + got := d.GenerateAtFixedLength(2) + want := []string{"a0", "a1", "a2"} + if !equalStringSlices(got, want) { + t.Fatalf("GenerateAtFixedLength(2) = %v, want %v", got, want) + } +} + +func TestGenerateAtFixedLengthWithLimit_HitsCap(t *testing.T) { + d := NewDankEncoder(smallRegex, 16) + got, err := d.GenerateAtFixedLengthWithLimit(2, 2) + if !errors.Is(err, ErrResultLimitReached) { + t.Fatalf("expected ErrResultLimitReached, got %v", err) + } + if len(got) != 2 { + t.Fatalf("expected exactly 2 results at the cap, got %d (%v)", len(got), got) + } +} + +func TestGenerateAtFixedLengthWithLimit_NoCap(t *testing.T) { + d := NewDankEncoder(smallRegex, 16) + got, err := d.GenerateAtFixedLengthWithLimit(2, 0) + if err != nil { + t.Fatalf("unexpected error with maxResults=0: %v", err) + } + if len(got) != 3 { + t.Fatalf("expected 3 results without cap, got %d (%v)", len(got), got) + } +} + +func TestGenerateAtFixedLengthWithContext_Cancellation(t *testing.T) { + d := NewDankEncoder(explodingRegex, 16) + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond) + defer cancel() + + start := time.Now() + got, err := d.GenerateAtFixedLengthWithContext(ctx, 6, 0) + elapsed := time.Since(start) + + if !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) { + t.Fatalf("expected context cancellation error, got %v", err) + } + if elapsed > 500*time.Millisecond { + t.Fatalf("DFS did not honour context deadline (took %s)", elapsed) + } + // Partial result slice should still be returned and sorted. + if !isSorted(got) { + t.Fatalf("partial results should be sorted, got %v", got[:min(10, len(got))]) + } +} + +func equalStringSlices(a, b []string) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +func isSorted(s []string) bool { + for i := 1; i < len(s); i++ { + if s[i-1] > s[i] { + return false + } + } + return true +} + +func min(a, b int) int { + if a < b { + return a + } + return b +} From de3cc933accc35e261c29c00b991d2ff2b488001 Mon Sep 17 00:00:00 2001 From: "Chris (ChrisJr404)" <11917633+ChrisJr404@users.noreply.github.com> Date: Sun, 3 May 2026 12:02:15 -0400 Subject: [PATCH 2/4] address review: clarify limit semantics + reject negative fixedLen - GenerateAtFixedLengthWithLimit doc-comment matched the runtime check (>=) rather than the original 'more than' wording. - Add an ErrInvalidFixedLength guard at the entry to generateAtFixedLength so dfsGenerateFixed cannot be called with a negative remaining counter. - New test exercises the negative-length path. --- internal/dank/dank.go | 15 ++++++++++++--- internal/dank/dank_test.go | 11 +++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/internal/dank/dank.go b/internal/dank/dank.go index bf4296a3..c04e224a 100644 --- a/internal/dank/dank.go +++ b/internal/dank/dank.go @@ -484,9 +484,10 @@ func (d *DankEncoder) GenerateAtFixedLength(fixedLen int) []string { } // GenerateAtFixedLengthWithLimit returns up to maxResults strings of exactly -// fixedLen. If the DFS produces more than maxResults results, generation -// stops early and ErrResultLimitReached is returned alongside the partial -// (still sorted) result slice. A maxResults <= 0 disables the cap. +// fixedLen. Generation stops once the result count reaches maxResults — i.e. +// the returned slice contains exactly maxResults entries when truncated — and +// ErrResultLimitReached is returned alongside the (sorted) partial result +// slice. A maxResults <= 0 disables the cap. func (d *DankEncoder) GenerateAtFixedLengthWithLimit(fixedLen, maxResults int) ([]string, error) { return d.generateAtFixedLength(context.Background(), fixedLen, maxResults) } @@ -502,10 +503,18 @@ func (d *DankEncoder) GenerateAtFixedLengthWithContext(ctx context.Context, fixe return d.generateAtFixedLength(ctx, fixedLen, maxResults) } +// ErrInvalidFixedLength is returned by the GenerateAtFixedLength* family when +// fixedLen is negative, which would otherwise cause the DFS to recurse +// indefinitely with an ever-decreasing remaining counter. +var ErrInvalidFixedLength = errors.New("dank: fixedLen must be >= 0") + // generateAtFixedLength is the shared implementation behind the three public // generators. ctx and maxResults can each be supplied independently // (background / 0) to recover the historical no-limit behaviour. func (d *DankEncoder) generateAtFixedLength(ctx context.Context, fixedLen, maxResults int) ([]string, error) { + if fixedLen < 0 { + return nil, ErrInvalidFixedLength + } var ( results []string err error diff --git a/internal/dank/dank_test.go b/internal/dank/dank_test.go index c6589785..331f9cdd 100644 --- a/internal/dank/dank_test.go +++ b/internal/dank/dank_test.go @@ -47,6 +47,17 @@ func TestGenerateAtFixedLengthWithLimit_NoCap(t *testing.T) { } } +func TestGenerateAtFixedLengthWithLimit_NegativeFixedLen(t *testing.T) { + d := NewDankEncoder(smallRegex, 16) + got, err := d.GenerateAtFixedLengthWithLimit(-1, 10) + if !errors.Is(err, ErrInvalidFixedLength) { + t.Fatalf("expected ErrInvalidFixedLength for negative fixedLen, got %v", err) + } + if len(got) != 0 { + t.Fatalf("expected empty slice on validation failure, got %v", got) + } +} + func TestGenerateAtFixedLengthWithContext_Cancellation(t *testing.T) { d := NewDankEncoder(explodingRegex, 16) ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond) From f47dc5a190ad69ea0dcfed7705d6cdc6755dca5f Mon Sep 17 00:00:00 2001 From: "Chris (ChrisJr404)" <11917633+ChrisJr404@users.noreply.github.com> Date: Sun, 3 May 2026 16:14:49 -0400 Subject: [PATCH 3/4] address review: nil-context guard + clarify legacy wrapper behaviour Two follow-ups from CodeRabbit on PR #291: 1. GenerateAtFixedLengthWithContext now normalises a nil ctx to context.Background() before recursing. Previously dfsGenerateFixed would call ctx.Err() unconditionally and panic on nil. New regression test (TestGenerateAtFixedLengthWithContext_NilContext) covers it. 2. GenerateAtFixedLength's docstring already pointed at the With* variants for limit handling but did not mention that a negative fixedLen returns an empty slice (the wrapper deliberately swallows ErrInvalidFixedLength to preserve the historical signature). Spell that out so callers know which entry point surfaces the sentinel error. --- internal/dank/dank.go | 8 ++++++++ internal/dank/dank_test.go | 14 ++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/internal/dank/dank.go b/internal/dank/dank.go index c04e224a..25d82875 100644 --- a/internal/dank/dank.go +++ b/internal/dank/dank.go @@ -478,6 +478,11 @@ func (d *DankEncoder) NumWords(minLen, maxLen int) int64 { // inputs that produce a very large language, prefer // GenerateAtFixedLengthWithLimit or GenerateAtFixedLengthWithContext. // See https://github.com/projectdiscovery/alterx/issues/285. +// +// A negative fixedLen returns an empty slice (rather than recursing +// forever); use GenerateAtFixedLengthWithLimit or +// GenerateAtFixedLengthWithContext if you need to surface +// ErrInvalidFixedLength explicitly. func (d *DankEncoder) GenerateAtFixedLength(fixedLen int) []string { results, _ := d.generateAtFixedLength(context.Background(), fixedLen, 0) return results @@ -500,6 +505,9 @@ func (d *DankEncoder) GenerateAtFixedLengthWithLimit(fixedLen, maxResults int) ( // ctx for cancellation. The partial result slice (sorted) is returned even // when ctx.Err() / ErrResultLimitReached fires. func (d *DankEncoder) GenerateAtFixedLengthWithContext(ctx context.Context, fixedLen, maxResults int) ([]string, error) { + if ctx == nil { + ctx = context.Background() + } return d.generateAtFixedLength(ctx, fixedLen, maxResults) } diff --git a/internal/dank/dank_test.go b/internal/dank/dank_test.go index 331f9cdd..54265425 100644 --- a/internal/dank/dank_test.go +++ b/internal/dank/dank_test.go @@ -58,6 +58,20 @@ func TestGenerateAtFixedLengthWithLimit_NegativeFixedLen(t *testing.T) { } } +func TestGenerateAtFixedLengthWithContext_NilContext(t *testing.T) { + d := NewDankEncoder(smallRegex, 16) + // Passing a nil context must not panic; the public entry point normalises + // it to context.Background() before reaching the recursive ctx.Err() call. + got, err := d.GenerateAtFixedLengthWithContext(nil, 2, 0) + if err != nil { + t.Fatalf("expected nil error with nil ctx, got %v", err) + } + want := []string{"a0", "a1", "a2"} + if !equalStringSlices(got, want) { + t.Fatalf("GenerateAtFixedLengthWithContext(nil, 2, 0) = %v, want %v", got, want) + } +} + func TestGenerateAtFixedLengthWithContext_Cancellation(t *testing.T) { d := NewDankEncoder(explodingRegex, 16) ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond) From da9b8f9ada9d17a32a432f95395ed2876a8fadb4 Mon Sep 17 00:00:00 2001 From: ChrisJr404 Date: Tue, 5 May 2026 18:25:34 -0400 Subject: [PATCH 4/4] chore(lint): silence staticcheck SA1012 in nil-context test The TestGenerateAtFixedLengthWithContext_NilContext case in #285 deliberately passes a nil ctx to confirm the public-API guard, but staticcheck flags the literal nil under SA1012 and fails the Lint Test workflow. Route the nil through a typed context.Context variable so the test still exercises the guard without tripping the linter. --- internal/dank/dank_test.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/internal/dank/dank_test.go b/internal/dank/dank_test.go index 54265425..2331d6bb 100644 --- a/internal/dank/dank_test.go +++ b/internal/dank/dank_test.go @@ -62,7 +62,10 @@ func TestGenerateAtFixedLengthWithContext_NilContext(t *testing.T) { d := NewDankEncoder(smallRegex, 16) // Passing a nil context must not panic; the public entry point normalises // it to context.Background() before reaching the recursive ctx.Err() call. - got, err := d.GenerateAtFixedLengthWithContext(nil, 2, 0) + // Assign through a typed var so staticcheck SA1012 doesn't flag the + // literal nil at the call site - we are explicitly exercising the guard. + var nilCtx context.Context + got, err := d.GenerateAtFixedLengthWithContext(nilCtx, 2, 0) if err != nil { t.Fatalf("expected nil error with nil ctx, got %v", err) }