diff --git a/pkg/massdns/process.go b/pkg/massdns/process.go index 4a2aafb..d68119d 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 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..0b25696 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,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 } } }