diff --git a/v2/go.mod b/v2/go.mod index e5e8d3aea..6338b968d 100644 --- a/v2/go.mod +++ b/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 diff --git a/v2/go.sum b/v2/go.sum index ad4f20252..0444bd2cb 100644 --- a/v2/go.sum +++ b/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= diff --git a/v2/pkg/executer/executer_http.go b/v2/pkg/executer/executer_http.go index c568116d6..3cc56414d 100644 --- a/v2/pkg/executer/executer_http.go +++ b/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{} diff --git a/v2/pkg/requests/bulk-http-request.go b/v2/pkg/requests/bulk-http-request.go index a5c169311..120ae069c 100644 --- a/v2/pkg/requests/bulk-http-request.go +++ b/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