Skip to content

Commit

Permalink
Merge pull request #5187 from projectdiscovery/feat-4808-planner
Browse files Browse the repository at this point in the history
Fixing issues with multi-thread execution
  • Loading branch information
Mzack9999 authored Jun 13, 2024
2 parents cfd7b6f + 1e85e70 commit 51b1d7c
Show file tree
Hide file tree
Showing 20 changed files with 149 additions and 135 deletions.
3 changes: 2 additions & 1 deletion cmd/integration-test/http.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down Expand Up @@ -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
}
Expand Down
3 changes: 2 additions & 1 deletion cmd/nuclei/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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"),
Expand Down
5 changes: 3 additions & 2 deletions internal/pdcp/writer.go
Original file line number Diff line number Diff line change
Expand Up @@ -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}$`
)

Expand Down
1 change: 1 addition & 0 deletions lib/multi.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
13 changes: 8 additions & 5 deletions lib/sdk.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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()
}
Expand All @@ -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()
}
Expand All @@ -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 {
Expand Down
7 changes: 6 additions & 1 deletion lib/sdk_private.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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) {
Expand Down Expand Up @@ -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)
})
Expand Down
137 changes: 74 additions & 63 deletions pkg/catalog/loader/loader.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"os"
"sort"
"strings"
"sync"

"github.com/logrusorgru/aurora"
"github.com/pkg/errors"
Expand All @@ -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"
)
Expand Down Expand Up @@ -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 == "" {
Expand All @@ -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
Expand Down
3 changes: 2 additions & 1 deletion pkg/output/output.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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)
Expand Down
3 changes: 2 additions & 1 deletion pkg/protocols/common/automaticscan/automaticscan.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions pkg/protocols/common/protocolinit/init.go
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,5 @@ func Init(options *types.Options) error {
}

func Close() {
protocolstate.Dialer.Close()
protocolstate.Dialer = nil
protocolstate.Close()
}
5 changes: 5 additions & 0 deletions pkg/protocols/common/protocolstate/state.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -212,5 +216,6 @@ func Close() {
Dialer.Close()
Dialer = nil
}
Dialer = nil
StopActiveMemGuardian()
}
Loading

0 comments on commit 51b1d7c

Please sign in to comment.