HTTP executor refactor + simplifying logics

v2/pkg/executer/executer_http.go

@@ -13,7 +13,6 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/corpix/uarand"
@@ -576,13 +575,3 @@ func (e *HTTPExecuter) setCustomHeaders(r *requests.HTTPRequest) {
 		}
 	}
 }
-
-type Result struct {
-	sync.Mutex
-	GotResults  bool
-	Meta        map[string]interface{}
-	Matches     map[string]interface{}
-	Extractions map[string]interface{}
-	historyData map[string]interface{}
-	Error       error
-}

v2/pkg/protocols/common/generators/generators.go

@@ -51,7 +51,8 @@ func New(payloads map[string]interface{}, Type Type) (*Generator, error) {
 			totalLength = len(v)
 		}
 	}
-	return &Generator{Type: Type, payloads: compiled}, nil
+	generator := &Generator{Type: Type, payloads: compiled}
+	return generator, nil
 }
 
 // Iterator is a single instance of an iterator for a generator structure
@@ -88,16 +89,17 @@ func (i *Iterator) Reset() {
 	}
 }
 
+// Remaining returns the amount of requests left for the generator.
+func (i *Iterator) Remaining() int {
+	return i.total - i.position
+}
+
 // Total returns the amount of input combinations available
 func (i *Iterator) Total() int {
 	count := 0
 	switch i.Type {
-	case Sniper:
+	case Sniper, PitchFork:
 		count = len(i.payloads[0].values)
-	case PitchFork:
-		for _, p := range i.payloads {
-			count = len(p.values)
-		}
 	case ClusterBomb:
 		count = 1
 		for _, p := range i.payloads {
@@ -131,6 +133,7 @@ func (i *Iterator) sniperValue() (map[string]interface{}, bool) {
 	}
 	values[payload.name] = payload.value()
 	payload.incrementPosition()
+	i.position++
 	return values, true
 }
 
@@ -145,6 +148,7 @@ func (i *Iterator) pitchforkValue() (map[string]interface{}, bool) {
 		values[p.name] = p.value()
 		p.incrementPosition()
 	}
+	i.position++
 	return values, true
 }
 

v2/pkg/protocols/http/build_request.go

@@ -18,6 +18,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/replacer"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/race"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/raw"
+	"github.com/projectdiscovery/rawhttp"
 	"github.com/projectdiscovery/retryablehttp-go"
 )
 
@@ -97,6 +98,7 @@ type generatedRequest struct {
 	original        *Request
 	rawRequest      *raw.Request
 	meta            map[string]interface{}
+	pipelinedClient *rawhttp.PipelineClient
 	request         *retryablehttp.Request
 }
 
@@ -128,6 +130,15 @@ func (r *requestGenerator) Make(baseURL string, dynamicValues map[string]interfa
 	return r.makeHTTPRequestFromModel(ctx, data, values)
 }
 
+// Remaining returns the remaining number of requests for the generator
+func (r *requestGenerator) Remaining() int {
+	if r.payloadIterator != nil {
+		payloadRemaining := r.payloadIterator.Remaining()
+		return (len(r.request.Raw) - r.currentIndex + 1) * payloadRemaining
+	}
+	return len(r.request.Path) - r.currentIndex + 1
+}
+
 // baseURLWithTemplatePrefs returns the url for BaseURL keeping
 // the template port and path preference
 func baseURLWithTemplatePrefs(data string, parsedURL *url.URL) string {

v2/pkg/protocols/http/http.go

@@ -61,6 +61,7 @@ type Request struct {
 
 	options       *protocols.ExecuterOptions
 	attackType    generators.Type
+	totalRequests int
 	generator     *generators.Generator // optional, only enabled when using payloads
 	httpClient    *retryablehttp.Client
 	rawhttpClient *rawhttp.Client
@@ -95,5 +96,15 @@ func (r *Request) Compile(options *protocols.ExecuterOptions) error {
 		}
 	}
 	r.options = options
+	r.totalRequests = r.Requests()
 	return nil
 }
+
+// Requests returns the total number of requests the YAML rule will perform
+func (r *Request) Requests() int {
+	if len(r.Payloads) > 0 {
+		payloadRequests := r.generator.NewIterator().Total()
+		return len(r.Raw) * payloadRequests
+	}
+	return len(r.Path)
+}

v2/pkg/protocols/http/request.go

@@ -1,128 +1,112 @@
 package http
 
-/*
-func (e *Request) ExecuteRaceRequest(reqURL string) *Result {
-	result := &Result{
-		Matches:     make(map[string]interface{}),
-		Extractions: make(map[string]interface{}),
-	}
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"os"
+	"strings"
+	"sync"
+	"time"
 
-	dynamicvalues := make(map[string]interface{})
+	"github.com/corpix/uarand"
+	"github.com/pkg/errors"
+	"github.com/projectdiscovery/nuclei/v2/pkg/matchers"
+	"github.com/projectdiscovery/nuclei/v2/pkg/output"
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/generators"
+	"github.com/projectdiscovery/nuclei/v2/pkg/requests"
+	"github.com/projectdiscovery/rawhttp"
+	"github.com/remeh/sizedwaitgroup"
+	"go.uber.org/multierr"
+)
 
-	// verify if the URL is already being processed
-	if e.HasGenerator(reqURL) {
-		return result
-	}
+const defaultMaxWorkers = 150
 
-	e.CreateGenerator(reqURL)
+// executeRaceRequest executes race condition request for a URL
+func (e *Request) executeRaceRequest(reqURL string, dynamicValues map[string]interface{}) ([]*output.InternalWrappedEvent, error) {
+	generator := e.newGenerator()
 
-	// Workers that keeps enqueuing new requests
 	maxWorkers := e.RaceNumberRequests
 	swg := sizedwaitgroup.New(maxWorkers)
-	for i := 0; i < e.RaceNumberRequests; i++ {
-		swg.Add()
-		// base request
-		result.Lock()
-		request, err := e.MakeHTTPRequest(reqURL, dynamicvalues, e.Current(reqURL))
-		payloads, _ := e.GetPayloadsValues(reqURL)
-		result.Unlock()
-		// ignore the error due to the base request having null paylods
-		if err == requests.ErrNoPayload {
-			// pass through
-		} else if err != nil {
-			result.Error = err
-		}
-		go func(httpRequest *requests.HTTPRequest) {
-			defer swg.Done()
 
-			// If the request was built correctly then execute it
-			err = e.handleHTTP(reqURL, httpRequest, dynamicvalues, result, payloads, "")
+	var requestErr error
+	var mutex *sync.Mutex
+	var outputs []*output.InternalWrappedEvent
+	for i := 0; i < e.RaceNumberRequests; i++ {
+		request, err := generator.Make(reqURL, nil)
 		if err != nil {
-				result.Error = errors.Wrap(err, "could not handle http request")
+			break
 		}
+
+		swg.Add()
+		go func(httpRequest *generatedRequest) {
+			output, err := e.executeRequest(reqURL, httpRequest, dynamicValues)
+			mutex.Lock()
+			if err != nil {
+				requestErr = multierr.Append(requestErr, err)
+			} else {
+				outputs = append(outputs, output...)
+			}
+			mutex.Unlock()
+			swg.Done()
 		}(request)
 	}
-
 	swg.Wait()
-
-	return result
+	return outputs, requestErr
 }
 
-func (e *Request) ExecuteParallelHTTP(p *progress.Progress, reqURL string) *Result {
-	result := &Result{
-		Matches:     make(map[string]interface{}),
-		Extractions: make(map[string]interface{}),
-	}
-
-	dynamicvalues := make(map[string]interface{})
-
-	// verify if the URL is already being processed
-	if e.HasGenerator(reqURL) {
-		return result
-	}
-
-	remaining := e.GetRequestCount()
-	e.CreateGenerator(reqURL)
+// executeRaceRequest executes race condition request for a URL
+func (e *Request) executeParallelHTTP(reqURL string, dynamicValues map[string]interface{}) ([]*output.InternalWrappedEvent, error) {
+	generator := e.newGenerator()
 
 	// Workers that keeps enqueuing new requests
 	maxWorkers := e.Threads
 	swg := sizedwaitgroup.New(maxWorkers)
-	for e.Next(reqURL) {
-		result.Lock()
-		request, err := e.MakeHTTPRequest(reqURL, dynamicvalues, e.Current(reqURL))
-		payloads, _ := e.GetPayloadsValues(reqURL)
-		result.Unlock()
-		// ignore the error due to the base request having null paylods
-		if err == requests.ErrNoPayload {
-			// pass through
-		} else if err != nil {
-			result.Error = err
-			p.Drop(remaining)
-		} else {
+
+	var requestErr error
+	var mutex *sync.Mutex
+	var outputs []*output.InternalWrappedEvent
+	for {
+		request, err := generator.Make(reqURL, dynamicValues)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			e.options.Progress.DecrementRequests(int64(generator.Remaining()))
+			return nil, err
+		}
 		swg.Add()
-			go func(httpRequest *requests.HTTPRequest) {
+		go func(httpRequest *generatedRequest) {
 			defer swg.Done()
 
-				e.ratelimiter.Take()
-
-				// If the request was built correctly then execute it
-				err = e.handleHTTP(reqURL, httpRequest, dynamicvalues, result, payloads, "")
+			e.options.RateLimiter.Take()
+			output, err := e.executeRequest(reqURL, httpRequest, dynamicValues)
+			mutex.Lock()
 			if err != nil {
-					e.traceLog.Request(e.template.ID, reqURL, "http", err)
-					result.Error = errors.Wrap(err, "could not handle http request")
-					p.Drop(remaining)
+				requestErr = multierr.Append(requestErr, err)
 			} else {
-					e.traceLog.Request(e.template.ID, reqURL, "http", nil)
+				outputs = append(outputs, output...)
 			}
+			mutex.Unlock()
 		}(request)
-		}
-		p.Update()
-		e.Increment(reqURL)
+		e.options.Progress.IncrementRequests()
 	}
 	swg.Wait()
-
-	return result
+	return outputs, requestErr
 }
 
-func (e *Request) ExecuteTurboHTTP(reqURL string) *Result {
-	result := &Result{
-		Matches:     make(map[string]interface{}),
-		Extractions: make(map[string]interface{}),
-	}
-
-	dynamicvalues := make(map[string]interface{})
-
-	// verify if the URL is already being processed
-	if e.HasGenerator(reqURL) {
-		return result
-	}
-
-	e.CreateGenerator(reqURL)
+// executeRaceRequest executes race condition request for a URL
+func (e *Request) executeTurboHTTP(reqURL string, dynamicValues map[string]interface{}) ([]*output.InternalWrappedEvent, error) {
+	generator := e.newGenerator()
 
 	// need to extract the target from the url
 	URL, err := url.Parse(reqURL)
 	if err != nil {
-		return result
+		return nil, err
 	}
 
 	pipeOptions := rawhttp.DefaultPipelineOptions
@@ -143,119 +127,90 @@ func (e *Request) ExecuteTurboHTTP(reqURL string) *Result {
 		maxWorkers = pipeOptions.MaxPendingRequests
 	}
 	swg := sizedwaitgroup.New(maxWorkers)
-	for e.Next(reqURL) {
-		result.Lock()
-		request, err := e.MakeHTTPRequest(reqURL, dynamicvalues, e.Current(reqURL))
-		payloads, _ := e.GetPayloadsValues(reqURL)
-		result.Unlock()
-		// ignore the error due to the base request having null paylods
-		if err == requests.ErrNoPayload {
-			// pass through
-		} else if err != nil {
-			result.Error = err
-		} else {
+
+	var requestErr error
+	var mutex *sync.Mutex
+	var outputs []*output.InternalWrappedEvent
+	for {
+		request, err := generator.Make(reqURL, dynamicValues)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			e.options.Progress.DecrementRequests(int64(generator.Remaining()))
+			return nil, err
+		}
+		request.pipelinedClient = pipeclient
+
 		swg.Add()
-			go func(httpRequest *requests.HTTPRequest) {
+		go func(httpRequest *generatedRequest) {
 			defer swg.Done()
 
-				// HTTP pipelining ignores rate limit
-				// If the request was built correctly then execute it
-				request.Pipeline = true
-				request.PipelineClient = pipeclient
-				err = e.handleHTTP(reqURL, httpRequest, dynamicvalues, result, payloads, "")
+			output, err := e.executeRequest(reqURL, httpRequest, dynamicValues)
+			mutex.Lock()
 			if err != nil {
-					e.traceLog.Request(e.template.ID, reqURL, "http", err)
-					result.Error = errors.Wrap(err, "could not handle http request")
+				requestErr = multierr.Append(requestErr, err)
 			} else {
-					e.traceLog.Request(e.template.ID, reqURL, "http", nil)
+				outputs = append(outputs, output...)
 			}
-				request.PipelineClient = nil
+			mutex.Unlock()
 		}(request)
-		}
-
-		e.Increment(reqURL)
+		e.options.Progress.IncrementRequests()
 	}
 	swg.Wait()
-	return result
+	return outputs, requestErr
 }
 
 // ExecuteHTTP executes the HTTP request on a URL
-func (e *Request) ExecuteHTTP(p *progress.Progress, reqURL string) *Result {
+func (e *Request) ExecuteHTTP(reqURL string, dynamicValues map[string]interface{}) ([]*output.InternalWrappedEvent, error) {
 	// verify if pipeline was requested
 	if e.Pipeline {
-		return e.ExecuteTurboHTTP(reqURL)
+		return e.executeTurboHTTP(reqURL, dynamicValues)
 	}
 
 	// verify if a basic race condition was requested
 	if e.Race && e.RaceNumberRequests > 0 {
-		return e.ExecuteRaceRequest(reqURL)
+		return e.executeRaceRequest(reqURL, dynamicValues)
 	}
 
 	// verify if parallel elaboration was requested
 	if e.Threads > 0 {
-		return e.ExecuteParallelHTTP(p, reqURL)
+		return e.executeParallelHTTP(reqURL, dynamicValues)
 	}
 
-	var requestNumber int
+	generator := e.newGenerator()
 
-	result := &Result{
-		Matches:     make(map[string]interface{}),
-		Extractions: make(map[string]interface{}),
-		historyData: make(map[string]interface{}),
-	}
-
-	dynamicvalues := make(map[string]interface{})
-
-	// verify if the URL is already being processed
-	if e.HasGenerator(reqURL) {
-		return result
-	}
-
-	remaining := e.GetRequestCount()
-	e.CreateGenerator(reqURL)
-
-	for e.Next(reqURL) {
-		requestNumber++
-		result.Lock()
-		httpRequest, err := e.MakeHTTPRequest(reqURL, dynamicvalues, e.Current(reqURL))
-		payloads, _ := e.GetPayloadsValues(reqURL)
-		result.Unlock()
-		// ignore the error due to the base request having null paylods
-		if err == requests.ErrNoPayload {
-			// pass through
-		} else if err != nil {
-			result.Error = err
-			p.Drop(remaining)
-		} else {
-			e.ratelimiter.Take()
-			// If the request was built correctly then execute it
-			format := "%s_" + strconv.Itoa(requestNumber)
-			err = e.handleHTTP(reqURL, httpRequest, dynamicvalues, result, payloads, format)
-			if err != nil {
-				result.Error = errors.Wrap(err, "could not handle http request")
-				p.Drop(remaining)
-				e.traceLog.Request(e.template.ID, reqURL, "http", err)
-			} else {
-				e.traceLog.Request(e.template.ID, reqURL, "http", nil)
-			}
-		}
-		p.Update()
-
-		// Check if has to stop processing at first valid result
-		if e.stopAtFirstMatch && result.GotResults {
-			p.Drop(remaining)
+	var requestErr error
+	var outputs []*output.InternalWrappedEvent
+	for {
+		request, err := generator.Make(reqURL, dynamicValues)
+		if err == io.EOF {
 			break
 		}
-
-		// move always forward with requests
-		e.Increment(reqURL)
-		remaining--
-	}
-	gologger.Verbosef("Sent for [%s] to %s\n", "http-request", e.template.ID, reqURL)
-	return result
+		if err != nil {
+			e.options.Progress.DecrementRequests(int64(generator.Remaining()))
+			return nil, err
 		}
 
-func (e *Request) handleHTTP(reqURL string, request *requests.HTTPRequest, dynamicvalues map[string]interface{}, result *Result, payloads map[string]interface{}, format string) error {
+		e.options.RateLimiter.Take()
+		output, err := e.executeRequest(reqURL, request, dynamicValues)
+		if err != nil {
+			requestErr = multierr.Append(requestErr, err)
+		} else {
+			outputs = append(outputs, output...)
+		}
+		e.options.Progress.IncrementRequests()
+
+		if request.original.options.Options.StopAtFirstMatch && len(output) > 0 {
+			e.options.Progress.DecrementRequests(int64(generator.Remaining()))
+			break
+		}
+	}
+	return outputs, requestErr
+}
+
+// executeRequest executes the actual generated request and returns error if occured
+func (e *Request) executeRequest(reqURL string, request *generatedRequest, dynamicvalues map[string]interface{}) ([]*output.InternalWrappedEvent, error) {
 	// Add User-Agent value randomly to the customHeaders slice if `random-agent` flag is given
 	if e.options.Options.RandomAgent {
 		// nolint:errcheck // ignoring error
@@ -285,7 +240,7 @@ func (e *Request) handleHTTP(reqURL string, request *requests.HTTPRequest, dynam
 
 	timeStart := time.Now()
 
-	if request.Pipeline {
+	if request.original.Pipeline {
 		resp, err = request.PipelineClient.DoRaw(request.RawRequest.Method, reqURL, request.RawRequest.Path, requests.ExpandMapValues(request.RawRequest.Headers), ioutil.NopCloser(strings.NewReader(request.RawRequest.Data)))
 		if err != nil {
 			if resp != nil {
@@ -295,10 +250,10 @@ func (e *Request) handleHTTP(reqURL string, request *requests.HTTPRequest, dynam
 			return err
 		}
 		e.traceLog.Request(e.template.ID, reqURL, "http", nil)
-	} else if request.Unsafe {
+	} else if request.original.Unsafe {
 		// rawhttp
 		// burp uses "\r\n" as new line character
-		request.RawRequest.Data = strings.ReplaceAll(request.RawRequest.Data, "\n", "\r\n")
+		request.rawRequest.Data = strings.ReplaceAll(request.RawRequest.Data, "\n", "\r\n")
 		options := e.rawHTTPClient.Options
 		options.AutomaticContentLength = request.AutomaticContentLengthHeader
 		options.AutomaticHostHeader = request.AutomaticHostHeader
@@ -474,6 +429,6 @@ func (e *Request) handleHTTP(reqURL string, request *requests.HTTPRequest, dynam
 		result.Unlock()
 	}
 
+	gologger.Verbosef("Sent for [%s] to %s\n", "http-request", e.template.ID, reqURL)
 	return nil
 }
-*/

v2/pkg/requests/doc.go

@@ -1,3 +0,0 @@
-// Package requests implements requests for templates that
-// will be sent to hosts.
-package requests

v2/pkg/requests/dump.go

@@ -1,21 +0,0 @@
-package requests
-
-import (
-	"bytes"
-	"io/ioutil"
-	"net/http/httputil"
-	"strings"
-
-	"github.com/projectdiscovery/rawhttp"
-)
-
-func Dump(req *HTTPRequest, reqURL string) ([]byte, error) {
-	if req.Request != nil {
-		// Create a copy on the fly of the request body - ignore errors
-		bodyBytes, _ := req.Request.BodyBytes()
-		req.Request.Request.Body = ioutil.NopCloser(bytes.NewReader(bodyBytes))
-		return httputil.DumpRequest(req.Request.Request, true)
-	}
-
-	return rawhttp.DumpRequestRaw(req.RawRequest.Method, reqURL, req.RawRequest.Path, ExpandMapValues(req.RawRequest.Headers), ioutil.NopCloser(strings.NewReader(req.RawRequest.Data)))
-}

v2/pkg/requests/util.go

@@ -1,67 +0,0 @@
-package requests
-
-import (
-	"bytes"
-	"compress/gzip"
-	"fmt"
-	"io/ioutil"
-	"strings"
-)
-
-func newReplacer(values map[string]interface{}) *strings.Replacer {
-	var replacerItems []string
-	for key, val := range values {
-		replacerItems = append(
-			replacerItems,
-			fmt.Sprintf("%s%s%s", markerParenthesisOpen, key, markerParenthesisClose),
-			fmt.Sprintf("%s", val),
-			fmt.Sprintf("%s%s%s", markerGeneral, key, markerGeneral),
-			fmt.Sprintf("%s", val),
-		)
-	}
-
-	return strings.NewReplacer(replacerItems...)
-}
-
-// HandleDecompression if the user specified a custom encoding (as golang transport doesn't do this automatically)
-func HandleDecompression(r *HTTPRequest, bodyOrig []byte) (bodyDec []byte, err error) {
-	if r.Request == nil {
-		return bodyOrig, nil
-	}
-
-	encodingHeader := strings.TrimSpace(strings.ToLower(r.Request.Header.Get("Accept-Encoding")))
-	if encodingHeader == "gzip" || encodingHeader == "gzip, deflate" {
-		gzipreader, err := gzip.NewReader(bytes.NewReader(bodyOrig))
-		if err != nil {
-			return bodyDec, err
-		}
-		defer gzipreader.Close()
-
-		bodyDec, err = ioutil.ReadAll(gzipreader)
-		if err != nil {
-			return bodyDec, err
-		}
-
-		return bodyDec, nil
-	}
-
-	return bodyOrig, nil
-}
-
-// ZipMapValues converts values from strings slices to flat string
-func ZipMapValues(m map[string][]string) (m1 map[string]string) {
-	m1 = make(map[string]string)
-	for k, v := range m {
-		m1[k] = strings.Join(v, "")
-	}
-	return
-}
-
-// ExpandMapValues converts values from flat string to strings slice
-func ExpandMapValues(m map[string]string) (m1 map[string][]string) {
-	m1 = make(map[string][]string)
-	for k, v := range m {
-		m1[k] = []string{v}
-	}
-	return
-}