From 173104bf6823632e2e7f926c2548f0a36d9c4e56 Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Tue, 16 Dec 2025 19:30:43 +0400 Subject: [PATCH 1/5] Memory + Speed improvements --- pkg/massdns/process.go | 18 +++++++----- pkg/parser/parser.go | 2 +- pkg/wildcards/resolver.go | 61 ++++++++++++++++++++++----------------- 3 files changed, 45 insertions(+), 36 deletions(-) diff --git a/pkg/massdns/process.go b/pkg/massdns/process.go index 4a2aafb..786176f 100644 --- a/pkg/massdns/process.go +++ b/pkg/massdns/process.go @@ -11,6 +11,7 @@ import ( "os/exec" "strconv" "strings" + "sync/atomic" "time" "github.com/projectdiscovery/dnsx/libs/dnsx" @@ -21,6 +22,7 @@ import ( "github.com/projectdiscovery/utils/batcher" fileutil "github.com/projectdiscovery/utils/file" ioutil "github.com/projectdiscovery/utils/io" + mapsutil "github.com/projectdiscovery/utils/maps" "github.com/remeh/sizedwaitgroup" "github.com/weppos/publicsuffix-go/publicsuffix" ) @@ -53,7 +55,7 @@ func (instance *Instance) RunWithContext(ctx context.Context) (stdout, stderr st // 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.Split(instance.options.MassDnsCmd, " ")...) + args = append(args, strings.Fields(instance.options.MassDnsCmd)...) } cmd := exec.CommandContext(ctx, instance.options.MassdnsPath, args...) @@ -178,7 +180,7 @@ func (instance *Instance) runChunk(ctx context.Context, chunkFile string) (stdou // 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.Split(instance.options.MassDnsCmd, " ")...) + args = append(args, strings.Fields(instance.options.MassDnsCmd)...) } cmd := exec.CommandContext(ctx, instance.options.MassdnsPath, args...) @@ -374,10 +376,10 @@ func (instance *Instance) writeOutput(store *store.Store) error { } } - uniqueMap := make(map[string]struct{}) + uniqueMap := mapsutil.NewSyncLockMap[string, struct{}]() // write count of resolved hosts - resolvedCount := 0 + var resolvedCount atomic.Int32 // if trusted resolvers are specified verify the results var dnsResolver *dnsx.DNSX @@ -400,10 +402,10 @@ func (instance *Instance) writeOutput(store *store.Store) error { store.Iterate(func(ip string, hostnames []string, counter int) { for _, hostname := range hostnames { // Skip if we already printed this subdomain once - if _, ok := uniqueMap[hostname]; ok { + if _, ok := uniqueMap.Has(hostname) { continue } - uniqueMap[hostname] = struct{}{} + uniqueMap.Set(hostname, struct{}{}) swg.Add() go func(hostname string) { @@ -451,14 +453,14 @@ func (instance *Instance) writeOutput(store *store.Store) error { _, _ = safeWriter.Write([]byte(data)) } gologger.Silent().Msgf("%s", data) - resolvedCount++ + resolvedCount.Add(1) }(hostname) } }) swg.Wait() - gologger.Info().Msgf("Total resolved: %d\n", resolvedCount) + gologger.Info().Msgf("Total resolved: %d\n", resolvedCount.Load()) // Close the files and return if output != nil { diff --git a/pkg/parser/parser.go b/pkg/parser/parser.go index 1554467..d5a7049 100644 --- a/pkg/parser/parser.go +++ b/pkg/parser/parser.go @@ -57,7 +57,7 @@ func ParseReader(reader io.Reader, onResult OnResultFN) error { } else { // Non empty line represents DNS answer section, we split on space, // iterate over all the parts, and write the answer to the struct. - parts := strings.Split(text, " ") + parts := strings.Fields(text) if len(parts) != 3 { continue diff --git a/pkg/wildcards/resolver.go b/pkg/wildcards/resolver.go index bdf7ae5..a854cc0 100644 --- a/pkg/wildcards/resolver.go +++ b/pkg/wildcards/resolver.go @@ -4,6 +4,7 @@ import ( "errors" "fmt" "strings" + "sync" "github.com/miekg/dns" "github.com/projectdiscovery/dnsx/libs/dnsx" @@ -63,25 +64,40 @@ func (w *Resolver) SetProbeCount(count int) { } } -// probeWildcardIPs probes the given wildcard pattern multiple times and returns all IPs found. +// probeWildcardIPs probes the given wildcard pattern multiple times concurrently and returns all IPs found. // Returns nil if the first probe returns NXDOMAIN (not a wildcard level). +// All queries are executed in parallel for better performance. func (w *Resolver) probeWildcardIPs(pattern string, count int) []string { - var ips []string + if count <= 0 { + return nil + } + + var wg sync.WaitGroup + ips := sliceutil.NewSyncSlice[string]() + + // Launch all queries concurrently for i := 0; i < count; i++ { - probeHost := strings.ReplaceAll(pattern, "*.", xid.New().String()+".") - in, err := w.client.QueryOne(probeHost) - if err != nil { - continue - } - if in == nil || in.StatusCodeRaw != dns.RcodeSuccess { - if i == 0 { - return nil + wg.Add(1) + go func() { + defer wg.Done() + + probeHost := strings.ReplaceAll(pattern, "*.", xid.New().String()+".") + in, err := w.client.QueryOne(probeHost) + + if err == nil && in != nil && in.StatusCodeRaw == dns.RcodeSuccess { + ips.Append(in.A...) } - break - } - ips = append(ips, in.A...) + }() } - return ips + + wg.Wait() + + // Check if first query failed (original behavior) + if ips.Len() == 0 { + return nil + } + + return sliceutil.Dedupe(ips.Slice) } // generateWildcardPermutations generates wildcard permutations for a given subdomain @@ -114,15 +130,6 @@ func generateWildcardPermutations(subdomain, domain string) []string { return hosts } -func getSyncLockMapValues(m *mapsutil.SyncLockMap[string, struct{}]) map[string]struct{} { - values := make(map[string]struct{}) - _ = m.Iterate(func(key string, value struct{}) error { - values[key] = value - return nil - }) - return values -} - // LookupHost returns wildcard IP addresses of a wildcard if it's a wildcard. // To determine, first we split the target host by dots, create permutation // of it's levels, check for wildcard on each one of them and if found any, @@ -168,7 +175,7 @@ func (w *Resolver) LookupHost(host string, knownIPs []string) (bool, map[string] if cachedValueOk { for _, knownIP := range knownIPs { if _, ipExists := cachedValue.IPS.Get(knownIP); ipExists { - return true, getSyncLockMapValues(cachedValue.IPS) + return true, cachedValue.IPS.Map } } // Cache hit but IP not found - re-probe to catch missed round-robin IPs @@ -180,7 +187,7 @@ func (w *Resolver) LookupHost(host string, knownIPs []string) (bool, map[string] _ = w.wildcardAnswersCache.Set(original, cachedValue) for _, knownIP := range knownIPs { if _, ipExists := cachedValue.IPS.Get(knownIP); ipExists { - return true, getSyncLockMapValues(cachedValue.IPS) + return true, cachedValue.IPS.Map } } } @@ -210,7 +217,7 @@ func (w *Resolver) LookupHost(host string, knownIPs []string) (bool, map[string] _ = w.wildcardAnswersCache.Set(original, cachedValue) for _, knownIP := range knownIPs { if _, ipExists := cachedValue.IPS.Get(knownIP); ipExists { - return true, getSyncLockMapValues(cachedValue.IPS) + return true, cachedValue.IPS.Map } } @@ -221,7 +228,7 @@ func (w *Resolver) LookupHost(host string, knownIPs []string) (bool, map[string] if err == nil && in != nil && in.StatusCodeRaw == dns.RcodeSuccess { for _, record := range in.A { if _, ipExists := cachedValue.IPS.Get(record); ipExists { - return true, getSyncLockMapValues(cachedValue.IPS) + return true, cachedValue.IPS.Map } } } From fa01af89bae20c5d6f170a568f4a8835da693960 Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Tue, 16 Dec 2025 19:37:46 +0400 Subject: [PATCH 2/5] lint --- pkg/massdns/process.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pkg/massdns/process.go b/pkg/massdns/process.go index 786176f..d68119d 100644 --- a/pkg/massdns/process.go +++ b/pkg/massdns/process.go @@ -402,10 +402,10 @@ func (instance *Instance) writeOutput(store *store.Store) error { store.Iterate(func(ip string, hostnames []string, counter int) { for _, hostname := range hostnames { // Skip if we already printed this subdomain once - if _, ok := uniqueMap.Has(hostname) { + if uniqueMap.Has(hostname) { continue } - uniqueMap.Set(hostname, struct{}{}) + _ = uniqueMap.Set(hostname, struct{}{}) swg.Add() go func(hostname string) { From f0c52039e498c4991f93e1177d9bbbcb63cb41ea Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Tue, 16 Dec 2025 19:53:55 +0400 Subject: [PATCH 3/5] lint --- pkg/wildcards/resolver.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pkg/wildcards/resolver.go b/pkg/wildcards/resolver.go index a854cc0..93cb611 100644 --- a/pkg/wildcards/resolver.go +++ b/pkg/wildcards/resolver.go @@ -76,8 +76,9 @@ func (w *Resolver) probeWildcardIPs(pattern string, count int) []string { ips := sliceutil.NewSyncSlice[string]() // Launch all queries concurrently - for i := 0; i < count; i++ { + for range count { wg.Add(1) + go func() { defer wg.Done() From 89943325f282e4814df10360443651d419ac37ce Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Tue, 16 Dec 2025 20:12:38 +0400 Subject: [PATCH 4/5] keep first query nil return --- pkg/wildcards/resolver.go | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/pkg/wildcards/resolver.go b/pkg/wildcards/resolver.go index 93cb611..3c25276 100644 --- a/pkg/wildcards/resolver.go +++ b/pkg/wildcards/resolver.go @@ -74,26 +74,42 @@ func (w *Resolver) probeWildcardIPs(pattern string, count int) []string { var wg sync.WaitGroup ips := sliceutil.NewSyncSlice[string]() + var firstQueryFailed bool + var firstQueryOnce sync.Once // Launch all queries concurrently - for range count { + for i := 0; i < count; i++ { wg.Add(1) - go func() { + go func(index int) { defer wg.Done() probeHost := strings.ReplaceAll(pattern, "*.", xid.New().String()+".") in, err := w.client.QueryOne(probeHost) + // Track first query (index 0) failure for early exit behavior + if index == 0 { + firstQueryOnce.Do(func() { + if err != nil || in == nil || in.StatusCodeRaw != dns.RcodeSuccess { + firstQueryFailed = true + } + }) + } + if err == nil && in != nil && in.StatusCodeRaw == dns.RcodeSuccess { ips.Append(in.A...) } - }() + }(i) } wg.Wait() - // Check if first query failed (original behavior) + // Check if first query failed (original behavior - return nil if first probe fails) + if firstQueryFailed { + return nil + } + + // If no IPs collected, return nil if ips.Len() == 0 { return nil } From a23053f8b4bc3fc6cce607d5ec9c2e3c3089040f Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Tue, 16 Dec 2025 20:19:45 +0400 Subject: [PATCH 5/5] keep early exit on first try --- pkg/wildcards/resolver.go | 61 ++++++++++++++++++++++----------------- 1 file changed, 34 insertions(+), 27 deletions(-) diff --git a/pkg/wildcards/resolver.go b/pkg/wildcards/resolver.go index 3c25276..0b25696 100644 --- a/pkg/wildcards/resolver.go +++ b/pkg/wildcards/resolver.go @@ -66,49 +66,56 @@ func (w *Resolver) SetProbeCount(count int) { // probeWildcardIPs probes the given wildcard pattern multiple times concurrently and returns all IPs found. // Returns nil if the first probe returns NXDOMAIN (not a wildcard level). -// All queries are executed in parallel for better performance. +// First query is executed sequentially for early exit, remaining queries run in parallel. func (w *Resolver) probeWildcardIPs(pattern string, count int) []string { if count <= 0 { return nil } - var wg sync.WaitGroup ips := sliceutil.NewSyncSlice[string]() - var firstQueryFailed bool - var firstQueryOnce sync.Once - // Launch all queries concurrently - for i := 0; i < count; i++ { - wg.Add(1) + probe := func() ([]string, bool) { + probeHost := strings.ReplaceAll(pattern, "*.", xid.New().String()+".") + in, err := w.client.QueryOne(probeHost) + if err != nil || in == nil || in.StatusCodeRaw != dns.RcodeSuccess { + return nil, false + } + return in.A, true + } - go func(index int) { - defer wg.Done() + // Execute first query sequentially for early exit behavior + resultIPs, success := probe() + if !success { + return nil + } + if len(resultIPs) > 0 { + ips.Append(resultIPs...) + } - probeHost := strings.ReplaceAll(pattern, "*.", xid.New().String()+".") - in, err := w.client.QueryOne(probeHost) + // If only one query requested, return now + if count == 1 { + if ips.Len() == 0 { + return nil + } + return sliceutil.Dedupe(ips.Slice) + } - // Track first query (index 0) failure for early exit behavior - if index == 0 { - firstQueryOnce.Do(func() { - if err != nil || in == nil || in.StatusCodeRaw != dns.RcodeSuccess { - firstQueryFailed = true - } - }) - } + // Launch remaining queries concurrently + var wg sync.WaitGroup + for i := 1; i < count; i++ { + wg.Add(1) + go func() { + defer wg.Done() - if err == nil && in != nil && in.StatusCodeRaw == dns.RcodeSuccess { - ips.Append(in.A...) + resultIPs, success := probe() + if success && len(resultIPs) > 0 { + ips.Append(resultIPs...) } - }(i) + }() } wg.Wait() - // Check if first query failed (original behavior - return nil if first probe fails) - if firstQueryFailed { - return nil - } - // If no IPs collected, return nil if ips.Len() == 0 { return nil