Feat conn (#3991)

* .

* introducing connection reuse with scan strategy

* bumping up to 500

probably should be set to a lower number to push connection reuse

* Removed debug panic

* merge conflict fix

* dep update

* removing useless comparison

* lint errs

* fmt

---------

Co-authored-by: Ice3man <nizamulrana@gmail.com>
Co-authored-by: sandeep <8293321+ehsandeep@users.noreply.github.com>

v2/pkg/protocols/common/automaticscan/automaticscan.go

@@ -16,6 +16,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/contextargs"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/httpclientpool"
+	httputil "github.com/projectdiscovery/nuclei/v2/pkg/protocols/utils/http"
 	"github.com/projectdiscovery/nuclei/v2/pkg/templates"
 	"github.com/projectdiscovery/nuclei/v2/pkg/templates/types"
 	"github.com/projectdiscovery/retryablehttp-go"
@@ -58,13 +59,13 @@ func New(opts Options) (*Service, error) {
 
 	var mappingData map[string]string
 	config := config.DefaultConfig
-	if err == nil {
-		mappingFile := filepath.Join(config.TemplatesDirectory, mappingFilename)
-		if file, err := os.Open(mappingFile); err == nil {
-			_ = yaml.NewDecoder(file).Decode(&mappingData)
-			file.Close()
-		}
+
+	mappingFile := filepath.Join(config.TemplatesDirectory, mappingFilename)
+	if file, err := os.Open(mappingFile); err == nil {
+		_ = yaml.NewDecoder(file).Decode(&mappingData)
+		file.Close()
 	}
+
 	if opts.ExecuterOpts.Options.Verbose {
 		gologger.Verbose().Msgf("Normalized mapping (%d): %v\n", len(mappingData), mappingData)
 	}
@@ -86,7 +87,9 @@ func New(opts Options) (*Service, error) {
 	childExecuter := opts.Engine.ChildExecuter()
 
 	httpclient, err := httpclientpool.Get(opts.ExecuterOpts.Options, &httpclientpool.Configuration{
-		Connection: &httpclientpool.ConnectionConfiguration{DisableKeepAlive: true},
+		Connection: &httpclientpool.ConnectionConfiguration{
+			DisableKeepAlive: httputil.ShouldDisableKeepAlive(opts.ExecuterOpts.Options),
+		},
 	})
 	if err != nil {
 		return nil, errors.Wrap(err, "could not get http client")

v2/pkg/protocols/http/build_request.go

@@ -20,6 +20,7 @@ import (
 	protocolutils "github.com/projectdiscovery/nuclei/v2/pkg/protocols/utils"
 	httputil "github.com/projectdiscovery/nuclei/v2/pkg/protocols/utils/http"
 	"github.com/projectdiscovery/nuclei/v2/pkg/types"
+	"github.com/projectdiscovery/nuclei/v2/pkg/types/scanstrategy"
 	"github.com/projectdiscovery/rawhttp"
 	"github.com/projectdiscovery/retryablehttp-go"
 	errorutil "github.com/projectdiscovery/utils/errors"
@@ -341,7 +342,7 @@ func (r *requestGenerator) fillRequest(req *retryablehttp.Request, values map[st
 	}
 
 	// In case of multiple threads the underlying connection should remain open to allow reuse
-	if r.request.Threads <= 0 && req.Header.Get("Connection") == "" {
+	if r.request.Threads <= 0 && req.Header.Get("Connection") == "" && r.options.Options.ScanStrategy != scanstrategy.HostSpray.String() {
 		req.Close = true
 	}
 

v2/pkg/protocols/http/http.go

@@ -14,6 +14,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/fuzz"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/generators"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/httpclientpool"
+	httputil "github.com/projectdiscovery/nuclei/v2/pkg/protocols/utils/http"
 	"github.com/projectdiscovery/rawhttp"
 	"github.com/projectdiscovery/retryablehttp-go"
 	fileutil "github.com/projectdiscovery/utils/file"
@@ -249,7 +250,9 @@ func (request *Request) Compile(options *protocols.ExecutorOptions) error {
 		MaxRedirects: request.MaxRedirects,
 		NoTimeout:    false,
 		CookieReuse:  request.CookieReuse,
-		Connection:   &httpclientpool.ConnectionConfiguration{DisableKeepAlive: true},
+		Connection: &httpclientpool.ConnectionConfiguration{
+			DisableKeepAlive: httputil.ShouldDisableKeepAlive(options.Options),
+		},
 		RedirectFlow: httpclientpool.DontFollowRedirect,
 	}
 

v2/pkg/protocols/http/httpclientpool/clientpool.go

@@ -21,8 +21,10 @@ import (
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/protocolstate"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/utils"
 	"github.com/projectdiscovery/nuclei/v2/pkg/types"
+	"github.com/projectdiscovery/nuclei/v2/pkg/types/scanstrategy"
 	"github.com/projectdiscovery/rawhttp"
 	"github.com/projectdiscovery/retryablehttp-go"
+	mapsutil "github.com/projectdiscovery/utils/maps"
 )
 
 var (
@@ -31,9 +33,8 @@ var (
 
 	rawHttpClient     *rawhttp.Client
 	forceMaxRedirects int
-	poolMutex         *sync.RWMutex
 	normalClient      *retryablehttp.Client
-	clientPool        map[string]*retryablehttp.Client
+	clientPool        *mapsutil.SyncLockMap[string, *retryablehttp.Client]
 )
 
 // Init initializes the clientpool implementation
@@ -45,8 +46,9 @@ func Init(options *types.Options) error {
 	if options.ShouldFollowHTTPRedirects() {
 		forceMaxRedirects = options.MaxRedirects
 	}
-	poolMutex = &sync.RWMutex{}
-	clientPool = make(map[string]*retryablehttp.Client)
+	clientPool = &mapsutil.SyncLockMap[string, *retryablehttp.Client]{
+		Map: make(mapsutil.Map[string, *retryablehttp.Client]),
+	}
 
 	client, err := wrappedGet(options, &Configuration{})
 	if err != nil {
@@ -160,12 +162,9 @@ func wrappedGet(options *types.Options, configuration *Configuration) (*retryabl
 	}
 
 	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()
 
 	// Multiple Host
 	retryableHttpOptions := retryablehttp.DefaultOptionsSpraying
@@ -174,7 +173,7 @@ func wrappedGet(options *types.Options, configuration *Configuration) (*retryabl
 	maxConnsPerHost := 0
 	maxIdleConnsPerHost := -1
 
-	if configuration.Threads > 0 {
+	if configuration.Threads > 0 || options.ScanStrategy == scanstrategy.HostSpray.String() {
 		// Single host
 		retryableHttpOptions = retryablehttp.DefaultOptionsSingle
 		disableKeepAlives = false
@@ -203,6 +202,7 @@ func wrappedGet(options *types.Options, configuration *Configuration) (*retryabl
 		redirectFlow = DontFollowRedirect
 		maxRedirects = 0
 	}
+
 	// override connection's settings if required
 	if configuration.Connection != nil {
 		disableKeepAlives = configuration.Connection.DisableKeepAlive
@@ -298,9 +298,9 @@ func wrappedGet(options *types.Options, configuration *Configuration) (*retryabl
 
 	// Only add to client pool if we don't have a cookie jar in place.
 	if jar == nil {
-		poolMutex.Lock()
-		clientPool[hash] = client
-		poolMutex.Unlock()
+		if err := clientPool.Set(hash, client); err != nil {
+			return nil, err
+		}
 	}
 	return client, nil
 }

v2/pkg/protocols/http/request.go

@@ -17,6 +17,7 @@ import (
 	"go.uber.org/multierr"
 	"moul.io/http2curl"
 
+	"github.com/projectdiscovery/fastdialer/fastdialer"
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/nuclei/v2/pkg/operators"
 	"github.com/projectdiscovery/nuclei/v2/pkg/output"
@@ -878,7 +879,7 @@ func (request *Request) pruneSignatureInternalValues(maps ...map[string]interfac
 
 func (request *Request) newContext(input *contextargs.Context) context.Context {
 	if input.MetaInput.CustomIP != "" {
-		return context.WithValue(context.Background(), "ip", input.MetaInput.CustomIP) //nolint
+		return context.WithValue(context.Background(), fastdialer.IP, input.MetaInput.CustomIP)
 	}
 	return context.Background()
 }

v2/pkg/protocols/utils/http/requtils.go

@@ -4,6 +4,8 @@ import (
 	"regexp"
 	"strings"
 
+	"github.com/projectdiscovery/nuclei/v2/pkg/types"
+	"github.com/projectdiscovery/nuclei/v2/pkg/types/scanstrategy"
 	"github.com/projectdiscovery/retryablehttp-go"
 	urlutil "github.com/projectdiscovery/utils/url"
 )
@@ -42,3 +44,9 @@ func SetHeader(req *retryablehttp.Request, name, value string) {
 		req.Host = value
 	}
 }
+
+// ShouldDisableKeepAlive depending on scan strategy
+func ShouldDisableKeepAlive(options *types.Options) bool {
+	// with host-spray strategy keep-alive must be enabled
+	return options.ScanStrategy != scanstrategy.HostSpray.String()
+}