Merge pull request #5187 from projectdiscovery/feat-4808-planner

Fixing issues with multi-thread execution
projectdiscovery · Jun 13, 2024 · 51b1d7c · 51b1d7c
2 parents cfd7b6f + 1e85e70
commit 51b1d7c
Show file tree

Hide file tree

Showing 20 changed files with 149 additions and 135 deletions.
diff --git a/cmd/integration-test/http.go b/cmd/integration-test/http.go
@@ -23,6 +23,7 @@ import (
 	logutil "github.com/projectdiscovery/utils/log"
 	sliceutil "github.com/projectdiscovery/utils/slice"
 	stringsutil "github.com/projectdiscovery/utils/strings"
+	unitutils "github.com/projectdiscovery/utils/unit"
 )
 
 var httpTestcases = []TestCaseInfo{
@@ -509,7 +510,7 @@ func (h *httpPostMultipartBody) Execute(filePath string) error {
 	var routerErr error
 
 	router.POST("/", func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
-		if err := r.ParseMultipartForm(1 * 1024); err != nil {
+		if err := r.ParseMultipartForm(unitutils.Mega); err != nil {
 			routerErr = err
 			return
 		}

diff --git a/cmd/nuclei/main.go b/cmd/nuclei/main.go
@@ -38,6 +38,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/utils/monitor"
 	errorutil "github.com/projectdiscovery/utils/errors"
 	fileutil "github.com/projectdiscovery/utils/file"
+	unitutils "github.com/projectdiscovery/utils/unit"
 	updateutils "github.com/projectdiscovery/utils/update"
 )
 
@@ -304,7 +305,7 @@ on extensive configurability, massive extensibility and ease of use.`)
 		flagSet.StringVarP(&options.AttackType, "attack-type", "at", "", "type of payload combinations to perform (batteringram,pitchfork,clusterbomb)"),
 		flagSet.StringVarP(&options.SourceIP, "source-ip", "sip", "", "source ip address to use for network scan"),
 		flagSet.IntVarP(&options.ResponseReadSize, "response-size-read", "rsr", 0, "max response size to read in bytes"),
-		flagSet.IntVarP(&options.ResponseSaveSize, "response-size-save", "rss", 1*1024*1024, "max response size to read in bytes"),
+		flagSet.IntVarP(&options.ResponseSaveSize, "response-size-save", "rss", unitutils.Mega, "max response size to read in bytes"),
 		flagSet.DurationVarP(&options.ResponseReadTimeout, "response-read-timeout", "rrt", time.Duration(5*time.Second), "response read timeout in seconds"),
 		flagSet.CallbackVar(resetCallback, "reset", "reset removes all nuclei configuration and data files (including nuclei-templates)"),
 		flagSet.BoolVarP(&options.TlsImpersonate, "tls-impersonate", "tlsi", false, "enable experimental client hello (ja3) tls randomization"),

diff --git a/internal/pdcp/writer.go b/internal/pdcp/writer.go
@@ -19,15 +19,16 @@ import (
 	"github.com/projectdiscovery/retryablehttp-go"
 	pdcpauth "github.com/projectdiscovery/utils/auth/pdcp"
 	errorutil "github.com/projectdiscovery/utils/errors"
+	unitutils "github.com/projectdiscovery/utils/unit"
 	updateutils "github.com/projectdiscovery/utils/update"
 	urlutil "github.com/projectdiscovery/utils/url"
 )
 
 const (
 	uploadEndpoint = "/v1/scans/import"
 	appendEndpoint = "/v1/scans/%s/import"
-	flushTimer     = time.Duration(1) * time.Minute
-	MaxChunkSize   = 1024 * 1024 * 4 // 4 MB
+	flushTimer     = time.Minute
+	MaxChunkSize   = 4 * unitutils.Mega // 4 MB
 	xidRe          = `^[a-z0-9]{20}$`
 )
 

diff --git a/lib/multi.go b/lib/multi.go
@@ -128,6 +128,7 @@ func (e *ThreadSafeNucleiEngine) ExecuteNucleiWithOptsCtx(ctx context.Context, t
 			return err
 		}
 	}
+
 	// create ephemeral nuclei objects/instances/types using base nuclei engine
 	unsafeOpts, err := createEphemeralObjects(ctx, e.eng, tmpEngine.opts)
 	if err != nil {

diff --git a/lib/sdk.go b/lib/sdk.go
@@ -18,7 +18,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/hosterrorscache"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/interactsh"
-	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate"
+	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolinit"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/headless/engine"
 	"github.com/projectdiscovery/nuclei/v3/pkg/reporting"
 	"github.com/projectdiscovery/nuclei/v3/pkg/templates"
@@ -183,8 +183,7 @@ func (e *NucleiEngine) SignTemplate(tmplSigner *signer.TemplateSigner, data []by
 	return buff.Bytes(), err
 }
 
-// Close all resources used by nuclei engine
-func (e *NucleiEngine) Close() {
+func (e *NucleiEngine) closeInternal() {
 	if e.interactshClient != nil {
 		e.interactshClient.Close()
 	}
@@ -206,8 +205,6 @@ func (e *NucleiEngine) Close() {
 	if e.rateLimiter != nil {
 		e.rateLimiter.Stop()
 	}
-	// close global shared resources
-	protocolstate.Close()
 	if e.inputProvider != nil {
 		e.inputProvider.Close()
 	}
@@ -219,6 +216,12 @@ func (e *NucleiEngine) Close() {
 	}
 }
 
+// Close all resources used by nuclei engine
+func (e *NucleiEngine) Close() {
+	e.closeInternal()
+	protocolinit.Close()
+}
+
 // ExecuteCallbackWithCtx executes templates on targets and calls callback on each result(only if results are found)
 // enable matcher-status option if you expect this callback to be called for all results regardless if it matched or not
 func (e *NucleiEngine) ExecuteCallbackWithCtx(ctx context.Context, callback ...func(event *output.ResultEvent)) error {

diff --git a/lib/sdk_private.go b/lib/sdk_private.go
@@ -25,6 +25,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/hosterrorscache"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/interactsh"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolinit"
+	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/http/httpclientpool"
 	"github.com/projectdiscovery/nuclei/v3/pkg/reporting"
 	"github.com/projectdiscovery/nuclei/v3/pkg/templates"
@@ -34,7 +35,7 @@ import (
 	"github.com/projectdiscovery/ratelimit"
 )
 
-var sharedInit sync.Once = sync.Once{}
+var sharedInit *sync.Once
 
 // applyRequiredDefaults to options
 func (e *NucleiEngine) applyRequiredDefaults(ctx context.Context) {
@@ -117,6 +118,10 @@ func (e *NucleiEngine) init(ctx context.Context) error {
 
 	e.parser = templates.NewParser()
 
+	if sharedInit == nil || protocolstate.ShouldInit() {
+		sharedInit = &sync.Once{}
+	}
+
 	sharedInit.Do(func() {
 		_ = protocolinit.Init(e.opts)
 	})

diff --git a/pkg/catalog/loader/loader.go b/pkg/catalog/loader/loader.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"sort"
 	"strings"
+	"sync"
 
 	"github.com/logrusorgru/aurora"
 	"github.com/pkg/errors"
@@ -23,6 +24,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/workflows"
 	"github.com/projectdiscovery/retryablehttp-go"
 	errorutil "github.com/projectdiscovery/utils/errors"
+	sliceutil "github.com/projectdiscovery/utils/slice"
 	stringsutil "github.com/projectdiscovery/utils/strings"
 	urlutil "github.com/projectdiscovery/utils/url"
 )
@@ -425,10 +427,10 @@ func (store *Store) LoadTemplatesWithTags(templatesList, tags []string) []*templ
 	store.logErroredTemplates(errs)
 	templatePathMap := store.pathFilter.Match(includedTemplates)
 
-	loadedTemplates := make([]*templates.Template, 0, len(templatePathMap))
+	loadedTemplates := sliceutil.NewSyncSlice[*templates.Template]()
 
 	loadTemplate := func(tmpl *templates.Template) {
-		loadedTemplates = append(loadedTemplates, tmpl)
+		loadedTemplates.Append(tmpl)
 		// increment signed/unsigned counters
 		if tmpl.Verified {
 			if tmpl.TemplateVerifier == "" {
@@ -441,80 +443,89 @@ func (store *Store) LoadTemplatesWithTags(templatesList, tags []string) []*templ
 		}
 	}
 
+	var wgLoadTemplates sync.WaitGroup
+
 	for templatePath := range templatePathMap {
-		loaded, err := store.config.ExecutorOptions.Parser.LoadTemplate(templatePath, store.tagFilter, tags, store.config.Catalog)
-		if loaded || store.pathFilter.MatchIncluded(templatePath) {
-			parsed, err := templates.Parse(templatePath, store.preprocessor, store.config.ExecutorOptions)
-			if err != nil {
-				// exclude templates not compatible with offline matching from total runtime warning stats
-				if !errors.Is(err, templates.ErrIncompatibleWithOfflineMatching) {
-					stats.Increment(templates.RuntimeWarningsStats)
-				}
-				gologger.Warning().Msgf("Could not parse template %s: %s\n", templatePath, err)
-			} else if parsed != nil {
-				if !parsed.Verified && store.config.ExecutorOptions.Options.DisableUnsignedTemplates {
-					// skip unverified templates when prompted to
-					stats.Increment(templates.SkippedUnsignedStats)
-					continue
-				}
-				// if template has request signature like aws then only signed and verified templates are allowed
-				if parsed.UsesRequestSignature() && !parsed.Verified {
-					stats.Increment(templates.SkippedRequestSignatureStats)
-					continue
-				}
-				// DAST only templates
-				if store.config.ExecutorOptions.Options.DAST {
-					// check if the template is a DAST template
-					if parsed.IsFuzzing() {
-						loadTemplate(parsed)
+		wgLoadTemplates.Add(1)
+		go func(templatePath string) {
+			defer wgLoadTemplates.Done()
+
+			loaded, err := store.config.ExecutorOptions.Parser.LoadTemplate(templatePath, store.tagFilter, tags, store.config.Catalog)
+			if loaded || store.pathFilter.MatchIncluded(templatePath) {
+				parsed, err := templates.Parse(templatePath, store.preprocessor, store.config.ExecutorOptions)
+				if err != nil {
+					// exclude templates not compatible with offline matching from total runtime warning stats
+					if !errors.Is(err, templates.ErrIncompatibleWithOfflineMatching) {
+						stats.Increment(templates.RuntimeWarningsStats)
 					}
-				} else if len(parsed.RequestsHeadless) > 0 && !store.config.ExecutorOptions.Options.Headless {
-					// donot include headless template in final list if headless flag is not set
-					stats.Increment(templates.ExcludedHeadlessTmplStats)
-					if config.DefaultConfig.LogAllEvents {
-						gologger.Print().Msgf("[%v] Headless flag is required for headless template '%s'.\n", aurora.Yellow("WRN").String(), templatePath)
+					gologger.Warning().Msgf("Could not parse template %s: %s\n", templatePath, err)
+				} else if parsed != nil {
+					if !parsed.Verified && store.config.ExecutorOptions.Options.DisableUnsignedTemplates {
+						// skip unverified templates when prompted to
+						stats.Increment(templates.SkippedUnsignedStats)
+						return
 					}
-				} else if len(parsed.RequestsCode) > 0 && !store.config.ExecutorOptions.Options.EnableCodeTemplates {
-					// donot include 'Code' protocol custom template in final list if code flag is not set
-					stats.Increment(templates.ExcludedCodeTmplStats)
-					if config.DefaultConfig.LogAllEvents {
-						gologger.Print().Msgf("[%v] Code flag is required for code protocol template '%s'.\n", aurora.Yellow("WRN").String(), templatePath)
+					// if template has request signature like aws then only signed and verified templates are allowed
+					if parsed.UsesRequestSignature() && !parsed.Verified {
+						stats.Increment(templates.SkippedRequestSignatureStats)
+						return
 					}
-				} else if len(parsed.RequestsCode) > 0 && !parsed.Verified && len(parsed.Workflows) == 0 {
-					// donot include unverified 'Code' protocol custom template in final list
-					stats.Increment(templates.SkippedCodeTmplTamperedStats)
-					// these will be skipped so increment skip counter
-					stats.Increment(templates.SkippedUnsignedStats)
-					if config.DefaultConfig.LogAllEvents {
-						gologger.Print().Msgf("[%v] Tampered/Unsigned template at %v.\n", aurora.Yellow("WRN").String(), templatePath)
-					}
-				} else if parsed.IsFuzzing() && !store.config.ExecutorOptions.Options.DAST {
-					stats.Increment(templates.ExludedDastTmplStats)
-					if config.DefaultConfig.LogAllEvents {
-						gologger.Print().Msgf("[%v] -dast flag is required for DAST template '%s'.\n", aurora.Yellow("WRN").String(), templatePath)
+					// DAST only templates
+					if store.config.ExecutorOptions.Options.DAST {
+						// check if the template is a DAST template
+						if parsed.IsFuzzing() {
+							loadTemplate(parsed)
+						}
+					} else if len(parsed.RequestsHeadless) > 0 && !store.config.ExecutorOptions.Options.Headless {
+						// donot include headless template in final list if headless flag is not set
+						stats.Increment(templates.ExcludedHeadlessTmplStats)
+						if config.DefaultConfig.LogAllEvents {
+							gologger.Print().Msgf("[%v] Headless flag is required for headless template '%s'.\n", aurora.Yellow("WRN").String(), templatePath)
+						}
+					} else if len(parsed.RequestsCode) > 0 && !store.config.ExecutorOptions.Options.EnableCodeTemplates {
+						// donot include 'Code' protocol custom template in final list if code flag is not set
+						stats.Increment(templates.ExcludedCodeTmplStats)
+						if config.DefaultConfig.LogAllEvents {
+							gologger.Print().Msgf("[%v] Code flag is required for code protocol template '%s'.\n", aurora.Yellow("WRN").String(), templatePath)
+						}
+					} else if len(parsed.RequestsCode) > 0 && !parsed.Verified && len(parsed.Workflows) == 0 {
+						// donot include unverified 'Code' protocol custom template in final list
+						stats.Increment(templates.SkippedCodeTmplTamperedStats)
+						// these will be skipped so increment skip counter
+						stats.Increment(templates.SkippedUnsignedStats)
+						if config.DefaultConfig.LogAllEvents {
+							gologger.Print().Msgf("[%v] Tampered/Unsigned template at %v.\n", aurora.Yellow("WRN").String(), templatePath)
+						}
+					} else if parsed.IsFuzzing() && !store.config.ExecutorOptions.Options.DAST {
+						stats.Increment(templates.ExludedDastTmplStats)
+						if config.DefaultConfig.LogAllEvents {
+							gologger.Print().Msgf("[%v] -dast flag is required for DAST template '%s'.\n", aurora.Yellow("WRN").String(), templatePath)
+						}
+					} else {
+						loadTemplate(parsed)
 					}
-				} else {
-					loadTemplate(parsed)
 				}
 			}
-		}
-		if err != nil {
-			if strings.Contains(err.Error(), templates.ErrExcluded.Error()) {
-				stats.Increment(templates.TemplatesExcludedStats)
-				if config.DefaultConfig.LogAllEvents {
-					gologger.Print().Msgf("[%v] %v\n", aurora.Yellow("WRN").String(), err.Error())
+			if err != nil {
+				if strings.Contains(err.Error(), templates.ErrExcluded.Error()) {
+					stats.Increment(templates.TemplatesExcludedStats)
+					if config.DefaultConfig.LogAllEvents {
+						gologger.Print().Msgf("[%v] %v\n", aurora.Yellow("WRN").String(), err.Error())
+					}
+					return
 				}
-				continue
+				gologger.Warning().Msg(err.Error())
 			}
-			gologger.Warning().Msg(err.Error())
-		}
+		}(templatePath)
 	}
 
-	sort.SliceStable(loadedTemplates, func(i, j int) bool {
-		return loadedTemplates[i].Path < loadedTemplates[j].Path
+	wgLoadTemplates.Wait()
+
+	sort.SliceStable(loadedTemplates.Slice, func(i, j int) bool {
+		return loadedTemplates.Slice[i].Path < loadedTemplates.Slice[j].Path
 	})
 
-	return loadedTemplates
+	return loadedTemplates.Slice
 }
 
 // IsHTTPBasedProtocolUsed returns true if http/headless protocol is being used for

diff --git a/pkg/output/output.go b/pkg/output/output.go
@@ -33,6 +33,7 @@ import (
 	"github.com/projectdiscovery/utils/errkit"
 	fileutil "github.com/projectdiscovery/utils/file"
 	osutils "github.com/projectdiscovery/utils/os"
+	unitutils "github.com/projectdiscovery/utils/unit"
 	urlutil "github.com/projectdiscovery/utils/url"
 )
 
@@ -449,7 +450,7 @@ func (w *StandardWriter) WriteFailure(wrappedEvent *InternalWrappedEvent) error
 	return w.Write(data)
 }
 
-var maxTemplateFileSizeForEncoding = 1024 * 1024
+var maxTemplateFileSizeForEncoding = unitutils.Mega
 
 func (w *StandardWriter) encodeTemplate(templatePath string) string {
 	data, err := os.ReadFile(templatePath)

diff --git a/pkg/protocols/common/automaticscan/automaticscan.go b/pkg/protocols/common/automaticscan/automaticscan.go
@@ -32,13 +32,14 @@ import (
 	sliceutil "github.com/projectdiscovery/utils/slice"
 	stringsutil "github.com/projectdiscovery/utils/strings"
 	syncutil "github.com/projectdiscovery/utils/sync"
+	unitutils "github.com/projectdiscovery/utils/unit"
 	wappalyzer "github.com/projectdiscovery/wappalyzergo"
 	"gopkg.in/yaml.v2"
 )
 
 const (
 	mappingFilename = "wappalyzer-mapping.yml"
-	maxDefaultBody  = 4 * 1024 * 1024 // 4MB
+	maxDefaultBody  = 4 * unitutils.Mega
 )
 
 // Options contains configuration options for automatic scan service

diff --git a/pkg/protocols/common/protocolinit/init.go b/pkg/protocols/common/protocolinit/init.go
@@ -38,6 +38,5 @@ func Init(options *types.Options) error {
 }
 
 func Close() {
-	protocolstate.Dialer.Close()
-	protocolstate.Dialer = nil
+	protocolstate.Close()
 }
diff --git a/pkg/protocols/common/protocolstate/state.go b/pkg/protocols/common/protocolstate/state.go
@@ -22,6 +22,10 @@ var (
 	Dialer *fastdialer.Dialer
 )
 
+func ShouldInit() bool {
+	return Dialer == nil
+}
+
 // Init creates the Dialer instance based on user configuration
 func Init(options *types.Options) error {
 	if Dialer != nil {
@@ -212,5 +216,6 @@ func Close() {
 		Dialer.Close()
 		Dialer = nil
 	}
+	Dialer = nil
 	StopActiveMemGuardian()
 }