diff --git a/pkg/passive/passive.go b/pkg/passive/passive.go index 81a7cf04..9b2661ee 100644 --- a/pkg/passive/passive.go +++ b/pkg/passive/passive.go @@ -85,14 +85,14 @@ func (a *Agent) EnumerateSubdomainsWithCtx(ctx context.Context, domain string, p func (a *Agent) buildMultiRateLimiter(ctx context.Context, globalRateLimit int, rateLimit *subscraping.CustomRateLimit) (*ratelimit.MultiLimiter, error) { var multiRateLimiter *ratelimit.MultiLimiter var err error + if rateLimit == nil { + rateLimit = &subscraping.CustomRateLimit{} + } for _, source := range a.sources { - var rl uint - if sourceRateLimit, ok := rateLimit.Custom.Get(strings.ToLower(source.Name())); ok { - rl = sourceRateLimitOrDefault(uint(globalRateLimit), sourceRateLimit) - } + rl, duration := resolveSourceRateLimit(globalRateLimit, rateLimit, source.Name()) if rl > 0 { - multiRateLimiter, err = addRateLimiter(ctx, multiRateLimiter, source.Name(), rl, time.Second) + multiRateLimiter, err = addRateLimiter(ctx, multiRateLimiter, source.Name(), rl, duration) } else { multiRateLimiter, err = addRateLimiter(ctx, multiRateLimiter, source.Name(), math.MaxUint32, time.Millisecond) } @@ -104,6 +104,30 @@ func (a *Agent) buildMultiRateLimiter(ctx context.Context, globalRateLimit int, return multiRateLimiter, err } +// resolveSourceRateLimit returns the effective rate limit and duration for a source. +// Priority: per-source custom limit > global -rl limit > unlimited. +// Duration comes from -rls (e.g. hackertarget=2/m → 2 per minute), defaulting to per-second. +func resolveSourceRateLimit(globalRateLimit int, rateLimit *subscraping.CustomRateLimit, sourceName string) (uint, time.Duration) { + duration := time.Second // default: requests per second + + sourceLower := strings.ToLower(sourceName) + if sourceRL, ok := rateLimit.Custom.Get(sourceLower); ok { + rl := sourceRateLimitOrDefault(uint(max(globalRateLimit, 0)), sourceRL) + // Use per-source duration from -rls if set (e.g. "2/m" → time.Minute) + if d, ok := rateLimit.CustomDuration.Get(sourceLower); ok && d > 0 { + duration = d + } + return rl, duration + } + + // No per-source limit: fall back to global -rl + if globalRateLimit > 0 { + return uint(globalRateLimit), duration + } + + return 0, duration +} + func sourceRateLimitOrDefault(defaultRateLimit uint, sourceRateLimit uint) uint { if sourceRateLimit > 0 { return sourceRateLimit diff --git a/pkg/passive/ratelimit_test.go b/pkg/passive/ratelimit_test.go new file mode 100644 index 00000000..9cce7b57 --- /dev/null +++ b/pkg/passive/ratelimit_test.go @@ -0,0 +1,125 @@ +package passive + +import ( + "context" + "math" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mapsutil "github.com/projectdiscovery/utils/maps" + + "github.com/projectdiscovery/subfinder/v2/pkg/subscraping" +) + +func newCustomRateLimit(counts map[string]uint, durations map[string]time.Duration) *subscraping.CustomRateLimit { + crl := &subscraping.CustomRateLimit{ + Custom: mapsutil.SyncLockMap[string, uint]{ + Map: make(map[string]uint), + }, + CustomDuration: mapsutil.SyncLockMap[string, time.Duration]{ + Map: make(map[string]time.Duration), + }, + } + for k, v := range counts { + _ = crl.Custom.Set(k, v) + } + for k, d := range durations { + _ = crl.CustomDuration.Set(k, d) + } + return crl +} + +func TestResolveSourceRateLimit_PerSourceOverride(t *testing.T) { + // Per-source limit should override global + crl := newCustomRateLimit( + map[string]uint{"hackertarget": 5}, + map[string]time.Duration{"hackertarget": time.Minute}, + ) + + rl, dur := resolveSourceRateLimit(10, crl, "hackertarget") + assert.Equal(t, uint(5), rl) + assert.Equal(t, time.Minute, dur) +} + +func TestResolveSourceRateLimit_GlobalFallback(t *testing.T) { + // Sources without per-source limit should use global -rl + crl := newCustomRateLimit(nil, nil) + + rl, dur := resolveSourceRateLimit(10, crl, "crtsh") + assert.Equal(t, uint(10), rl) + assert.Equal(t, time.Second, dur) +} + +func TestResolveSourceRateLimit_NoLimit(t *testing.T) { + // No global and no per-source → unlimited (0) + crl := newCustomRateLimit(nil, nil) + + rl, _ := resolveSourceRateLimit(0, crl, "crtsh") + assert.Equal(t, uint(0), rl) +} + +func TestResolveSourceRateLimit_NilRateLimit(t *testing.T) { + // nil CustomRateLimit should not panic + rl, dur := resolveSourceRateLimit(5, &subscraping.CustomRateLimit{}, "crtsh") + assert.Equal(t, uint(5), rl) + assert.Equal(t, time.Second, dur) +} + +func TestResolveSourceRateLimit_PerSourceDefaultDuration(t *testing.T) { + // Per-source limit without explicit duration defaults to per-second + crl := newCustomRateLimit( + map[string]uint{"hackertarget": 3}, + nil, + ) + + rl, dur := resolveSourceRateLimit(0, crl, "hackertarget") + assert.Equal(t, uint(3), rl) + assert.Equal(t, time.Second, dur) +} + +func TestBuildMultiRateLimiter_GlobalAppliedToAll(t *testing.T) { + // With -rl 5 and no per-source overrides, all sources should be rate-limited + agent := New([]string{"hackertarget", "crtsh"}, []string{}, false, false) + crl := newCustomRateLimit(nil, nil) + + limiter, err := agent.buildMultiRateLimiter(context.Background(), 5, crl) + require.NoError(t, err) + require.NotNil(t, limiter) +} + +func TestBuildMultiRateLimiter_NilRateLimit(t *testing.T) { + // nil rateLimit should not panic (guard in buildMultiRateLimiter) + agent := New([]string{"hackertarget"}, []string{}, false, false) + + limiter, err := agent.buildMultiRateLimiter(context.Background(), 0, nil) + require.NoError(t, err) + require.NotNil(t, limiter) +} + +func TestBuildMultiRateLimiter_PerSourceWithDuration(t *testing.T) { + // Per-source limit with custom duration (e.g. -rls hackertarget=2/m) + agent := New([]string{"hackertarget", "crtsh"}, []string{}, false, false) + crl := newCustomRateLimit( + map[string]uint{"hackertarget": 2}, + map[string]time.Duration{"hackertarget": time.Minute}, + ) + + limiter, err := agent.buildMultiRateLimiter(context.Background(), 0, crl) + require.NoError(t, err) + require.NotNil(t, limiter) +} + +func TestBuildMultiRateLimiter_UnlimitedWhenNoLimits(t *testing.T) { + // Without any rate limits, sources should get MaxUint32 (effectively unlimited) + agent := New([]string{"hackertarget"}, []string{}, false, false) + crl := newCustomRateLimit(nil, nil) + + limiter, err := agent.buildMultiRateLimiter(context.Background(), 0, crl) + require.NoError(t, err) + require.NotNil(t, limiter) + // No assertion on internal state - just verify it doesn't error + _ = math.MaxUint32 // referenced for clarity +} diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go index f00ce641..b9dd844d 100644 --- a/pkg/runner/runner.go +++ b/pkg/runner/runner.go @@ -10,6 +10,7 @@ import ( "regexp" "strconv" "strings" + "time" "github.com/projectdiscovery/gologger" contextutil "github.com/projectdiscovery/utils/context" @@ -61,11 +62,17 @@ func NewRunner(options *Options) (*Runner, error) { Custom: mapsutil.SyncLockMap[string, uint]{ Map: make(map[string]uint), }, + CustomDuration: mapsutil.SyncLockMap[string, time.Duration]{ + Map: make(map[string]time.Duration), + }, } for source, sourceRateLimit := range options.RateLimits.AsMap() { if sourceRateLimit.MaxCount > 0 && sourceRateLimit.MaxCount <= math.MaxUint { _ = runner.rateLimit.Custom.Set(source, sourceRateLimit.MaxCount) + if sourceRateLimit.Duration > 0 { + _ = runner.rateLimit.CustomDuration.Set(source, sourceRateLimit.Duration) + } } } diff --git a/pkg/subscraping/types.go b/pkg/subscraping/types.go index 30d0f407..38118717 100644 --- a/pkg/subscraping/types.go +++ b/pkg/subscraping/types.go @@ -16,7 +16,8 @@ const ( ) type CustomRateLimit struct { - Custom mapsutil.SyncLockMap[string, uint] + Custom mapsutil.SyncLockMap[string, uint] + CustomDuration mapsutil.SyncLockMap[string, time.Duration] } // BasicAuth request's Authorization header