echo: protect profiles by IPecho/v0.16.3

2 files changed, 64 insertions, 34 deletions

diff --git a/echo/echo_default.go b/echo/echo_default.go
index dadeb90..5cd4b75 100644
--- a/echo/echo_default.go
+++ b/echo/echo_default.go
@@ -82,6 +82,7 @@ type EchoWrapper struct {
         autocert          autocert.PrimingCertProvider
         templateFS        fs.FS
         errorHandler      ErrorHandler
+        netboxIpFilter    *glmw.NetboxIPFilter
 }
 
 // NewEchoWrapper creates a new instance of Echo and wraps it in an
@@ -163,10 +164,6 @@ func (w *EchoWrapper) AddErrorHandler(h ContentErrorHandler, mime ...string) {
 }
 
 func (w *EchoWrapper) Configure(c EchoConfig) (err error) {
-        if w.Debug || c.EnableTrace {
-                glmw.RegisterPprof(w.Echo)
-        }
-
         w.configureAutocert(&c)
 
         if err := w.configureIpExtractor(&c); err != nil {
@@ -197,7 +194,11 @@ func (w *EchoWrapper) Configure(c EchoConfig) (err error) {
         w.configureCORS(&c)
         w.configureCSP(&c)
         w.configureServerHeader(&c)
+
+        // These all depend on w.netboxIpFilter
+        w.configureNetboxIpFilter(&c)
         w.configurePrometheus(&c)
+        w.configurePprof(&c)
 
         return nil
 }
@@ -404,19 +405,21 @@ func (w *EchoWrapper) configureServerHeader(c *EchoConfig) {
         }
 }
 
-func (w *EchoWrapper) configurePrometheus(c *EchoConfig) {
-        if !c.DisablePrometheus && c.NetboxClient != nil {
-                // TODO: Should constrain this by site probably but monitoring happens
-                // across sites so those prefixes need to be included
-                f := &glmw.NetboxIPFilter{
-                        NetboxClient:     c.NetboxClient,
-                        Tag:              "management",
-                        IncludeLocalhost: true,
-                        Logger:           w.Logger,
-                }
-                w.runner.AddInitJob(f.Init)
-                w.runner.AddJob(f.RunRefresh)
+func (w *EchoWrapper) configureNetboxIpFilter(c *EchoConfig) {
+        // TODO: Should constrain this by site probably but monitoring happens
+        // across sites so those prefixes need to be included
+        w.netboxIpFilter = &glmw.NetboxIPFilter{
+                NetboxClient:     c.NetboxClient,
+                Tag:              "management",
+                IncludeLocalhost: true,
+                Logger:           w.Logger,
+        }
+        w.runner.AddInitJob(w.netboxIpFilter.Init)
+        w.runner.AddJob(w.netboxIpFilter.RunRefresh)
+}
 
+func (w *EchoWrapper) configurePrometheus(c *EchoConfig) {
+        if !c.DisablePrometheus && w.netboxIpFilter != nil {
                 var prom *prometheus.Prometheus
                 if c.PrometheusConfig != nil {
                         prom = prometheus.NewPrometheusWithConfig(c.PrometheusConfig)
@@ -425,6 +428,16 @@ func (w *EchoWrapper) configurePrometheus(c *EchoConfig) {
                 }
 
                 w.Use(prom.MiddlewareHandler)
-                w.GET(prom.Config.MetricsPath, prom.MetricsHandler, f.Middleware)
+                w.GET(prom.Config.MetricsPath, prom.MetricsHandler, w.netboxIpFilter.Middleware)
+        }
+}
+
+func (w *EchoWrapper) configurePprof(c *EchoConfig) {
+        if (w.Debug || c.EnableTrace) && w.netboxIpFilter != nil {
+                glmw.RegisterPprofWithConfig(
+                        w.Echo,
+                        glmw.DefaultPprofConfig,
+                        w.netboxIpFilter.Middleware,
+                )
         }
 }
diff --git a/echo/middleware/pprof.go b/echo/middleware/pprof.go
index cc69505..8c7fe1b 100644
--- a/echo/middleware/pprof.go
+++ b/echo/middleware/pprof.go
@@ -5,28 +5,45 @@ import (
         "net/http/pprof"
 
         "github.com/labstack/echo/v4"
+        "github.com/labstack/echo/v4/middleware"
 )
 
-func RegisterPprof(e *echo.Echo, prefixOptions ...string) {
-        prefix := "/debug/pprof"
-        if len(prefixOptions) > 0 {
-                prefix = prefixOptions[0]
+type PprofConfig struct {
+        Skipper middleware.Skipper
+        Prefix  string
+}
+
+var DefaultPprofConfig = PprofConfig{
+        Skipper: middleware.DefaultSkipper,
+        Prefix:  "/debug/pprof",
+}
+
+func RegisterPprof(e *echo.Echo) {
+        RegisterPprofWithConfig(e, DefaultPprofConfig)
+}
+
+func RegisterPprofWithConfig(e *echo.Echo, config PprofConfig, mw ...echo.MiddlewareFunc) {
+        if config.Skipper == nil {
+                config.Skipper = DefaultPprofConfig.Skipper
+        }
+        if config.Prefix == "" {
+                config.Prefix = DefaultPprofConfig.Prefix
         }
 
-        prefixRouter := e.Group(prefix)
+        g := e.Group(config.Prefix)
         {
-                prefixRouter.GET("/", pprofAdapter(pprof.Index))
-                prefixRouter.GET("/allocs", pprofAdapter(pprof.Handler("allocs").ServeHTTP))
-                prefixRouter.GET("/block", pprofAdapter(pprof.Handler("block").ServeHTTP))
-                prefixRouter.GET("/cmdline", pprofAdapter(pprof.Cmdline))
-                prefixRouter.GET("/goroutine", pprofAdapter(pprof.Handler("goroutine").ServeHTTP))
-                prefixRouter.GET("/heap", pprofAdapter(pprof.Handler("heap").ServeHTTP))
-                prefixRouter.GET("/mutex", pprofAdapter(pprof.Handler("mutex").ServeHTTP))
-                prefixRouter.GET("/profile", pprofAdapter(pprof.Profile))
-                prefixRouter.POST("/symbol", pprofAdapter(pprof.Symbol))
-                prefixRouter.GET("/symbol", pprofAdapter(pprof.Symbol))
-                prefixRouter.GET("/threadcreate", pprofAdapter(pprof.Handler("threadcreate").ServeHTTP))
-                prefixRouter.GET("/trace", pprofAdapter(pprof.Trace))
+                g.GET("/", pprofAdapter(pprof.Index), mw...)
+                g.GET("/allocs", pprofAdapter(pprof.Handler("allocs").ServeHTTP), mw...)
+                g.GET("/block", pprofAdapter(pprof.Handler("block").ServeHTTP), mw...)
+                g.GET("/cmdline", pprofAdapter(pprof.Cmdline), mw...)
+                g.GET("/goroutine", pprofAdapter(pprof.Handler("goroutine").ServeHTTP), mw...)
+                g.GET("/heap", pprofAdapter(pprof.Handler("heap").ServeHTTP), mw...)
+                g.GET("/mutex", pprofAdapter(pprof.Handler("mutex").ServeHTTP), mw...)
+                g.GET("/profile", pprofAdapter(pprof.Profile), mw...)
+                g.POST("/symbol", pprofAdapter(pprof.Symbol), mw...)
+                g.GET("/symbol", pprofAdapter(pprof.Symbol), mw...)
+                g.GET("/threadcreate", pprofAdapter(pprof.Handler("threadcreate").ServeHTTP), mw...)
+                g.GET("/trace", pprofAdapter(pprof.Trace), mw...)
         }
 }