From d075a4e2172a66dd384cfda1cc02ef250cafc178 Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Tue, 19 Aug 2025 20:39:16 +0200 Subject: [PATCH] per-target rate limit --- go.mod | 2 +- go.sum | 6 ++++++ internal/runner/runner.go | 12 +++++++++--- internal/server/nuclei_sdk.go | 2 +- pkg/protocols/dns/request.go | 2 +- pkg/protocols/http/request.go | 4 ++-- pkg/protocols/http/request_fuzz.go | 2 +- pkg/protocols/protocols.go | 11 +++++++---- pkg/testutils/testutils.go | 7 ++++++- 9 files changed, 34 insertions(+), 14 deletions(-) diff --git a/go.mod b/go.mod index 2fd4ee3b5..79417a871 100644 --- a/go.mod +++ b/go.mod @@ -99,7 +99,7 @@ require ( github.com/projectdiscovery/mapcidr v1.1.34 github.com/projectdiscovery/n3iwf v0.0.0-20230523120440-b8cd232ff1f5 github.com/projectdiscovery/networkpolicy v0.1.18 - github.com/projectdiscovery/ratelimit v0.0.81 + github.com/projectdiscovery/ratelimit v0.0.82-0.20250819165347-d56e782c656d github.com/projectdiscovery/rdap v0.9.0 github.com/projectdiscovery/sarif v0.0.1 github.com/projectdiscovery/tlsx v1.1.9 diff --git a/go.sum b/go.sum index 5fa540e33..3c04783ce 100644 --- a/go.sum +++ b/go.sum @@ -805,6 +805,12 @@ github.com/projectdiscovery/networkpolicy v0.1.18 h1:DAeP73SvcuT4evaohNS7BPELw+V github.com/projectdiscovery/networkpolicy v0.1.18/go.mod h1:2yWanKsU2oBZ75ch94IsEQy6hByFp+3oTiSyC6ew3TE= github.com/projectdiscovery/ratelimit v0.0.81 h1:u6lW+rAhS/UO0amHTYmYLipPK8NEotA9521hdojBtgI= github.com/projectdiscovery/ratelimit v0.0.81/go.mod h1:tK04WXHuC4i6AsFkByInODSNf45gd9sfaMHzmy2bAsA= +github.com/projectdiscovery/ratelimit v0.0.82-0.20250819014026-b6ac32486d12 h1:oS/0Oe6tv7SG2zoFBofwvoB3mfDb5Yk76np4zU0Mcg8= +github.com/projectdiscovery/ratelimit v0.0.82-0.20250819014026-b6ac32486d12/go.mod h1:z076BrLkBb5yS7uhHNoCTf8X/BvFSGRxwQ8EzEL9afM= +github.com/projectdiscovery/ratelimit v0.0.82-0.20250819164203-5e0837ee74a5 h1:tulP7aVzZsWlXbWo6W78nVFqJ05c6zSS8dsos/cBaSY= +github.com/projectdiscovery/ratelimit v0.0.82-0.20250819164203-5e0837ee74a5/go.mod h1:z076BrLkBb5yS7uhHNoCTf8X/BvFSGRxwQ8EzEL9afM= +github.com/projectdiscovery/ratelimit v0.0.82-0.20250819165347-d56e782c656d h1:zUptoE5gh4gyJsbcHfNUEKgdoieyAHW2hCfWXZ/Zye4= +github.com/projectdiscovery/ratelimit v0.0.82-0.20250819165347-d56e782c656d/go.mod h1:z076BrLkBb5yS7uhHNoCTf8X/BvFSGRxwQ8EzEL9afM= github.com/projectdiscovery/rawhttp v0.1.90 h1:LOSZ6PUH08tnKmWsIwvwv1Z/4zkiYKYOSZ6n+8RFKtw= github.com/projectdiscovery/rawhttp v0.1.90/go.mod h1:VZYAM25UI/wVB3URZ95ZaftgOnsbphxyAw/XnQRRz4Y= github.com/projectdiscovery/rdap v0.9.0 h1:wPhHx5pQ2QI+WGhyNb2PjhTl0NtB39Nk7YFZ9cp8ZGA= diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 2a3c2c470..cf1bcecfa 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -87,7 +87,7 @@ type Runner struct { colorizer aurora.Aurora issuesClient reporting.Client browser *engine.Browser - rateLimiter *ratelimit.Limiter + rateLimiter *ratelimit.AutoLimiter hostErrors hosterrorscache.CacheInterface resumeCfg *types.ResumeCfg pprofServer *pprofutil.PprofServer @@ -384,11 +384,17 @@ func New(options *types.Options) (*Runner, error) { if options.RateLimit > 0 && options.RateLimitDuration == 0 { options.RateLimitDuration = time.Second } + var ratelimiter *ratelimit.AutoLimiter if options.RateLimit == 0 && options.RateLimitDuration == 0 { - runner.rateLimiter = ratelimit.NewUnlimited(context.Background()) + ratelimiter = ratelimit.NewAutoLimiter(context.Background(), ratelimit.WithUnlimited()) } else { - runner.rateLimiter = ratelimit.New(context.Background(), uint(options.RateLimit), options.RateLimitDuration) + ratelimiter = ratelimit.NewAutoLimiter( + context.Background(), + ratelimit.WithMaxCount(uint(options.RateLimit)), + ratelimit.WithDuration(options.RateLimitDuration), + ) } + runner.rateLimiter = ratelimiter if tmpDir, err := os.MkdirTemp("", "nuclei-tmp-*"); err == nil { runner.tmpDir = tmpDir diff --git a/internal/server/nuclei_sdk.go b/internal/server/nuclei_sdk.go index 022d9ab9b..29173bf4f 100644 --- a/internal/server/nuclei_sdk.go +++ b/internal/server/nuclei_sdk.go @@ -50,7 +50,7 @@ type NucleiExecutorOptions struct { Progress progress.Progress Catalog catalog.Catalog IssuesClient reporting.Client - RateLimiter *ratelimit.Limiter + RateLimiter *ratelimit.AutoLimiter Interactsh *interactsh.Client ProjectFile *projectfile.ProjectFile Browser *browserEngine.Browser diff --git a/pkg/protocols/dns/request.go b/pkg/protocols/dns/request.go index 3cc0cd715..df7360c34 100644 --- a/pkg/protocols/dns/request.go +++ b/pkg/protocols/dns/request.go @@ -151,7 +151,7 @@ func (request *Request) execute(input *contextargs.Context, domain string, metad } } - request.options.RateLimitTake() + request.options.RateLimitTake(input) // Send the request to the target servers response, err := dnsClient.Do(compiledRequest) diff --git a/pkg/protocols/http/request.go b/pkg/protocols/http/request.go index 381f4d385..29a1ccdde 100644 --- a/pkg/protocols/http/request.go +++ b/pkg/protocols/http/request.go @@ -291,7 +291,7 @@ func (request *Request) executeParallelHTTP(input *contextargs.Context, dynamicV return } // putting ratelimiter here prevents any unnecessary waiting if any - request.options.RateLimitTake() + request.options.RateLimitTake(input) // after ratelimit take, check if we need to stop if spmHandler.FoundFirstMatch() || request.isUnresponsiveAddress(updatedInput) || spmHandler.Cancelled() { @@ -470,7 +470,7 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa executeFunc := func(data string, payloads, dynamicValue map[string]interface{}) (bool, error) { hasInteractMatchers := interactsh.HasMatchers(request.CompiledOperators) - request.options.RateLimitTake() + request.options.RateLimitTake(input) ctx := request.newContext(input) ctxWithTimeout, cancel := context.WithTimeoutCause(ctx, request.options.Options.GetTimeouts().HttpTimeout, ErrHttpEngineRequestDeadline) diff --git a/pkg/protocols/http/request_fuzz.go b/pkg/protocols/http/request_fuzz.go index 3a7e2cc74..2c71b83bc 100644 --- a/pkg/protocols/http/request_fuzz.go +++ b/pkg/protocols/http/request_fuzz.go @@ -181,7 +181,7 @@ func (request *Request) executeGeneratedFuzzingRequest(gr fuzz.GeneratedRequest, if request.options.HostErrorsCache != nil && request.options.HostErrorsCache.Check(request.options.ProtocolType.String(), input) { return false } - request.options.RateLimitTake() + request.options.RateLimitTake(input) req := &generatedRequest{ request: gr.Request, dynamicValues: gr.DynamicValues, diff --git a/pkg/protocols/protocols.go b/pkg/protocols/protocols.go index 702aa5729..1a777ee02 100644 --- a/pkg/protocols/protocols.go +++ b/pkg/protocols/protocols.go @@ -78,7 +78,7 @@ type ExecutorOptions struct { // Progress is a progress client for scan reporting Progress progress.Progress // RateLimiter is a rate-limiter for limiting sent number of requests. - RateLimiter *ratelimit.Limiter + RateLimiter *ratelimit.AutoLimiter // Catalog is a template catalog implementation for nuclei Catalog catalog.Catalog // ProjectFile is the project file for nuclei @@ -143,7 +143,7 @@ type ExecutorOptions struct { // todo: centralizing components is not feasible with current clogged architecture // a possible approach could be an internal event bus with pub-subs? This would be less invasive than // reworking dep injection from scratch -func (e *ExecutorOptions) RateLimitTake() { +func (e *ExecutorOptions) RateLimitTake(ctx *contextargs.Context) { // The code below can race and there isn't a great way to fix this without adding an idempotent // function to the rate limiter implementation. For now, stick with whatever rate is already set. /* @@ -152,8 +152,11 @@ func (e *ExecutorOptions) RateLimitTake() { e.RateLimiter.SetDuration(e.Options.RateLimitDuration) } */ - if e.RateLimiter != nil { - e.RateLimiter.Take() + if hostErrorCache, ok := e.HostErrorsCache.(*hosterrorscache.Cache); ok { + key := hostErrorCache.GetKeyFromContext(ctx, nil) + if e.RateLimiter != nil { + e.RateLimiter.Take(key) + } } } diff --git a/pkg/testutils/testutils.go b/pkg/testutils/testutils.go index 5f791c2c1..1df5e3c11 100644 --- a/pkg/testutils/testutils.go +++ b/pkg/testutils/testutils.go @@ -88,6 +88,11 @@ type TemplateInfo struct { // NewMockExecuterOptions creates a new mock executeroptions struct func NewMockExecuterOptions(options *types.Options, info *TemplateInfo) *protocols.ExecutorOptions { progressImpl, _ := progress.NewStatsTicker(0, false, false, false, 0) + rateLimiter := ratelimit.NewAutoLimiter( + context.Background(), + ratelimit.WithMaxCount(uint(options.RateLimit)), + ratelimit.WithDuration(time.Second), + ) executerOpts := &protocols.ExecutorOptions{ TemplateID: info.ID, TemplateInfo: info.Info, @@ -99,7 +104,7 @@ func NewMockExecuterOptions(options *types.Options, info *TemplateInfo) *protoco IssuesClient: nil, Browser: nil, Catalog: disk.NewCatalog(config.DefaultConfig.TemplatesDirectory), - RateLimiter: ratelimit.New(context.Background(), uint(options.RateLimit), time.Second), + RateLimiter: rateLimiter, } executerOpts.CreateTemplateCtxStore() return executerOpts