fix(flow): bound runtime execution

pkg/protocols/protocols.go

@@ -143,6 +143,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)
@@ -392,6 +395,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,
@@ -622,6 +626,7 @@ func (e *ExecutorOptions) ApplyNewEngineOptions(n *ExecutorOptions) {
 	e.Logger = n.Logger
 	e.CustomFastdialer = n.CustomFastdialer
 	e.ClusterMappings = n.ClusterMappings.Copy()
+	e.MaxFlowProtocolRequests = n.MaxFlowProtocolRequests
 }
 
 // GetOptions returns scan options, or nil for an unbound executor.

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 {
@@ -55,7 +67,8 @@ type FlowExecutor struct {
 	flattenKeys []string
 
 	// executions counts how many times the flow ran each request
-	executions map[protocols.Request]*atomic.Int64
+	executions           map[protocols.Request]*atomic.Int64
+	protocolRequestCount atomic.Int64
 }
 
 // NewFlowExecutor creates a new flow executor from a list of requests
@@ -232,12 +245,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")
@@ -245,16 +254,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()
 		}
 	}()
 
@@ -288,11 +295,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 {
@@ -303,6 +318,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
+		}
+	}
+}
+
 // reconcileProgress settles progress, whose total assumed each request runs
 // once: requests the flow never ran count as processed, as protocols do for
 // skipped requests, and extra runs of a request are added to the total

pkg/tmplexec/flow/flow_executor_test.go

@@ -13,6 +13,7 @@ import (
 	"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"
@@ -21,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"
 )
@@ -51,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()) {
@@ -72,6 +75,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()) {
@@ -93,6 +99,32 @@ func (f *FlowExecutor) requestExecutor(runtime *goja.Runtime, reqMap mapsutil.Ma
 	return matcherStatus.Load()
 }
 
+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) {