diff --git a/env/env.go b/env/env.go index 76358b0..69ed094 100644 --- a/env/env.go +++ b/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 +} diff --git a/env/env_test.go b/env/env_test.go index 80c5cd6..dbc31c5 100644 --- a/env/env_test.go +++ b/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) + } +} diff --git a/pprof/README.md b/pprof/README.md new file mode 100644 index 0000000..783597e --- /dev/null +++ b/pprof/README.md @@ -0,0 +1,52 @@ +## PProfiling Usage Guide + +### Environment Variables + +- `PPROF`: Enable or disable profiling. Set to 1 to enable. +- `MEM_PROFILE_DIR`: Directory to write memory profiles to. +- `CPU_PROFILE_DIR`: Directory to write CPU profiles to. +- `PPROF_TIME`: Polling time for CPU and memory profiles (with unit ex: 10s). +- `MEM_PROFILE_RATE`: Memory profiling rate (default 4096). + + +## How to Use + +1. Set the environment variables as per your requirements. + +```bash +export PPROF=1 +export MEM_PROFILE_DIR=/path/to/memprofile +export CPU_PROFILE_DIR=/path/to/cpuprofile +export PPROF_TIME=10s +export MEM_PROFILE_RATE=4096 +``` + +2. Run your Go application. The profiler will start automatically if PPROF is set to 1. + +**Output** + +- Memory profiles will be written to the directory specified by MEM_PROFILE_DIR. +- CPU profiles will be written to the directory specified by CPU_PROFILE_DIR. +- Profiles will be written at intervals specified by PPROF_TIME. +- Memory profiling rate is controlled by MEM_PROFILE_RATE. + +### Example + +```bash +[+] GOOS: linux +[+] GOARCH: amd64 +[+] Command: /path/to/your/app +Available PPROF Config Options: +MEM_PROFILE_DIR - directory to write memory profiles to +CPU_PROFILE_DIR - directory to write cpu profiles to +PPROF_TIME - polling time for cpu and memory profiles (with unit ex: 10s) +MEM_PROFILE_RATE - memory profiling rate (default 4096) +profile: memory profiling enabled (rate 4096), /path/to/memprofile +profile: ticker enabled (rate 10s) +profile: cpu profiling enabled (ticker 10s) +``` + +### Note + +- The polling time (PPROF_TIME) should be set according to your application's performance and profiling needs. +- The memory profiling rate (MEM_PROFILE_RATE) controls the granularity of the memory profiling. Higher values provide more detail but consume more resources. \ No newline at end of file diff --git a/pprof/pprof.go b/pprof/pprof.go new file mode 100644 index 0000000..653b9ad --- /dev/null +++ b/pprof/pprof.go @@ -0,0 +1,88 @@ +package pprof + +import ( + "bytes" + "log" + "os" + "path/filepath" + "runtime" + "runtime/pprof" + "strconv" + "strings" + "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.Printf("[+] Command: %v\n", strings.Join(os.Args, " ")) + 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++ + } + }() + } +}