From 9e862d07763b47943591b51331dd77ce045ef024 Mon Sep 17 00:00:00 2001 From: Ice3man Date: Tue, 16 Dec 2025 19:25:24 +0530 Subject: [PATCH 1/2] feat: pass all ips at once in wildcard resolver (improvements) --- pkg/massdns/process.go | 85 ++++++++++++++++++---------------- pkg/wildcards/resolver.go | 28 +++++++---- pkg/wildcards/resolver_test.go | 2 +- 3 files changed, 63 insertions(+), 52 deletions(-) diff --git a/pkg/massdns/process.go b/pkg/massdns/process.go index ef9a9ea..574fc84 100644 --- a/pkg/massdns/process.go +++ b/pkg/massdns/process.go @@ -277,66 +277,69 @@ func (instance *Instance) autoExtractRootDomains(store *store.Store) error { } func (instance *Instance) filterWildcards(st *store.Store) error { + // Build hostname -> IPs map to avoid redundant DNS queries + hostnameToIPs := make(map[string][]string) + hostnameCounters := make(map[string]int) + + st.Iterate(func(ip string, hostnames []string, counter int) { + for _, hostname := range hostnames { + hostnameToIPs[hostname] = append(hostnameToIPs[hostname], ip) + if counter > hostnameCounters[hostname] { + hostnameCounters[hostname] = counter + } + } + }) + // Start to work in parallel on wildcards wildcardWg := sizedwaitgroup.New(instance.options.WildcardsThreads) - var allCancelFunc []context.CancelFunc - - st.Iterate(func(ip string, hostnames []string, counter int) { - ipCtx, ipCancelFunc := context.WithCancel(context.Background()) - allCancelFunc = append(allCancelFunc, ipCancelFunc) - // We've stumbled upon a wildcard, just ignore it. - if instance.wildcardStore.Has(ip) { - return + for hostname, ips := range hostnameToIPs { + // Skip if any IP is already marked as wildcard + hasWildcardIP := false + for _, ip := range ips { + if instance.wildcardStore.Has(ip) { + hasWildcardIP = true + break + } + } + if hasWildcardIP { + continue } - // Perform wildcard detection on the ip, if an IP is found in the wildcard - // we add it to the wildcard map so that further runs don't require such filtering again. + counter := hostnameCounters[hostname] + // Perform wildcard detection on the hostname if counter >= 5 or strict mode if counter >= 5 || instance.options.StrictWildcard { - for _, hostname := range hostnames { - wildcardWg.Add() - go func(ctx context.Context, ipCancelFunc context.CancelFunc, IP string, hostname string) { - defer wildcardWg.Done() + wildcardWg.Add() + go func(hostname string, ips []string) { + defer wildcardWg.Done() - gologger.Info().Msgf("Started filtering wildcards for %s\n", hostname) + gologger.Info().Msgf("Started filtering wildcards for %s (with %d IPs)\n", hostname, len(ips)) - select { - case <-ctx.Done(): - return - default: - } - - isWildcard, ips := instance.wildcardResolver.LookupHost(hostname, IP) - if len(ips) > 0 { - for ip := range ips { - // we add the single ip to the wildcard list - if err := instance.wildcardStore.Set(ip); err != nil { - gologger.Error().Msgf("could not set wildcard ip: %s", err) - } - gologger.Info().Msgf("Removing wildcard %s\n", ip) + isWildcard, wildcardIPs := instance.wildcardResolver.LookupHost(hostname, ips) + if len(wildcardIPs) > 0 { + for ip := range wildcardIPs { + if err := instance.wildcardStore.Set(ip); err != nil { + gologger.Error().Msgf("could not set wildcard ip: %s", err) } + gologger.Info().Msgf("Removing wildcard %s\n", ip) } + } - if isWildcard { - // we also mark the original ip as wildcard, since at least once it resolved to this host - if err := instance.wildcardStore.Set(IP); err != nil { + if isWildcard { + for _, ip := range ips { + if err := instance.wildcardStore.Set(ip); err != nil { gologger.Error().Msgf("could not set wildcard ip: %s", err) } - ipCancelFunc() - gologger.Info().Msgf("Removed wildcard %s\n", IP) } + gologger.Info().Msgf("Removed wildcard hostname %s with %d IPs\n", hostname, len(ips)) + } - }(ipCtx, ipCancelFunc, ip, hostname) - } + }(hostname, ips) } - }) + } wildcardWg.Wait() - for _, cancelFunc := range allCancelFunc { - cancelFunc() - } - // Do a second pass as well and remove all the wildcards // from the store that we have found so that everything is covered allWildcardIPs := instance.wildcardResolver.GetAllWildcardIPs() diff --git a/pkg/wildcards/resolver.go b/pkg/wildcards/resolver.go index d198b59..bdf7ae5 100644 --- a/pkg/wildcards/resolver.go +++ b/pkg/wildcards/resolver.go @@ -127,7 +127,7 @@ func getSyncLockMapValues(m *mapsutil.SyncLockMap[string, struct{}]) map[string] // 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, // we remove all the hosts that have this IP from the map. -func (w *Resolver) LookupHost(host string, ip string) (bool, map[string]struct{}) { +func (w *Resolver) LookupHost(host string, knownIPs []string) (bool, map[string]struct{}) { wildcards := make(map[string]struct{}) var domain string @@ -166,8 +166,10 @@ func (w *Resolver) LookupHost(host string, ip string) (bool, map[string]struct{} // and it is used always for resolutions in future. cachedValue, cachedValueOk := w.wildcardAnswersCache.Get(original) if cachedValueOk { - if _, ipExists := cachedValue.IPS.Get(ip); ipExists { - return true, getSyncLockMapValues(cachedValue.IPS) + for _, knownIP := range knownIPs { + if _, ipExists := cachedValue.IPS.Get(knownIP); ipExists { + return true, getSyncLockMapValues(cachedValue.IPS) + } } // Cache hit but IP not found - re-probe to catch missed round-robin IPs if extraIPs := w.probeWildcardIPs(original, reProbeCount); len(extraIPs) > 0 { @@ -176,8 +178,10 @@ func (w *Resolver) LookupHost(host string, ip string) (bool, map[string]struct{} _ = cachedValue.IPS.Set(record, struct{}{}) } _ = w.wildcardAnswersCache.Set(original, cachedValue) - if _, ipExists := cachedValue.IPS.Get(ip); ipExists { - return true, getSyncLockMapValues(cachedValue.IPS) + for _, knownIP := range knownIPs { + if _, ipExists := cachedValue.IPS.Get(knownIP); ipExists { + return true, getSyncLockMapValues(cachedValue.IPS) + } } } } @@ -204,8 +208,10 @@ func (w *Resolver) LookupHost(host string, ip string) (bool, map[string]struct{} _ = cachedValue.IPS.Set(record, struct{}{}) } _ = w.wildcardAnswersCache.Set(original, cachedValue) - if _, ipExists := cachedValue.IPS.Get(ip); ipExists { - return true, getSyncLockMapValues(cachedValue.IPS) + for _, knownIP := range knownIPs { + if _, ipExists := cachedValue.IPS.Get(knownIP); ipExists { + return true, getSyncLockMapValues(cachedValue.IPS) + } } // Resolve actual host multiple times to catch round-robin IPs @@ -223,9 +229,11 @@ func (w *Resolver) LookupHost(host string, ip string) (bool, map[string]struct{} } } - // check if original ip are among wildcards - if _, ok := wildcards[ip]; ok { - return true, wildcards + // check if any of the knownIPs are among wildcards + for _, knownIP := range knownIPs { + if _, ok := wildcards[knownIP]; ok { + return true, wildcards + } } return false, wildcards diff --git a/pkg/wildcards/resolver_test.go b/pkg/wildcards/resolver_test.go index 1ab3490..a001ca8 100644 --- a/pkg/wildcards/resolver_test.go +++ b/pkg/wildcards/resolver_test.go @@ -44,7 +44,7 @@ func Test_Resolver_LookupHost(t *testing.T) { require.NoError(t, err) require.NotEmpty(t, ips) - return resolver.LookupHost(subdomain, ips[0]) + return resolver.LookupHost(subdomain, ips) } t.Run("normal", func(t *testing.T) { isWildcard, wildcards := lookupAndResolve("www.google.com", resolver) From ab67b7a1f8f6b75e9333fd4bd3d285d4efc8acb9 Mon Sep 17 00:00:00 2001 From: Ice3man Date: Tue, 16 Dec 2025 19:35:57 +0530 Subject: [PATCH 2/2] revert misc --- pkg/massdns/process.go | 85 ++++++++++++++++++++---------------------- 1 file changed, 41 insertions(+), 44 deletions(-) diff --git a/pkg/massdns/process.go b/pkg/massdns/process.go index 574fc84..62fd208 100644 --- a/pkg/massdns/process.go +++ b/pkg/massdns/process.go @@ -277,69 +277,66 @@ func (instance *Instance) autoExtractRootDomains(store *store.Store) error { } func (instance *Instance) filterWildcards(st *store.Store) error { - // Build hostname -> IPs map to avoid redundant DNS queries - hostnameToIPs := make(map[string][]string) - hostnameCounters := make(map[string]int) - - st.Iterate(func(ip string, hostnames []string, counter int) { - for _, hostname := range hostnames { - hostnameToIPs[hostname] = append(hostnameToIPs[hostname], ip) - if counter > hostnameCounters[hostname] { - hostnameCounters[hostname] = counter - } - } - }) - // Start to work in parallel on wildcards wildcardWg := sizedwaitgroup.New(instance.options.WildcardsThreads) - for hostname, ips := range hostnameToIPs { - // Skip if any IP is already marked as wildcard - hasWildcardIP := false - for _, ip := range ips { - if instance.wildcardStore.Has(ip) { - hasWildcardIP = true - break - } - } - if hasWildcardIP { - continue + var allCancelFunc []context.CancelFunc + + st.Iterate(func(ip string, hostnames []string, counter int) { + ipCtx, ipCancelFunc := context.WithCancel(context.Background()) + allCancelFunc = append(allCancelFunc, ipCancelFunc) + // We've stumbled upon a wildcard, just ignore it. + if instance.wildcardStore.Has(ip) { + return } - counter := hostnameCounters[hostname] - // Perform wildcard detection on the hostname if counter >= 5 or strict mode + // Perform wildcard detection on the ip, if an IP is found in the wildcard + // we add it to the wildcard map so that further runs don't require such filtering again. if counter >= 5 || instance.options.StrictWildcard { - wildcardWg.Add() - go func(hostname string, ips []string) { - defer wildcardWg.Done() + for _, hostname := range hostnames { + wildcardWg.Add() + go func(ctx context.Context, ipCancelFunc context.CancelFunc, IP string, hostname string) { + defer wildcardWg.Done() - gologger.Info().Msgf("Started filtering wildcards for %s (with %d IPs)\n", hostname, len(ips)) + gologger.Info().Msgf("Started filtering wildcards for %s\n", hostname) - isWildcard, wildcardIPs := instance.wildcardResolver.LookupHost(hostname, ips) - if len(wildcardIPs) > 0 { - for ip := range wildcardIPs { - if err := instance.wildcardStore.Set(ip); err != nil { - gologger.Error().Msgf("could not set wildcard ip: %s", err) + select { + case <-ctx.Done(): + return + default: + } + + isWildcard, ips := instance.wildcardResolver.LookupHost(hostname, []string{IP}) + if len(ips) > 0 { + for ip := range ips { + // we add the single ip to the wildcard list + if err := instance.wildcardStore.Set(ip); err != nil { + gologger.Error().Msgf("could not set wildcard ip: %s", err) + } + gologger.Info().Msgf("Removing wildcard %s\n", ip) } - gologger.Info().Msgf("Removing wildcard %s\n", ip) } - } - if isWildcard { - for _, ip := range ips { - if err := instance.wildcardStore.Set(ip); err != nil { + if isWildcard { + // we also mark the original ip as wildcard, since at least once it resolved to this host + if err := instance.wildcardStore.Set(IP); err != nil { gologger.Error().Msgf("could not set wildcard ip: %s", err) } + ipCancelFunc() + gologger.Info().Msgf("Removed wildcard %s\n", IP) } - gologger.Info().Msgf("Removed wildcard hostname %s with %d IPs\n", hostname, len(ips)) - } - }(hostname, ips) + }(ipCtx, ipCancelFunc, ip, hostname) + } } - } + }) wildcardWg.Wait() + for _, cancelFunc := range allCancelFunc { + cancelFunc() + } + // Do a second pass as well and remove all the wildcards // from the store that we have found so that everything is covered allWildcardIPs := instance.wildcardResolver.GetAllWildcardIPs()