Improvements

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 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 {

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

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,64 @@ 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).
+// First query is executed sequentially for early exit, remaining queries run in parallel.
 func (w *Resolver) probeWildcardIPs(pattern string, count int) []string {
-	var ips []string
-	for i := 0; i < count; i++ {
+	if count <= 0 {
+		return nil
+	}
+
+	ips := sliceutil.NewSyncSlice[string]()
+
+	probe := func() ([]string, bool) {
 		probeHost := strings.ReplaceAll(pattern, "*.", xid.New().String()+".")
 		in, err := w.client.QueryOne(probeHost)
-		if err != nil {
-			continue
+		if err != nil || in == nil || in.StatusCodeRaw != dns.RcodeSuccess {
+			return nil, false
 		}
-		if in == nil || in.StatusCodeRaw != dns.RcodeSuccess {
-			if i == 0 {
-				return nil
-			}
-			break
+		return in.A, true
+	}
+
+	// Execute first query sequentially for early exit behavior
+	resultIPs, success := probe()
+	if !success {
+		return nil
+	}
+	if len(resultIPs) > 0 {
+		ips.Append(resultIPs...)
+	}
+
+	// If only one query requested, return now
+	if count == 1 {
+		if ips.Len() == 0 {
+			return nil
 		}
-		ips = append(ips, in.A...)
+		return sliceutil.Dedupe(ips.Slice)
 	}
-	return ips
+
+	// Launch remaining queries concurrently
+	var wg sync.WaitGroup
+	for i := 1; i < count; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			resultIPs, success := probe()
+			if success && len(resultIPs) > 0 {
+				ips.Append(resultIPs...)
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	// If no IPs collected, return nil
+	if ips.Len() == 0 {
+		return nil
+	}
+
+	return sliceutil.Dedupe(ips.Slice)
 }
 
 // generateWildcardPermutations generates wildcard permutations for a given subdomain
@@ -114,15 +154,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 +199,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 +211,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 +241,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 +252,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
 						}
 					}
 				}