add rate limit per source option

.github/workflows/build-test.yml

@@ -18,7 +18,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v4
         with:
-          go-version: 1.19
+          go-version: 1.20.x
 
       - name: Check out code
         uses: actions/checkout@v3

.github/workflows/lint-test.yml

@@ -15,7 +15,7 @@ jobs:
       - uses: actions/checkout@v3
       - uses: actions/setup-go@v4
         with: 
-          go-version: 1.19
+          go-version: 1.20.x
       - name: Run golangci-lint
         uses: golangci/[email protected]
         with:

.github/workflows/release-binary.yml

@@ -18,7 +18,7 @@ jobs:
       - name: "Set up Go"
         uses: actions/setup-go@v4
         with: 
-          go-version: 1.19
+          go-version: 1.20.x
 
       - name: "Create release on GitHub"
         uses: goreleaser/goreleaser-action@v3

README.md

@@ -74,6 +74,7 @@ FILTER:
 
 RATE-LIMIT:
   -rl, -rate-limit int  maximum number of http requests to send per second
+  -rls value            maximum number of http requests to send per second four providers in key=value format (-rls "hackertarget=10/s,shodan=15/s")
   -t int                number of concurrent goroutines for resolving (-active only) (default 10)
 
 UPDATE:
@@ -110,7 +111,7 @@ OPTIMIZATION:
 
 # Installation
 
-`subfinder` requires **go1.19** to install successfully. Run the following command to install the latest version:
+`subfinder` requires **go1.20** to install successfully. Run the following command to install the latest version:
 
 ```sh
 go install -v github.com/projectdiscovery/subfinder/v2/cmd/subfinder@latest

v2/go.mod

@@ -1,6 +1,6 @@
 module github.com/projectdiscovery/subfinder/v2
 
-go 1.18
+go 1.20
 
 require (
 	github.com/corpix/uarand v0.2.0

v2/pkg/passive/passive.go

@@ -3,25 +3,53 @@ package passive
 import (
 	"context"
 	"fmt"
+	"math"
 	"sort"
+	"strings"
 	"sync"
 	"time"
 
+	"github.com/projectdiscovery/ratelimit"
 	"github.com/projectdiscovery/subfinder/v2/pkg/subscraping"
 )
 
+type EnumerationOptions struct {
+	customRateLimiter *subscraping.CustomRateLimit
+}
+
+type EnumerateOption func(opts *EnumerationOptions)
+
+func WithCustomRateLimit(crl *subscraping.CustomRateLimit) EnumerateOption {
+	return func(opts *EnumerationOptions) {
+		opts.customRateLimiter = crl
+	}
+}
+
 // EnumerateSubdomains wraps EnumerateSubdomainsWithCtx with an empty context
-func (a *Agent) EnumerateSubdomains(domain string, proxy string, rateLimit, timeout int, maxEnumTime time.Duration) chan subscraping.Result {
-	return a.EnumerateSubdomainsWithCtx(context.Background(), domain, proxy, rateLimit, timeout, maxEnumTime)
+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...)
 }
 
 // EnumerateSubdomainsWithCtx enumerates all the subdomains for a given domain
-func (a *Agent) EnumerateSubdomainsWithCtx(ctx context.Context, domain string, proxy string, rateLimit, timeout int, maxEnumTime time.Duration) chan subscraping.Result {
+func (a *Agent) EnumerateSubdomainsWithCtx(ctx context.Context, domain string, proxy string, rateLimit int, timeout int, maxEnumTime time.Duration, options ...EnumerateOption) chan subscraping.Result {
 	results := make(chan subscraping.Result)
+
 	go func() {
 		defer close(results)
 
-		session, err := subscraping.NewSession(domain, proxy, rateLimit, timeout)
+		var enumerateOptions EnumerationOptions
+		for _, enumerateOption := range options {
+			enumerateOption(&enumerateOptions)
+		}
+
+		multiRateLimiter, err := a.buildMultiRateLimiter(ctx, rateLimit, enumerateOptions.customRateLimiter)
+		if err != nil {
+			results <- subscraping.Result{
+				Type: subscraping.Error, Error: fmt.Errorf("could not init multi rate limiter for %s: %s", domain, err),
+			}
+			return
+		}
+		session, err := subscraping.NewSession(domain, proxy, multiRateLimiter, timeout)
 		if err != nil {
 			results <- subscraping.Result{
 				Type: subscraping.Error, Error: fmt.Errorf("could not init passive session for %s: %s", domain, err),
@@ -36,9 +64,9 @@ func (a *Agent) EnumerateSubdomainsWithCtx(ctx context.Context, domain string, p
 		// Run each source in parallel on the target domain
 		for _, runner := range a.sources {
 			wg.Add(1)
-
 			go func(source subscraping.Source) {
-				for resp := range source.Run(ctx, domain, session) {
+				ctxWithValue := context.WithValue(ctx, subscraping.CtxSourceArg, source.Name())
+				for resp := range source.Run(ctxWithValue, domain, session) {
 					results <- resp
 				}
 				wg.Done()
@@ -50,6 +78,54 @@ func (a *Agent) EnumerateSubdomainsWithCtx(ctx context.Context, domain string, p
 	return results
 }
 
+func (a *Agent) buildMultiRateLimiter(ctx context.Context, globalRateLimit int, rateLimit *subscraping.CustomRateLimit) (*ratelimit.MultiLimiter, error) {
+	var multiRateLimiter *ratelimit.MultiLimiter
+	var err error
+	for _, source := range a.sources {
+		var rl uint
+		if sourceRateLimit, ok := rateLimit.Custom.Get(strings.ToLower(source.Name())); ok {
+			rl = sourceRateLimitOrDefault(uint(globalRateLimit), sourceRateLimit)
+		}
+
+		if rl > 0 {
+			multiRateLimiter, err = addRateLimiter(ctx, multiRateLimiter, source.Name(), rl, time.Second)
+		} else {
+			multiRateLimiter, err = addRateLimiter(ctx, multiRateLimiter, source.Name(), math.MaxUint32, time.Millisecond)
+		}
+
+		if err != nil {
+			break
+		}
+	}
+	return multiRateLimiter, err
+}
+
+func sourceRateLimitOrDefault(defaultRateLimit uint, sourceRateLimit uint) uint {
+	if sourceRateLimit > 0 {
+		return sourceRateLimit
+	}
+	return defaultRateLimit
+}
+
+func addRateLimiter(ctx context.Context, multiRateLimiter *ratelimit.MultiLimiter, key string, maxCount uint, duration time.Duration) (*ratelimit.MultiLimiter, error) {
+	if multiRateLimiter == nil {
+		mrl, err := ratelimit.NewMultiLimiter(ctx, &ratelimit.Options{
+			Key:         key,
+			IsUnlimited: maxCount == math.MaxUint32,
+			MaxCount:    maxCount,
+			Duration:    duration,
+		})
+		return mrl, err
+	}
+	err := multiRateLimiter.Add(&ratelimit.Options{
+		Key:         key,
+		IsUnlimited: maxCount == math.MaxUint32,
+		MaxCount:    maxCount,
+		Duration:    duration,
+	})
+	return multiRateLimiter, err
+}
+
 func (a *Agent) GetStatistics() map[string]subscraping.Statistics {
 	stats := make(map[string]subscraping.Statistics)
 	sort.Slice(a.sources, func(i, j int) bool {

v2/pkg/passive/sources_w_auth_test.go

@@ -3,15 +3,18 @@ package passive
 import (
 	"context"
 	"fmt"
+	"math"
 	"os"
 	"reflect"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/gologger/levels"
+	"github.com/projectdiscovery/ratelimit"
 	"github.com/projectdiscovery/subfinder/v2/pkg/subscraping"
 )
 
@@ -21,8 +24,16 @@ func TestSourcesWithKeys(t *testing.T) {
 
 	gologger.DefaultLogger.SetMaxLevel(levels.LevelDebug)
 
-	ctx := context.Background()
-	session, err := subscraping.NewSession(domain, "", 0, timeout)
+	ctxParent := context.Background()
+	var multiRateLimiter *ratelimit.MultiLimiter
+	for _, source := range AllSources {
+		if !source.NeedsKey() {
+			continue
+		}
+		multiRateLimiter, _ = addRateLimiter(ctxParent, multiRateLimiter, source.Name(), math.MaxInt32, time.Millisecond)
+	}
+
+	session, err := subscraping.NewSession(domain, "", multiRateLimiter, timeout)
 	assert.Nil(t, err)
 
 	var expected = subscraping.Result{Type: subscraping.Subdomain, Value: domain, Error: nil}
@@ -42,7 +53,8 @@ func TestSourcesWithKeys(t *testing.T) {
 		t.Run(source.Name(), func(t *testing.T) {
 			var results []subscraping.Result
 
-			for result := range source.Run(ctx, domain, session) {
+			ctxWithValue := context.WithValue(ctxParent, subscraping.CtxSourceArg, source.Name())
+			for result := range source.Run(ctxWithValue, domain, session) {
 				results = append(results, result)
 
 				assert.Equal(t, source.Name(), result.Source, "wrong source name")

v2/pkg/passive/sources_wo_auth_test.go

@@ -3,32 +3,45 @@ package passive
 import (
 	"context"
 	"fmt"
+	"math"
 	"reflect"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/exp/slices"
 
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/gologger/levels"
+	"github.com/projectdiscovery/ratelimit"
 	"github.com/projectdiscovery/subfinder/v2/pkg/subscraping"
 )
 
 func TestSourcesWithoutKeys(t *testing.T) {
 	ignoredSources := []string{
-		"commoncrawl", // commoncrawl is under resourced and will likely time-out so step over it for this test https://groups.google.com/u/2/g/common-crawl/c/3QmQjFA_3y4/m/vTbhGqIBBQAJ
-		"riddler",     // Fails with 403: There might be too much traffic or a configuration error
-		"crtsh",       // Fails in GH Action (possibly IP-based ban) causing a timeout.
+		"commoncrawl",  // commoncrawl is under resourced and will likely time-out so step over it for this test https://groups.google.com/u/2/g/common-crawl/c/3QmQjFA_3y4/m/vTbhGqIBBQAJ
+		"riddler",      // Fails with 403: There might be too much traffic or a configuration error
+		"crtsh",        // Fails in GH Action (possibly IP-based ban) causing a timeout.
+		"hackertarget", // Fails in GH Action (possibly IP-based ban) but works locally
 	}
 
 	domain := "hackerone.com"
 	timeout := 60
 
 	gologger.DefaultLogger.SetMaxLevel(levels.LevelDebug)
 
-	ctx := context.Background()
-	session, err := subscraping.NewSession(domain, "", 0, timeout)
+	ctxParent := context.Background()
+
+	var multiRateLimiter *ratelimit.MultiLimiter
+	for _, source := range AllSources {
+		if source.NeedsKey() || slices.Contains(ignoredSources, source.Name()) {
+			continue
+		}
+		multiRateLimiter, _ = addRateLimiter(ctxParent, multiRateLimiter, source.Name(), math.MaxInt32, time.Millisecond)
+	}
+
+	session, err := subscraping.NewSession(domain, "", multiRateLimiter, timeout)
 	assert.Nil(t, err)
 
 	var expected = subscraping.Result{Type: subscraping.Subdomain, Value: domain, Error: nil}
@@ -41,7 +54,8 @@ func TestSourcesWithoutKeys(t *testing.T) {
 		t.Run(source.Name(), func(t *testing.T) {
 			var results []subscraping.Result
 
-			for result := range source.Run(ctx, domain, session) {
+			ctxWithValue := context.WithValue(ctxParent, subscraping.CtxSourceArg, source.Name())
+			for result := range source.Run(ctxWithValue, domain, session) {
 				results = append(results, result)
 
 				assert.Equal(t, source.Name(), result.Source, "wrong source name")

v2/pkg/runner/enumerate.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/projectdiscovery/gologger"
 
+	"github.com/projectdiscovery/subfinder/v2/pkg/passive"
 	"github.com/projectdiscovery/subfinder/v2/pkg/resolve"
 	"github.com/projectdiscovery/subfinder/v2/pkg/subscraping"
 )
@@ -40,7 +41,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)
+	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))
 
 	wg := &sync.WaitGroup{}
 	wg.Add(1)