Put pprof behind token gate

Author: sandyydk

cmd/cloudinfo/config.go

@@ -183,6 +183,10 @@ type configuration struct {
 	ServiceLoader loader.Config
 
 	Store cistore.Config
+
+	Pprof struct {
+		SecretToken string `mapstructure:"secret_token"`
+	}
 }
 
 // Validate validates the configuration.
@@ -344,6 +348,8 @@ func configure(v *viper.Viper, p *pflag.FlagSet) {
 
 	_ = v.BindEnv("provider.digitalocean.accessToken", "DIGITALOCEAN_ACCESS_TOKEN")
 
+	_ = v.BindEnv("pprof.secret_token", "CLOUDINFO_PPROF_SECRET_TOKEN")
+
 	// Management
 	v.SetDefault("management.enabled", true)
 	v.SetDefault("management.address", ":8001")

cmd/cloudinfo/main.go

@@ -234,7 +234,7 @@ func main() {
 		errorHandler,
 	)
 
-	routeHandler := api.NewRouteHandler(prodInfo, buildInfo, graphqlHandler, cloudInfoLogger)
+	routeHandler := api.NewRouteHandler(prodInfo, buildInfo, graphqlHandler, cloudInfoLogger, config.Pprof.SecretToken)
 
 	// new default gin engine (recovery, logger middleware)
 	router := gin.Default()

config.toml.dist

@@ -130,6 +130,10 @@ enabled = false
 enabled = true
 address = ":8001"
 
+[pprof]
+# Required in the Token header to access /debug/pprof/*. Empty disables pprof.
+secret_token = ""
+
 [serviceloader]
 serviceConfigLocation = "./configs"
 serviceConfigName = "services"

internal/app/cloudinfo/api/pprof.go

@@ -0,0 +1,72 @@
+// Copyright © 2018 Banzai Cloud
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package api
+
+import (
+	"crypto/subtle"
+	"net/http"
+	"net/http/pprof"
+
+	"github.com/gin-gonic/gin"
+)
+
+type pprofRoute struct {
+	pattern string
+	handler http.Handler
+}
+
+func pprofRoutes() []pprofRoute {
+	return []pprofRoute{
+		{pattern: "/debug/pprof/", handler: http.HandlerFunc(pprof.Index)},
+		{pattern: "/debug/pprof/cmdline", handler: http.HandlerFunc(pprof.Cmdline)},
+		{pattern: "/debug/pprof/profile", handler: http.HandlerFunc(pprof.Profile)},
+		{pattern: "/debug/pprof/symbol", handler: http.HandlerFunc(pprof.Symbol)},
+		{pattern: "/debug/pprof/trace", handler: http.HandlerFunc(pprof.Trace)},
+		{pattern: "/debug/pprof/allocs", handler: pprof.Handler("allocs")},
+		{pattern: "/debug/pprof/block", handler: pprof.Handler("block")},
+		{pattern: "/debug/pprof/goroutine", handler: pprof.Handler("goroutine")},
+		{pattern: "/debug/pprof/heap", handler: pprof.Handler("heap")},
+		{pattern: "/debug/pprof/mutex", handler: pprof.Handler("mutex")},
+		{pattern: "/debug/pprof/threadcreate", handler: pprof.Handler("threadcreate")},
+	}
+}
+
+func pprofAuth(secretToken string) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		provided := c.GetHeader("Token")
+		if !tokenEquals(provided, secretToken) {
+			c.AbortWithStatus(http.StatusForbidden)
+			return
+		}
+		c.Next()
+	}
+}
+
+func attachPprof(router gin.IRouter, secretToken string) {
+	if secretToken == "" {
+		return
+	}
+	auth := pprofAuth(secretToken)
+	for _, route := range pprofRoutes() {
+		router.Any(route.pattern, auth, gin.WrapH(route.handler))
+	}
+}
+
+func tokenEquals(provided, expected string) bool {
+	if len(provided) != len(expected) {
+		return false
+	}
+	return subtle.ConstantTimeCompare([]byte(provided), []byte(expected)) == 1
+}

internal/app/cloudinfo/api/pprof_test.go

@@ -0,0 +1,127 @@
+// Copyright © 2018 Banzai Cloud
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package api
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMain(m *testing.M) {
+	gin.SetMode(gin.TestMode)
+	m.Run()
+}
+
+func TestPprofAuth_RejectsMissingAndInvalidToken(t *testing.T) {
+	router := gin.New()
+	router.GET("/debug/pprof/heap", pprofAuth("secret-token"), func(c *gin.Context) {
+		c.Status(http.StatusOK)
+	})
+
+	t.Run("missing token", func(t *testing.T) {
+		req := httptest.NewRequest(http.MethodGet, "/debug/pprof/heap", nil)
+		rec := httptest.NewRecorder()
+		router.ServeHTTP(rec, req)
+		assert.Equal(t, http.StatusForbidden, rec.Code)
+	})
+
+	t.Run("wrong token", func(t *testing.T) {
+		req := httptest.NewRequest(http.MethodGet, "/debug/pprof/heap", nil)
+		req.Header.Set("Token", "not-the-secret")
+		rec := httptest.NewRecorder()
+		router.ServeHTTP(rec, req)
+		assert.Equal(t, http.StatusForbidden, rec.Code)
+	})
+}
+
+func TestPprofAuth_AllowsMatchingToken(t *testing.T) {
+	router := gin.New()
+	router.GET("/debug/pprof/heap", pprofAuth("secret-token"), func(c *gin.Context) {
+		c.String(http.StatusOK, "ok")
+	})
+
+	req := httptest.NewRequest(http.MethodGet, "/debug/pprof/heap", nil)
+	req.Header.Set("Token", "secret-token")
+	rec := httptest.NewRecorder()
+	router.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "ok", rec.Body.String())
+}
+
+func TestAttachPprof_ProtectsHeapEndpoint(t *testing.T) {
+	router := gin.New()
+	attachPprof(router, "secret-token")
+
+	t.Run("unauthenticated", func(t *testing.T) {
+		req := httptest.NewRequest(http.MethodGet, "/debug/pprof/heap", nil)
+		rec := httptest.NewRecorder()
+		router.ServeHTTP(rec, req)
+		assert.Equal(t, http.StatusForbidden, rec.Code)
+	})
+
+	t.Run("authenticated", func(t *testing.T) {
+		req := httptest.NewRequest(http.MethodGet, "/debug/pprof/heap", nil)
+		req.Header.Set("Token", "secret-token")
+		rec := httptest.NewRecorder()
+		router.ServeHTTP(rec, req)
+		require.Equal(t, http.StatusOK, rec.Code)
+		assert.NotEmpty(t, rec.Body.Bytes())
+	})
+}
+
+func TestAttachPprof_SkipsHandlersWhenTokenUnset(t *testing.T) {
+	router := gin.New()
+	attachPprof(router, "")
+
+	for _, route := range router.Routes() {
+		assert.NotContains(t, route.Path, "/debug/pprof")
+	}
+
+	req := httptest.NewRequest(http.MethodGet, "/debug/pprof/heap", nil)
+	req.Header.Set("Token", "anything")
+	rec := httptest.NewRecorder()
+	router.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusNotFound, rec.Code)
+}
+
+func TestPprofRoutes(t *testing.T) {
+	routes := pprofRoutes()
+	require.NotEmpty(t, routes)
+
+	seen := make(map[string]bool, len(routes))
+	for _, route := range routes {
+		assert.NotEmpty(t, route.pattern)
+		assert.NotNil(t, route.handler)
+		assert.False(t, seen[route.pattern], "duplicate pprof route %s", route.pattern)
+		seen[route.pattern] = true
+	}
+
+	assert.True(t, seen["/debug/pprof/heap"])
+	assert.True(t, seen["/debug/pprof/goroutine"])
+}
+
+func TestTokenEquals(t *testing.T) {
+	assert.True(t, tokenEquals("abc", "abc"))
+	assert.False(t, tokenEquals("abc", "abd"))
+	assert.False(t, tokenEquals("abc", "abcd"))
+	assert.False(t, tokenEquals("", "a"))
+}

internal/app/cloudinfo/api/routes.go

@@ -19,7 +19,6 @@ import (
 	"io/fs"
 	"io/ioutil"
 	"net/http"
-	"net/http/pprof"
 	"strings"
 
 	"emperror.dev/emperror"
@@ -28,7 +27,6 @@ import (
 	"github.com/gin-contrib/cors"
 	"github.com/gin-contrib/static"
 	"github.com/gin-gonic/gin"
-	"github.com/google/uuid"
 
 	"github.com/banzaicloud/cloudinfo/internal/cloudinfo"
 	"github.com/banzaicloud/cloudinfo/internal/cloudinfo/metrics"
@@ -40,21 +38,23 @@ import (
 
 // RouteHandler configures the REST API routes in the gin router
 type RouteHandler struct {
-	log            cloudinfo.Logger
-	prod           types.CloudInfo
-	buildInfo      buildinfo.BuildInfo
-	errorResponder Responder
-	graphqlHandler http.Handler
+	log              cloudinfo.Logger
+	prod             types.CloudInfo
+	buildInfo        buildinfo.BuildInfo
+	errorResponder   Responder
+	graphqlHandler   http.Handler
+	pprofSecretToken string
 }
 
 // NewRouteHandler creates a new RouteHandler and returns a reference to it
-func NewRouteHandler(p types.CloudInfo, bi buildinfo.BuildInfo, graphqlHandler http.Handler, log cloudinfo.Logger) *RouteHandler {
+func NewRouteHandler(p types.CloudInfo, bi buildinfo.BuildInfo, graphqlHandler http.Handler, log cloudinfo.Logger, pprofSecretToken string) *RouteHandler {
 	return &RouteHandler{
-		prod:           p,
-		buildInfo:      bi,
-		errorResponder: NewErrorResponder(),
-		graphqlHandler: graphqlHandler,
-		log:            log,
+		prod:             p,
+		buildInfo:        bi,
+		errorResponder:   NewErrorResponder(),
+		graphqlHandler:   graphqlHandler,
+		log:              log,
+		pprofSecretToken: pprofSecretToken,
 	}
 }
 
@@ -124,13 +124,7 @@ func (r *RouteHandler) ConfigureRoutes(router *gin.Engine, basePath string) {
 
 	base.POST("/graphql", r.query())
 
-	heapRoutePath := fmt.Sprintf("/heap/%s", uuid.New().String())
-	r.log.Info("Heap pprof path", map[string]interface{}{"path": heapRoutePath})
-	router.GET(heapRoutePath, gin.WrapF(pprof.Handler("heap").ServeHTTP))
-
-	goRoutineRoutePath := fmt.Sprintf("/goroutine/%s", uuid.New().String())
-	r.log.Info("Goroutine pprof path", map[string]interface{}{"path": goRoutineRoutePath})
-	router.GET(goRoutineRoutePath, gin.WrapF(pprof.Handler("goroutine").ServeHTTP))
+	attachPprof(router, r.pprofSecretToken)
 }
 
 func (r *RouteHandler) signalStatus(c *gin.Context) {