adding connection pool within same template

v2/go.mod

@@ -18,6 +18,7 @@ require (
 	github.com/projectdiscovery/retryablehttp-go v1.0.1
 	github.com/remeh/sizedwaitgroup v1.0.0
 	github.com/vbauerster/mpb/v5 v5.3.0
+	go.uber.org/ratelimit v0.1.0
 	golang.org/x/net v0.0.0-20201006153459-a7d1128ccaa0
 	golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e
 	gopkg.in/yaml.v2 v2.3.0

v2/go.sum

@@ -58,6 +58,8 @@ github.com/vbauerster/mpb v1.1.3 h1:IRgic8VFaURXkW0VxDLkNOiNaAgtw0okB2YIaVvJDI4=
 github.com/vbauerster/mpb v3.4.0+incompatible h1:mfiiYw87ARaeRW6x5gWwYRUawxaW1tLAD8IceomUCNw=
 github.com/vbauerster/mpb/v5 v5.3.0 h1:vgrEJjUzHaSZKDRRxul5Oh4C72Yy/5VEMb0em+9M0mQ=
 github.com/vbauerster/mpb/v5 v5.3.0/go.mod h1:4yTkvAb8Cm4eylAp6t0JRq6pXDkFJ4krUlDqWYkakAs=
+go.uber.org/ratelimit v0.1.0 h1:U2AruXqeTb4Eh9sYQSTrMhH8Cb7M0Ian2ibBOnBcnAw=
+go.uber.org/ratelimit v0.1.0/go.mod h1:2X8KaoNd1J0lZV+PxJk/5+DGbO/tpwLR1m++a7FnB/Y=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 h1:psW17arqaxU48Z5kZ0CQnkZWQJsqcURM6tKiBApRjXI=

v2/pkg/executer/executer_http.go

@@ -14,6 +14,7 @@ import (
 	"os"
 	"regexp"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/pkg/errors"
@@ -27,6 +28,7 @@ import (
 	"github.com/projectdiscovery/rawhttp"
 	"github.com/projectdiscovery/retryablehttp-go"
 	"github.com/remeh/sizedwaitgroup"
+	"go.uber.org/ratelimit"
 	"golang.org/x/net/proxy"
 )
 
@@ -127,10 +129,58 @@ func NewHTTPExecuter(options *HTTPOptions) (*HTTPExecuter, error) {
 	return executer, nil
 }
 
-func (e *HTTPExecuter) ExecuteTurboHTTP(ctx context.Context, p progress.IProgress, reqURL string) (result Result) {
+func (e *HTTPExecuter) ExecuteParallelHTTP(p progress.IProgress, reqURL string) (result Result) {
 	result.Matches = make(map[string]interface{})
 	result.Extractions = make(map[string]interface{})
+	dynamicvalues := make(map[string]interface{})
 
+	// verify if the URL is already being processed
+	if e.bulkHTTPRequest.HasGenerator(reqURL) {
+		return
+	}
+
+	var rateLimit ratelimit.Limiter
+	if e.bulkHTTPRequest.RateLimit > 0 {
+		rateLimit = ratelimit.New(e.bulkHTTPRequest.RateLimit)
+	} else {
+		rateLimit = ratelimit.NewUnlimited()
+	}
+
+	remaining := e.bulkHTTPRequest.GetRequestCount()
+	e.bulkHTTPRequest.CreateGenerator(reqURL)
+
+	// Workers that keeps enqueuing new requests
+	maxWorkers := e.bulkHTTPRequest.Threads
+	swg := sizedwaitgroup.New(maxWorkers)
+	for e.bulkHTTPRequest.Next(reqURL) && !result.Done {
+		request, err := e.bulkHTTPRequest.MakeHTTPRequest(context.Background(), reqURL, dynamicvalues, e.bulkHTTPRequest.Current(reqURL))
+		if err != nil {
+			result.Error = err
+			p.Drop(remaining)
+		} else {
+			swg.Add()
+			go func(httpRequest *requests.HTTPRequest) {
+				defer swg.Done()
+
+				rateLimit.Take()
+
+				// If the request was built correctly then execute it
+				err = e.handleHTTP(reqURL, httpRequest, dynamicvalues, &result)
+				if err != nil {
+					result.Error = errors.Wrap(err, "could not handle http request")
+					p.Drop(remaining)
+				}
+			}(request)
+		}
+		e.bulkHTTPRequest.Increment(reqURL)
+	}
+
+	swg.Wait()
+
+	return result
+}
+
+func (e *HTTPExecuter) ExecuteTurboHTTP(p progress.IProgress, reqURL string) (result Result) {
 	result.Matches = make(map[string]interface{})
 	result.Extractions = make(map[string]interface{})
 	dynamicvalues := make(map[string]interface{})
@@ -165,7 +215,7 @@ func (e *HTTPExecuter) ExecuteTurboHTTP(ctx context.Context, p progress.IProgres
 
 	swg := sizedwaitgroup.New(maxWorkers)
 	for e.bulkHTTPRequest.Next(reqURL) && !result.Done {
-		request, err := e.bulkHTTPRequest.MakeHTTPRequest(ctx, reqURL, dynamicvalues, e.bulkHTTPRequest.Current(reqURL))
+		request, err := e.bulkHTTPRequest.MakeHTTPRequest(context.Background(), reqURL, dynamicvalues, e.bulkHTTPRequest.Current(reqURL))
 		if err != nil {
 			result.Error = err
 			p.Drop(remaining)
@@ -198,7 +248,11 @@ func (e *HTTPExecuter) ExecuteTurboHTTP(ctx context.Context, p progress.IProgres
 func (e *HTTPExecuter) ExecuteHTTP(ctx context.Context, p progress.IProgress, reqURL string) (result Result) {
 	// verify if pipeline was requested
 	if e.bulkHTTPRequest.Pipeline {
-		return e.ExecuteTurboHTTP(ctx, p, reqURL)
+		return e.ExecuteTurboHTTP(p, reqURL)
+	}
+
+	if e.bulkHTTPRequest.Threads > 0 {
+		return e.ExecuteParallelHTTP(p, reqURL)
 	}
 
 	result.Matches = make(map[string]interface{})
@@ -334,11 +388,13 @@ func (e *HTTPExecuter) handleHTTP(reqURL string, request *requests.HTTPRequest,
 			// If the matcher has matched, and its an OR
 			// write the first output then move to next matcher.
 			if matcherCondition == matchers.ORCondition {
+				result.Lock()
 				result.Matches[matcher.Name] = nil
 				// probably redundant but ensures we snapshot current payload values when matchers are valid
 				result.Meta = request.Meta
-				e.writeOutputHTTP(request, resp, body, matcher, nil)
 				result.GotResults = true
+				result.Unlock()
+				e.writeOutputHTTP(request, resp, body, matcher, nil)
 			}
 		}
 	}
@@ -360,16 +416,19 @@ func (e *HTTPExecuter) handleHTTP(reqURL string, request *requests.HTTPRequest,
 			}
 		}
 		// probably redundant but ensures we snapshot current payload values when extractors are valid
+		result.Lock()
 		result.Meta = request.Meta
 		result.Extractions[extractor.Name] = extractorResults
+		result.Unlock()
 	}
 
 	// Write a final string of output if matcher type is
 	// AND or if we have extractors for the mechanism too.
 	if len(outputExtractorResults) > 0 || matcherCondition == matchers.ANDCondition {
 		e.writeOutputHTTP(request, resp, body, nil, outputExtractorResults)
-
+		result.Lock()
 		result.GotResults = true
+		result.Unlock()
 	}
 
 	return nil
@@ -380,7 +439,21 @@ func (e *HTTPExecuter) Close() {}
 
 // makeHTTPClient creates a http client
 func makeHTTPClient(proxyURL *url.URL, options *HTTPOptions) *retryablehttp.Client {
+	// Multiple Host
 	retryablehttpOptions := retryablehttp.DefaultOptionsSpraying
+	disableKeepAlives := true
+	maxIdleConns := 0
+	maxConnsPerHost := 0
+	maxIdleConnsPerHost := -1
+
+	if options.BulkHTTPRequest.Threads > 0 {
+		// Single host
+		retryablehttpOptions = retryablehttp.DefaultOptionsSingle
+		disableKeepAlives = false
+		maxIdleConnsPerHost = 500
+		maxConnsPerHost = 500
+	}
+
 	retryablehttpOptions.RetryWaitMax = 10 * time.Second
 	retryablehttpOptions.RetryMax = options.Retries
 	followRedirects := options.BulkHTTPRequest.Redirects
@@ -391,12 +464,14 @@ func makeHTTPClient(proxyURL *url.URL, options *HTTPOptions) *retryablehttp.Clie
 			Timeout:   30 * time.Second,
 			KeepAlive: 30 * time.Second,
 		}).DialContext,
-		MaxIdleConnsPerHost: -1,
+		MaxIdleConns:        maxIdleConns,
+		MaxIdleConnsPerHost: maxIdleConnsPerHost,
+		MaxConnsPerHost:     maxConnsPerHost,
 		TLSClientConfig: &tls.Config{
 			Renegotiation:      tls.RenegotiateOnceAsClient,
 			InsecureSkipVerify: true,
 		},
-		DisableKeepAlives: true,
+		DisableKeepAlives: disableKeepAlives,
 	}
 
 	// Attempts to overwrite the dial function with the socks proxied version
@@ -479,6 +554,7 @@ func (e *HTTPExecuter) setCustomHeaders(r *requests.HTTPRequest) {
 }
 
 type Result struct {
+	sync.Mutex
 	GotResults  bool
 	Done        bool
 	Meta        map[string]interface{}

v2/pkg/requests/bulk-http-request.go

@@ -75,6 +75,9 @@ type BulkHTTPRequest struct {
 	DisableAutoHostname bool `yaml:"disable-automatic-host-header,omitempty"`
 	// DisableAutoContentLength Enable/Disable Content-Length header for unsafe raw requests
 	DisableAutoContentLength bool `yaml:"disable-automatic-content-length-header,omitempty"`
+	Threads                  int  `yaml:"threads,omitempty"`
+	RateLimit                int  `yaml:"rate-limit,omitempty"`
+
 	// Internal Finite State Machine keeping track of scan process
 	gsfm *GeneratorFSM
 }
@@ -244,8 +247,12 @@ func (r *BulkHTTPRequest) handleRawWithPaylods(ctx context.Context, raw, baseURL
 }
 
 func (r *BulkHTTPRequest) fillRequest(req *http.Request, values map[string]interface{}) (*retryablehttp.Request, error) {
-	setHeader(req, "Connection", "close")
-	req.Close = true
+	// In case of multiple threads the underlying connection should remain open to allow reuse
+	if r.Threads <= 0 {
+		setHeader(req, "Connection", "close")
+		req.Close = true
+	}
+
 	replacer := newReplacer(values)
 
 	// Check if the user requested a request body