Added hang monitor for goroutine dumping (#1949)

* Added hang monitor for goroutine dumping

* misc

* Made hang monitor optional with flag

* Added stack comparison for monitoring + misc

* Removed debug statements

* misc update

Co-authored-by: sandeep <sandeep@projectdiscovery.io>

v2/cmd/nuclei/main.go

@@ -17,6 +17,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http"
 	templateTypes "github.com/projectdiscovery/nuclei/v2/pkg/templates/types"
 	"github.com/projectdiscovery/nuclei/v2/pkg/types"
+	"github.com/projectdiscovery/nuclei/v2/pkg/utils/monitor"
 )
 
 var (
@@ -33,6 +34,11 @@ func main() {
 
 	runner.ParseOptions(options)
 
+	if options.HangMonitor {
+		cancel := monitor.NewStackMonitor(10 * time.Second)
+		defer cancel()
+	}
+
 	nucleiRunner, err := runner.New(options)
 	if err != nil {
 		gologger.Fatal().Msgf("Could not create runner: %s\n", err)
@@ -198,6 +204,7 @@ on extensive configurability, massive extensibility and ease of use.`)
 		flagSet.StringVarP(&options.TraceLogFile, "trace-log", "tlog", "", "file to write sent requests trace log"),
 		flagSet.StringVarP(&options.ErrorLogFile, "error-log", "elog", "", "file to write sent requests error log"),
 		flagSet.BoolVar(&options.Version, "version", false, "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.BoolVar(&options.VerboseVerbose, "vv", false, "display templates loaded for scan"),
 		flagSet.BoolVarP(&options.EnablePprof, "enable-pprof", "ep", false, "enable pprof debugging server"),

v2/go.mod

@@ -86,6 +86,7 @@ require (
 
 require (
 	git.mills.io/prologic/smtpd v0.0.0-20210710122116-a525b76c287a // indirect
+	github.com/DataDog/gostackparse v0.5.0 // indirect
 	github.com/Mzack9999/go-http-digest-auth-client v0.6.1-0.20220414142836-eb8883508809 // indirect
 	github.com/Mzack9999/ldapserver v1.0.2-0.20211229000134-b44a0d6ad0dd // indirect
 	github.com/PuerkitoBio/goquery v1.6.0 // indirect

v2/go.sum

@@ -36,6 +36,8 @@ git.mills.io/prologic/smtpd v0.0.0-20210710122116-a525b76c287a/go.mod h1:C7hXLmF
 github.com/AndreasBriese/bbloom v0.0.0-20190825152654-46b345b51c96/go.mod h1:bOvUY6CB00SOBii9/FifXqc0awNKxLFCL/+pkDPuyl8=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
+github.com/DataDog/gostackparse v0.5.0 h1:jb72P6GFHPHz2W0onsN51cS3FkaMDcjb0QzgxxA4gDk=
+github.com/DataDog/gostackparse v0.5.0/go.mod h1:lTfqcJKqS9KnXQGnyQMCugq3u1FP6UZMfWR0aitKFMM=
 github.com/Ice3man543/nvd v1.0.8/go.mod h1:0DxLJk6revOcJKiZxa2K+rNF/HO1zJO97lqQtXhXfSc=
 github.com/Knetic/govaluate v3.0.0+incompatible/go.mod h1:r7JcOSlj0wfOMncg0iLm8Leh48TZaKVeNIfJntJ2wa0=
 github.com/Knetic/govaluate v3.0.1-0.20171022003610-9aa49832a739+incompatible h1:1G1pk05UrOh0NlF1oeaaix1x8XzrfjIDK47TY0Zehcw=

v2/pkg/types/types.go

@@ -174,6 +174,8 @@ type Options struct {
 	TemplatesVersion bool
 	// TemplateList lists available templates
 	TemplateList bool
+	// HangMonitor enables nuclei hang monitoring
+	HangMonitor bool
 	// Stdin specifies whether stdin input was given to the process
 	Stdin bool
 	// StopAtFirstMatch stops processing template at first full match (this may break chained requests)

v2/pkg/utils/monitor/monitor.go

@@ -0,0 +1,139 @@
+// Package monitor implements a goroutine based monitoring for
+// detecting stuck scanner processes and dumping stack and other
+// relevant information for investigation.
+package monitor
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"runtime"
+	"strings"
+	"time"
+
+	"github.com/DataDog/gostackparse"
+	"github.com/projectdiscovery/gologger"
+	"github.com/rs/xid"
+)
+
+// Agent is an agent for monitoring hanging programs
+type Agent struct {
+	cancel    context.CancelFunc
+	lastStack []string
+
+	goroutineCount   int
+	currentIteration int // number of times we've checked hang
+}
+
+const defaultMonitorIteration = 6
+
+// NewStackMonitor returns a new stack monitor instance
+func NewStackMonitor(interval time.Duration) context.CancelFunc {
+	ctx, cancel := context.WithCancel(context.Background())
+	ticker := time.NewTicker(interval)
+
+	monitor := &Agent{cancel: cancel}
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				ticker.Stop()
+			case <-ticker.C:
+				monitor.monitorWorker()
+			default:
+				continue
+			}
+		}
+	}()
+	return cancel
+}
+
+// monitorWorker is a worker for monitoring running goroutines
+func (s *Agent) monitorWorker() {
+	current := runtime.NumGoroutine()
+	if current != s.goroutineCount {
+		s.goroutineCount = current
+		s.currentIteration = 0
+		return
+	}
+	s.currentIteration++
+
+	if s.currentIteration == defaultMonitorIteration-1 {
+		lastStackTrace := generateStackTraceSlice(getStack(true))
+		s.lastStack = lastStackTrace
+		return
+	}
+
+	// cancel the monitoring goroutine if we discover
+	// we've been stuck for some iterations.
+	if s.currentIteration == defaultMonitorIteration {
+		currentStack := getStack(true)
+
+		// Bail out if the stacks don't match from previous iteration
+		newStack := generateStackTraceSlice(currentStack)
+		if !compareStringSliceEqual(s.lastStack, newStack) {
+			s.currentIteration = 0
+			return
+		}
+		s.cancel()
+		stackTraceFile := fmt.Sprintf("nuclei-stacktrace-%s.dump", xid.New().String())
+		gologger.Error().Msgf("Detected hanging goroutine (count=%d/%d) = %s\n", current, s.goroutineCount, stackTraceFile)
+		if err := ioutil.WriteFile(stackTraceFile, currentStack, os.ModePerm); err != nil {
+			gologger.Error().Msgf("Could not write stack trace for goroutines: %s\n", err)
+		}
+		os.Exit(1) // exit forcefully if we've been stuck
+	}
+}
+
+// getStack returns full stack trace of the program
+var getStack = func(all bool) []byte {
+	for i := 1024 * 1024; ; i *= 2 {
+		buf := make([]byte, i)
+		if n := runtime.Stack(buf, all); n < i {
+			return buf[:n-1]
+		}
+	}
+}
+
+// generateStackTraceSlice returns a list of current stack in string slice format
+func generateStackTraceSlice(stack []byte) []string {
+	goroutines, _ := gostackparse.Parse(bytes.NewReader(stack))
+
+	var builder strings.Builder
+	var stackList []string
+	for _, goroutine := range goroutines {
+		builder.WriteString(goroutine.State)
+		builder.WriteString("|")
+
+		for _, frame := range goroutine.Stack {
+			builder.WriteString(frame.Func)
+			builder.WriteString(";")
+		}
+		stackList = append(stackList, builder.String())
+		builder.Reset()
+	}
+	return stackList
+}
+
+// compareStringSliceEqual compares two string slices for equality without order
+func compareStringSliceEqual(first, second []string) bool {
+	if len(first) != len(second) {
+		return false
+	}
+	diff := make(map[string]int, len(first))
+	for _, x := range first {
+		diff[x]++
+	}
+	for _, y := range second {
+		if _, ok := diff[y]; !ok {
+			return false
+		}
+		diff[y] -= 1
+		if diff[y] == 0 {
+			delete(diff, y)
+		}
+	}
+	return len(diff) == 0
+}

v2/pkg/utils/monitor/monitor_test.go

@@ -0,0 +1,15 @@
+package monitor
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMonitorCompareStringSliceEqual(t *testing.T) {
+	value := compareStringSliceEqual([]string{"a", "b"}, []string{"b", "a"})
+	require.True(t, value, "could not get correct value")
+
+	value = compareStringSliceEqual([]string{"a", "c"}, []string{"b", "a"})
+	require.False(t, value, "could get incorrect value")
+}