From 5f9317c7435cc1f46a4773976caf39c627943947 Mon Sep 17 00:00:00 2001 From: Sebastian Kurfuerst Date: Thu, 29 Sep 2022 09:57:20 +0200 Subject: [PATCH] FEATURE: allow pprof profiling in prunner --- app/app.go | 7 ++++ server/server.go | 39 ++++++++++++-------- server/server_test.go | 83 ++++++++++++++++++++++++++++++++++++++++--- 3 files changed, 110 insertions(+), 19 deletions(-) diff --git a/app/app.go b/app/app.go index 1be89ec..55fd3e5 100644 --- 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 index 8593558..c69be7f 100644 --- 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 index 0b19a50..9e85f49 100644 --- 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) +}