fix: address PR review comments

go.mod

@@ -35,6 +35,7 @@ require (
 	golang.org/x/net v0.58.0
 	golang.org/x/oauth2 v0.36.0
 	golang.org/x/text v0.41.0
+	golang.org/x/time v0.15.0
 )
 
 require (
@@ -465,7 +466,6 @@ require (
 	golang.org/x/exp v0.0.0-20260527015227-08cc5374adb3
 	golang.org/x/mod v0.39.0 // indirect
 	golang.org/x/sys v0.47.0
-	golang.org/x/time v0.15.0 // indirect
 	golang.org/x/tools v0.49.0
 	google.golang.org/protobuf v1.36.11
 	gopkg.in/alecthomas/kingpin.v2 v2.2.6 // indirect

pkg/protocols/headless/engine/instance.go

@@ -21,6 +21,8 @@ type Instance struct {
 	requestLog map[string]string // contains actual request that was sent
 }
 
+const browserCleanupTimeout = 5 * time.Second
+
 // NewInstance creates a new instance for the current browser.
 //
 // The login process is repeated only once for a browser, and the created
@@ -29,7 +31,31 @@ type Instance struct {
 // Users can also choose to run the login->actions process again
 // which uses a new incognito browser instance to run actions.
 func (b *Browser) NewInstance() (*Instance, error) {
-	browser, err := b.engine.Incognito()
+	return b.newInstance(b.engine)
+}
+
+// NewInstanceWithContext bounds browser-context creation and all subsequent instance calls.
+func (b *Browser) NewInstanceWithContext(ctx context.Context, timeout time.Duration) (*Instance, error) {
+	operationCtx := ctx
+	cancel := func() {}
+	if timeout > 0 {
+		operationCtx, cancel = context.WithTimeout(ctx, timeout)
+	}
+	browser, err := b.newInstanceBrowser(b.engine.Context(operationCtx))
+	cancel()
+	if err != nil {
+		return nil, err
+	}
+	browser.engine = browser.engine.Context(ctx)
+	return browser, nil
+}
+
+func (b *Browser) newInstance(engine *rod.Browser) (*Instance, error) {
+	return b.newInstanceBrowser(engine)
+}
+
+func (b *Browser) newInstanceBrowser(engine *rod.Browser) (*Instance, error) {
+	browser, err := engine.Incognito()
 	if err != nil {
 		return nil, err
 	}
@@ -49,7 +75,9 @@ func (i *Instance) GetRequestLog() map[string]string {
 
 // Close closes all the tabs and pages for a browser instance
 func (i *Instance) Close() error {
-	return i.engine.Close()
+	ctx, cancel := context.WithTimeout(context.Background(), browserCleanupTimeout)
+	defer cancel()
+	return i.engine.Context(ctx).Close()
 }
 
 // SetInteractsh client

pkg/protocols/headless/engine/instance_test.go

@@ -0,0 +1,56 @@
+package engine
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/go-rod/rod"
+	"github.com/go-rod/rod/lib/cdp"
+	"github.com/go-rod/rod/lib/proto"
+	"github.com/stretchr/testify/require"
+)
+
+type cleanupCDPClient struct {
+	events      chan *cdp.Event
+	contextErrs chan error
+}
+
+func (c *cleanupCDPClient) Event() <-chan *cdp.Event {
+	return c.events
+}
+
+func (c *cleanupCDPClient) Call(ctx context.Context, _ string, method string, _ interface{}) ([]byte, error) {
+	if method == "Target.disposeBrowserContext" {
+		c.contextErrs <- ctx.Err()
+	}
+	return []byte("{}"), nil
+}
+
+func TestNewInstanceWithContextCancelsBlockedBrowserContextCreation(t *testing.T) {
+	client := newBlockingCDPClient("Target.createBrowserContext", 500*time.Millisecond)
+	browser := &Browser{engine: rod.New().Client(client)}
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	defer cancel()
+
+	started := time.Now()
+	_, err := browser.NewInstanceWithContext(ctx, time.Second)
+
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.Less(t, time.Since(started), 200*time.Millisecond)
+}
+
+func TestInstanceCloseUsesLiveCleanupContext(t *testing.T) {
+	client := &cleanupCDPClient{
+		events:      make(chan *cdp.Event),
+		contextErrs: make(chan error, 1),
+	}
+	canceledCtx, cancel := context.WithCancel(context.Background())
+	cancel()
+	browser := rod.New().Client(client).Context(canceledCtx)
+	browser.BrowserContextID = proto.BrowserBrowserContextID("context-id")
+	instance := &Instance{engine: browser}
+
+	require.NoError(t, instance.Close())
+	require.NoError(t, <-client.contextErrs)
+}

pkg/protocols/headless/engine/page.go

@@ -2,6 +2,7 @@ package engine
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -40,6 +41,7 @@ type Page struct {
 	payloads           map[string]interface{}
 	variables          map[string]interface{}
 	lastActionNavigate *Action
+	cancel             context.CancelFunc
 }
 
 // HistoryData contains the page request/response pairs
@@ -61,22 +63,33 @@ func (i *Instance) Run(ctx *contextargs.Context, actions []*Action, payloads map
 		return nil, nil, err
 	}
 
-	page, err := i.engine.Page(proto.TargetCreateTarget{})
+	pageContext := ctx.Context()
+	cancelPageCreation := func() {}
+	if options.Timeout > 0 {
+		pageContext, cancelPageCreation = context.WithTimeout(pageContext, options.Timeout)
+	}
+	page, err := i.engine.Context(pageContext).Page(proto.TargetCreateTarget{})
 	if err != nil {
+		cancelPageCreation()
 		return nil, nil, err
 	}
+	pageSetupComplete := false
+	defer func() {
+		if !pageSetupComplete {
+			_ = page.Close()
+			cancelPageCreation()
+		}
+	}()
 	// A Rod page inherits the browser process context by default. Bind it to the
 	// request context so scan cancellation interrupts in-flight headless actions.
-	page = page.Context(ctx.Context()).Timeout(options.Timeout)
+	page = page.Context(pageContext)
 
 	if err = i.browser.applyDefaultHeaders(page); err != nil {
-		_ = page.Close()
 		return nil, nil, err
 	}
 
 	if i.browser.customAgent != "" {
 		if userAgentErr := page.SetUserAgent(&proto.NetworkSetUserAgentOverride{UserAgent: i.browser.customAgent}); userAgentErr != nil {
-			_ = page.Close()
 			return nil, nil, userAgentErr
 		}
 	}
@@ -88,7 +101,6 @@ func (i *Instance) Run(ctx *contextargs.Context, actions []*Action, payloads map
 	target := ctx.MetaInput.Input
 	input, err := urlutil.Parse(target)
 	if err != nil {
-		_ = page.Close()
 		return nil, nil, errkit.Wrapf(err, "could not parse URL %s", target)
 	}
 
@@ -109,7 +121,9 @@ func (i *Instance) Run(ctx *contextargs.Context, actions []*Action, payloads map
 		payloads:  payloads,
 		variables: variables,
 		inputURL:  input,
+		cancel:    cancelPageCreation,
 	}
+	pageSetupComplete = true
 
 	successfulPageCreation := false
 	defer func() {
@@ -245,7 +259,12 @@ func (p *Page) Close() {
 	if p.hijackNative != nil {
 		_ = p.hijackNative.Stop()
 	}
-	_ = p.page.Close()
+	cleanupCtx, cancelCleanup := context.WithTimeout(context.Background(), browserCleanupTimeout)
+	_ = p.page.Context(cleanupCtx).Close()
+	cancelCleanup()
+	if p.cancel != nil {
+		p.cancel()
+	}
 }
 
 // Page returns the current page for the actions

pkg/protocols/headless/engine/page_test.go

@@ -0,0 +1,63 @@
+package engine
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/go-rod/rod"
+	"github.com/go-rod/rod/lib/cdp"
+	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
+	"github.com/stretchr/testify/require"
+)
+
+var errCDPCallBlocked = errors.New("cdp call remained blocked")
+
+type blockingCDPClient struct {
+	method string
+	wait   time.Duration
+	events chan *cdp.Event
+}
+
+func newBlockingCDPClient(method string, wait time.Duration) *blockingCDPClient {
+	return &blockingCDPClient{method: method, wait: wait, events: make(chan *cdp.Event)}
+}
+
+func (c *blockingCDPClient) Event() <-chan *cdp.Event {
+	return c.events
+}
+
+func (c *blockingCDPClient) Call(ctx context.Context, _ string, method string, _ interface{}) ([]byte, error) {
+	if method != c.method {
+		return []byte("{}"), nil
+	}
+
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case <-time.After(c.wait):
+		return nil, errCDPCallBlocked
+	}
+}
+
+func TestRunCancelsBlockedPageCreation(t *testing.T) {
+	client := newBlockingCDPClient("Target.createTarget", 500*time.Millisecond)
+	instance := &Instance{
+		browser: &Browser{},
+		engine:  rod.New().Client(client),
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	defer cancel()
+
+	started := time.Now()
+	_, _, err := instance.Run(
+		contextargs.NewWithInput(ctx, "http://example.com"),
+		nil,
+		nil,
+		&Options{Timeout: time.Second},
+	)
+
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.Less(t, time.Since(started), 200*time.Millisecond)
+}

pkg/protocols/headless/request.go

@@ -132,7 +132,8 @@ func extractBaseURLFromActions(steps []*engine.Action) (string, error) {
 }
 
 func (request *Request) executeRequestWithPayloads(input *contextargs.Context, payloads map[string]interface{}, previous output.InternalEvent, interactshURLs []string, callback protocols.OutputEventCallback) error {
-	instance, err := request.options.Browser.NewInstance()
+	pageTimeout := time.Duration(request.options.Options.PageTimeout) * time.Second
+	instance, err := request.options.Browser.NewInstanceWithContext(input.Context(), pageTimeout)
 	if err != nil {
 		request.options.Output.Request(request.options.TemplatePath, input.MetaInput.Input, request.Type().String(), err)
 		request.options.Progress.IncrementFailedRequestsBy(1)
@@ -150,7 +151,7 @@ func (request *Request) executeRequestWithPayloads(input *contextargs.Context, p
 		return errors.Wrap(err, errCouldNotGetHtmlElement)
 	}
 	options := &engine.Options{
-		Timeout:       time.Duration(request.options.Options.PageTimeout) * time.Second,
+		Timeout:       pageTimeout,
 		DisableCookie: request.DisableCookie,
 		Options:       request.options.Options,
 	}

pkg/protocols/http/httpclientpool/clientpool.go

@@ -25,7 +25,6 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/http/httpcache"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/utils"
 	"github.com/projectdiscovery/nuclei/v3/pkg/types"
-	"github.com/projectdiscovery/ratelimit"
 	"github.com/projectdiscovery/rawhttp"
 	"github.com/projectdiscovery/retryablehttp-go"
 	urlutil "github.com/projectdiscovery/utils/url"
@@ -662,7 +661,7 @@ func isURLEncoded(s string) bool {
 
 // GetPerHostRateLimiter gets or creates a rate limiter for a specific host
 // Returns nil if per-host rate limiting is not enabled
-func GetPerHostRateLimiter(options *types.Options, hostname string) (*ratelimit.Limiter, error) {
+func GetPerHostRateLimiter(options *types.Options, hostname string) (*PerHostRateLimiter, error) {
 	if !options.PerHostRateLimit {
 		return nil, nil
 	}

pkg/protocols/http/httpclientpool/perhost_ratelimit_pool.go

@@ -10,8 +10,7 @@ import (
 	"github.com/hashicorp/golang-lru/v2/expirable"
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/nuclei/v3/pkg/types"
-	"github.com/projectdiscovery/nuclei/v3/pkg/utils"
-	"github.com/projectdiscovery/ratelimit"
+	"golang.org/x/time/rate"
 )
 
 type PerHostRateLimitPool struct {
@@ -31,7 +30,7 @@ type PerHostRateLimitPool struct {
 const ppsWindowSize = 100
 
 type rateLimitEntry struct {
-	limiter        *ratelimit.Limiter
+	limiter        *PerHostRateLimiter
 	createdAt      time.Time
 	accessCount    atomic.Uint64
 	requestCount   atomic.Uint64
@@ -46,6 +45,63 @@ type rateLimitEntry struct {
 	requestMu         sync.Mutex
 }
 
+type PerHostRateLimiter struct {
+	limiter   *rate.Limiter
+	ctx       context.Context
+	cancel    context.CancelFunc
+	unlimited bool
+}
+
+func (l *PerHostRateLimiter) Wait(ctx context.Context) error {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	if err := ctx.Err(); err != nil {
+		return context.Cause(ctx)
+	}
+	if err := l.ctx.Err(); err != nil {
+		return context.Cause(l.ctx)
+	}
+	if l.unlimited {
+		return nil
+	}
+
+	now := time.Now()
+	reservation := l.limiter.ReserveN(now, 1)
+	if !reservation.OK() {
+		return fmt.Errorf("per-host rate limiter rejected reservation")
+	}
+	delay := reservation.DelayFrom(now)
+	if delay <= 0 {
+		return nil
+	}
+
+	timer := time.NewTimer(delay)
+	defer timer.Stop()
+	select {
+	case <-timer.C:
+		return nil
+	case <-ctx.Done():
+		reservation.Cancel()
+		return context.Cause(ctx)
+	case <-l.ctx.Done():
+		reservation.Cancel()
+		return context.Cause(l.ctx)
+	}
+}
+
+func (l *PerHostRateLimiter) Take() {
+	_ = l.Wait(context.Background())
+}
+
+func (l *PerHostRateLimiter) CanTake() bool {
+	return l.unlimited || l.limiter.Tokens() > 0
+}
+
+func (l *PerHostRateLimiter) Stop() {
+	l.cancel()
+}
+
 func NewPerHostRateLimitPool(size int, maxIdleTime, maxLifetime time.Duration, options *types.Options) *PerHostRateLimitPool {
 	if size < 0 {
 		size = 0
@@ -88,7 +144,7 @@ func NewPerHostRateLimitPool(size int, maxIdleTime, maxLifetime time.Duration, o
 
 func (p *PerHostRateLimitPool) GetOrCreate(
 	host string,
-) (*ratelimit.Limiter, error) {
+) (*PerHostRateLimiter, error) {
 	normalizedHost := normalizeHostPort(host)
 
 	// Try to get entry (this refreshes TTL in expirable LRU)
@@ -175,20 +231,28 @@ func (p *PerHostRateLimitPool) GetOrCreate(
 // of starting one fixed-window timer goroutine per host. Apart from producing
 // steadier target pressure, this keeps a large host pool cheap and prevents
 // several concurrent engines from consuming the whole window at once.
-func newPerHostRateLimiter(options *types.Options) (*ratelimit.Limiter, error) {
+func newPerHostRateLimiter(options *types.Options) (*PerHostRateLimiter, error) {
+	ctx, cancel := context.WithCancel(context.Background())
 	if options == nil {
-		return utils.GetRateLimiter(context.Background(), 0, 0), nil
+		return &PerHostRateLimiter{ctx: ctx, cancel: cancel, unlimited: true}, nil
 	}
 	if options.RateLimit < 0 {
+		cancel()
 		return nil, fmt.Errorf("rate limit must not be negative: %d", options.RateLimit)
 	}
 	if options.RateLimitDuration < 0 {
+		cancel()
 		return nil, fmt.Errorf("rate limit duration must not be negative: %s", options.RateLimitDuration)
 	}
 	if options.RateLimit == 0 || options.RateLimitDuration == 0 {
-		return utils.GetRateLimiter(context.Background(), 0, 0), nil
+		return &PerHostRateLimiter{ctx: ctx, cancel: cancel, unlimited: true}, nil
 	}
-	return ratelimit.NewLeakyBucket(context.Background(), uint(options.RateLimit), options.RateLimitDuration), nil
+	requestsPerSecond := rate.Limit(float64(options.RateLimit) / options.RateLimitDuration.Seconds())
+	return &PerHostRateLimiter{
+		limiter: rate.NewLimiter(requestsPerSecond, options.RateLimit),
+		ctx:     ctx,
+		cancel:  cancel,
+	}, nil
 }
 
 func (p *PerHostRateLimitPool) EvictHost(host string) bool {
@@ -243,7 +307,7 @@ type RateLimitPoolStats struct {
 	Size      int
 }
 
-func (p *PerHostRateLimitPool) GetLimiterForHost(host string) (*ratelimit.Limiter, bool) {
+func (p *PerHostRateLimitPool) GetLimiterForHost(host string) (*PerHostRateLimiter, bool) {
 	normalizedHost := normalizeHostPort(host)
 
 	if entry, ok := p.cache.Peek(normalizedHost); ok {

pkg/protocols/http/httpclientpool/perhost_ratelimit_pool_close_test.go

@@ -1,6 +1,7 @@
 package httpclientpool
 
 import (
+	"context"
 	"fmt"
 	"runtime"
 	"testing"
@@ -48,6 +49,27 @@ func TestPerHostRateLimitPoolSmoothsLowRates(t *testing.T) {
 		"the first request beyond the initial bucket should be continuously refilled, not wait for a fixed window")
 }
 
+func TestPerHostRateLimitCancellationReleasesReservation(t *testing.T) {
+	opts := &types.Options{RateLimit: 1, RateLimitDuration: time.Second}
+	pool := NewPerHostRateLimitPool(1, time.Hour, time.Hour, opts)
+	t.Cleanup(pool.Close)
+
+	limiter, err := pool.GetOrCreate("https://cancel.example.com")
+	require.NoError(t, err)
+	require.NoError(t, limiter.Wait(context.Background()))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	result := make(chan error, 1)
+	go func() { result <- limiter.Wait(ctx) }()
+	time.Sleep(20 * time.Millisecond)
+	cancel()
+	require.ErrorIs(t, <-result, context.Canceled)
+
+	started := time.Now()
+	require.NoError(t, limiter.Wait(context.Background()))
+	require.Less(t, time.Since(started), 1500*time.Millisecond)
+}
+
 func TestUnboundedPerHostRateLimitPoolRetainsBudgetsBeyondDefaultCapacity(t *testing.T) {
 	opts := &types.Options{RateLimit: 10, RateLimitDuration: time.Second}
 	pool := NewPerHostRateLimitPool(0, time.Hour, time.Hour, opts)

pkg/protocols/http/request.go

@@ -73,7 +73,7 @@ func (request *Request) Type() templateTypes.ProtocolType {
 }
 
 // rateLimitTake handles rate limiting, using per-host rate limiter if enabled, otherwise global
-func (request *Request) rateLimitTake(hostname string) error {
+func (request *Request) rateLimitTake(ctx context.Context, hostname string) error {
 	if request.options.Options.PerHostRateLimit && hostname != "" {
 		// Use per-host rate limiter
 		limiter, err := httpclientpool.GetPerHostRateLimiter(request.options.Options, hostname)
@@ -81,7 +81,9 @@ func (request *Request) rateLimitTake(hostname string) error {
 			return err
 		}
 		if limiter != nil {
-			limiter.Take()
+			if err := limiter.Wait(ctx); err != nil {
+				return err
+			}
 			// Record request for pps stats
 			httpclientpool.RecordPerHostRateLimitRequest(request.options.Options, hostname)
 			return nil
@@ -281,7 +283,7 @@ func (request *Request) executeParallelHTTP(input *contextargs.Context, dynamicV
 					// Extract from request URL if available
 					hostname = t.req.request.Request.URL.String()
 				}
-				if err := request.rateLimitTake(hostname); err != nil {
+				if err := request.rateLimitTake(t.updatedInput.Context(), hostname); err != nil {
 					select {
 					case <-spmHandler.Done():
 						spmHandler.Release()
@@ -581,7 +583,7 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa
 				// Use the generated URL directly - the normalization function will extract host:port correctly
 				hostname = generatedHttpRequest.URL()
 			}
-			if err := request.rateLimitTake(hostname); err != nil {
+			if err := request.rateLimitTake(ctxWithTimeout, hostname); err != nil {
 				return true, err
 			}
 

pkg/protocols/http/request_fuzz.go

@@ -196,7 +196,7 @@ func (request *Request) executeGeneratedFuzzingRequest(gr fuzz.GeneratedRequest,
 	if gr.Request != nil && gr.Request.Request != nil && gr.Request.Request.URL != nil {
 		hostname = gr.Request.Request.URL.String()
 	}
-	if err := request.rateLimitTake(hostname); err != nil {
+	if err := request.rateLimitTake(input.Context(), hostname); err != nil {
 		return false, err
 	}
 	req := &generatedRequest{

pkg/protocols/http/request_test.go

@@ -727,3 +727,59 @@ func TestExecuteParallelHTTP_GoroutineLeaks(t *testing.T) {
 		require.Equal(t, context.Canceled, err)
 	})
 }
+
+func TestPerHostRateLimitHonorsRequestContextCancellation(t *testing.T) {
+	options := testutils.DefaultOptions.Copy()
+	options.ExecutionId = t.Name()
+	options.PerHostRateLimit = true
+	options.PerHostRateLimitPoolSize = -1
+	options.RateLimit = 1
+	options.RateLimitDuration = 2 * time.Second
+	options.RestrictLocalNetworkAccess = false
+	testutils.Init(options)
+	t.Cleanup(func() { testutils.Cleanup(options) })
+
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
+		_, _ = fmt.Fprint(writer, "ok")
+	}))
+	defer server.Close()
+
+	request := &Request{
+		ID:     "per-host-rate-limit-cancellation",
+		Method: HTTPMethodTypeHolder{MethodType: HTTPGet},
+		Path:   []string{"{{BaseURL}}/"},
+		Operators: operators.Operators{Matchers: []*matchers.Matcher{{
+			Part:  "body",
+			Type:  matchers.MatcherTypeHolder{MatcherType: matchers.WordsMatcher},
+			Words: []string{"ok"},
+		}}},
+	}
+	executerOptions := testutils.NewMockExecuterOptions(options, &testutils.TemplateInfo{
+		ID:   request.ID,
+		Info: model.Info{SeverityHolder: severity.Holder{Severity: severity.Low}, Name: "test"},
+	})
+	require.NoError(t, request.Compile(executerOptions))
+
+	execute := func(ctx context.Context) error {
+		return request.ExecuteWithResults(
+			contextargs.NewWithInput(ctx, server.URL),
+			make(output.InternalEvent),
+			make(output.InternalEvent),
+			func(*output.InternalWrappedEvent) {},
+		)
+	}
+	require.NoError(t, execute(context.Background()))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	result := make(chan error, 1)
+	go func() { result <- execute(ctx) }()
+	time.Sleep(100 * time.Millisecond)
+	cancelledAt := time.Now()
+	cancel()
+
+	err := <-result
+	cancellationDelay := time.Since(cancelledAt)
+	t.Logf("request returned %s after its context was cancelled", cancellationDelay.Round(time.Millisecond))
+	require.ErrorIs(t, err, context.Canceled)
+	require.Less(t, cancellationDelay, 250*time.Millisecond)
+}