feat(nuclei): generate trace file when using profile-mem (#5690)

* feat(nuclei): generate trace file when using `profile-mem`

Signed-off-by: Dwi Siswanto <[email protected]>

* docs(DESIGN): dynamically grep mod path

Signed-off-by: Dwi Siswanto <[email protected]>

---------

Signed-off-by: Dwi Siswanto <[email protected]>

Author: dwisiswant0

.gitignore

@@ -42,4 +42,9 @@ pkg/protocols/common/helpers/deserialization/testdata/ValueObject2.ser
 vendor
 
 # Headless `screenshot` action
-*.png
+*.png
+
+# Profiling & tracing
+*.prof
+*.pprof
+*.trace

DESIGN.md

@@ -457,26 +457,49 @@ func (template *Template) compileProtocolRequests(options protocols.ExecuterOpti
 That's it, you've added a new protocol to Nuclei. The next good step would be to write integration tests which are described in `integration-tests` and `cmd/integration-tests` directories.
 
 
-## Profiling Instructions
+## Profiling and Tracing
 
-To enable dumping of Memory profiling data, `-profile-mem` flag can be used along with path to a file. This writes a pprof formatted file which can be used for investigate resource usage with `pprof` tool.
+To analyze Nuclei's performance and resource usage, you can generate memory profiles and trace files using the `-profile-mem` flag:
 
-```console
-$ nuclei -t nuclei-templates/ -u https://example.com -profile-mem mem.pprof
+```bash
+nuclei -t nuclei-templates/ -u https://example.com -profile-mem=nuclei-$(git describe --tags)
 ```
 
-To view profile data in pprof, first install pprof. Then run the below command -
+This command creates two files:
 
-```console
-$ go tool pprof mem.pprof
+* `nuclei.prof`: Memory (heap) profile
+* `nuclei.trace`: Execution trace
+
+### Analyzing the Memory Profile
+
+1. View the profile in the terminal:
+
+```bash
+go tool pprof nuclei.prof
+```
+
+2. Display top memory consumers:
+
+```bash
+go tool pprof -top nuclei.prof | grep "$(go list -m)" | head -10
 ```
 
-To open a web UI on a port to visualize debug data, the below command can be used.
+3. Visualize the profile in a web browser:
 
-```console
-$ go tool pprof -http=:8081 mem.pprof
+```bash
+go tool pprof -http=:$(shuf -i 1000-99999 -n 1) nuclei.prof
 ```
 
+### Analyzing the Trace File
+
+To examine the execution trace:
+
+```bash
+go tool trace nuclei.trace
+```
+
+These tools help identify performance bottlenecks and memory leaks, allowing for targeted optimizations of Nuclei's codebase.
+
 ## Project Structure
 
 - [pkg/reporting](./pkg/reporting) - Reporting modules for nuclei.

cmd/nuclei/main.go

@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"runtime/pprof"
+	"runtime/trace"
 	"strings"
 	"time"
 
@@ -103,21 +104,40 @@ func main() {
 		return
 	}
 
-	// Profiling related code
+	// Profiling & tracing related code
 	if memProfile != "" {
-		f, err := os.Create(memProfile)
+		memProfile = strings.TrimSuffix(memProfile, filepath.Ext(memProfile)) + ".prof"
+		memProfileFile, err := os.Create(memProfile)
 		if err != nil {
-			gologger.Fatal().Msgf("profile: could not create memory profile %q: %v", memProfile, err)
+			gologger.Fatal().Msgf("profile: could not create memory profile %q file: %v", memProfile, err)
 		}
-		old := runtime.MemProfileRate
+
+		traceFilepath := strings.TrimSuffix(memProfile, filepath.Ext(memProfile)) + ".trace"
+		traceFile, err := os.Create(traceFilepath)
+		if err != nil {
+			gologger.Fatal().Msgf("profile: could not create trace %q file: %v", traceFilepath, err)
+		}
+
+		oldMemProfileRate := runtime.MemProfileRate
 		runtime.MemProfileRate = 4096
-		gologger.Print().Msgf("profile: memory profiling enabled (rate %d), %s", runtime.MemProfileRate, memProfile)
+
+		// Start tracing
+		if err := trace.Start(traceFile); err != nil {
+			gologger.Fatal().Msgf("profile: could not start trace: %v", err)
+		}
 
 		defer func() {
-			_ = pprof.Lookup("heap").WriteTo(f, 0)
-			f.Close()
-			runtime.MemProfileRate = old
-			gologger.Print().Msgf("profile: memory profiling disabled, %s", memProfile)
+			// Start CPU profiling
+			if err := pprof.WriteHeapProfile(memProfileFile); err != nil {
+				gologger.Fatal().Msgf("profile: could not start CPU profile: %v", err)
+			}
+			memProfileFile.Close()
+			traceFile.Close()
+			trace.Stop()
+			runtime.MemProfileRate = oldMemProfileRate
+
+			gologger.Info().Msgf("Memory profile saved at %q", memProfile)
+			gologger.Info().Msgf("Traced at %q", traceFilepath)
 		}()
 	}
 
@@ -402,7 +422,7 @@ on extensive configurability, massive extensibility and ease of use.`)
 		flagSet.CallbackVar(printVersion, "version", "show nuclei version"),
 		flagSet.BoolVarP(&options.HangMonitor, "hang-monitor", "hm", false, "enable nuclei hang monitoring"),
 		flagSet.BoolVarP(&options.Verbose, "verbose", "v", false, "show verbose output"),
-		flagSet.StringVar(&memProfile, "profile-mem", "", "optional nuclei memory profile dump file"),
+		flagSet.StringVar(&memProfile, "profile-mem", "", "generate memory (heap) profile & trace files"),
 		flagSet.BoolVar(&options.VerboseVerbose, "vv", false, "display templates loaded for scan"),
 		flagSet.BoolVarP(&options.ShowVarDump, "show-var-dump", "svd", false, "show variables dump for debugging"),
 		flagSet.BoolVarP(&options.EnablePprof, "enable-pprof", "ep", false, "enable pprof debugging server"),