Support concurrent Nuclei engines in the same process

lib/config.go

@@ -537,7 +537,7 @@ func WithResumeFile(file string) NucleiSDKOptions {
 	}
 }
 
-// WithLogger allows setting gologger instance
+// WithLogger allows setting a shared gologger instance
 func WithLogger(logger *gologger.Logger) NucleiSDKOptions {
 	return func(e *NucleiEngine) error {
 		e.Logger = logger

lib/sdk_private.go

@@ -231,6 +231,25 @@ func (e *NucleiEngine) init(ctx context.Context) error {
 		}
 	}
 
+	// Handle the case where the user passed an existing parser that we can use as a cache
+	if e.opts.Parser != nil {
+		if cachedParser, ok := e.opts.Parser.(*templates.Parser); ok {
+			e.parser = cachedParser
+			e.opts.Parser = cachedParser
+			e.executerOpts.Parser = cachedParser
+			e.executerOpts.Options.Parser = cachedParser
+		}
+	}
+
+	// Create a new parser if necessary
+	if e.parser == nil {
+		op := templates.NewParser()
+		e.parser = op
+		e.opts.Parser = op
+		e.executerOpts.Parser = op
+		e.executerOpts.Options.Parser = op
+	}
+
 	e.engine = core.New(e.opts)
 	e.engine.SetExecuterOptions(e.executerOpts)
 

pkg/external/customtemplates/github_test.go

@@ -1,23 +1,25 @@
 package customtemplates
 
 import (
+	"bytes"
 	"context"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/projectdiscovery/gologger"
+	"github.com/projectdiscovery/gologger/levels"
 	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
 	"github.com/projectdiscovery/nuclei/v3/pkg/testutils"
-	osutils "github.com/projectdiscovery/utils/os"
+	"github.com/projectdiscovery/nuclei/v3/pkg/utils"
 	"github.com/stretchr/testify/require"
 )
 
 func TestDownloadCustomTemplatesFromGitHub(t *testing.T) {
-	if osutils.IsOSX() {
-		t.Skip("skipping on macos due to unknown failure (works locally)")
-	}
-
-	gologger.DefaultLogger.SetWriter(&testutils.NoopWriter{})
+	// Capture output to check for rate limit errors
+	outputBuffer := &bytes.Buffer{}
+	gologger.DefaultLogger.SetWriter(&utils.CaptureWriter{Buffer: outputBuffer})
+	gologger.DefaultLogger.SetMaxLevel(levels.LevelDebug)
 
 	templatesDirectory := t.TempDir()
 	config.DefaultConfig.SetTemplatesDir(templatesDirectory)
@@ -29,5 +31,12 @@ func TestDownloadCustomTemplatesFromGitHub(t *testing.T) {
 	require.Nil(t, err, "could not create custom templates manager")
 
 	ctm.Download(context.Background())
+
+	// Check if output contains rate limit error and skip test if so
+	output := outputBuffer.String()
+	if strings.Contains(output, "API rate limit exceeded") {
+		t.Skip("GitHub API rate limit exceeded, skipping test")
+	}
+
 	require.DirExists(t, filepath.Join(templatesDirectory, "github", "projectdiscovery", "nuclei-templates-test"), "cloned directory does not exists")
 }

pkg/installer/template.go

@@ -53,7 +53,7 @@ func (t *templateUpdateResults) String() string {
 		},
 	}
 	table := tablewriter.NewWriter(&buff)
-	table.Header("Total", "Added", "Modified", "Removed")
+	table.Header([]string{"Total", "Added", "Modified", "Removed"})
 	for _, v := range data {
 		_ = table.Append(v)
 	}

pkg/js/libs/fs/fs.go

@@ -1,6 +1,7 @@
 package fs
 
 import (
+	"context"
 	"os"
 
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate"
@@ -27,8 +28,9 @@ import (
 // // when no itemType is provided, it will return both files and directories
 // const items = fs.ListDir('/tmp');
 // ```
-func ListDir(path string, itemType string) ([]string, error) {
-	finalPath, err := protocolstate.NormalizePath(path)
+func ListDir(ctx context.Context, path string, itemType string) ([]string, error) {
+	executionId := ctx.Value("executionId").(string)
+	finalPath, err := protocolstate.NormalizePathWithExecutionId(executionId, path)
 	if err != nil {
 		return nil, err
 	}
@@ -57,8 +59,9 @@ func ListDir(path string, itemType string) ([]string, error) {
 // // here permitted directories are $HOME/nuclei-templates/*
 // const content = fs.ReadFile('helpers/usernames.txt');
 // ```
-func ReadFile(path string) ([]byte, error) {
-	finalPath, err := protocolstate.NormalizePath(path)
+func ReadFile(ctx context.Context, path string) ([]byte, error) {
+	executionId := ctx.Value("executionId").(string)
+	finalPath, err := protocolstate.NormalizePathWithExecutionId(executionId, path)
 	if err != nil {
 		return nil, err
 	}
@@ -74,8 +77,8 @@ func ReadFile(path string) ([]byte, error) {
 // // here permitted directories are $HOME/nuclei-templates/*
 // const content = fs.ReadFileAsString('helpers/usernames.txt');
 // ```
-func ReadFileAsString(path string) (string, error) {
-	bin, err := ReadFile(path)
+func ReadFileAsString(ctx context.Context, path string) (string, error) {
+	bin, err := ReadFile(ctx, path)
 	if err != nil {
 		return "", err
 	}
@@ -91,14 +94,14 @@ func ReadFileAsString(path string) (string, error) {
 // const contents = fs.ReadFilesFromDir('helpers/ssh-keys');
 // log(contents);
 // ```
-func ReadFilesFromDir(dir string) ([]string, error) {
-	files, err := ListDir(dir, "file")
+func ReadFilesFromDir(ctx context.Context, dir string) ([]string, error) {
+	files, err := ListDir(ctx, dir, "file")
 	if err != nil {
 		return nil, err
 	}
 	var results []string
 	for _, file := range files {
-		content, err := ReadFileAsString(dir + "/" + file)
+		content, err := ReadFileAsString(ctx, dir+"/"+file)
 		if err != nil {
 			return nil, err
 		}

pkg/protocols/common/protocolstate/file.go

@@ -4,22 +4,65 @@ import (
 	"strings"
 
 	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
+	"github.com/projectdiscovery/nuclei/v3/pkg/types"
 	errorutil "github.com/projectdiscovery/utils/errors"
 	fileutil "github.com/projectdiscovery/utils/file"
+	mapsutil "github.com/projectdiscovery/utils/maps"
 )
 
 var (
 	// LfaAllowed means local file access is allowed
-	LfaAllowed bool
+	LfaAllowed *mapsutil.SyncLockMap[string, bool]
 )
 
+func init() {
+	LfaAllowed = mapsutil.NewSyncLockMap[string, bool]()
+}
+
+// IsLfaAllowed returns whether local file access is allowed
+func IsLfaAllowed(options *types.Options) bool {
+	if GetLfaAllowed(options) {
+		return true
+	}
+
+	// Otherwise look into dialers
+	dialers, ok := dialers.Get(options.ExecutionId)
+	if ok && dialers != nil {
+		dialers.Lock()
+		defer dialers.Unlock()
+
+		return dialers.LocalFileAccessAllowed
+	}
+
+	// otherwise just return option value
+	return options.AllowLocalFileAccess
+}
+
+func SetLfaAllowed(options *types.Options) {
+	_ = LfaAllowed.Set(options.ExecutionId, options.AllowLocalFileAccess)
+}
+
+func GetLfaAllowed(options *types.Options) bool {
+	allowed, ok := LfaAllowed.Get(options.ExecutionId)
+
+	return ok && allowed
+}
+
+func NormalizePathWithExecutionId(executionId string, filePath string) (string, error) {
+	options := &types.Options{
+		ExecutionId: executionId,
+	}
+	return NormalizePath(options, filePath)
+}
+
 // Normalizepath normalizes path and returns absolute path
 // it returns error if path is not allowed
 // this respects the sandbox rules and only loads files from
 // allowed directories
-func NormalizePath(filePath string) (string, error) {
-	// TODO: this should be tied to executionID
-	if LfaAllowed {
+func NormalizePath(options *types.Options, filePath string) (string, error) {
+	// TODO: this should be tied to executionID using *types.Options
+	if IsLfaAllowed(options) {
+		// if local file access is allowed, we can return the absolute path
 		return filePath, nil
 	}
 	cleaned, err := fileutil.ResolveNClean(filePath, config.DefaultConfig.GetTemplateDir())

pkg/protocols/common/protocolstate/headless.go

@@ -74,18 +74,6 @@ func InitHeadless(options *types.Options) {
 	}
 }
 
-// AllowLocalFileAccess returns whether local file access is allowed
-func IsLfaAllowed(options *types.Options) bool {
-	dialers, ok := dialers.Get(options.ExecutionId)
-	if ok && dialers != nil {
-		dialers.Lock()
-		defer dialers.Unlock()
-
-		return dialers.LocalFileAccessAllowed
-	}
-	return false
-}
-
 func IsRestrictLocalNetworkAccess(options *types.Options) bool {
 	dialers, ok := dialers.Get(options.ExecutionId)
 	if ok && dialers != nil {

pkg/protocols/common/protocolstate/state.go

@@ -200,9 +200,7 @@ func initDialers(options *types.Options) error {
 
 	StartActiveMemGuardian(context.Background())
 
-	// TODO: this should be tied to executionID
-	// overidde global settings with latest options
-	LfaAllowed = options.AllowLocalFileAccess
+	SetLfaAllowed(options)
 
 	return nil
 }

pkg/protocols/dns/dnsclientpool/clientpool.go

@@ -8,12 +8,14 @@ import (
 	"github.com/pkg/errors"
 	"github.com/projectdiscovery/nuclei/v3/pkg/types"
 	"github.com/projectdiscovery/retryabledns"
+	mapsutil "github.com/projectdiscovery/utils/maps"
 )
 
 var (
-	poolMutex    *sync.RWMutex
+	clientPool *mapsutil.SyncLockMap[string, *retryabledns.Client]
+
 	normalClient *retryabledns.Client
-	clientPool   map[string]*retryabledns.Client
+	m            sync.Mutex
 )
 
 // defaultResolvers contains the list of resolvers known to be trusted.
@@ -26,12 +28,14 @@ var defaultResolvers = []string{
 
 // Init initializes the client pool implementation
 func Init(options *types.Options) error {
+	m.Lock()
+	defer m.Unlock()
+
 	// Don't create clients if already created in the past.
 	if normalClient != nil {
 		return nil
 	}
-	poolMutex = &sync.RWMutex{}
-	clientPool = make(map[string]*retryabledns.Client)
+	clientPool = mapsutil.NewSyncLockMap[string, *retryabledns.Client]()
 
 	resolvers := defaultResolvers
 	if len(options.InternalResolversList) > 0 {
@@ -45,6 +49,12 @@ func Init(options *types.Options) error {
 	return nil
 }
 
+func getNormalClient() *retryabledns.Client {
+	m.Lock()
+	defer m.Unlock()
+	return normalClient
+}
+
 // Configuration contains the custom configuration options for a client
 type Configuration struct {
 	// Retries contains the retries for the dns client
@@ -71,15 +81,12 @@ func (c *Configuration) Hash() string {
 // Get creates or gets a client for the protocol based on custom configuration
 func Get(options *types.Options, configuration *Configuration) (*retryabledns.Client, error) {
 	if (configuration.Retries <= 1) && len(configuration.Resolvers) == 0 {
-		return normalClient, nil
+		return getNormalClient(), nil
 	}
 	hash := configuration.Hash()
-	poolMutex.RLock()
-	if client, ok := clientPool[hash]; ok {
-		poolMutex.RUnlock()
+	if client, ok := clientPool.Get(hash); ok {
 		return client, nil
 	}
-	poolMutex.RUnlock()
 
 	resolvers := defaultResolvers
 	if len(options.InternalResolversList) > 0 {
@@ -95,9 +102,7 @@ func Get(options *types.Options, configuration *Configuration) (*retryabledns.Cl
 	if err != nil {
 		return nil, errors.Wrap(err, "could not create dns client")
 	}
+	_ = clientPool.Set(hash, client)
 
-	poolMutex.Lock()
-	clientPool[hash] = client
-	poolMutex.Unlock()
 	return client, nil
 }

pkg/protocols/http/httpclientpool/clientpool.go

@@ -154,16 +154,16 @@ func GetRawHTTP(options *protocols.ExecutorOptions) *rawhttp.Client {
 		return dialers.RawHTTPClient
 	}
 
-	rawHttpOptions := rawhttp.DefaultOptions
+	rawHttpOptionsCopy := *rawhttp.DefaultOptions
 	if options.Options.AliveHttpProxy != "" {
-		rawHttpOptions.Proxy = options.Options.AliveHttpProxy
+		rawHttpOptionsCopy.Proxy = options.Options.AliveHttpProxy
 	} else if options.Options.AliveSocksProxy != "" {
-		rawHttpOptions.Proxy = options.Options.AliveSocksProxy
+		rawHttpOptionsCopy.Proxy = options.Options.AliveSocksProxy
 	} else if dialers.Fastdialer != nil {
-		rawHttpOptions.FastDialer = dialers.Fastdialer
+		rawHttpOptionsCopy.FastDialer = dialers.Fastdialer
 	}
-	rawHttpOptions.Timeout = options.Options.GetTimeouts().HttpTimeout
-	dialers.RawHTTPClient = rawhttp.NewClient(rawHttpOptions)
+	rawHttpOptionsCopy.Timeout = options.Options.GetTimeouts().HttpTimeout
+	dialers.RawHTTPClient = rawhttp.NewClient(&rawHttpOptionsCopy)
 	return dialers.RawHTTPClient
 }