diff --git a/README.md b/README.md index ecdd3d52..c5d66bc1 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,23 @@ Supported firewalls: - ipset only (IPv4 :heavy_check_mark: / IPv6 :heavy_check_mark: ) - pf (IPV4 :heavy_check_mark: / IPV6 :heavy_check_mark: ) +## Profiling + +On-demand Go **pprof** endpoints and optional automatic heap dumps are controlled only by environment variables (no YAML changes). The pprof server uses a separate listener from Prometheus so routes are not mixed with `/metrics`. + +| Variable | Meaning | +|----------|---------| +| `CS_PROFILING_ENABLED` | When set to `true`, starts the pprof HTTP server. | +| `CS_PROFILING_ADDR` | Listen address for pprof (default `:6060`). | +| `CS_PROFILING_HEAP_DUMP_DIR` | If set to a non-empty path, runs a background watcher that can write heap profiles to this directory when memory is high. Independent of `CS_PROFILING_ENABLED`. | +| `CS_PROFILING_HEAP_DUMP_THRESHOLD_MB` | Heap allocation threshold in mebibytes before a dump (default `200`). | +| `CS_PROFILING_HEAP_POLL_INTERVAL` | How often heap use is checked, as a Go duration (default `30s`). | +| `CS_PROFILING_HEAP_DUMP_COOLDOWN` | Minimum time between successful heap dumps, as a Go duration (default `5m`). | + +For garbage-collection tracing, set **`GODEBUG=gctrace=1`** in the container environment before the process starts (the Go runtime reads this at startup). + +Protect the pprof port with network policy or bind to loopback only (`127.0.0.1:6060`) where appropriate; profiling endpoints expose sensitive in-process data. + # Installation Please follow the [official documentation](https://doc.crowdsec.net/docs/bouncers/firewall). diff --git a/cmd/root.go b/cmd/root.go index 9a7a48b5..6fe1df48 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -28,6 +28,7 @@ import ( "github.com/crowdsecurity/cs-firewall-bouncer/pkg/backend" "github.com/crowdsecurity/cs-firewall-bouncer/pkg/cfg" "github.com/crowdsecurity/cs-firewall-bouncer/pkg/metrics" + "github.com/crowdsecurity/cs-firewall-bouncer/pkg/profiling" ) const bouncerType = "crowdsec-firewall-bouncer" @@ -250,6 +251,9 @@ func Execute() error { }() } + profiling.StartPprofServerIfEnabled() + profiling.StartHeapWatcherIfEnabled(ctx) + g.Go(func() error { log.Infof("Processing new and deleted decisions . . .") diff --git a/go.mod b/go.mod index a5acf8ac..614d2b77 100644 --- a/go.mod +++ b/go.mod @@ -7,6 +7,7 @@ require ( github.com/crowdsecurity/go-cs-bouncer v0.0.21 github.com/crowdsecurity/go-cs-lib v0.0.25 github.com/google/nftables v0.3.0 + github.com/google/pprof v0.0.0-20260507013755-92041b743c96 github.com/prometheus/client_golang v1.23.2 github.com/prometheus/client_model v0.6.2 github.com/sirupsen/logrus v1.9.4 diff --git a/go.sum b/go.sum index fec0a302..6d125393 100644 --- a/go.sum +++ b/go.sum @@ -51,6 +51,8 @@ github.com/google/go-querystring v1.1.0 h1:AnCroh3fv4ZBgVIf1Iwtovgjaw/GiKJo8M8yD github.com/google/go-querystring v1.1.0/go.mod h1:Kcdr2DB4koayq7X8pmAG4sNG59So17icRSOU623lUBU= github.com/google/nftables v0.3.0 h1:bkyZ0cbpVeMHXOrtlFc8ISmfVqq5gPJukoYieyVmITg= github.com/google/nftables v0.3.0/go.mod h1:BCp9FsrbF1Fn/Yu6CLUc9GGZFw/+hsxfluNXXmxBfRM= +github.com/google/pprof v0.0.0-20260507013755-92041b743c96 h1:YDDnaZ9afWajDboPMt9Vikqca/yWAX7KAxVzb4lJU1M= +github.com/google/pprof v0.0.0-20260507013755-92041b743c96/go.mod h1:MxpfABSjhmINe3F1It9d+8exIHFvUqtLIRCdOGNXqiI= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= diff --git a/pkg/profiling/profiling.go b/pkg/profiling/profiling.go new file mode 100644 index 00000000..446536b9 --- /dev/null +++ b/pkg/profiling/profiling.go @@ -0,0 +1,209 @@ +// Package profiling exposes an optional HTTP pprof server and heap auto-dump +// tooling, gated by CS_PROFILING_* environment variables. +// +// Pprof routes are registered only on a dedicated [http.ServeMux] (not via +// `_ "net/http/pprof"`), so handlers are not attached to [http.DefaultServeMux] +// alongside the Prometheus `/metrics` server. +package profiling + +import ( + "compress/gzip" + "context" + "fmt" + "net" + "net/http" + urlpprof "net/http/pprof" + "os" + "path/filepath" + "runtime" + runtimepprof "runtime/pprof" + "strconv" + "strings" + "time" + + log "github.com/sirupsen/logrus" +) + +const ( + envProfilingEnabled = "CS_PROFILING_ENABLED" + envProfilingAddr = "CS_PROFILING_ADDR" + envHeapDumpDir = "CS_PROFILING_HEAP_DUMP_DIR" + envHeapDumpThresholdMB = "CS_PROFILING_HEAP_DUMP_THRESHOLD_MB" + envHeapPollInterval = "CS_PROFILING_HEAP_POLL_INTERVAL" + envHeapDumpCooldown = "CS_PROFILING_HEAP_DUMP_COOLDOWN" + defaultProfilingAddr = ":6060" + defaultHeapThresholdMiB uint64 = 200 + defaultHeapPollInterval = 30 * time.Second + defaultHeapCooldown = 5 * time.Minute +) + +func registerPprofHandlers(mux *http.ServeMux) { + mux.HandleFunc("/debug/pprof/cmdline", urlpprof.Cmdline) + mux.HandleFunc("/debug/pprof/profile", urlpprof.Profile) + mux.HandleFunc("/debug/pprof/symbol", urlpprof.Symbol) + mux.HandleFunc("/debug/pprof/trace", urlpprof.Trace) + mux.HandleFunc("/debug/pprof/", urlpprof.Index) +} + +type heapConfig struct { + thresholdBytes uint64 + pollInterval time.Duration + cooldown time.Duration +} + +func parseHeapConfig() heapConfig { + var thresholdMiB uint64 = defaultHeapThresholdMiB + if v := strings.TrimSpace(os.Getenv(envHeapDumpThresholdMB)); v != "" { + if parsed, err := strconv.ParseUint(v, 10, 64); err == nil && parsed > 0 { + thresholdMiB = parsed + } else if err != nil { + log.Warningf("heap watcher: invalid %s=%q (using default %d MiB): %v", envHeapDumpThresholdMB, v, defaultHeapThresholdMiB, err) + } + } + thresholdBytes := thresholdMiB * 1024 * 1024 + + pollInterval := defaultHeapPollInterval + if v := strings.TrimSpace(os.Getenv(envHeapPollInterval)); v != "" { + if d, err := time.ParseDuration(v); err == nil && d > 0 { + pollInterval = d + } else if err != nil { + log.Warningf("heap watcher: invalid %s=%q (using default %s): %v", envHeapPollInterval, v, defaultHeapPollInterval, err) + } + } + + cooldown := defaultHeapCooldown + if v := strings.TrimSpace(os.Getenv(envHeapDumpCooldown)); v != "" { + if d, err := time.ParseDuration(v); err == nil && d > 0 { + cooldown = d + } else if err != nil { + log.Warningf("heap watcher: invalid %s=%q (using default %s): %v", envHeapDumpCooldown, v, defaultHeapCooldown, err) + } + } + + return heapConfig{ + thresholdBytes: thresholdBytes, + pollInterval: pollInterval, + cooldown: cooldown, + } +} + +// Start binds a dedicated HTTP server for /debug/pprof/* on addr and runs it in +// a background goroutine. Listen errors are logged and do not stop the +// process. Returns nil after the listener is accepted (or after logging a bind +// failure). +func Start(addr string) error { + mux := http.NewServeMux() + registerPprofHandlers(mux) + + ln, err := net.Listen("tcp", addr) + if err != nil { + log.Errorf("pprof server: failed to listen on %s: %v", addr, err) + return nil + } + + startOnListener(ln, mux) + return nil +} + +func startOnListener(ln net.Listener, mux *http.ServeMux) *http.Server { + log.Infof("pprof server listening on %s (set GODEBUG=gctrace=1 in the environment for GC trace output from the runtime)", ln.Addr().String()) + + srv := &http.Server{Handler: mux} + go func() { + if serveErr := srv.Serve(ln); serveErr != nil && serveErr != http.ErrServerClosed { + log.Errorf("pprof server: %v", serveErr) + } + }() + return srv +} + +// StartPprofServerIfEnabled starts the pprof server when CS_PROFILING_ENABLED is "true". +func StartPprofServerIfEnabled() { + if !strings.EqualFold(strings.TrimSpace(os.Getenv(envProfilingEnabled)), "true") { + return + } + addr := strings.TrimSpace(os.Getenv(envProfilingAddr)) + if addr == "" { + addr = defaultProfilingAddr + } + _ = Start(addr) +} + +// StartHeapWatcher runs a poll loop that writes heap profiles when HeapAlloc crosses +// the configured threshold (subject to cooldown). Reads configuration from the +// environment on each invocation; does nothing when CS_PROFILING_HEAP_DUMP_DIR is empty. +func StartHeapWatcher(ctx context.Context) { + dir := strings.TrimSpace(os.Getenv(envHeapDumpDir)) + if dir == "" { + return + } + + cfg := parseHeapConfig() + go heapWatcherLoop(ctx, dir, cfg.thresholdBytes, cfg.pollInterval, cfg.cooldown, nil) +} + +func heapWatcherLoop(ctx context.Context, dir string, thresholdBytes uint64, pollInterval, cooldown time.Duration, nowFn func() time.Time) { + if nowFn == nil { + nowFn = time.Now + } + + ticker := time.NewTicker(pollInterval) + defer ticker.Stop() + + var lastDump time.Time + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + var ms runtime.MemStats + runtime.ReadMemStats(&ms) + + if ms.HeapAlloc < thresholdBytes { + continue + } + if !lastDump.IsZero() && nowFn().Sub(lastDump) < cooldown { + continue + } + + ts := strings.ReplaceAll(nowFn().UTC().Format(time.RFC3339), ":", "-") + filename := fmt.Sprintf("heap-%s.pb.gz", ts) + fullPath := filepath.Join(dir, filename) + + if err := writeHeapProfileGZ(fullPath); err != nil { + log.Errorf("heap watcher: failed to write heap profile to %s: %v", fullPath, err) + continue + } + + lastDump = nowFn() + log.Infof("heap watcher: wrote heap profile to %s (HeapAlloc=%d bytes, threshold=%d bytes)", fullPath, ms.HeapAlloc, thresholdBytes) + } + } +} + +func writeHeapProfileGZ(path string) error { + if err := os.MkdirAll(filepath.Dir(path), 0o750); err != nil { + return err + } + f, err := os.Create(path) + if err != nil { + return err + } + defer f.Close() + + gw := gzip.NewWriter(f) + if err := runtimepprof.WriteHeapProfile(gw); err != nil { + _ = gw.Close() + return err + } + return gw.Close() +} + +// StartHeapWatcherIfEnabled starts the heap watcher when CS_PROFILING_HEAP_DUMP_DIR is non-empty. +func StartHeapWatcherIfEnabled(ctx context.Context) { + if strings.TrimSpace(os.Getenv(envHeapDumpDir)) == "" { + return + } + StartHeapWatcher(ctx) +} diff --git a/pkg/profiling/profiling_test.go b/pkg/profiling/profiling_test.go new file mode 100644 index 00000000..35c5c43c --- /dev/null +++ b/pkg/profiling/profiling_test.go @@ -0,0 +1,204 @@ +package profiling + +import ( + "compress/gzip" + "context" + "net" + "net/http" + "os" + "path/filepath" + "testing" + "time" + + "github.com/google/pprof/profile" + "github.com/stretchr/testify/require" +) + +func TestWriteHeapProfileGZ_CreatesValidFile(t *testing.T) { + t.Parallel() + dir := t.TempDir() + path := filepath.Join(dir, "heap.pb.gz") + require.NoError(t, writeHeapProfileGZ(path)) + + _, err := os.Stat(path) + require.NoError(t, err) + + f, err := os.Open(path) + require.NoError(t, err) + defer f.Close() + + gz, err := gzip.NewReader(f) + require.NoError(t, err) + defer gz.Close() + + _, err = profile.Parse(gz) + require.NoError(t, err) +} + +func TestWriteHeapProfileGZ_CreatesIntermediateDirs(t *testing.T) { + t.Parallel() + dir := t.TempDir() + path := filepath.Join(dir, "nested", "sub", "heap.pb.gz") + require.NoError(t, writeHeapProfileGZ(path)) + + _, err := os.Stat(path) + require.NoError(t, err) +} + +func TestStartPprofServerIfEnabled(t *testing.T) { + t.Run("does nothing when disabled", func(t *testing.T) { + t.Setenv(envProfilingEnabled, "") + t.Setenv(envProfilingAddr, "") + + StartPprofServerIfEnabled() + + ln, err := net.Listen("tcp", "127.0.0.1:6060") + require.NoError(t, err) + require.NoError(t, ln.Close()) + }) + + t.Run("default listen address", func(t *testing.T) { + t.Setenv(envProfilingEnabled, "true") + t.Setenv(envProfilingAddr, "") + + StartPprofServerIfEnabled() + + conn, err := net.DialTimeout("tcp", "127.0.0.1:6060", 2*time.Second) + require.NoError(t, err) + require.NoError(t, conn.Close()) + }) + + t.Run("custom listen address", func(t *testing.T) { + free, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + addr := free.Addr().String() + require.NoError(t, free.Close()) + + t.Setenv(envProfilingEnabled, "true") + t.Setenv(envProfilingAddr, addr) + + StartPprofServerIfEnabled() + + conn, err := net.DialTimeout("tcp", addr, 2*time.Second) + require.NoError(t, err) + require.NoError(t, conn.Close()) + }) +} + +func TestPprofServer_ServesHeapEndpoint(t *testing.T) { + mux := http.NewServeMux() + registerPprofHandlers(mux) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + srv := startOnListener(ln, mux) + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _ = srv.Shutdown(ctx) + }) + + client := &http.Client{Timeout: 5 * time.Second} + resp, err := client.Get("http://" + ln.Addr().String() + "/debug/pprof/heap") + require.NoError(t, err) + t.Cleanup(func() { _ = resp.Body.Close() }) + require.Equal(t, http.StatusOK, resp.StatusCode) +} + +func TestParseHeapConfig_Defaults(t *testing.T) { + t.Setenv(envHeapDumpThresholdMB, "") + t.Setenv(envHeapPollInterval, "") + t.Setenv(envHeapDumpCooldown, "") + + cfg := parseHeapConfig() + require.Equal(t, defaultHeapThresholdMiB*1024*1024, cfg.thresholdBytes) + require.Equal(t, defaultHeapPollInterval, cfg.pollInterval) + require.Equal(t, defaultHeapCooldown, cfg.cooldown) +} + +func TestParseHeapConfig_InvalidValuesFallBackToDefaults(t *testing.T) { + t.Setenv(envHeapDumpThresholdMB, "not-a-number") + t.Setenv(envHeapPollInterval, "30-smurfs") + t.Setenv(envHeapDumpCooldown, "5-forever") + + cfg := parseHeapConfig() + require.Equal(t, defaultHeapThresholdMiB*1024*1024, cfg.thresholdBytes) + require.Equal(t, defaultHeapPollInterval, cfg.pollInterval) + require.Equal(t, defaultHeapCooldown, cfg.cooldown) +} + +func TestHeapWatcherLoop_DumpsWhenThresholdExceeded(t *testing.T) { + dir := t.TempDir() + hold := make([]byte, 4<<20) + _ = hold + + ctx, cancel := context.WithTimeout(context.Background(), 800*time.Millisecond) + defer cancel() + + fixed := time.Date(2021, 3, 4, 5, 6, 7, 0, time.UTC) + nowFn := func() time.Time { return fixed } + + go heapWatcherLoop(ctx, dir, 1, 10*time.Millisecond, time.Hour, nowFn) + + require.Eventually(t, func() bool { + entries, err := os.ReadDir(dir) + return err == nil && len(entries) > 0 + }, 600*time.Millisecond, 20*time.Millisecond, "expected a heap dump file") +} + +func TestHeapWatcherLoop_CooldownPreventsRepeatDump(t *testing.T) { + dir := t.TempDir() + hold := make([]byte, 4<<20) + _ = hold + + ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) + defer cancel() + + t0 := time.Date(2022, 4, 5, 6, 7, 8, 0, time.UTC) + nowFn := func() time.Time { return t0 } + + go heapWatcherLoop(ctx, dir, 1, 10*time.Millisecond, time.Hour, nowFn) + + require.Eventually(t, func() bool { + entries, err := os.ReadDir(dir) + return err == nil && len(entries) >= 1 + }, 250*time.Millisecond, 20*time.Millisecond) + + time.Sleep(100 * time.Millisecond) + + entries, err := os.ReadDir(dir) + require.NoError(t, err) + require.Len(t, entries, 1) +} + +func TestHeapWatcherLoop_ExitsOnContextCancel(t *testing.T) { + dir := t.TempDir() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + done := make(chan struct{}) + go func() { + heapWatcherLoop(ctx, dir, 1, time.Second, time.Hour, time.Now) + close(done) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("heap watcher loop did not exit after context cancel") + } + + entries, err := os.ReadDir(dir) + require.NoError(t, err) + require.Empty(t, entries) +} + +func TestStartHeapWatcherIfEnabled_DoesNothingWithoutDir(t *testing.T) { + t.Setenv(envHeapDumpDir, "") + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + StartHeapWatcherIfEnabled(ctx) +}