add rate limit per source option

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)
   -t int                number of concurrent goroutines for resolving (-active only) (default 10)
 
 UPDATE:

v2/go.mod

@@ -11,7 +11,7 @@ require (
 	github.com/projectdiscovery/dnsx v1.1.4
 	github.com/projectdiscovery/fdmax v0.0.4
 	github.com/projectdiscovery/gologger v1.1.10
-	github.com/projectdiscovery/ratelimit v0.0.8
+	github.com/projectdiscovery/ratelimit v0.0.9
 	github.com/projectdiscovery/utils v0.0.39
 	github.com/rs/xid v1.5.0
 	github.com/stretchr/testify v1.8.4

v2/go.sum

@@ -126,6 +126,8 @@ github.com/projectdiscovery/gologger v1.1.10 h1:XNRdtzLTdxiFGuK9gutoL752mykzXDoi
 github.com/projectdiscovery/gologger v1.1.10/go.mod h1:VqANHK7qcEq3i6/vV5HNWwdyv2aFPSrlaVDU4Ogrc6U=
 github.com/projectdiscovery/ratelimit v0.0.8 h1:K6S/DCr48xNxTXHRmU82wl1mj7j0VrXnAKr8sKTacHI=
 github.com/projectdiscovery/ratelimit v0.0.8/go.mod h1:JJAtj8Rd5DNqN5FgwyMHWIi4BHivOw1+8gDrpsBf8Ic=
+github.com/projectdiscovery/ratelimit v0.0.9 h1:28t2xDHUnyss1irzqPG3Oxz5hkRjl+3Q2I/aes7nau8=
+github.com/projectdiscovery/ratelimit v0.0.9/go.mod h1:f98UxLsHt0dWrHTbRDxos4+RvOLE0UFpyECfrfKBz1I=
 github.com/projectdiscovery/retryabledns v1.0.23 h1:8X6HECevl6n7K7kxqfHLTcsCVXDdz9HR+1hS6fto7sw=
 github.com/projectdiscovery/retryabledns v1.0.23/go.mod h1:tj8BEe0jWCB4m3aAhHq4S2dqfpod0h+BiCZmXkXAGac=
 github.com/projectdiscovery/retryablehttp-go v1.0.15 h1:kP9x9f++QimRwb8ABqnI1dhEymvnZXS2Wp2Zs4rWk/c=

v2/pkg/passive/passive.go

@@ -3,25 +3,37 @@ package passive
 import (
 	"context"
 	"fmt"
+	"math"
 	"sort"
+	"strconv"
+	"strings"
 	"sync"
 	"time"
 
+	"github.com/projectdiscovery/ratelimit"
 	"github.com/projectdiscovery/subfinder/v2/pkg/subscraping"
 )
 
 // 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, rateLimits map[string]interface{}, timeout int, maxEnumTime time.Duration) chan subscraping.Result {
+	return a.EnumerateSubdomainsWithCtx(context.Background(), domain, proxy, rateLimit, rateLimits, timeout, maxEnumTime)
 }
 
 // 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, rateLimits map[string]interface{}, timeout int, maxEnumTime time.Duration) chan subscraping.Result {
 	results := make(chan subscraping.Result)
+
 	go func() {
 		defer close(results)
 
-		session, err := subscraping.NewSession(domain, proxy, rateLimit, timeout)
+		multiRateLimiter, err := buildMultiRateLimiter(ctx, a, rateLimit, rateLimits)
+		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),
@@ -38,7 +50,9 @@ func (a *Agent) EnumerateSubdomainsWithCtx(ctx context.Context, domain string, p
 			wg.Add(1)
 
 			go func(source subscraping.Source) {
-				for resp := range source.Run(ctx, domain, session) {
+				//lint:ignore SA1029 reason
+				ctxWithValue := context.WithValue(ctx, "source", source.Name())
+				for resp := range source.Run(ctxWithValue, domain, session) {
 					results <- resp
 				}
 				wg.Done()
@@ -50,6 +64,57 @@ func (a *Agent) EnumerateSubdomainsWithCtx(ctx context.Context, domain string, p
 	return results
 }
 
+func buildMultiRateLimiter(ctx context.Context, a *Agent, rateLimit int, rateLimits map[string]interface{}) (*ratelimit.MultiLimiter, error) {
+	var multiRateLimiter *ratelimit.MultiLimiter
+	var err error
+	for _, source := range a.sources {
+		var rl uint
+		if sourceRateLimit, ok := rateLimits[strings.ToLower(source.Name())]; ok {
+			rl = sourceRateLimitOrGlobal(rateLimit, 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 sourceRateLimitOrGlobal(globalRateLimit int, sourceRateLimit interface{}) uint {
+	if sourceRateLimitStr, ok := sourceRateLimit.(string); ok {
+		sourceRateLimitUint, err := strconv.ParseUint(sourceRateLimitStr, 10, 64)
+		if err == nil {
+			return uint(sourceRateLimitUint)
+		}
+	}
+	return uint(globalRateLimit)
+}
+
+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,18 @@ func TestSourcesWithKeys(t *testing.T) {
 
 	gologger.DefaultLogger.SetMaxLevel(levels.LevelDebug)
 
-	ctx := context.Background()
-	session, err := subscraping.NewSession(domain, "", 0, timeout)
+	ctxParent := context.Background()
+	//lint:ignore SA1029 reason
+	ctx := context.WithValue(ctxParent, "default", "default")
+	multiRateLimiter, err := ratelimit.NewMultiLimiter(ctx, &ratelimit.Options{
+		Key:         "default",
+		IsUnlimited: true,
+		MaxCount:    math.MaxUint32,
+		Duration:    time.Millisecond,
+	})
+	assert.Nil(t, err)
+
+	session, err := subscraping.NewSession(domain, "", multiRateLimiter, timeout)
 	assert.Nil(t, err)
 
 	var expected = subscraping.Result{Type: subscraping.Subdomain, Value: domain, Error: nil}
@@ -42,7 +55,7 @@ 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) {
+			for result := range source.Run(ctxParent, 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,15 +3,18 @@ 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"
 )
 
@@ -27,8 +30,19 @@ func TestSourcesWithoutKeys(t *testing.T) {
 
 	gologger.DefaultLogger.SetMaxLevel(levels.LevelDebug)
 
-	ctx := context.Background()
-	session, err := subscraping.NewSession(domain, "", 0, timeout)
+	ctxParent := context.Background()
+
+	//lint:ignore SA1029 reason
+	ctx := context.WithValue(ctxParent, "default", "default")
+	multiRateLimiter, err := ratelimit.NewMultiLimiter(ctx, &ratelimit.Options{
+		Key:         "default",
+		IsUnlimited: true,
+		MaxCount:    math.MaxUint32,
+		Duration:    time.Millisecond,
+	})
+	assert.Nil(t, err)
+
+	session, err := subscraping.NewSession(domain, "", multiRateLimiter, timeout)
 	assert.Nil(t, err)
 
 	var expected = subscraping.Result{Type: subscraping.Subdomain, Value: domain, Error: nil}

v2/pkg/runner/enumerate.go

@@ -40,7 +40,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.RateLimits.AsMap(), r.options.Timeout, time.Duration(r.options.MaxEnumerationTime)*time.Minute)
 
 	wg := &sync.WaitGroup{}
 	wg.Add(1)

v2/pkg/runner/options.go

@@ -59,7 +59,8 @@ type Options struct {
 	Config             string              // Config contains the location of the config file
 	ProviderConfig     string              // ProviderConfig contains the location of the provider config file
 	Proxy              string              // HTTP proxy
-	RateLimit          int                 // Maximum number of HTTP requests to send per second
+	RateLimit          int                 // Global maximum number of HTTP requests to send per second
+	RateLimits         goflags.RuntimeMap  // Maximum number of HTTP requests to send per second
 	ExcludeIps         bool
 	Match              goflags.StringSlice
 	Filter             goflags.StringSlice
@@ -114,7 +115,8 @@ func ParseOptions() *Options {
 	)
 
 	createGroup(flagSet, "rate-limit", "Rate-limit",
-		flagSet.IntVarP(&options.RateLimit, "rate-limit", "rl", 0, "maximum number of http requests to send per second"),
+		flagSet.IntVarP(&options.RateLimit, "rate-limit", "rl", 0, "maximum number of http requests to send per second (global)"),
+		flagSet.RuntimeMapVar(&options.RateLimits, "rls", nil, "maximum number of http requests to send per second four providers in key=value format (-rls hackertarget=10)"),
 		flagSet.IntVar(&options.Threads, "t", 10, "number of concurrent goroutines for resolving (-active only)"),
 	)
 

v2/pkg/subscraping/agent.go

@@ -18,7 +18,7 @@ import (
 )
 
 // NewSession creates a new session object for a domain
-func NewSession(domain string, proxy string, rateLimit, timeout int) (*Session, error) {
+func NewSession(domain string, proxy string, multiRateLimiter *ratelimit.MultiLimiter, timeout int) (*Session, error) {
 	Transport := &http.Transport{
 		MaxIdleConns:        100,
 		MaxIdleConnsPerHost: 100,
@@ -49,11 +49,7 @@ func NewSession(domain string, proxy string, rateLimit, timeout int) (*Session,
 	session := &Session{Client: client}
 
 	// Initiate rate limit instance
-	if rateLimit > 0 {
-		session.RateLimiter = ratelimit.New(context.Background(), uint(rateLimit), time.Second)
-	} else {
-		session.RateLimiter = ratelimit.NewUnlimited(context.Background())
-	}
+	session.MultiRateLimiter = multiRateLimiter
 
 	// Create a new extractor object for the current domain
 	extractor, err := NewSubdomainExtractor(domain)
@@ -106,7 +102,8 @@ func (s *Session) HTTPRequest(ctx context.Context, method, requestURL, cookies s
 		req.Header.Set(key, value)
 	}
 
-	s.RateLimiter.Take()
+	sourceName := ctx.Value("source").(string)
+	s.MultiRateLimiter.Take(sourceName)
 
 	return httpRequestWrapper(s.Client, req)
 }
@@ -125,7 +122,7 @@ func (s *Session) DiscardHTTPResponse(response *http.Response) {
 
 // Close the session
 func (s *Session) Close() {
-	s.RateLimiter.Stop()
+	s.MultiRateLimiter.Stop()
 	s.Client.CloseIdleConnections()
 }
 

v2/pkg/subscraping/types.go

@@ -59,7 +59,7 @@ type Session struct {
 	// Client is the current http client
 	Client *http.Client
 	// Rate limit instance
-	RateLimiter *ratelimit.Limiter
+	MultiRateLimiter *ratelimit.MultiLimiter
 }
 
 // Result is a result structure returned by a source