update instances and add codeexectimeout

projectdiscovery · Jun 7, 2024 · 32a73bf · 32a73bf
1 parent b3eb134
commit 32a73bf
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 13 deletions.
diff --git a/pkg/js/compiler/compiler.go b/pkg/js/compiler/compiler.go
@@ -4,11 +4,11 @@ package compiler
 import (
 	"context"
 	"fmt"
-	"time"
 
 	"github.com/dop251/goja"
 
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/generators"
+	"github.com/projectdiscovery/nuclei/v3/pkg/types"
 	contextutil "github.com/projectdiscovery/utils/context"
 	"github.com/projectdiscovery/utils/errkit"
 	stringsutil "github.com/projectdiscovery/utils/strings"
@@ -45,6 +45,8 @@ type ExecuteOptions struct {
 
 	Context context.Context
 
+	TimeoutVariants types.TimeoutVariants
+
 	// Manually exported objects
 	exports map[string]interface{}
 }
@@ -113,7 +115,8 @@ func (c *Compiler) ExecuteWithOptions(program *goja.Program, args *ExecuteArgs,
 	}
 
 	// execute with context and timeout
-	ctx, cancel := context.WithTimeoutCause(opts.Context, time.Duration(opts.Timeout)*time.Second, ErrJSExecDeadline)
+
+	ctx, cancel := context.WithTimeoutCause(opts.Context, opts.TimeoutVariants.JsCompilerExecutionTimeout, ErrJSExecDeadline)
 	defer cancel()
 	// execute the script
 	results, err := contextutil.ExecFuncWithTwoReturns(ctx, func() (val goja.Value, err error) {

diff --git a/pkg/protocols/code/code.go b/pkg/protocols/code/code.go
@@ -199,11 +199,12 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa
 
 		result, err := request.options.JsCompiler.ExecuteWithOptions(request.preConditionCompiled, args,
 			&compiler.ExecuteOptions{
-				Timeout:  timeout,
-				Source:   &request.PreCondition,
-				Callback: registerPreConditionFunctions,
-				Cleanup:  cleanUpPreConditionFunctions,
-				Context:  input.Context(),
+				Timeout:         timeout,
+				TimeoutVariants: request.options.TimeoutVariants,
+				Source:          &request.PreCondition,
+				Callback:        registerPreConditionFunctions,
+				Cleanup:         cleanUpPreConditionFunctions,
+				Context:         input.Context(),
 			})
 		if err != nil {
 			return errorutil.NewWithTag(request.TemplateID, "could not execute pre-condition: %s", err)
@@ -218,7 +219,7 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa
 		}
 	}
 
-	ctx, cancel := context.WithTimeoutCause(input.Context(), time.Duration(timeout)*time.Second, ErrCodeExecutionDeadline)
+	ctx, cancel := context.WithTimeoutCause(input.Context(), request.options.TimeoutVariants.CodeExecutionTimeout, ErrCodeExecutionDeadline)
 	defer cancel()
 	// Note: we use contextutil despite the fact that gozero accepts context as argument
 	gOutput, err := contextutil.ExecFuncWithTwoReturns(ctx, func() (*gozerotypes.Result, error) {

diff --git a/pkg/protocols/javascript/js.go b/pkg/protocols/javascript/js.go
@@ -153,9 +153,10 @@ func (request *Request) Compile(options *protocols.ExecutorOptions) error {
 		}
 
 		opts := &compiler.ExecuteOptions{
-			Timeout: request.Timeout,
-			Source:  &request.Init,
-			Context: context.Background(),
+			Timeout:         request.Timeout,
+			TimeoutVariants: request.options.TimeoutVariants,
+			Source:          &request.Init,
+			Context:         context.Background(),
 		}
 		// register 'export' function to export variables from init code
 		// these are saved in args and are available in pre-condition and request code
@@ -345,7 +346,11 @@ func (request *Request) ExecuteWithResults(target *contextargs.Context, dynamicV
 		argsCopy.TemplateCtx = templateCtx.GetAll()
 
 		result, err := request.options.JsCompiler.ExecuteWithOptions(request.preConditionCompiled, argsCopy,
-			&compiler.ExecuteOptions{Timeout: request.Timeout, Source: &request.PreCondition, Context: target.Context()})
+			&compiler.ExecuteOptions{
+				Timeout:         request.Timeout,
+				TimeoutVariants: requestOptions.TimeoutVariants,
+				Source:          &request.PreCondition, Context: target.Context(),
+			})
 		if err != nil {
 			return errorutil.NewWithTag(request.TemplateID, "could not execute pre-condition: %s", err)
 		}
@@ -500,7 +505,12 @@ func (request *Request) executeRequestWithPayloads(hostPort string, input *conte
 	}
 
 	results, err := request.options.JsCompiler.ExecuteWithOptions(request.scriptCompiled, argsCopy,
-		&compiler.ExecuteOptions{Timeout: request.Timeout, Source: &request.Code, Context: input.Context()})
+		&compiler.ExecuteOptions{
+			Timeout:         request.Timeout,
+			TimeoutVariants: requestOptions.TimeoutVariants,
+			Source:          &request.Code,
+			Context:         input.Context(),
+		})
 	if err != nil {
 		// shouldn't fail even if it returned error instead create a failure event
 		results = compiler.ExecuteResult{"success": false, "error": err.Error()}

diff --git a/pkg/types/types.go b/pkg/types/types.go
@@ -401,6 +401,7 @@ type TimeoutVariants struct {
 	MaxResponseHeaderTimeout   time.Duration
 	ResponseReadTimeout        time.Duration
 	JsCompilerExecutionTimeout time.Duration
+	CodeExecutionTimeout       time.Duration
 	HttpTimeout                time.Duration
 	DialTimeout                time.Duration
 }
@@ -414,6 +415,7 @@ func (options *Options) BuildTimeoutVariants() TimeoutVariants {
 		//response read timeout in seconds
 		ResponseReadTimeout:        time.Second * 5,
 		JsCompilerExecutionTimeout: time.Second * time.Duration(int(float64(options.Timeout)*1.5)),
+		CodeExecutionTimeout:       time.Second * 3,
 		//http timeout for the client
 		HttpTimeout: time.Second * time.Duration(options.Timeout*3),
 		//timeout for network requests