diff --git a/v2/internal/runner/processor.go b/v2/internal/runner/processor.go index 1edd9f8f3..50db69784 100644 --- a/v2/internal/runner/processor.go +++ b/v2/internal/runner/processor.go @@ -44,10 +44,7 @@ func (r *Runner) processWorkflowWithList(template *templates.Template) bool { wg.Add() go func(URL string) { defer wg.Done() - match, err := template.CompiledWorkflow.RunWorkflow(URL) - if err != nil { - gologger.Warning().Msgf("[%s] Could not execute step: %s\n", r.colorizer.BrightBlue(template.ID), err) - } + match := template.CompiledWorkflow.RunWorkflow(URL) results.CAS(false, match) }(URL) return nil diff --git a/v2/pkg/workflows/compile.go b/v2/pkg/workflows/compile.go index 3e973d8b7..d7392dbf6 100644 --- a/v2/pkg/workflows/compile.go +++ b/v2/pkg/workflows/compile.go @@ -4,12 +4,13 @@ import ( "os" "github.com/pkg/errors" + "github.com/projectdiscovery/nuclei/v2/pkg/protocols" "gopkg.in/yaml.v2" ) // Parse a yaml workflow file -func Parse(file string) (*Workflow, error) { - workflow := &Workflow{} +func Parse(file string, options *protocols.ExecuterOptions) (*Workflow, error) { + workflow := &Workflow{options: options} f, err := os.Open(file) if err != nil { diff --git a/v2/pkg/workflows/execute.go b/v2/pkg/workflows/execute.go index 223beb633..c22ccb029 100644 --- a/v2/pkg/workflows/execute.go +++ b/v2/pkg/workflows/execute.go @@ -3,25 +3,32 @@ package workflows import ( "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/nuclei/v2/pkg/output" + "github.com/remeh/sizedwaitgroup" "go.uber.org/atomic" ) // RunWorkflow runs a workflow on an input and returns true or false -func (w *Workflow) RunWorkflow(input string) (bool, error) { +func (w *Workflow) RunWorkflow(input string) bool { results := &atomic.Bool{} + swg := sizedwaitgroup.New(w.options.Options.TemplateThreads) for _, template := range w.Workflows { - err := w.runWorkflowStep(template, input, results) - if err != nil { - return results.Load(), err - } + swg.Add() + func(template *WorkflowTemplate) { + err := w.runWorkflowStep(template, input, results, &swg) + if err != nil { + gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, err) + } + swg.Done() + }(template) } - return results.Load(), nil + swg.Wait() + return results.Load() } // runWorkflowStep runs a workflow step for the workflow. It executes the workflow // in a recursive manner running all subtemplates and matchers. -func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, results *atomic.Bool) error { +func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, results *atomic.Bool, swg *sizedwaitgroup.SizedWaitGroup) error { var firstMatched bool var mainErr error @@ -46,9 +53,6 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res } if len(template.Matchers) > 0 { - var executionErr error - var mainErr error - for _, executer := range template.Executers { executer.Options.Progress.AddToTotal(int64(executer.Executer.Requests())) @@ -65,10 +69,14 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res } for _, subtemplate := range matcher.Subtemplates { - if err := w.runWorkflowStep(subtemplate, input, results); err != nil { - executionErr = err - break - } + swg.Add() + + go func(subtemplate *WorkflowTemplate) { + if err := w.runWorkflowStep(subtemplate, input, results, swg); err != nil { + gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", subtemplate.Template, err) + } + swg.Done() + }(subtemplate) } } }) @@ -80,21 +88,20 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res } continue } - if executionErr != nil { - if len(template.Executers) == 1 { - mainErr = executionErr - } else { - gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, executionErr) - } - } } return mainErr } if len(template.Subtemplates) > 0 && firstMatched { for _, subtemplate := range template.Subtemplates { - if err := w.runWorkflowStep(subtemplate, input, results); err != nil { - return err - } + swg.Add() + + func(template *WorkflowTemplate) { + err := w.runWorkflowStep(template, input, results, swg) + if err != nil { + gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, err) + } + swg.Done() + }(subtemplate) } } return mainErr diff --git a/v2/pkg/workflows/execute_test.go b/v2/pkg/workflows/execute_test.go index 394b6e9b1..6a94cefc8 100644 --- a/v2/pkg/workflows/execute_test.go +++ b/v2/pkg/workflows/execute_test.go @@ -7,20 +7,20 @@ import ( "github.com/projectdiscovery/nuclei/v2/pkg/operators" "github.com/projectdiscovery/nuclei/v2/pkg/output" "github.com/projectdiscovery/nuclei/v2/pkg/protocols" + "github.com/projectdiscovery/nuclei/v2/pkg/types" "github.com/stretchr/testify/require" ) func TestWorkflowsSimple(t *testing.T) { progress, _ := progress.NewProgress(false, false, 0) - workflow := &Workflow{Workflows: []*WorkflowTemplate{ + workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ {Executers: []*ProtocolExecuterPair{{ Executer: &mockExecuter{result: true}, Options: &protocols.ExecuterOptions{Progress: progress}}, }}, }} - matched, err := workflow.RunWorkflow("https://test.com") - require.Nil(t, err, "could not run workflow") + matched := workflow.RunWorkflow("https://test.com") require.True(t, matched, "could not get correct match value") } @@ -28,7 +28,7 @@ func TestWorkflowsSimpleMultiple(t *testing.T) { progress, _ := progress.NewProgress(false, false, 0) var firstInput, secondInput string - workflow := &Workflow{Workflows: []*WorkflowTemplate{ + workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ {Executers: []*ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { firstInput = input @@ -41,8 +41,7 @@ func TestWorkflowsSimpleMultiple(t *testing.T) { }}, }} - matched, err := workflow.RunWorkflow("https://test.com") - require.Nil(t, err, "could not run workflow") + matched := workflow.RunWorkflow("https://test.com") require.True(t, matched, "could not get correct match value") require.Equal(t, "https://test.com", firstInput, "could not get correct first input") @@ -53,7 +52,7 @@ func TestWorkflowsSubtemplates(t *testing.T) { progress, _ := progress.NewProgress(false, false, 0) var firstInput, secondInput string - workflow := &Workflow{Workflows: []*WorkflowTemplate{ + workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ {Executers: []*ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { firstInput = input @@ -65,8 +64,7 @@ func TestWorkflowsSubtemplates(t *testing.T) { }}}}, }} - matched, err := workflow.RunWorkflow("https://test.com") - require.Nil(t, err, "could not run workflow") + matched := workflow.RunWorkflow("https://test.com") require.True(t, matched, "could not get correct match value") require.Equal(t, "https://test.com", firstInput, "could not get correct first input") @@ -77,7 +75,7 @@ func TestWorkflowsSubtemplatesNoMatch(t *testing.T) { progress, _ := progress.NewProgress(false, false, 0) var firstInput, secondInput string - workflow := &Workflow{Workflows: []*WorkflowTemplate{ + workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ {Executers: []*ProtocolExecuterPair{{ Executer: &mockExecuter{result: false, executeHook: func(input string) { firstInput = input @@ -89,8 +87,7 @@ func TestWorkflowsSubtemplatesNoMatch(t *testing.T) { }}}}, }} - matched, err := workflow.RunWorkflow("https://test.com") - require.Nil(t, err, "could not run workflow") + matched := workflow.RunWorkflow("https://test.com") require.False(t, matched, "could not get correct match value") require.Equal(t, "https://test.com", firstInput, "could not get correct first input") @@ -101,7 +98,7 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) { progress, _ := progress.NewProgress(false, false, 0) var firstInput, secondInput string - workflow := &Workflow{Workflows: []*WorkflowTemplate{ + workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ {Executers: []*ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { firstInput = input @@ -118,8 +115,7 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) { }}}}}}, }} - matched, err := workflow.RunWorkflow("https://test.com") - require.Nil(t, err, "could not run workflow") + matched := workflow.RunWorkflow("https://test.com") require.True(t, matched, "could not get correct match value") require.Equal(t, "https://test.com", firstInput, "could not get correct first input") @@ -130,7 +126,7 @@ func TestWorkflowsSubtemplatesWithMatcherNoMatch(t *testing.T) { progress, _ := progress.NewProgress(false, false, 0) var firstInput, secondInput string - workflow := &Workflow{Workflows: []*WorkflowTemplate{ + workflow := &Workflow{options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ {Executers: []*ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { firstInput = input @@ -147,8 +143,7 @@ func TestWorkflowsSubtemplatesWithMatcherNoMatch(t *testing.T) { }}}}}}, }} - matched, err := workflow.RunWorkflow("https://test.com") - require.Nil(t, err, "could not run workflow") + matched := workflow.RunWorkflow("https://test.com") require.False(t, matched, "could not get correct match value") require.Equal(t, "https://test.com", firstInput, "could not get correct first input") diff --git a/v2/pkg/workflows/workflows.go b/v2/pkg/workflows/workflows.go index 860b1fdcb..0ac29f677 100644 --- a/v2/pkg/workflows/workflows.go +++ b/v2/pkg/workflows/workflows.go @@ -6,6 +6,8 @@ import "github.com/projectdiscovery/nuclei/v2/pkg/protocols" type Workflow struct { // Workflows is a yaml based workflow declaration code. Workflows []*WorkflowTemplate `yaml:"workflows,omitempty"` + + options *protocols.ExecuterOptions } // WorkflowTemplate is a template to be ran as part of a workflow