code refactor + request dump fix

Fixes indirectly #844

v2/pkg/protocols/common/protocolinit/init.go

@@ -5,6 +5,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/protocolstate"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/dns/dnsclientpool"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/httpclientpool"
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/signerpool"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/network/networkclientpool"
 	"github.com/projectdiscovery/nuclei/v2/pkg/types"
 )
@@ -22,6 +23,9 @@ func Init(options *types.Options) error {
 	if err := httpclientpool.Init(options); err != nil {
 		return err
 	}
+	if err := signerpool.Init(options); err != nil {
+		return err
+	}
 	return networkclientpool.Init(options)
 }
 

v2/pkg/protocols/http/request.go

@@ -28,6 +28,8 @@ import (
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/interactsh"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/tostring"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/httpclientpool"
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/signer"
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/signerpool"
 	templateTypes "github.com/projectdiscovery/nuclei/v2/pkg/templates/types"
 	"github.com/projectdiscovery/nuclei/v2/pkg/types"
 	"github.com/projectdiscovery/rawhttp"
@@ -293,6 +295,7 @@ func (request *Request) executeRequest(reqURL string, generatedRequest *generate
 		err           error
 	)
 
+	// Dump request for variables checks
 	// For race conditions we can't dump the request body at this point as it's already waiting the open-gate event, already handled with a similar code within the race function
 	if !generatedRequest.original.Race {
 		var dumpError error
@@ -308,11 +311,6 @@ func (request *Request) executeRequest(reqURL string, generatedRequest *generate
 			gologger.Warning().Msgf("[%s] Could not make http request for %s: %v\n", request.options.TemplateID, reqURL, varErr)
 			return errStopExecution
 		}
-
-		if request.options.Options.Debug || request.options.Options.DebugRequests {
-			gologger.Info().Msgf("[%s] Dumped HTTP request for %s\n\n", request.options.TemplateID, reqURL)
-			gologger.Print().Msgf("%s", dumpedRequestString)
-		}
 	}
 
 	var formedURL string
@@ -352,39 +350,24 @@ func (request *Request) executeRequest(reqURL string, generatedRequest *generate
 		if resp == nil {
 			switch request.Signature.Value {
 			case AWSSignature:
-				var awsSigner *AwsSigner
+				var awsSigner signer.Signer
 				payloads := request.options.Options.Vars.AsMap()
 				awsAccessKeyId := types.ToString(payloads["aws-id"])
 				awsSecretAccessKey := types.ToString(payloads["aws-secret"])
-				if awsAccessKeyId != "" && awsSecretAccessKey != "" {
-					awsSigner, err = NewAwsSigner(awsAccessKeyId, awsSecretAccessKey)
-				} else {
-					// env variables
-					awsSigner, err = NewAwsSignerFromEnv()
-					if err != nil {
-						// $HOME/.aws/credentials
-						awsSigner, err = NewAwsSignerFromFile()
-						if err != nil {
-							return err
-						}
-					}
+				awsSignerArgs := signer.AwsSignerArgs{AwsId: awsAccessKeyId, AwsSecretToken: awsSecretAccessKey}
+				service := types.ToString(payloads["service"])
+				region := types.ToString(payloads["region"])
+				awsSignatureArguments := signer.AwsSignatureArguments{
+					Service: types.ToString(service),
+					Region:  types.ToString(region),
+					Time:    time.Now(),
 				}
+
+				awsSigner, err := signerpool.Get(request.options.Options, &signerpool.Configuration{SignerArgs: awsSignerArgs})
 				if err != nil {
 					return err
 				}
-
-				service := types.ToString(payloads["service"])
-				region := types.ToString(payloads["region"])
-				if service == "" || region == "" {
-					return errors.New("service and region are mandatory")
-				}
-
-				args := SignArguments{
-					Service: types.ToString(payloads["service"]),
-					Region:  types.ToString(payloads["region"]),
-					Time:    time.Now(),
-				}
-				err = awsSigner.SignHTTP(generatedRequest.request.Request, args)
+				err = awsSigner.SignHTTP(generatedRequest.request.Request, awsSignatureArguments)
 				if err != nil {
 					return err
 				}
@@ -392,6 +375,21 @@ func (request *Request) executeRequest(reqURL string, generatedRequest *generate
 			resp, err = request.httpClient.Do(generatedRequest.request)
 		}
 	}
+
+	// Dump the requests containing all headers
+	if !generatedRequest.original.Race {
+		var dumpError error
+		dumpedRequest, dumpError = dump(generatedRequest, reqURL)
+		if dumpError != nil {
+			return dumpError
+		}
+		dumpedRequestString := string(dumpedRequest)
+		if request.options.Options.Debug || request.options.Options.DebugRequests {
+			gologger.Info().Msgf("[%s] Dumped HTTP request for %s\n\n", request.options.TemplateID, reqURL)
+			gologger.Print().Msgf("%s", dumpedRequestString)
+		}
+	}
+
 	if err != nil {
 		// rawhttp doesn't support draining response bodies.
 		if resp != nil && resp.Body != nil && generatedRequest.rawRequest == nil {

v2/pkg/protocols/http/signer/aws.go

@@ -1,4 +1,4 @@
-package http
+package signer
 
 import (
 	"bytes"
@@ -17,21 +17,44 @@ type AwsSigner struct {
 	signer *v4.Signer
 }
 
-type SignArguments struct {
+type AwsSignerArgs struct {
+	AwsId          string
+	AwsSecretToken string
+}
+
+func (awsSignerArgs AwsSignerArgs) Validate() error {
+	if awsSignerArgs.AwsId == "" {
+		return errors.New("empty id")
+	}
+	if awsSignerArgs.AwsSecretToken == "" {
+		return errors.New("empty token")
+	}
+
+	return nil
+}
+
+type AwsSignatureArguments struct {
 	Service string
 	Region  string
 	Time    time.Time
 }
 
-func NewAwsSigner(awsId, awsSecretToken string) (*AwsSigner, error) {
-	if awsId == "" {
-		return nil, errors.New("empty id")
+func (awsSignatureArguments AwsSignatureArguments) Validate() error {
+	if awsSignatureArguments.Region == "" {
+		return errors.New("empty region")
 	}
-	if awsSecretToken == "" {
-		return nil, errors.New("empty token")
+	if awsSignatureArguments.Service == "" {
+		return errors.New("empty service")
 	}
 
-	creds := credentials.NewStaticCredentials(awsId, awsSecretToken, "")
+	return nil
+}
+
+func NewAwsSigner(args AwsSignerArgs) (*AwsSigner, error) {
+	if err := args.Validate(); err != nil {
+		return nil, err
+	}
+	creds := credentials.NewStaticCredentials(args.AwsId, args.AwsSecretToken, "")
 	if creds == nil {
 		return nil, errors.New("couldn't create the credentials structure")
 	}
@@ -57,7 +80,12 @@ func NewAwsSignerFromFile() (*AwsSigner, error) {
 	return &AwsSigner{creds: creds}, nil
 }
 
-func (awsSigner *AwsSigner) SignHTTP(request *http.Request, args SignArguments) error {
+func (awsSigner *AwsSigner) SignHTTP(request *http.Request, args interface{}) error {
+	signatureArgs, err := awsSigner.checkSignatureArgs(args)
+	if err != nil {
+		return err
+	}
+
 	awsSigner.prepareRequest(request)
 	var body *bytes.Reader
 	if request.Body != nil {
@@ -69,16 +97,21 @@ func (awsSigner *AwsSigner) SignHTTP(request *http.Request, args SignArguments)
 		body = bytes.NewReader(bodyBytes)
 	}
 
-	if _, err := awsSigner.signer.Sign(request, body, args.Service, args.Region, args.Time); err != nil {
+	if _, err := awsSigner.signer.Sign(request, body, signatureArgs.Service, signatureArgs.Region, signatureArgs.Time); err != nil {
 		return err
 	}
 	return nil
 }
 
-func (awsSigner *AwsSigner) CalculateHTTPHeaders(request *http.Request, args SignArguments) (map[string]string, error) {
+func (awsSigner *AwsSigner) CalculateHTTPHeaders(request *http.Request, args interface{}) (map[string]string, error) {
+	signatureArgs, err := awsSigner.checkSignatureArgs(args)
+	if err != nil {
+		return nil, err
+	}
+
 	reqClone := request.Clone(context.Background())
 	awsSigner.prepareRequest(reqClone)
-	err := awsSigner.SignHTTP(reqClone, args)
+	err = awsSigner.SignHTTP(reqClone, signatureArgs)
 	if err != nil {
 		return nil, err
 	}
@@ -88,6 +121,13 @@ func (awsSigner *AwsSigner) CalculateHTTPHeaders(request *http.Request, args Sig
 	return headers, nil
 }
 
+func (awsSigner *AwsSigner) checkSignatureArgs(args interface{}) (AwsSignatureArguments, error) {
+	if signatureArgs, ok := args.(AwsSignatureArguments); ok {
+		return signatureArgs, signatureArgs.Validate()
+	}
+	return AwsSignatureArguments{}, errors.New("wrong signature type")
+}
+
 func (awsSigner *AwsSigner) prepareRequest(request *http.Request) {
 	request.Header.Del("Host")
 }

v2/pkg/protocols/http/signer/signer.go

@@ -0,0 +1,40 @@
+package signer
+
+import (
+	"errors"
+	"net/http"
+)
+
+type Signer interface {
+	SignHTTP(request *http.Request, args interface{}) error
+	CalculateHTTPHeaders(request *http.Request, args interface{}) (map[string]string, error)
+}
+
+type SignerArgs interface {
+	Validate() error
+}
+
+type SignatureArguments interface {
+	Validate() error
+}
+
+func NewSigner(args SignerArgs) (signer Signer, err error) {
+	switch signerArgs := args.(type) {
+	case AwsSignerArgs:
+		awsSigner, err := NewAwsSigner(signerArgs)
+		if err != nil {
+			// env variables
+			awsSigner, err = NewAwsSignerFromEnv()
+			if err != nil {
+				// $HOME/.aws/credentials
+				awsSigner, err = NewAwsSignerFromFile()
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+		return awsSigner, err
+	default:
+		return nil, errors.New("unknown signature arguments type")
+	}
+}

v2/pkg/protocols/http/signerpool/signerpool.go

@@ -0,0 +1,57 @@
+package signerpool
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/signer"
+
+	"github.com/projectdiscovery/nuclei/v2/pkg/types"
+)
+
+var (
+	poolMutex  *sync.RWMutex
+	clientPool map[string]signer.Signer
+)
+
+// Init initializes the clientpool implementation
+func Init(options *types.Options) error {
+	poolMutex = &sync.RWMutex{}
+	clientPool = make(map[string]signer.Signer)
+	return nil
+}
+
+// Configuration contains the custom configuration options for a client
+type Configuration struct {
+	SignerArgs signer.SignerArgs
+}
+
+// Hash returns the hash of the configuration to allow client pooling
+func (c *Configuration) Hash() string {
+	builder := &strings.Builder{}
+	builder.WriteString(fmt.Sprintf("%v", c.SignerArgs))
+	hash := builder.String()
+	return hash
+}
+
+// Get creates or gets a client for the protocol based on custom configuration
+func Get(options *types.Options, configuration *Configuration) (signer.Signer, error) {
+	hash := configuration.Hash()
+	poolMutex.RLock()
+	if client, ok := clientPool[hash]; ok {
+		poolMutex.RUnlock()
+		return client, nil
+	}
+	poolMutex.RUnlock()
+
+	client, err := signer.NewSigner(configuration.SignerArgs)
+	if err != nil {
+		return nil, err
+	}
+
+	poolMutex.Lock()
+	clientPool[hash] = client
+	poolMutex.Unlock()
+	return client, nil
+}