From d68af67e6ea08c9df13f44af8e9fbd9dc378faf2 Mon Sep 17 00:00:00 2001 From: Dwi Siswanto Date: Mon, 14 Oct 2024 16:23:36 +0700 Subject: [PATCH] feat(nuclei): generate trace file when using `profile-mem` (#5690) * feat(nuclei): generate trace file when using `profile-mem` Signed-off-by: Dwi Siswanto * docs(DESIGN): dynamically grep mod path Signed-off-by: Dwi Siswanto --------- Signed-off-by: Dwi Siswanto --- .gitignore | 7 ++++++- DESIGN.md | 43 +++++++++++++++++++++++++++++++++---------- cmd/nuclei/main.go | 40 ++++++++++++++++++++++++++++++---------- 3 files changed, 69 insertions(+), 21 deletions(-) diff --git a/.gitignore b/.gitignore index 146b0a892c..6145103342 100644 --- a/.gitignore +++ b/.gitignore @@ -42,4 +42,9 @@ pkg/protocols/common/helpers/deserialization/testdata/ValueObject2.ser vendor # Headless `screenshot` action -*.png \ No newline at end of file +*.png + +# Profiling & tracing +*.prof +*.pprof +*.trace \ No newline at end of file diff --git a/DESIGN.md b/DESIGN.md index 9d92e28f8d..af0e27baf9 100644 --- a/DESIGN.md +++ b/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. diff --git a/cmd/nuclei/main.go b/cmd/nuclei/main.go index 8fe5b29573..caac7e5546 100644 --- a/cmd/nuclei/main.go +++ b/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"),