Merge pull request #38 from Flowpack/pprof-profiling-for-prunner

FEATURE: allow pprof profiling in prunner
Flowpack · Sep 29, 2022 · ed281a2 · ed281a2
2 parents 1976713 + 5f9317c
commit ed281a2
Show file tree

Hide file tree

Showing 3 changed files with 110 additions and 19 deletions.
diff --git a/app/app.go b/app/app.go
@@ -57,6 +57,12 @@ func New(info Info) *cli.App {
 			Value:   false,
 			EnvVars: []string{"PRUNNER_VERBOSE"},
 		},
+		&cli.BoolFlag{
+			Name:    "enable-profiling",
+			Usage:   "Enable the Profiling endpoints underneath /debug/pprof",
+			Value:   false,
+			EnvVars: []string{"PRUNNER_ENABLE_PROFILING"},
+		},
 		&cli.BoolFlag{
 			Name:    "disable-ansi",
 			Usage:   "Force disable ANSI log output and output log in logfmt format",
@@ -231,6 +237,7 @@ func appAction(c *cli.Context) error {
 		outputStore,
 		middleware.RequestLogger(createLogFormatter(c)),
 		tokenAuth,
+		c.Bool("enable-profiling"),
 	)
 
 	// Set up a simple REST API for listing jobs and scheduling pipelines

diff --git a/server/server.go b/server/server.go
@@ -33,7 +33,7 @@ type server struct {
 	outputStore taskctl.OutputStore
 }
 
-func NewServer(pRunner *prunner.PipelineRunner, outputStore taskctl.OutputStore, logger func(http.Handler) http.Handler, tokenAuth *jwtauth.JWTAuth) *server {
+func NewServer(pRunner *prunner.PipelineRunner, outputStore taskctl.OutputStore, logger func(http.Handler) http.Handler, tokenAuth *jwtauth.JWTAuth, enableProfiling bool) *server {
 	srv := &server{
 		pRunner:     pRunner,
 		outputStore: outputStore,
@@ -42,22 +42,31 @@ func NewServer(pRunner *prunner.PipelineRunner, outputStore taskctl.OutputStore,
 	r := chi.NewRouter()
 	r.Use(logger)
 	r.Use(middleware.Recoverer)
-	// Seek, verify and validate JWT tokens
-	r.Use(jwtauth.Verifier(tokenAuth))
-	// Handle valid / invalid tokens
-	r.Use(jwtauth.Authenticator)
-
-	r.Route("/pipelines", func(r chi.Router) {
-		r.Get("/", srv.pipelines)
-		r.Get("/jobs", srv.pipelinesJobs)
-		r.Post("/schedule", srv.pipelinesSchedule)
-	})
-	r.Route("/job", func(r chi.Router) {
-		r.Get("/detail", srv.jobDetail)
-		r.Get("/logs", srv.jobLogs)
-		r.Post("/cancel", srv.jobCancel)
+
+	// we do not want JWT authentication for /debug/pprof,
+	// that's why we need to create a new handler group (to scope the authentication middlewares)
+	r.Group(func(r chi.Router) {
+		// Seek, verify and validate JWT tokens
+		r.Use(jwtauth.Verifier(tokenAuth))
+		// Handle valid / invalid tokens
+		r.Use(jwtauth.Authenticator)
+
+		r.Route("/pipelines", func(r chi.Router) {
+			r.Get("/", srv.pipelines)
+			r.Get("/jobs", srv.pipelinesJobs)
+			r.Post("/schedule", srv.pipelinesSchedule)
+		})
+		r.Route("/job", func(r chi.Router) {
+			r.Get("/detail", srv.jobDetail)
+			r.Get("/logs", srv.jobLogs)
+			r.Post("/cancel", srv.jobCancel)
+		})
 	})
 
+	if enableProfiling {
+		r.Mount("/debug", middleware.Profiler())
+	}
+
 	srv.handler = r
 
 	return srv

diff --git a/server/server_test.go b/server/server_test.go
@@ -69,7 +69,7 @@ func TestServer_Pipelines(t *testing.T) {
 
 	tokenAuth := jwtauth.New("HS256", []byte("not-very-secret"), nil)
 	noopMiddleware := func(next http.Handler) http.Handler { return next }
-	srv := NewServer(pRunner, outputStore, noopMiddleware, tokenAuth)
+	srv := NewServer(pRunner, outputStore, noopMiddleware, tokenAuth, false)
 
 	claims := make(map[string]interface{})
 	jwtauth.SetIssuedNow(claims)
@@ -91,6 +91,34 @@ func TestServer_Pipelines(t *testing.T) {
 	}`, rec.Body.String())
 }
 
+func TestServer_PipelinesCanNotBeAccessedWithWrongJwtToken(t *testing.T) {
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+
+	pRunner, err := prunner.NewPipelineRunner(ctx, defs, func(j *prunner.PipelineJob) taskctl.Runner {
+		return &test.MockRunner{}
+	}, nil, nil)
+	require.NoError(t, err)
+
+	outputStore := test.NewMockOutputStore()
+
+	tokenAuth := jwtauth.New("HS256", []byte("not-very-secret"), nil)
+	noopMiddleware := func(next http.Handler) http.Handler { return next }
+	srv := NewServer(pRunner, outputStore, noopMiddleware, tokenAuth, false)
+
+	wrongTokenAuthClient := jwtauth.New("HS256", []byte("THIS-IS-WRONG"), nil)
+	claims := make(map[string]interface{})
+	jwtauth.SetIssuedNow(claims)
+	_, tokenString, _ := wrongTokenAuthClient.Encode(claims)
+
+	req := httptest.NewRequest("GET", "/pipelines", nil)
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", tokenString))
+	rec := httptest.NewRecorder()
+	srv.ServeHTTP(rec, req)
+
+	require.Equal(t, http.StatusUnauthorized, rec.Code)
+}
+
 func TestServer_PipelinesSchedule(t *testing.T) {
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	defer cancelFunc()
@@ -104,7 +132,7 @@ func TestServer_PipelinesSchedule(t *testing.T) {
 
 	tokenAuth := jwtauth.New("HS256", []byte("not-very-secret"), nil)
 	noopMiddleware := func(next http.Handler) http.Handler { return next }
-	srv := NewServer(pRunner, outputStore, noopMiddleware, tokenAuth)
+	srv := NewServer(pRunner, outputStore, noopMiddleware, tokenAuth, false)
 
 	claims := make(map[string]interface{})
 	jwtauth.SetIssuedNow(claims)
@@ -152,7 +180,7 @@ func TestServer_JobCreationTimeIsRoundedForPhpCompatibility(t *testing.T) {
 
 	tokenAuth := jwtauth.New("HS256", []byte("not-very-secret"), nil)
 	noopMiddleware := func(next http.Handler) http.Handler { return next }
-	srv := NewServer(pRunner, outputStore, noopMiddleware, tokenAuth)
+	srv := NewServer(pRunner, outputStore, noopMiddleware, tokenAuth, false)
 
 	claims := make(map[string]interface{})
 	jwtauth.SetIssuedNow(claims)
@@ -220,7 +248,7 @@ func TestServer_JobCancel(t *testing.T) {
 
 	tokenAuth := jwtauth.New("HS256", []byte("not-very-secret"), nil)
 	noopMiddleware := func(next http.Handler) http.Handler { return next }
-	srv := NewServer(pRunner, outputStore, noopMiddleware, tokenAuth)
+	srv := NewServer(pRunner, outputStore, noopMiddleware, tokenAuth, false)
 
 	claims := make(map[string]interface{})
 	jwtauth.SetIssuedNow(claims)
@@ -268,3 +296,50 @@ func TestServer_JobCancel(t *testing.T) {
 
 	}, 50*time.Millisecond, "job exists and was completed")
 }
+
+func TestServer_NoAccessToProfilingRoutesIfDisabled(t *testing.T) {
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+
+	pRunner, err := prunner.NewPipelineRunner(ctx, defs, func(j *prunner.PipelineJob) taskctl.Runner {
+		return &test.MockRunner{}
+	}, nil, nil)
+	require.NoError(t, err)
+
+	outputStore := test.NewMockOutputStore()
+
+	tokenAuth := jwtauth.New("HS256", []byte("not-very-secret"), nil)
+	noopMiddleware := func(next http.Handler) http.Handler { return next }
+	srv := NewServer(pRunner, outputStore, noopMiddleware, tokenAuth, false)
+
+	req := httptest.NewRequest("GET", "/debug/pprof", nil)
+
+	rec := httptest.NewRecorder()
+	srv.ServeHTTP(rec, req)
+
+	require.Equal(t, http.StatusNotFound, rec.Code)
+}
+
+func TestServer_AccessToProfilingWorksIfEnabledWithoutToken(t *testing.T) {
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+
+	pRunner, err := prunner.NewPipelineRunner(ctx, defs, func(j *prunner.PipelineJob) taskctl.Runner {
+		return &test.MockRunner{}
+	}, nil, nil)
+	require.NoError(t, err)
+
+	outputStore := test.NewMockOutputStore()
+
+	tokenAuth := jwtauth.New("HS256", []byte("not-very-secret"), nil)
+	noopMiddleware := func(next http.Handler) http.Handler { return next }
+	// !!! the last parameter is "enableProfiling: true"
+	srv := NewServer(pRunner, outputStore, noopMiddleware, tokenAuth, true)
+
+	req := httptest.NewRequest("GET", "/debug/pprof/", nil)
+
+	rec := httptest.NewRecorder()
+	srv.ServeHTTP(rec, req)
+
+	require.Equal(t, http.StatusOK, rec.Code)
+}