From a1cc52c3ff958310e4a48cdd21d419650fe215e6 Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Mon, 12 Oct 2020 20:15:21 +0200 Subject: [PATCH] adding bulk-size --- v2/internal/runner/options.go | 2 ++ v2/internal/runner/processor.go | 10 +++++----- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/v2/internal/runner/options.go b/v2/internal/runner/options.go index 0ac8145b..423161d0 100644 --- a/v2/internal/runner/options.go +++ b/v2/internal/runner/options.go @@ -40,6 +40,7 @@ type Options struct { TemplatesDirectory string // TemplatesDirectory is the directory to use for storing templates RateLimit int // Rate-Limit of requests per specified target StopAtFirstMatch bool // Stop processing template at first full match (this may break chained requests) + BulkSize int // Number of targets analyzed in parallel for each template } type multiStringFlag []string @@ -81,6 +82,7 @@ func ParseOptions() *Options { flag.BoolVar(&options.TemplateList, "tl", false, "List available templates") flag.IntVar(&options.RateLimit, "rate-limit", -1, "Per Target Rate-Limit") flag.BoolVar(&options.StopAtFirstMatch, "stop-at-first-match", false, "Stop processing http requests at first match (this may break template/workflow logic)") + flag.IntVar(&options.BulkSize, "bulk-size", 150, "Number of hosts analyzed in parallel per template") flag.Parse() diff --git a/v2/internal/runner/processor.go b/v2/internal/runner/processor.go index ebc4ab17..8b3d0656 100644 --- a/v2/internal/runner/processor.go +++ b/v2/internal/runner/processor.go @@ -9,7 +9,6 @@ import ( "path" "path/filepath" "strings" - "sync" tengo "github.com/d5/tengo/v2" "github.com/d5/tengo/v2/stdlib" @@ -21,6 +20,7 @@ import ( "github.com/projectdiscovery/nuclei/v2/pkg/requests" "github.com/projectdiscovery/nuclei/v2/pkg/templates" "github.com/projectdiscovery/nuclei/v2/pkg/workflows" + "github.com/remeh/sizedwaitgroup" ) // workflowTemplates contains the initialized workflow templates per template group @@ -79,12 +79,12 @@ func (r *Runner) processTemplateWithList(p progress.IProgress, template *templat var globalresult atomicboolean.AtomBool - var wg sync.WaitGroup + wg := sizedwaitgroup.New(r.options.BulkSize) scanner := bufio.NewScanner(strings.NewReader(r.input)) for scanner.Scan() { URL := scanner.Text() - wg.Add(1) + wg.Add() go func(URL string) { defer wg.Done() @@ -125,12 +125,12 @@ func (r *Runner) processWorkflowWithList(p progress.IProgress, workflow *workflo logicBytes := []byte(workflow.Logic) - var wg sync.WaitGroup + wg := sizedwaitgroup.New(r.options.BulkSize) scanner := bufio.NewScanner(strings.NewReader(r.input)) for scanner.Scan() { targetURL := scanner.Text() - wg.Add(1) + wg.Add() go func(targetURL string) { defer wg.Done()