fix(sources/virustotal): cap pagination + clearer 429 message

README.md

@@ -96,6 +96,7 @@ CONFIGURATION:
   -nW, -active                  display active subdomains only
   -proxy string                 http proxy to use with subfinder
   -ei, -exclude-ip              exclude IPs from the list of domains
+  -mr, -max-results int         limit the number of results per source (0 = unlimited; honored by paginating sources such as virustotal)
 
 DEBUG:
   -silent             show only subdomains in output

pkg/passive/passive.go

@@ -16,6 +16,7 @@ import (
 
 type EnumerationOptions struct {
 	customRateLimiter *subscraping.CustomRateLimit
+	maxResults        int
 }
 
 type EnumerateOption func(opts *EnumerationOptions)
@@ -26,6 +27,15 @@ func WithCustomRateLimit(crl *subscraping.CustomRateLimit) EnumerateOption {
 	}
 }
 
+// WithMaxResults sets a per-source limit on the number of results emitted.
+// A value of 0 (the default) means no limit. Sources that paginate can honor
+// this to avoid unnecessary requests (e.g. to stay within API quotas).
+func WithMaxResults(maxResults int) EnumerateOption {
+	return func(opts *EnumerationOptions) {
+		opts.maxResults = maxResults
+	}
+}
+
 // EnumerateSubdomains wraps EnumerateSubdomainsWithCtx with an empty context
 func (a *Agent) EnumerateSubdomains(domain string, proxy string, rateLimit int, timeout int, maxEnumTime time.Duration, options ...EnumerateOption) chan subscraping.Result {
 	return a.EnumerateSubdomainsWithCtx(context.Background(), domain, proxy, rateLimit, timeout, maxEnumTime, options...)
@@ -58,6 +68,7 @@ func (a *Agent) EnumerateSubdomainsWithCtx(ctx context.Context, domain string, p
 			return
 		}
 		defer session.Close()
+		session.MaxResults = enumerateOptions.maxResults
 
 		ctx, cancel := context.WithTimeout(ctx, maxEnumTime)
 

pkg/runner/enumerate.go

@@ -50,7 +50,7 @@ func (r *Runner) EnumerateSingleDomainWithCtx(ctx context.Context, domain string
 
 	// Run the passive subdomain enumeration
 	now := time.Now()
-	passiveResults := r.passiveAgent.EnumerateSubdomainsWithCtx(ctx, domain, r.options.Proxy, r.options.RateLimit, r.options.Timeout, time.Duration(r.options.MaxEnumerationTime)*time.Minute, passive.WithCustomRateLimit(r.rateLimit))
+	passiveResults := r.passiveAgent.EnumerateSubdomainsWithCtx(ctx, domain, r.options.Proxy, r.options.RateLimit, r.options.Timeout, time.Duration(r.options.MaxEnumerationTime)*time.Minute, passive.WithCustomRateLimit(r.rateLimit), passive.WithMaxResults(r.options.MaxResults))
 
 	wg := &sync.WaitGroup{}
 	wg.Add(1)

pkg/runner/options.go

@@ -69,6 +69,11 @@ type Options struct {
 	filterRegexes      []*regexp.Regexp
 	ResultCallback     OnResultCallback // OnResult callback
 	DisableUpdateCheck bool             // DisableUpdateCheck disable update checking
+
+	// MaxResults limits the number of results requested per source.
+	// A value of 0 (default) means no limit. Only sources that paginate
+	// honor this (currently virustotal), to help stay within API quotas.
+	MaxResults int `yaml:"max-results,omitempty"`
 }
 
 // OnResultCallback (hostResult)
@@ -128,6 +133,7 @@ func ParseOptions() *Options {
 		flagSet.BoolVarP(&options.RemoveWildcard, "active", "nW", false, "display active subdomains only"),
 		flagSet.StringVar(&options.Proxy, "proxy", "", "http proxy to use with subfinder"),
 		flagSet.BoolVarP(&options.ExcludeIps, "exclude-ip", "ei", false, "exclude IPs from the list of domains"),
+		flagSet.IntVarP(&options.MaxResults, "max-results", "mr", 0, "limit the number of results per source (0 = unlimited; honored by paginating sources such as virustotal)"),
 	)
 
 	flagSet.CreateGroup("debug", "Debug",

pkg/runner/validate.go

@@ -40,6 +40,11 @@ func (options *Options) validateOptions() error {
 		return errors.New("hostip flag must be used with RemoveWildcard option")
 	}
 
+	// The per-source results limit cannot be negative.
+	if options.MaxResults < 0 {
+		return fmt.Errorf("max-results cannot be negative")
+	}
+
 	if options.Match != nil {
 		options.matchRegexes = make([]*regexp.Regexp, len(options.Match))
 		var err error

pkg/subscraping/sources/virustotal/virustotal.go

@@ -4,6 +4,7 @@ package virustotal
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"time"
 
 	jsoniter "github.com/json-iterator/go"
@@ -47,6 +48,10 @@ func (s *Source) Run(ctx context.Context, domain string, session *subscraping.Se
 			close(results)
 		}(time.Now())
 
+		// Honor an optional per-source result limit to avoid unnecessary
+		// pagination requests (e.g. to stay within API quotas). 0 = no limit.
+		maxResults := session.MaxResults
+
 		randomApiKey := subscraping.PickRandom(s.apiKeys, s.Name())
 		if randomApiKey == "" {
 			return
@@ -65,20 +70,27 @@ func (s *Source) Run(ctx context.Context, domain string, session *subscraping.Se
 			s.requests++
 			resp, err := session.Get(ctx, url, "", map[string]string{"x-apikey": randomApiKey})
 			if err != nil {
+				// The free tier grants 500 requests/day; once it is exhausted every
+				// call returns HTTP 429. Surface an actionable message instead of the
+				// generic "unexpected status code 429" so operators know to supply an
+				// enterprise key or lower the scope (see #1718).
+				if resp != nil && resp.StatusCode == http.StatusTooManyRequests {
+					err = fmt.Errorf("virustotal quota exhausted (HTTP 429); some subdomains for %s may be missing", domain)
+				}
 				results <- subscraping.Result{Source: s.Name(), Type: subscraping.Error, Error: err}
 				s.errors++
 				session.DiscardHTTPResponse(resp)
 				return
 			}
-			defer func() {
-				if err := resp.Body.Close(); err != nil {
-					results <- subscraping.Result{Source: s.Name(), Type: subscraping.Error, Error: err}
-					s.errors++
-				}
-			}()
 
 			var data response
 			err = jsoniter.NewDecoder(resp.Body).Decode(&data)
+			// Close the body per iteration; deferring inside the loop would keep
+			// every page's body (and its connection) open until the goroutine exits.
+			if closeErr := resp.Body.Close(); closeErr != nil {
+				results <- subscraping.Result{Source: s.Name(), Type: subscraping.Error, Error: closeErr}
+				s.errors++
+			}
 			if err != nil {
 				results <- subscraping.Result{Source: s.Name(), Type: subscraping.Error, Error: err}
 				s.errors++
@@ -92,6 +104,9 @@ func (s *Source) Run(ctx context.Context, domain string, session *subscraping.Se
 				case results <- subscraping.Result{Source: s.Name(), Type: subscraping.Subdomain, Value: subdomain.Id}:
 					s.results++
 				}
+				if maxResults > 0 && s.results >= maxResults {
+					return
+				}
 			}
 			cursor = data.Meta.Cursor
 			if cursor == "" {

pkg/subscraping/sources/virustotal/virustotal_test.go

@@ -0,0 +1,193 @@
+package virustotal
+
+import (
+	"context"
+	"fmt"
+	"math"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/projectdiscovery/ratelimit"
+	"github.com/projectdiscovery/subfinder/v2/pkg/subscraping"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// rewriteTransport sends every request to the test server's host while
+// preserving the original path/query, so the source's hardcoded
+// www.virustotal.com URL still resolves to our handler.
+type rewriteTransport struct {
+	target *url.URL
+}
+
+func (r *rewriteTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	req.URL.Scheme = r.target.Scheme
+	req.URL.Host = r.target.Host
+	return http.DefaultTransport.RoundTrip(req)
+}
+
+func newTestSession(t *testing.T, server *httptest.Server, maxResults int) *subscraping.Session {
+	t.Helper()
+	target, err := url.Parse(server.URL)
+	require.NoError(t, err)
+
+	mrl, err := ratelimit.NewMultiLimiter(context.Background(), &ratelimit.Options{
+		Key:         "virustotal",
+		IsUnlimited: false,
+		MaxCount:    math.MaxInt32,
+		Duration:    time.Millisecond,
+	})
+	require.NoError(t, err)
+
+	return &subscraping.Session{
+		Client:           &http.Client{Transport: &rewriteTransport{target: target}, Timeout: 5 * time.Second},
+		MultiRateLimiter: mrl,
+		MaxResults:       maxResults,
+	}
+}
+
+func runSource(t *testing.T, source *Source, session *subscraping.Session, domain string) (subs []string, errs []error) {
+	t.Helper()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	ctxWithValue := context.WithValue(ctx, subscraping.CtxSourceArg, "virustotal")
+
+	for r := range source.Run(ctxWithValue, domain, session) {
+		switch r.Type {
+		case subscraping.Subdomain:
+			subs = append(subs, r.Value)
+		case subscraping.Error:
+			errs = append(errs, r.Error)
+		}
+	}
+	return
+}
+
+func TestVirustotalSource_Metadata(t *testing.T) {
+	source := &Source{}
+	assert.Equal(t, "virustotal", source.Name())
+	assert.True(t, source.IsDefault())
+	assert.True(t, source.HasRecursiveSupport())
+	assert.True(t, source.NeedsKey())
+}
+
+func TestVirustotalSource_NoApiKey(t *testing.T) {
+	source := &Source{}
+
+	ctx := context.Background()
+	mrl, _ := ratelimit.NewMultiLimiter(ctx, &ratelimit.Options{
+		Key:         "virustotal",
+		IsUnlimited: false,
+		MaxCount:    math.MaxInt32,
+		Duration:    time.Millisecond,
+	})
+	session := &subscraping.Session{Client: http.DefaultClient, MultiRateLimiter: mrl}
+	ctxWithValue := context.WithValue(ctx, subscraping.CtxSourceArg, "virustotal")
+
+	var count int
+	for range source.Run(ctxWithValue, "example.com", session) {
+		count++
+	}
+	assert.Equal(t, 0, count, "expected no results without an api key")
+	assert.Equal(t, 0, source.Statistics().Requests)
+}
+
+// With no per-source limit the source must follow the cursor to the end and
+// return every subdomain, i.e. the default behaviour stays unbounded so paid
+// keys are never silently truncated.
+func TestVirustotalSource_FollowsCursorWhenUnlimited(t *testing.T) {
+	var hits int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		n := atomic.AddInt32(&hits, 1)
+		w.Header().Set("Content-Type", "application/json")
+		// Three pages of one subdomain each, then the cursor terminates.
+		if n < 3 {
+			_, _ = fmt.Fprintf(w, `{"data":[{"id":"sub%d.example.com"}],"meta":{"cursor":"next-%d"}}`, n, n)
+			return
+		}
+		_, _ = fmt.Fprintf(w, `{"data":[{"id":"sub%d.example.com"}],"meta":{"cursor":""}}`, n)
+	}))
+	defer server.Close()
+
+	source := &Source{}
+	source.AddApiKeys([]string{"test-key"})
+
+	subs, errs := runSource(t, source, newTestSession(t, server, 0), "example.com")
+
+	assert.Empty(t, errs)
+	assert.ElementsMatch(t, []string{"sub1.example.com", "sub2.example.com", "sub3.example.com"}, subs)
+	assert.Equal(t, 3, source.Statistics().Requests)
+}
+
+// When -max-results (session.MaxResults) is set, pagination must stop as soon
+// as the limit is reached so a single domain can't drain an API quota (#1718).
+func TestVirustotalSource_MaxResultsCapsPagination(t *testing.T) {
+	var hits int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		n := atomic.AddInt32(&hits, 1)
+		w.Header().Set("Content-Type", "application/json")
+		// Always return one subdomain plus a non-empty cursor so the source
+		// would loop forever without the max-results cap.
+		_, _ = fmt.Fprintf(w, `{"data":[{"id":"sub%d.example.com"}],"meta":{"cursor":"next-%d"}}`, n, n)
+	}))
+	defer server.Close()
+
+	const limit = 3
+	source := &Source{}
+	source.AddApiKeys([]string{"test-key"})
+
+	subs, errs := runSource(t, source, newTestSession(t, server, limit), "example.com")
+
+	assert.Empty(t, errs, "no errors expected on the happy path")
+	assert.Len(t, subs, limit, "should yield exactly max-results subdomains")
+	assert.Equal(t, limit, source.Statistics().Requests, "request count must be capped by max-results")
+	assert.Equal(t, int32(limit), atomic.LoadInt32(&hits))
+}
+
+// Cursor exhaustion must not waste an extra request.
+func TestVirustotalSource_StopsWhenCursorEmpty(t *testing.T) {
+	var hits int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		n := atomic.AddInt32(&hits, 1)
+		w.Header().Set("Content-Type", "application/json")
+		if n == 1 {
+			_, _ = fmt.Fprint(w, `{"data":[{"id":"a.example.com"}],"meta":{"cursor":"more"}}`)
+			return
+		}
+		_, _ = fmt.Fprint(w, `{"data":[{"id":"b.example.com"}],"meta":{"cursor":""}}`)
+	}))
+	defer server.Close()
+
+	source := &Source{}
+	source.AddApiKeys([]string{"test-key"})
+
+	subs, errs := runSource(t, source, newTestSession(t, server, 0), "example.com")
+
+	assert.Empty(t, errs)
+	assert.ElementsMatch(t, []string{"a.example.com", "b.example.com"}, subs)
+	assert.Equal(t, 2, source.Statistics().Requests)
+}
+
+// On HTTP 429 (quota exhausted) the source must surface a human-readable error
+// mentioning the quota, not a generic "unexpected status code 429" string.
+func TestVirustotalSource_RateLimitedReturnsActionableError(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusTooManyRequests)
+	}))
+	defer server.Close()
+
+	source := &Source{}
+	source.AddApiKeys([]string{"test-key"})
+
+	_, errs := runSource(t, source, newTestSession(t, server, 0), "example.com")
+
+	require.Len(t, errs, 1, "exactly one error expected on 429")
+	assert.Contains(t, errs[0].Error(), "quota exhausted")
+	assert.Contains(t, errs[0].Error(), "429")
+	assert.Equal(t, 1, source.Statistics().Requests, "must not retry past the first 429")
+	assert.Equal(t, 1, source.Statistics().Errors)
+}

pkg/subscraping/types.go

@@ -92,6 +92,11 @@ type Session struct {
 	MultiRateLimiter *ratelimit.MultiLimiter
 	// Timeout is the timeout in seconds for requests
 	Timeout int
+	// MaxResults is the maximum number of results a source should emit
+	// before stopping. A value of 0 means no limit (default behavior).
+	// Sources that paginate can honor this to avoid unnecessary requests
+	// (e.g. to stay within API quotas).
+	MaxResults int
 }
 
 // Result is a result structure returned by a source