diff --git a/v2/internal/runner/options.go b/v2/internal/runner/options.go index 85b6d35c0..8dfa7f190 100644 --- a/v2/internal/runner/options.go +++ b/v2/internal/runner/options.go @@ -14,7 +14,8 @@ import ( // the template requesting process. // nolint // false positive, options are allocated once and are necessary as is type Options struct { - Sandbox bool // Sandbox mode allows users to run isolated workflows with system commands disabled + MaxWorkflowDuration int // MaxWorkflowDuration is the maximum time a workflow can run for a URL + Sandbox bool // Sandbox mode allows users to run isolated workflows with system commands disabled Debug bool // Debug mode allows debugging request/responses for the engine Silent bool // Silent suppresses any extra text and only writes found URLs on screen. Version bool // Version specifies if we should just show version and exit @@ -68,6 +69,7 @@ func ParseOptions() *Options { options := &Options{} flag.BoolVar(&options.Sandbox, "sandbox", false, "Run workflows in isolated sandbox mode") + flag.IntVar(&options.MaxWorkflowDuration, "workflow-duration", 10, "Max time for workflow run on single URL in minutes") flag.StringVar(&options.Target, "target", "", "Target is a single target to scan using template") flag.Var(&options.Templates, "t", "Template input dir/file/files to run on host. Can be used multiple times. Supports globbing.") flag.Var(&options.ExcludedTemplates, "exclude", "Template input dir/file/files to exclude. Can be used multiple times. Supports globbing.") diff --git a/v2/internal/runner/processor.go b/v2/internal/runner/processor.go index 9bfd30ae3..2cd288613 100644 --- a/v2/internal/runner/processor.go +++ b/v2/internal/runner/processor.go @@ -8,6 +8,7 @@ import ( "path" "path/filepath" "strings" + "time" tengo "github.com/d5/tengo/v2" "github.com/d5/tengo/v2/stdlib" @@ -132,11 +133,9 @@ func (r *Runner) processWorkflowWithList(p *progress.Progress, workflow *workflo gologger.Warningf("Could not preload templates for workflow %s: %s\n", workflow.ID, err) return false } - logicBytes := []byte(workflow.Logic) wg := sizedwaitgroup.New(r.options.BulkSize) - r.hm.Scan(func(k, _ []byte) error { targetURL := string(k) wg.Add() @@ -163,7 +162,10 @@ func (r *Runner) processWorkflowWithList(p *progress.Progress, workflow *workflo variables[name] = variable } - _, err := script.RunContext(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), time.Duration(r.options.MaxWorkflowDuration)*time.Minute) + defer cancel() + + _, err := script.RunContext(ctx) if err != nil { gologger.Errorf("Could not execute workflow '%s': %s\n", workflow.ID, err) }