test(flow): fix compatibility imports

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)
+}

pkg/protocols/protocols.go

@@ -137,6 +137,9 @@ type ExecutorOptions struct {
 	ProtocolType templateTypes.ProtocolType
 	// Flow is execution flow for the template (written in javascript)
 	Flow string
+	// MaxFlowProtocolRequests bounds protocol requests issued by one Flow execution.
+	// Zero leaves the execution unbounded for callers that trust their templates.
+	MaxFlowProtocolRequests int64
 	// IsMultiProtocol is true if template has more than one protocol
 	IsMultiProtocol bool
 	// templateStore is a map which contains template context for each scan  (i.e input * template-id pair)
@@ -348,6 +351,7 @@ func (e *ExecutorOptions) Copy() *ExecutorOptions {
 		ResumeCfg:                    e.ResumeCfg,
 		ProtocolType:                 e.ProtocolType,
 		Flow:                         e.Flow,
+		MaxFlowProtocolRequests:      e.MaxFlowProtocolRequests,
 		IsMultiProtocol:              e.IsMultiProtocol,
 		JsCompiler:                   e.JsCompiler,
 		AuthProvider:                 e.AuthProvider,
@@ -532,4 +536,5 @@ func (e *ExecutorOptions) ApplyNewEngineOptions(n *ExecutorOptions) {
 	e.Logger = n.Logger
 	e.CustomFastdialer = n.CustomFastdialer
 	e.ClusterMappings = n.ClusterMappings.Copy()
+	e.MaxFlowProtocolRequests = n.MaxFlowProtocolRequests
 }

pkg/protocols/protocols_test.go

@@ -13,9 +13,10 @@ func TestExecutorOptionsCopyPreservesInteractshScope(t *testing.T) {
 	require.NoError(t, err)
 	scope := client.NewRequestScope()
 
-	copy := (&ExecutorOptions{Interactsh: client, InteractshScope: scope}).Copy()
+	copy := (&ExecutorOptions{Interactsh: client, InteractshScope: scope, MaxFlowProtocolRequests: 1000}).Copy()
 	require.Same(t, client, copy.Interactsh)
 	require.Same(t, scope, copy.InteractshScope)
+	require.Equal(t, int64(1000), copy.MaxFlowProtocolRequests)
 }
 
 func TestExecutorOptionsApplyNewEngineOptionsReplacesInteractshScope(t *testing.T) {
@@ -26,10 +27,12 @@ func TestExecutorOptionsApplyNewEngineOptionsReplacesInteractshScope(t *testing.
 
 	reused := &ExecutorOptions{Interactsh: client, InteractshScope: previous}
 	reused.ApplyNewEngineOptions(&ExecutorOptions{
-		Options:         &types.Options{},
-		Interactsh:      client,
-		InteractshScope: next,
+		Options:                 &types.Options{},
+		Interactsh:              client,
+		InteractshScope:         next,
+		MaxFlowProtocolRequests: 1000,
 	})
 
 	require.Same(t, next, reused.InteractshScope)
+	require.Equal(t, int64(1000), reused.MaxFlowProtocolRequests)
 }

pkg/tmplexec/flow/flow_executor.go

@@ -2,11 +2,13 @@ package flow
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"io"
 	"strconv"
 	"strings"
 	"sync/atomic"
+	"time"
 
 	"github.com/projectdiscovery/goja"
 	"github.com/projectdiscovery/nuclei/v3/pkg/js/compiler"
@@ -26,8 +28,18 @@ import (
 var (
 	// ErrInvalidRequestID is a request id error
 	ErrInvalidRequestID = errkit.New("invalid request id provided")
+	// ErrFlowRuntimeTerminationTimeout is returned when a flow VM does not stop after cancellation.
+	ErrFlowRuntimeTerminationTimeout = errkit.New("flow runtime failed to terminate after cancellation")
+	// ErrFlowProtocolRequestLimit is returned when Flow exceeds its configured protocol request budget.
+	ErrFlowProtocolRequestLimit = errkit.New("flow protocol request limit exceeded")
 )
 
+const flowRuntimeTerminationGrace = time.Second
+
+type flowRunResult struct {
+	err error
+}
+
 // ProtoOptions are options that can be passed to flow protocol callback
 // ex: dns(protoOptions) <- protoOptions are optional and can be anything
 type ProtoOptions struct {
@@ -54,7 +66,8 @@ type FlowExecutor struct {
 	// a request ex: if dynamic extractor returns ["value"] it will be converted to "value"
 	flattenKeys []string
 
-	executed *mapsutil.SyncLockMap[string, struct{}]
+	executed             *mapsutil.SyncLockMap[string, struct{}]
+	protocolRequestCount atomic.Int64
 }
 
 // NewFlowExecutor creates a new flow executor from a list of requests
@@ -229,12 +242,8 @@ func (f *FlowExecutor) ExecuteWithResults(ctx *scan.ScanContext) error {
 
 	// get a new runtime from pool
 	runtime := GetJSRuntime(f.options.Options)
-	defer func() {
-		// whether to reuse or not depends on the whether script modifies
-		// global scope or not,
-		PutJSRuntime(runtime, compiler.CanRunAsIIFE(f.options.Flow))
-	}()
-	defer func() {
+	cleanupRuntime := func() {
+		runtime.ClearInterrupt()
 		// remove set builtin
 		_ = runtime.GlobalObject().Delete("set")
 		_ = runtime.GlobalObject().Delete("template")
@@ -242,16 +251,14 @@ func (f *FlowExecutor) ExecuteWithResults(ctx *scan.ScanContext) error {
 			_ = runtime.GlobalObject().Delete(proto)
 		}
 		runtime.RemoveContextValue("executionId")
-	}()
-
-	// TODO(dwisiswant0): remove this once we get the RCA.
+		// whether to reuse or not depends on the whether script modifies
+		// global scope or not,
+		PutJSRuntime(runtime, compiler.CanRunAsIIFE(f.options.Flow))
+	}
+	cleanupOnReturn := true
 	defer func() {
-		if ci.IsCI() {
-			return
-		}
-
-		if r := recover(); r != nil {
-			f.ctx.LogError(fmt.Errorf("panic occurred while executing flow: %v", r))
+		if cleanupOnReturn {
+			cleanupRuntime()
 		}
 	}()
 
@@ -285,11 +292,19 @@ func (f *FlowExecutor) ExecuteWithResults(ctx *scan.ScanContext) error {
 	runtime.SetContextValue("executionId", f.options.Options.ExecutionId)
 
 	// pass flow and execute the js vm and handle errors
-	_, err := runtime.RunProgram(f.program)
+	runResult, abandoned := runFlowProgram(ctx.Context(), runtime, f.program)
+	if abandoned != nil {
+		cleanupOnReturn = false
+		go func() {
+			<-abandoned
+			cleanupRuntime()
+		}()
+		return errkit.Wrapf(ErrFlowRuntimeTerminationTimeout, "%v", ctx.Context().Err())
+	}
 	f.reconcileProgress()
-	if err != nil {
-		ctx.LogError(err)
-		return errkit.Wrapf(err, "failed to execute flow\n%v\n", f.options.Flow)
+	if runResult.err != nil {
+		ctx.LogError(runResult.err)
+		return errkit.Wrapf(runResult.err, "failed to execute flow\n%v\n", f.options.Flow)
 	}
 	runtimeErr := f.GetRuntimeErrors()
 	if runtimeErr != nil {
@@ -300,6 +315,38 @@ func (f *FlowExecutor) ExecuteWithResults(ctx *scan.ScanContext) error {
 	return nil
 }
 
+func runFlowProgram(ctx context.Context, runtime *goja.Runtime, program *goja.Program) (flowRunResult, <-chan flowRunResult) {
+	result := make(chan flowRunResult, 1)
+	go func() {
+		outcome := flowRunResult{}
+		defer func() {
+			if recovered := recover(); recovered != nil {
+				if ci.IsCI() {
+					panic(recovered)
+				}
+				outcome.err = fmt.Errorf("panic occurred while executing flow: %v", recovered)
+			}
+			result <- outcome
+		}()
+		_, outcome.err = runtime.RunProgram(program)
+	}()
+
+	select {
+	case outcome := <-result:
+		return outcome, nil
+	case <-ctx.Done():
+		runtime.Interrupt(ctx.Err())
+		timer := time.NewTimer(flowRuntimeTerminationGrace)
+		defer timer.Stop()
+		select {
+		case outcome := <-result:
+			return outcome, nil
+		case <-timer.C:
+			return flowRunResult{}, result
+		}
+	}
+}
+
 func (f *FlowExecutor) reconcileProgress() {
 	for proto, list := range f.allProtocols {
 		for idx, req := range list {

pkg/tmplexec/flow/flow_executor_test.go

@@ -6,10 +6,14 @@ import (
 	"log"
 	"net/http"
 	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/julienschmidt/httprouter"
+	"github.com/projectdiscovery/nuclei/v3/internal/tests/testutils"
 	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
 	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/disk"
 	"github.com/projectdiscovery/nuclei/v3/pkg/loader/workflow"
@@ -18,7 +22,6 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
 	"github.com/projectdiscovery/nuclei/v3/pkg/scan"
 	"github.com/projectdiscovery/nuclei/v3/pkg/templates"
-	"github.com/projectdiscovery/nuclei/v3/internal/tests/testutils"
 	"github.com/projectdiscovery/ratelimit"
 	"github.com/stretchr/testify/require"
 )
@@ -48,6 +51,79 @@ func setup() {
 	executerOpts.WorkflowLoader = workflowLoader
 }
 
+func TestFlowStopsWhenContextCanceled(t *testing.T) {
+	setup()
+
+	templatePath := filepath.Join(t.TempDir(), "infinite-flow.yaml")
+	require.NoError(t, os.WriteFile(templatePath, []byte(`id: infinite-flow
+info:
+  name: Infinite Flow
+  author: pdteam
+  severity: info
+flow: while (true) {}
+http:
+  - method: GET
+    path:
+      - "{{BaseURL}}"
+`), 0o600))
+
+	tmpl, err := templates.Parse(templatePath, nil, executerOpts)
+	require.NoError(t, err)
+	require.NoError(t, tmpl.Executer.Compile())
+
+	executionContext, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+	input := contextargs.NewWithInput(executionContext, "https://example.com")
+	scanContext := scan.NewScanContext(executionContext, input)
+	result := make(chan error, 1)
+	go func() {
+		_, executeErr := tmpl.Executer.Execute(scanContext)
+		result <- executeErr
+	}()
+
+	select {
+	case executeErr := <-result:
+		require.ErrorIs(t, executeErr, context.DeadlineExceeded)
+	case <-time.After(time.Second):
+		t.Fatal("flow execution did not stop after context cancellation")
+	}
+}
+
+func TestFlowProtocolRequestLimit(t *testing.T) {
+	setup()
+	executerOpts.MaxFlowProtocolRequests = 2
+
+	var requests atomic.Int64
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		requests.Add(1)
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer server.Close()
+
+	templatePath := filepath.Join(t.TempDir(), "bounded-flow.yaml")
+	require.NoError(t, os.WriteFile(templatePath, []byte(`id: bounded-flow
+info:
+  name: Bounded Flow
+  author: pdteam
+  severity: info
+flow: for (let i = 0; i < 3; i++) { http(); }
+http:
+  - method: GET
+    path:
+      - "{{BaseURL}}"
+`), 0o600))
+
+	tmpl, err := templates.Parse(templatePath, nil, executerOpts)
+	require.NoError(t, err)
+	require.NoError(t, tmpl.Executer.Compile())
+
+	input := contextargs.NewWithInput(context.Background(), server.URL)
+	scanContext := scan.NewScanContext(context.Background(), input)
+	_, err = tmpl.Executer.Execute(scanContext)
+	require.ErrorContains(t, err, "flow protocol request limit exceeded")
+	require.Equal(t, int64(2), requests.Load())
+}
+
 func TestFlowTemplateWithIndex(t *testing.T) {
 	// test
 	setup()

pkg/tmplexec/flow/flow_internal.go

@@ -34,6 +34,9 @@ func (f *FlowExecutor) requestExecutor(runtime *goja.Runtime, reqMap mapsutil.Ma
 		// execution logic for http()/dns() etc
 		for index := range f.allProtocols[opts.protoName] {
 			req := f.allProtocols[opts.protoName][index]
+			if !f.reserveProtocolRequests(runtime, req) {
+				return matcherStatus.Load()
+			}
 			// transform input if required
 			inputItem := f.ctx.Input.Clone()
 			if f.options.InputHelper != nil && f.ctx.Input.MetaInput.Input != "" && !utils.HasOfflineHTTPResponse(f.ctx.Input.MetaInput, req.Type()) {
@@ -71,6 +74,9 @@ func (f *FlowExecutor) requestExecutor(runtime *goja.Runtime, reqMap mapsutil.Ma
 			}
 			return matcherStatus.Load()
 		}
+		if !f.reserveProtocolRequests(runtime, req) {
+			return matcherStatus.Load()
+		}
 		// transform input if required
 		inputItem := f.ctx.Input.Clone()
 		if f.options.InputHelper != nil && f.ctx.Input.MetaInput.Input != "" && !utils.HasOfflineHTTPResponse(f.ctx.Input.MetaInput, req.Type()) {
@@ -100,6 +106,32 @@ func requestKey(proto string, req protocols.Request, id string) string {
 	return proto + ":" + id
 }
 
+func (f *FlowExecutor) reserveProtocolRequests(runtime *goja.Runtime, req protocols.Request) bool {
+	limit := f.options.MaxFlowProtocolRequests
+	if limit <= 0 {
+		return true
+	}
+
+	cost := int64(req.Requests())
+	if cost < 1 {
+		cost = 1
+	}
+	for {
+		current := f.protocolRequestCount.Load()
+		if cost > limit-current {
+			err := errkit.Wrapf(ErrFlowProtocolRequestLimit, "maximum is %d", limit)
+			if storeErr := f.allErrs.Set("flow:request-limit", err); storeErr != nil {
+				f.ctx.LogError(fmt.Errorf("failed to store flow runtime errors got %v", storeErr))
+			}
+			runtime.Interrupt(err)
+			return false
+		}
+		if f.protocolRequestCount.CompareAndSwap(current, current+cost) {
+			return true
+		}
+	}
+}
+
 // protocolResultCallback returns a callback that is executed
 // after execution of each protocol request
 func (f *FlowExecutor) protocolResultCallback(req protocols.Request, matcherStatus *atomic.Bool, _ *ProtoOptions) func(result *output.InternalWrappedEvent) {