adding bulk-size

projectdiscovery · Oct 12, 2020 · a1cc52c · a1cc52c
1 parent 34656d1
commit a1cc52c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git 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
@@ -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()