misc improvements in js protocol execution (#4643)

* js protocol  timeout using -timeout flag

* fix zgrab smb hang

* fix lint error

* custom timeout field in js protocol

* minor update: bound checking

* add 6 * -timeout in code protocol by default

pkg/js/compiler/compiler.go

@@ -2,7 +2,9 @@
 package compiler
 
 import (
+	"context"
 	"runtime/debug"
+	"time"
 
 	"github.com/dop251/goja"
 	"github.com/dop251/goja/parser"
@@ -36,6 +38,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/js/libs/goconsole"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/generators"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate"
+	contextutil "github.com/projectdiscovery/utils/context"
 )
 
 // Compiler provides a runtime to execute goja runtime
@@ -71,6 +74,9 @@ type ExecuteOptions struct {
 	// Callback can be used to register new runtime helper functions
 	// ex: export etc
 	Callback func(runtime *goja.Runtime) error
+
+	/// Timeout for this script execution
+	Timeout int
 }
 
 // ExecuteArgs is the arguments to pass to the script.
@@ -151,7 +157,19 @@ func (c *Compiler) ExecuteWithOptions(code string, args *ExecuteArgs, opts *Exec
 	args.TemplateCtx = generators.MergeMaps(args.TemplateCtx, args.Args)
 	_ = runtime.Set("template", args.TemplateCtx)
 
-	results, err := runtime.RunString(code)
+	if opts.Timeout <= 0 || opts.Timeout > 180 {
+		// some js scripts can take longer time so allow configuring timeout
+		// from template but keep it within sane limits (180s)
+		opts.Timeout = JsProtocolTimeout
+	}
+
+	// execute with context and timeout
+	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(opts.Timeout)*time.Second)
+	defer cancel()
+	// execute the script
+	results, err := contextutil.ExecFuncWithTwoReturns(ctx, func() (goja.Value, error) {
+		return runtime.RunString(code)
+	})
 	if err != nil {
 		return nil, err
 	}

pkg/js/compiler/init.go

@@ -0,0 +1,20 @@
+package compiler
+
+import "github.com/projectdiscovery/nuclei/v3/pkg/types"
+
+// jsprotocolInit
+
+var (
+	// Per Execution Javascript timeout in seconds
+	JsProtocolTimeout = 10
+)
+
+// Init initializes the javascript protocol
+func Init(opts *types.Options) error {
+	if opts.Timeout < 10 {
+		// keep existing 10s timeout
+		return nil
+	}
+	JsProtocolTimeout = opts.Timeout
+	return nil
+}

pkg/js/libs/smb/smb.go

@@ -3,7 +3,6 @@ package smb
 import (
 	"context"
 	"fmt"
-	"net"
 	"time"
 
 	"github.com/hirochachacha/go-smb2"
@@ -24,26 +23,30 @@ type SMBClient struct{}
 // Returns handshake log and error. If error is not nil,
 // state will be false
 func (c *SMBClient) ConnectSMBInfoMode(host string, port int) (*smb.SMBLog, error) {
+	if !protocolstate.IsHostAllowed(host) {
+		// host is not valid according to network policy
+		return nil, protocolstate.ErrHostDenied.Msgf(host)
+	}
 	conn, err := protocolstate.Dialer.Dial(context.TODO(), "tcp", fmt.Sprintf("%s:%d", host, port))
 	if err != nil {
 		return nil, err
 	}
-	defer conn.Close()
+	// try to get SMBv2/v3 info
+	result, err := c.getSMBInfo(conn, true, false)
+	_ = conn.Close() // close regardless of error
+	if err == nil {
+		return result, nil
+	}
 
-	_ = conn.SetDeadline(time.Now().Add(10 * time.Second))
+	// try to negotiate SMBv1
-	setupSession := true
+	conn, err = protocolstate.Dialer.Dial(context.TODO(), "tcp", fmt.Sprintf("%s:%d", host, port))
-
-	result, err := smb.GetSMBLog(conn, setupSession, false, false)
 	if err != nil {
-		conn.Close()
+		return nil, err
-		conn, err = net.DialTimeout("tcp", fmt.Sprintf("%s:%d", host, port), 10*time.Second)
+	}
-		if err != nil {
+	defer conn.Close()
-			return nil, err
+	result, err = c.getSMBInfo(conn, true, true)
-		}
+	if err != nil {
-		result, err = smb.GetSMBLog(conn, setupSession, true, false)
+		return result, nil
-		if err != nil {
-			return nil, err
-		}
 	}
 	return result, nil
 }
@@ -67,6 +70,10 @@ func (c *SMBClient) ListSMBv2Metadata(host string, port int) (*plugins.ServiceSM
 // Credentials cannot be blank. guest or anonymous credentials
 // can be used by providing empty password.
 func (c *SMBClient) ListShares(host string, port int, user, password string) ([]string, error) {
+	if !protocolstate.IsHostAllowed(host) {
+		// host is not valid according to network policy
+		return nil, protocolstate.ErrHostDenied.Msgf(host)
+	}
 	conn, err := protocolstate.Dialer.Dial(context.TODO(), "tcp", fmt.Sprintf("%s:%d", host, port))
 	if err != nil {
 		return nil, err

pkg/js/libs/smb/smb_private.go

@@ -9,8 +9,11 @@ import (
 	"github.com/praetorian-inc/fingerprintx/pkg/plugins"
 	"github.com/praetorian-inc/fingerprintx/pkg/plugins/services/smb"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate"
+	zgrabsmb "github.com/zmap/zgrab2/lib/smb/smb"
 )
 
+// ==== private helper functions/methods ====
+
 // collectSMBv2Metadata collects metadata for SMBv2 services.
 func collectSMBv2Metadata(host string, port int, timeout time.Duration) (*plugins.ServiceSMB, error) {
 	if timeout == 0 {
@@ -28,3 +31,17 @@ func collectSMBv2Metadata(host string, port int, timeout time.Duration) (*plugin
 	}
 	return metadata, nil
 }
+
+// getSMBInfo
+func (c *SMBClient) getSMBInfo(conn net.Conn, setupSession, v1 bool) (*zgrabsmb.SMBLog, error) {
+	_ = conn.SetDeadline(time.Now().Add(10 * time.Second))
+	defer func() {
+		_ = conn.SetDeadline(time.Time{})
+	}()
+
+	result, err := zgrabsmb.GetSMBLog(conn, setupSession, v1, false)
+	if err != nil {
+		return nil, err
+	}
+	return result, nil
+}

pkg/js/libs/smb/smbghost.go

@@ -20,6 +20,10 @@ const (
 // DetectSMBGhost tries to detect SMBGhost vulnerability
 // by using SMBv3 compression feature.
 func (c *SMBClient) DetectSMBGhost(host string, port int) (bool, error) {
+	if !protocolstate.IsHostAllowed(host) {
+		// host is not valid according to network policy
+		return false, protocolstate.ErrHostDenied.Msgf(host)
+	}
 	addr := net.JoinHostPort(host, strconv.Itoa(port))
 	conn, err := protocolstate.Dialer.Dial(context.TODO(), "tcp", addr)
 	if err != nil {

pkg/protocols/code/code.go

@@ -1,6 +1,7 @@
 package code
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"regexp"
@@ -26,11 +27,13 @@ import (
 	protocolutils "github.com/projectdiscovery/nuclei/v3/pkg/protocols/utils"
 	templateTypes "github.com/projectdiscovery/nuclei/v3/pkg/templates/types"
 	"github.com/projectdiscovery/nuclei/v3/pkg/types"
+	contextutil "github.com/projectdiscovery/utils/context"
 	errorutil "github.com/projectdiscovery/utils/errors"
 )
 
 const (
-	pythonEnvRegex = `os\.getenv\(['"]([^'"]+)['"]\)`
+	pythonEnvRegex    = `os\.getenv\(['"]([^'"]+)['"]\)`
+	TimeoutMultiplier = 6 // timeout multiplier for code protocol
 )
 
 var (
@@ -121,12 +124,17 @@ func (request *Request) GetID() string {
 }
 
 // ExecuteWithResults executes the protocol requests and returns results instead of writing them.
-func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicValues, previous output.InternalEvent, callback protocols.OutputEventCallback) error {
+func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicValues, previous output.InternalEvent, callback protocols.OutputEventCallback) (err error) {
 	metaSrc, err := gozero.NewSourceWithString(input.MetaInput.Input, "")
 	if err != nil {
 		return err
 	}
 	defer func() {
+		// catch any panics just in case
+		if r := recover(); r != nil {
+			gologger.Error().Msgf("[%s] Panic occurred in code protocol: %s\n", request.options.TemplateID, r)
+			err = fmt.Errorf("panic occurred: %s", r)
+		}
 		if err := metaSrc.Cleanup(); err != nil {
 			gologger.Warning().Msgf("%s\n", err)
 		}
@@ -150,9 +158,24 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa
 		allvars[name] = v
 		metaSrc.AddVariable(gozerotypes.Variable{Name: name, Value: v})
 	}
-	gOutput, err := request.gozero.Eval(context.Background(), request.src, metaSrc)
+	timeout := TimeoutMultiplier * request.options.Options.Timeout
-	if err != nil && gOutput == nil {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
-		return errorutil.NewWithErr(err).Msgf("[%s] Could not execute code on local machine %v", request.options.TemplateID, input.MetaInput.Input)
+	defer cancel()
+	// Note: we use contextutil despite the fact that gozero accepts context as argument
+	gOutput, err := contextutil.ExecFuncWithTwoReturns(ctx, func() (*gozerotypes.Result, error) {
+		return request.gozero.Eval(ctx, request.src, metaSrc)
+	})
+	if gOutput == nil {
+		// write error to stderr buff
+		var buff bytes.Buffer
+		if err != nil {
+			buff.WriteString(err.Error())
+		} else {
+			buff.WriteString("no output something went wrong")
+		}
+		gOutput = &gozerotypes.Result{
+			Stderr: buff,
+		}
 	}
 	gologger.Verbose().Msgf("[%s] Executed code on local machine %v", request.options.TemplateID, input.MetaInput.Input)
 

pkg/protocols/common/protocolinit/init.go

@@ -3,6 +3,7 @@ package protocolinit
 import (
 	"github.com/corpix/uarand"
 
+	"github.com/projectdiscovery/nuclei/v3/pkg/js/compiler"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/dns/dnsclientpool"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/http/httpclientpool"
@@ -34,6 +35,9 @@ func Init(options *types.Options) error {
 	if err := rdapclientpool.Init(options); err != nil {
 		return err
 	}
+	if err := compiler.Init(options); err != nil {
+		return err
+	}
 	return nil
 }
 

pkg/protocols/javascript/js.go

@@ -61,7 +61,9 @@ type Request struct {
 	// description: |
 	//   Code contains code to execute for the javascript request.
 	Code string `yaml:"code,omitempty" json:"code,omitempty" jsonschema:"title=code to execute in javascript,description=Executes inline javascript code for the request"`
-
+	// description: |
+	//   Timeout in seconds is optional timeout for each  javascript script execution (i.e init, pre-condition, code)
+	Timeout int `yaml:"timeout,omitempty" json:"timeout,omitempty" jsonschema:"title=timeout for javascript execution,description=Timeout in seconds is optional timeout for entire javascript script execution"`
 	// description: |
 	//   StopAtFirstMatch stops processing the request at first match.
 	StopAtFirstMatch bool `yaml:"stop-at-first-match,omitempty" json:"stop-at-first-match,omitempty" jsonschema:"title=stop at first match,description=Stop the execution after a match is found"`
@@ -141,7 +143,9 @@ func (request *Request) Compile(options *protocols.ExecutorOptions) error {
 			prettyPrint(request.TemplateID, buff.String())
 		}
 
-		opts := &compiler.ExecuteOptions{}
+		opts := &compiler.ExecuteOptions{
+			Timeout: request.Timeout,
+		}
 		// register 'export' function to export variables from init code
 		// these are saved in args and are available in pre-condition and request code
 		opts.Callback = func(runtime *goja.Runtime) error {
@@ -303,7 +307,7 @@ func (request *Request) ExecuteWithResults(target *contextargs.Context, dynamicV
 		}
 		argsCopy.TemplateCtx = templateCtx.GetAll()
 
-		result, err := request.options.JsCompiler.ExecuteWithOptions(request.PreCondition, argsCopy, nil)
+		result, err := request.options.JsCompiler.ExecuteWithOptions(request.PreCondition, argsCopy, &compiler.ExecuteOptions{Timeout: request.Timeout})
 		if err != nil {
 			return errorutil.NewWithTag(request.TemplateID, "could not execute pre-condition: %s", err)
 		}
@@ -426,7 +430,8 @@ func (request *Request) executeRequestWithPayloads(hostPort string, input *conte
 	}
 
 	results, err := request.options.JsCompiler.ExecuteWithOptions(string(requestData), argsCopy, &compiler.ExecuteOptions{
-		Pool: false,
+		Pool:    false,
+		Timeout: request.Timeout,
 	})
 	if err != nil {
 		// shouldn't fail even if it returned error instead create a failure event