memory leak fixes and optimizations (#4680)

* feat http response memory optimization + reuse buffers

* update nuclei version

* feat: reuse js vm's and compile to programs

* fix failing http integration test

* remove dead code + add -jsc

* feat reuse js vms in pool with concurrency

* update comments as per review

* bug fix+ update interactsh test to look for dns interaction

* try enabling all interactsh integration tests

---------

Co-authored-by: mzack <[email protected]>

cmd/integration-test/interactsh.go

@@ -1,11 +1,9 @@
 package main
 
-import osutils "github.com/projectdiscovery/utils/os"
-
 // All Interactsh related testcases
 var interactshTestCases = []TestCaseInfo{
-	{Path: "protocols/http/interactsh.yaml", TestCase: &httpInteractshRequest{}, DisableOn: func() bool { return osutils.IsWindows() || osutils.IsOSX() }},
-	{Path: "protocols/http/interactsh-stop-at-first-match.yaml", TestCase: &httpInteractshStopAtFirstMatchRequest{}, DisableOn: func() bool { return true }},
-	{Path: "protocols/http/default-matcher-condition.yaml", TestCase: &httpDefaultMatcherCondition{}, DisableOn: func() bool { return true }}, // disable this test for now
+	{Path: "protocols/http/interactsh.yaml", TestCase: &httpInteractshRequest{}, DisableOn: func() bool { return false }},
+	{Path: "protocols/http/interactsh-stop-at-first-match.yaml", TestCase: &httpInteractshStopAtFirstMatchRequest{}, DisableOn: func() bool { return false }}, // disable this test for now
+	{Path: "protocols/http/default-matcher-condition.yaml", TestCase: &httpDefaultMatcherCondition{}, DisableOn: func() bool { return false }},
 	{Path: "protocols/http/interactsh-requests-mc-and.yaml", TestCase: &httpInteractshRequestsWithMCAnd{}},
 }

cmd/nuclei/main.go

@@ -319,6 +319,7 @@ on extensive configurability, massive extensibility and ease of use.`)
 		flagSet.IntVarP(&options.TemplateThreads, "concurrency", "c", 25, "maximum number of templates to be executed in parallel"),
 		flagSet.IntVarP(&options.HeadlessBulkSize, "headless-bulk-size", "hbs", 10, "maximum number of headless hosts to be analyzed in parallel per template"),
 		flagSet.IntVarP(&options.HeadlessTemplateThreads, "headless-concurrency", "headc", 10, "maximum number of headless templates to be executed in parallel"),
+		flagSet.IntVarP(&options.JsConcurrency, "js-concurrency", "jsc", 120, "maximum number of javascript runtimes to be executed in parallel"),
 	)
 	flagSet.CreateGroup("optimization", "Optimizations",
 		flagSet.IntVar(&options.Timeout, "timeout", 10, "time to wait in seconds before timeout"),

integration_tests/protocols/http/default-matcher-condition.yaml

@@ -15,7 +15,7 @@ requests:
       - type: word
         part: interactsh_protocol
         words:
-          - "http"
+          - "dns"
 
       - type: status
         status:

integration_tests/protocols/http/interactsh-stop-at-first-match.yaml

@@ -24,6 +24,6 @@ requests:
 
     matchers:
       - type: word
-        part: interactsh_protocol # Confirms the HTTP Interaction
+        part: interactsh_protocol # Confirms DNS Interaction
         words:
-          - "http"
+          - "dns"

integration_tests/protocols/http/interactsh.yaml

@@ -16,4 +16,4 @@ requests:
       - type: word
         part: interactsh_protocol # Confirms the HTTP Interaction
         words:
-          - "http"
+          - "dns"

pkg/catalog/config/constants.go

@@ -17,7 +17,7 @@ const (
 	CLIConfigFileName               = "config.yaml"
 	ReportingConfigFilename         = "reporting-config.yaml"
 	// Version is the current version of nuclei
-	Version = `v3.1.7`
+	Version = `v3.1.8-dev`
 	// Directory Names of custom templates
 	CustomS3TemplatesDirName     = "s3"
 	CustomGitHubTemplatesDirName = "github"

pkg/js/compiler/compiler.go

@@ -3,78 +3,34 @@ package compiler
 
 import (
 	"context"
-	"runtime/debug"
+	"fmt"
 	"time"
 
 	"github.com/dop251/goja"
-	"github.com/dop251/goja/parser"
-	"github.com/dop251/goja_nodejs/console"
-	"github.com/dop251/goja_nodejs/require"
-	jsoniter "github.com/json-iterator/go"
-	"github.com/pkg/errors"
-
-	"github.com/projectdiscovery/gologger"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libbytes"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libfs"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libikev2"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libkerberos"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libldap"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libmssql"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libmysql"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libnet"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/liboracle"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libpop3"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libpostgres"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/librdp"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libredis"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/librsync"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libsmb"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libsmtp"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libssh"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libstructs"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libtelnet"
-	_ "github.com/projectdiscovery/nuclei/v3/pkg/js/generated/go/libvnc"
-	"github.com/projectdiscovery/nuclei/v3/pkg/js/global"
-	"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
 // based javascript scripts efficiently while also
 // providing them access to custom modules defined in libs/.
-type Compiler struct {
-	registry *require.Registry
-}
+type Compiler struct{}
 
 // New creates a new compiler for the goja runtime.
 func New() *Compiler {
-	registry := new(require.Registry) // this can be shared by multiple runtimes
-	// autoregister console node module with default printer it uses gologger backend
-	require.RegisterNativeModule(console.ModuleName, console.RequireWithPrinter(goconsole.NewGoConsolePrinter()))
-	return &Compiler{registry: registry}
+	return &Compiler{}
 }
 
 // ExecuteOptions provides options for executing a script.
 type ExecuteOptions struct {
-	// Pool specifies whether to use a pool of goja runtimes
-	// Can be used to speedup execution but requires
-	// the script to not make any global changes.
-	Pool bool
-
-	// CaptureOutput specifies whether to capture the output
-	// of the script execution.
-	CaptureOutput bool
-
-	// CaptureVariables specifies the variables to capture
-	// from the script execution.
-	CaptureVariables []string
-
 	// Callback can be used to register new runtime helper functions
 	// ex: export etc
 	Callback func(runtime *goja.Runtime) error
 
+	// Cleanup is extra cleanup function to be called after execution
+	Cleanup func(runtime *goja.Runtime)
+
 	/// Timeout for this script execution
 	Timeout int
 }
@@ -111,51 +67,30 @@ func (e ExecuteResult) GetSuccess() bool {
 
 // Execute executes a script with the default options.
 func (c *Compiler) Execute(code string, args *ExecuteArgs) (ExecuteResult, error) {
-	return c.ExecuteWithOptions(code, args, &ExecuteOptions{})
-}
-
-// VM returns a new goja runtime for the compiler.
-func (c *Compiler) VM() *goja.Runtime {
-	runtime := c.newRuntime(false)
-	runtime.SetParserOptions(parser.WithDisableSourceMaps)
-	c.registerHelpersForVM(runtime)
-	return runtime
+	p, err := goja.Compile("", code, false)
+	if err != nil {
+		return nil, err
+	}
+	return c.ExecuteWithOptions(p, args, &ExecuteOptions{})
 }
 
 // ExecuteWithOptions executes a script with the provided options.
-func (c *Compiler) ExecuteWithOptions(code string, args *ExecuteArgs, opts *ExecuteOptions) (ExecuteResult, error) {
-	defer func() {
-		if err := recover(); err != nil {
-			gologger.Error().Msgf("Recovered panic %s %v: %v", code, args, err)
-			gologger.Verbose().Msgf("%s", debug.Stack())
-			return
-		}
-	}()
+func (c *Compiler) ExecuteWithOptions(program *goja.Program, args *ExecuteArgs, opts *ExecuteOptions) (ExecuteResult, error) {
 	if opts == nil {
 		opts = &ExecuteOptions{}
 	}
-	runtime := c.newRuntime(opts.Pool)
-	c.registerHelpersForVM(runtime)
-
-	// register runtime functions if any
-	if opts.Callback != nil {
-		if err := opts.Callback(runtime); err != nil {
-			return nil, err
-		}
-	}
-
 	if args == nil {
 		args = NewExecuteArgs()
 	}
-	for k, v := range args.Args {
-		_ = runtime.Set(k, v)
-	}
+	// handle nil maps
 	if args.TemplateCtx == nil {
 		args.TemplateCtx = make(map[string]interface{})
 	}
+	if args.Args == nil {
+		args.Args = make(map[string]interface{})
+	}
 	// merge all args into templatectx
 	args.TemplateCtx = generators.MergeMaps(args.TemplateCtx, args.Args)
-	_ = runtime.Set("template", args.TemplateCtx)
 
 	if opts.Timeout <= 0 || opts.Timeout > 180 {
 		// some js scripts can take longer time so allow configuring timeout
@@ -170,72 +105,13 @@ func (c *Compiler) ExecuteWithOptions(code string, args *ExecuteArgs, opts *Exec
 	results, err := contextutil.ExecFuncWithTwoReturns(ctx, func() (val goja.Value, err error) {
 		defer func() {
 			if r := recover(); r != nil {
-				err = errors.Errorf("panic: %v", r)
+				err = fmt.Errorf("panic: %v", r)
 			}
 		}()
-		return runtime.RunString(code)
+		return executeProgram(program, args, opts)
 	})
 	if err != nil {
 		return nil, err
 	}
-	captured := results.Export()
-
-	if opts.CaptureOutput {
-		return convertOutputToResult(captured)
-	}
-	if len(opts.CaptureVariables) > 0 {
-		return c.captureVariables(runtime, opts.CaptureVariables)
-	}
-	// success is true by default . since js throws errors on failure
-	// hence output result is always success
-	return ExecuteResult{"response": captured, "success": results.ToBoolean()}, nil
-}
-
-// captureVariables captures the variables from the runtime.
-func (c *Compiler) captureVariables(runtime *goja.Runtime, variables []string) (ExecuteResult, error) {
-	results := make(ExecuteResult, len(variables))
-	for _, variable := range variables {
-		value := runtime.Get(variable)
-		if value == nil {
-			continue
-		}
-		results[variable] = value.Export()
-	}
-	return results, nil
-}
-
-func convertOutputToResult(output interface{}) (ExecuteResult, error) {
-	marshalled, err := jsoniter.Marshal(output)
-	if err != nil {
-		return nil, errors.Wrap(err, "could not marshal output")
-	}
-
-	var outputMap map[string]interface{}
-	if err := jsoniter.Unmarshal(marshalled, &outputMap); err != nil {
-		var v interface{}
-		if unmarshalErr := jsoniter.Unmarshal(marshalled, &v); unmarshalErr != nil {
-			return nil, unmarshalErr
-		}
-		outputMap = map[string]interface{}{"output": v}
-		return outputMap, nil
-	}
-	return outputMap, nil
-}
-
-// newRuntime creates a new goja runtime
-// TODO: Add support for runtime reuse for helper functions
-func (c *Compiler) newRuntime(reuse bool) *goja.Runtime {
-	return protocolstate.NewJSRuntime()
-}
-
-// registerHelpersForVM registers all the helper functions for the goja runtime.
-func (c *Compiler) registerHelpersForVM(runtime *goja.Runtime) {
-	_ = c.registry.Enable(runtime)
-	// by default import below modules every time
-	_ = runtime.Set("console", require.Require(runtime, console.ModuleName))
-
-	// Register embedded scripts
-	if err := global.RegisterNativeScripts(runtime); err != nil {
-		gologger.Error().Msgf("Could not register scripts: %s\n", err)
-	}
+	return ExecuteResult{"response": results.Export(), "success": results.ToBoolean()}, nil
 }

pkg/js/compiler/compiler_test.go

@@ -38,36 +38,6 @@ func TestExecuteResultGetSuccess(t *testing.T) {
 	}
 }
 
-func TestCompilerCaptureVariables(t *testing.T) {
-	compiler := New()
-	result, err := compiler.ExecuteWithOptions("var a = 1;", NewExecuteArgs(), &ExecuteOptions{CaptureVariables: []string{"a"}})
-	if err != nil {
-		t.Fatal(err)
-	}
-	gotValue, ok := result["a"]
-	if !ok {
-		t.Fatalf("expected a to be present in the result")
-	}
-	if gotValue.(int64) != 1 {
-		t.Fatalf("expected a to be 1, got=%v", gotValue)
-	}
-}
-
-func TestCompilerCaptureOutput(t *testing.T) {
-	compiler := New()
-	result, err := compiler.ExecuteWithOptions("let obj = {'a':'b'}; obj", NewExecuteArgs(), &ExecuteOptions{CaptureOutput: true})
-	if err != nil {
-		t.Fatal(err)
-	}
-	gotValue, ok := result["a"]
-	if !ok {
-		t.Fatalf("expected a to be present in the result")
-	}
-	if gotValue.(string) != "b" {
-		t.Fatalf("expected a to be b, got=%v", gotValue)
-	}
-}
-
 type noopWriter struct {
 	Callback func(data []byte, level levels.Level)
 }

pkg/js/compiler/init.go

@@ -7,6 +7,7 @@ import "github.com/projectdiscovery/nuclei/v3/pkg/types"
 var (
 	// Per Execution Javascript timeout in seconds
 	JsProtocolTimeout = 10
+	JsVmConcurrency   = 500
 )
 
 // Init initializes the javascript protocol
@@ -15,6 +16,11 @@ func Init(opts *types.Options) error {
 		// keep existing 10s timeout
 		return nil
 	}
+	if opts.JsConcurrency < 100 {
+		// 100 is reasonable default
+		opts.JsConcurrency = 100
+	}
 	JsProtocolTimeout = opts.Timeout
+	JsVmConcurrency = opts.JsConcurrency
 	return nil
 }