add pprof utils projectdiscovery/nuclei#4212

env/env.go

@@ -2,7 +2,9 @@ package env
 
 import (
 	"os"
+	"strconv"
 	"strings"
+	"time"
 )
 
 var (
@@ -20,3 +22,46 @@ func ExpandWithEnv(variables ...*string) {
 		*variable = os.Getenv(strings.TrimPrefix(*variable, "$"))
 	}
 }
+
+// EnvType is a type that can be used as a type for environment variables.
+type EnvType interface {
+	~string | ~int | ~bool | ~float64 | time.Duration
+}
+
+// GetEnvOrDefault returns the value of the environment variable or the default value if the variable is not set.
+// in requested type.
+func GetEnvOrDefault[T EnvType](key string, defaultValue T) T {
+	value := os.Getenv(key)
+	if value == "" {
+		return defaultValue
+	}
+	switch any(defaultValue).(type) {
+	case string:
+		return any(value).(T)
+	case int:
+		intVal, err := strconv.Atoi(value)
+		if err != nil || value == "" {
+			return defaultValue
+		}
+		return any(intVal).(T)
+	case bool:
+		boolVal, err := strconv.ParseBool(value)
+		if err != nil || value == "" {
+			return defaultValue
+		}
+		return any(boolVal).(T)
+	case float64:
+		floatVal, err := strconv.ParseFloat(value, 64)
+		if err != nil || value == "" {
+			return defaultValue
+		}
+		return any(floatVal).(T)
+	case time.Duration:
+		durationVal, err := time.ParseDuration(value)
+		if err != nil || value == "" {
+			return defaultValue
+		}
+		return any(durationVal).(T)
+	}
+	return defaultValue
+}

env/env_test.go

@@ -3,6 +3,7 @@ package env
 import (
 	"os"
 	"testing"
+	"time"
 )
 
 func TestExpandWithEnv(t *testing.T) {
@@ -42,3 +43,46 @@ func TestExpandWithEnvNilInput(t *testing.T) {
 	var nilVar *string = nil
 	ExpandWithEnv(nilVar)
 }
+
+func TestGetEnvOrDefault(t *testing.T) {
+	// Test for string
+	os.Setenv("TEST_STRING", "test")
+	resultString := GetEnvOrDefault("TEST_STRING", "default")
+	if resultString != "test" {
+		t.Errorf("Expected 'test', got %s", resultString)
+	}
+
+	// Test for int
+	os.Setenv("TEST_INT", "123")
+	resultInt := GetEnvOrDefault("TEST_INT", 0)
+	if resultInt != 123 {
+		t.Errorf("Expected 123, got %d", resultInt)
+	}
+
+	// Test for bool
+	os.Setenv("TEST_BOOL", "true")
+	resultBool := GetEnvOrDefault("TEST_BOOL", false)
+	if resultBool != true {
+		t.Errorf("Expected true, got %t", resultBool)
+	}
+
+	// Test for float64
+	os.Setenv("TEST_FLOAT", "1.23")
+	resultFloat := GetEnvOrDefault("TEST_FLOAT", 0.0)
+	if resultFloat != 1.23 {
+		t.Errorf("Expected 1.23, got %f", resultFloat)
+	}
+
+	// Test for time.Duration
+	os.Setenv("TEST_DURATION", "1h")
+	resultDuration := GetEnvOrDefault("TEST_DURATION", time.Duration(0))
+	if resultDuration != time.Hour {
+		t.Errorf("Expected 1h, got %s", resultDuration)
+	}
+
+	// Test for default value
+	resultDefault := GetEnvOrDefault("NON_EXISTING", "default")
+	if resultDefault != "default" {
+		t.Errorf("Expected 'default', got %s", resultDefault)
+	}
+}

pprof/pprof.go

@@ -0,0 +1,86 @@
+package pprof
+
+import (
+	"bytes"
+	"log"
+	"os"
+	"path/filepath"
+	"runtime"
+	"runtime/pprof"
+	"strconv"
+	"time"
+
+	"github.com/projectdiscovery/utils/env"
+)
+
+const (
+	PPROFSwitchENV = "PPROF"
+	MemProfileENV  = "MEM_PROFILE_DIR"
+	CPUProfileENV  = "CPU_PROFILE_DIR"
+	PPROFTimeENV   = "PPROF_TIME"
+	MemProfileRate = "MEM_PROFILE_RATE"
+)
+
+func init() {
+	if env.GetEnvOrDefault(PPROFSwitchENV, 0) == 1 {
+		log.Printf("GOOS: %v\n", runtime.GOOS)
+		log.Printf("GOARCH: %v\n", runtime.GOARCH)
+		log.Println("Available PPROF Config Options:")
+		log.Printf("%-16v - directory to write memory profiles to\n", MemProfileENV)
+		log.Printf("%-16v - directory to write cpu profiles to\n", CPUProfileENV)
+		log.Printf("%-16v - polling time for cpu and memory profiles (with unit ex: 10s)\n", PPROFTimeENV)
+		log.Printf("%-16v - memory profiling rate (default 4096)\n", MemProfileRate)
+
+		memProfilesDir := env.GetEnvOrDefault(MemProfileENV, "memdump")
+		cpuProfilesDir := env.GetEnvOrDefault(CPUProfileENV, "cpuprofile")
+		pprofTimeDuration := env.GetEnvOrDefault(PPROFTimeENV, time.Duration(3)*time.Second)
+		pprofRate := env.GetEnvOrDefault(MemProfileRate, 4096)
+
+		_ = os.MkdirAll(memProfilesDir, 0755)
+		_ = os.MkdirAll(cpuProfilesDir, 0755)
+
+		runtime.MemProfileRate = pprofRate
+		log.Printf("profile: memory profiling enabled (rate %d), %s\n", runtime.MemProfileRate, memProfilesDir)
+		log.Printf("profile: ticker enabled (rate %s)\n", pprofTimeDuration)
+
+		// cpu ticker and profiler
+		go func() {
+			ticker := time.NewTicker(pprofTimeDuration)
+			count := 0
+			buff := bytes.Buffer{}
+			log.Printf("profile: cpu profiling enabled (ticker %s)\n", pprofTimeDuration)
+			for {
+				err := pprof.StartCPUProfile(&buff)
+				if err != nil {
+					log.Fatalf("profile: could not start cpu profile: %s\n", err)
+				}
+				<-ticker.C
+				pprof.StopCPUProfile()
+				if err := os.WriteFile(filepath.Join(cpuProfilesDir, "cpuprofile-t"+strconv.Itoa(count)+".out"), buff.Bytes(), 0755); err != nil {
+					log.Fatalf("profile: could not write cpu profile: %s\n", err)
+				}
+				buff.Reset()
+				count++
+			}
+		}()
+
+		// memory ticker and profiler
+		go func() {
+			ticker := time.NewTicker(pprofTimeDuration)
+			count := 0
+			log.Printf("profile: memory profiling enabled (ticker %s)\n", pprofTimeDuration)
+			for {
+				<-ticker.C
+				var buff bytes.Buffer
+				if err := pprof.WriteHeapProfile(&buff); err != nil {
+					log.Printf("profile: could not write memory profile: %s\n", err)
+				}
+				err := os.WriteFile(filepath.ToSlash(filepath.Join(memProfilesDir, "memprofile-t"+strconv.Itoa(count)+".out")), buff.Bytes(), 0755)
+				if err != nil {
+					log.Printf("profile: could not write memory profile: %s\n", err)
+				}
+				count++
+			}
+		}()
+	}
+}