From 367241488835506bf263120cc9482ae6b3c16689 Mon Sep 17 00:00:00 2001 From: Drew Bailey <2614075+drewbailey@users.noreply.github.com> Date: Wed, 11 Dec 2019 14:58:41 -0500 Subject: [PATCH] test pprof headers and profile methods tidy up, add comments clean up seconds param assignment --- client/agent_endpoint.go | 10 ++- client/structs/structs.go | 2 + command/agent/agent_endpoint.go | 15 ++-- command/agent/http.go | 50 ++++++++++- command/agent/profile/pprof.go | 52 +++++++---- command/agent/profile/pprof_test.go | 134 ++++++++++++++++++++++++++++ nomad/client_agent_endpoint.go | 12 ++- 7 files changed, 243 insertions(+), 32 deletions(-) create mode 100644 command/agent/profile/pprof_test.go diff --git a/client/agent_endpoint.go b/client/agent_endpoint.go index 76fde73c4..c269091d0 100644 --- a/client/agent_endpoint.go +++ b/client/agent_endpoint.go @@ -39,18 +39,19 @@ func (a *Agent) Profile(args *cstructs.AgentPprofRequest, reply *cstructs.AgentP var resp []byte var err error + var headers map[string]string // Determine which profile to run // and generate profile. Blocks for args.Seconds switch args.ReqType { case profile.CPUReq: - resp, err = profile.CPUProfile(context.TODO(), args.Seconds) + resp, headers, err = profile.CPUProfile(context.TODO(), args.Seconds) case profile.CmdReq: - resp, err = profile.Cmdline() + resp, headers, err = profile.Cmdline() case profile.LookupReq: - resp, err = profile.Profile(args.Profile, args.Debug) + resp, headers, err = profile.Profile(args.Profile, args.Debug) case profile.TraceReq: - resp, err = profile.Trace(context.TODO(), args.Seconds) + resp, headers, err = profile.Trace(context.TODO(), args.Seconds) } if err != nil { @@ -63,6 +64,7 @@ func (a *Agent) Profile(args *cstructs.AgentPprofRequest, reply *cstructs.AgentP // Copy profile response to reply reply.Payload = resp reply.AgentID = a.c.NodeID() + reply.HTTPHeaders = headers return nil } diff --git a/client/structs/structs.go b/client/structs/structs.go index 10c9ad239..ae56406c4 100644 --- a/client/structs/structs.go +++ b/client/structs/structs.go @@ -85,6 +85,8 @@ type AgentPprofResponse struct { // Payload is the generated pprof profile Payload []byte + + HTTPHeaders map[string]string } // AllocFileInfo holds information about a file inside the AllocDir diff --git a/command/agent/agent_endpoint.go b/command/agent/agent_endpoint.go index 7eddfb3f5..aa4ac6829 100644 --- a/command/agent/agent_endpoint.go +++ b/command/agent/agent_endpoint.go @@ -334,7 +334,7 @@ func (s *HTTPServer) AgentForceLeaveRequest(resp http.ResponseWriter, req *http. return nil, err } -func (s *HTTPServer) AgentPprofRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) { +func (s *HTTPServer) AgentPprofRequest(resp http.ResponseWriter, req *http.Request) ([]byte, error) { path := strings.TrimPrefix(req.URL.Path, "/v1/agent/pprof/") switch { case path == "": @@ -357,7 +357,7 @@ func (s *HTTPServer) AgentPprofRequest(resp http.ResponseWriter, req *http.Reque } } -func (s *HTTPServer) agentPprof(reqType profile.ReqType, resp http.ResponseWriter, req *http.Request) (interface{}, error) { +func (s *HTTPServer) agentPprof(reqType profile.ReqType, resp http.ResponseWriter, req *http.Request) ([]byte, error) { var secret string s.parseToken(req, &secret) @@ -386,11 +386,10 @@ func (s *HTTPServer) agentPprof(reqType profile.ReqType, resp http.ResponseWrite if secondsParam == "" { seconds = 1 } else { - sec, err := strconv.Atoi(secondsParam) + seconds, err = strconv.Atoi(secondsParam) if err != nil { return nil, CodedError(400, err.Error()) } - seconds = sec } // Create the request @@ -435,11 +434,11 @@ func (s *HTTPServer) agentPprof(reqType profile.ReqType, resp http.ResponseWrite return nil, CodedError(code, msg) } - // Pprof cmdline is not a typical pprof - // so just return string instead of bytes - if args.ReqType == profile.CmdReq { - return string(reply.Payload), nil + // Set headers from profile request + for k, v := range reply.HTTPHeaders { + resp.Header().Set(k, v) } + return reply.Payload, nil } diff --git a/command/agent/http.go b/command/agent/http.go index 9710e1e47..f62b4b7ba 100644 --- a/command/agent/http.go +++ b/command/agent/http.go @@ -186,7 +186,7 @@ func (s *HTTPServer) registerHandlers(enableDebug bool) { s.mux.HandleFunc("/v1/agent/health", s.wrap(s.HealthRequest)) s.mux.HandleFunc("/v1/agent/monitor", s.wrap(s.AgentMonitor)) - s.mux.HandleFunc("/v1/agent/pprof/", s.wrap(s.AgentPprofRequest)) + s.mux.HandleFunc("/v1/agent/pprof/", s.wrapNonJSON(s.AgentPprofRequest)) s.mux.HandleFunc("/v1/metrics", s.wrap(s.MetricsRequest)) @@ -358,6 +358,54 @@ func (s *HTTPServer) wrap(handler func(resp http.ResponseWriter, req *http.Reque return f } +// wrapNonJON is used to wrap functions returning non JSON +// serializeable data to make them more convenient. It is primarily +// responsible for setting nomad headers and logging. +func (s *HTTPServer) wrapNonJSON(handler func(resp http.ResponseWriter, req *http.Request) ([]byte, error)) func(resp http.ResponseWriter, req *http.Request) { + f := func(resp http.ResponseWriter, req *http.Request) { + setHeaders(resp, s.agent.config.HTTPAPIResponseHeaders) + // Invoke the handler + reqURL := req.URL.String() + start := time.Now() + defer func() { + s.logger.Debug("request complete", "method", req.Method, "path", reqURL, "duration", time.Now().Sub(start)) + }() + obj, err := handler(resp, req) + + // Check for an error + if err != nil { + code := 500 + errMsg := err.Error() + if http, ok := err.(HTTPCodedError); ok { + code = http.Code() + } else if ecode, emsg, ok := structs.CodeFromRPCCodedErr(err); ok { + code = ecode + errMsg = emsg + } else { + // RPC errors get wrapped, so manually unwrap by only looking at their suffix + if strings.HasSuffix(errMsg, structs.ErrPermissionDenied.Error()) { + errMsg = structs.ErrPermissionDenied.Error() + code = 403 + } else if strings.HasSuffix(errMsg, structs.ErrTokenNotFound.Error()) { + errMsg = structs.ErrTokenNotFound.Error() + code = 403 + } + } + + resp.WriteHeader(code) + resp.Write([]byte(errMsg)) + s.logger.Error("request failed", "method", req.Method, "path", reqURL, "error", err, "code", code) + return + } + + // write response + if obj != nil { + resp.Write(obj) + } + } + return f +} + // decodeBody is used to decode a JSON request body func decodeBody(req *http.Request, out interface{}) error { dec := json.NewDecoder(req.Body) diff --git a/command/agent/profile/pprof.go b/command/agent/profile/pprof.go index d3285b1a7..eafb40dec 100644 --- a/command/agent/profile/pprof.go +++ b/command/agent/profile/pprof.go @@ -42,64 +42,86 @@ func IsErrProfileNotFound(err error) bool { // Cmdline responds with the running program's // command line, with arguments separated by NUL bytes. -func Cmdline() ([]byte, error) { +func Cmdline() ([]byte, map[string]string, error) { var buf bytes.Buffer fmt.Fprintf(&buf, strings.Join(os.Args, "\x00")) - return buf.Bytes(), nil + + return buf.Bytes(), + map[string]string{ + "X-Content-Type-Options": "nosniff", + "Content-Type": "text/plain; charset=utf-8", + }, nil } // Profile generates a pprof.Profile report for the given profile name // see runtime/pprof/pprof.go for available profiles. -func Profile(profile string, debug int) ([]byte, error) { +func Profile(profile string, debug int) ([]byte, map[string]string, error) { p := pprof.Lookup(profile) if p == nil { - return nil, NewErrProfileNotFound(profile) + return nil, nil, NewErrProfileNotFound(profile) } var buf bytes.Buffer if err := p.WriteTo(&buf, debug); err != nil { - return nil, err + return nil, nil, err } - return buf.Bytes(), nil + headers := map[string]string{ + "X-Content-Type-Options": "nosniff", + } + if debug != 0 { + headers["Content-Type"] = "text/plain; charset=utf-8" + } else { + headers["Content-Type"] = "application/octet-stream" + headers["Content-Disposition"] = fmt.Sprintf(`attachment; filename="%s"`, profile) + } + return buf.Bytes(), headers, nil } // CPUProfile generates a CPU Profile for a given duration -func CPUProfile(ctx context.Context, sec int) ([]byte, error) { +func CPUProfile(ctx context.Context, sec int) ([]byte, map[string]string, error) { if sec <= 0 { sec = 1 } var buf bytes.Buffer if err := pprof.StartCPUProfile(&buf); err != nil { - // trace.Start failed, no writes yet - return nil, err + return nil, nil, err } sleep(ctx, time.Duration(sec)*time.Second) pprof.StopCPUProfile() - return buf.Bytes(), nil + return buf.Bytes(), + map[string]string{ + "X-Content-Type-Options": "nosniff", + "Content-Type": "application/octet-stream", + "Content-Disposition": `attachment; filename="profile"`, + }, nil } // Trace runs a trace profile for a given duration -func Trace(ctx context.Context, sec int) ([]byte, error) { +func Trace(ctx context.Context, sec int) ([]byte, map[string]string, error) { if sec <= 0 { sec = 1 } var buf bytes.Buffer if err := trace.Start(&buf); err != nil { - // trace.Start failed, no writes yet - return nil, err + return nil, nil, err } - sleep(context.TODO(), time.Duration(sec)*time.Second) + sleep(ctx, time.Duration(sec)*time.Second) trace.Stop() - return buf.Bytes(), nil + return buf.Bytes(), + map[string]string{ + "X-Content-Type-Options": "nosniff", + "Content-Type": "application/octet-stream", + "Content-Disposition": `attachment; filename="trace"`, + }, nil } func sleep(ctx context.Context, d time.Duration) { diff --git a/command/agent/profile/pprof_test.go b/command/agent/profile/pprof_test.go new file mode 100644 index 000000000..103ee373b --- /dev/null +++ b/command/agent/profile/pprof_test.go @@ -0,0 +1,134 @@ +package profile + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestProfile(t *testing.T) { + cases := []struct { + desc string + profile string + debug int + expectedHeaders map[string]string + expectedErr error + }{ + { + desc: "profile that exists", + profile: "goroutine", + expectedHeaders: map[string]string{ + "X-Content-Type-Options": "nosniff", + "Content-Type": "application/octet-stream", + "Content-Disposition": `attachment; filename="goroutine"`, + }, + }, + { + desc: "profile that does not exist", + profile: "nonexistent", + expectedErr: NewErrProfileNotFound("nonexistent"), + expectedHeaders: nil, + }, + { + desc: "profile with debug enabled", + profile: "allocs", + debug: 1, + expectedHeaders: map[string]string{ + "X-Content-Type-Options": "nosniff", + "Content-Type": "text/plain; charset=utf-8", + }, + }, + } + + for _, tc := range cases { + t.Run(tc.desc, func(t *testing.T) { + resp, headers, err := Profile(tc.profile, tc.debug) + require.Equal(t, tc.expectedHeaders, headers) + + if tc.expectedErr != nil { + require.Nil(t, resp) + require.Equal(t, err, tc.expectedErr) + } else { + require.NotNil(t, resp) + } + }) + } +} + +func TestCPUProfile(t *testing.T) { + cases := []struct { + desc string + expectedHeaders map[string]string + }{ + { + desc: "successful cpu profile", + expectedHeaders: map[string]string{ + "X-Content-Type-Options": "nosniff", + "Content-Type": "application/octet-stream", + "Content-Disposition": `attachment; filename="profile"`, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.desc, func(t *testing.T) { + resp, headers, err := CPUProfile(context.Background(), 0) + require.NoError(t, err) + require.Equal(t, tc.expectedHeaders, headers) + + require.NotNil(t, resp) + }) + } +} + +func TestTrace(t *testing.T) { + cases := []struct { + desc string + expectedHeaders map[string]string + }{ + { + desc: "successful trace profile", + expectedHeaders: map[string]string{ + "X-Content-Type-Options": "nosniff", + "Content-Type": "application/octet-stream", + "Content-Disposition": `attachment; filename="trace"`, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.desc, func(t *testing.T) { + resp, headers, err := Trace(context.Background(), 0) + require.NoError(t, err) + require.Equal(t, tc.expectedHeaders, headers) + + require.NotNil(t, resp) + }) + } +} + +func TestCmdline(t *testing.T) { + cases := []struct { + desc string + expectedHeaders map[string]string + }{ + { + desc: "successful cmdline request", + expectedHeaders: map[string]string{ + "X-Content-Type-Options": "nosniff", + "Content-Type": "text/plain; charset=utf-8", + }, + }, + } + + for _, tc := range cases { + t.Run(tc.desc, func(t *testing.T) { + resp, headers, err := Cmdline() + require.NoError(t, err) + require.Equal(t, tc.expectedHeaders, headers) + + require.NotNil(t, resp) + }) + } +} diff --git a/nomad/client_agent_endpoint.go b/nomad/client_agent_endpoint.go index ba52afb1a..c93b0d059 100644 --- a/nomad/client_agent_endpoint.go +++ b/nomad/client_agent_endpoint.go @@ -62,21 +62,24 @@ func (a *Agent) Profile(args *cstructs.AgentPprofRequest, reply *cstructs.AgentP // Process the request on this server var resp []byte var err error + var headers map[string]string // Mark which server fulfilled the request reply.AgentID = a.srv.serf.LocalMember().Name // Determine which profile to run // and generate profile. Blocks for args.Seconds + // Our RPC endpoints currently don't support context + // or request cancellation so stubbing with TODO switch args.ReqType { case profile.CPUReq: - resp, err = profile.CPUProfile(context.TODO(), args.Seconds) + resp, headers, err = profile.CPUProfile(context.TODO(), args.Seconds) case profile.CmdReq: - resp, err = profile.Cmdline() + resp, headers, err = profile.Cmdline() case profile.LookupReq: - resp, err = profile.Profile(args.Profile, args.Debug) + resp, headers, err = profile.Profile(args.Profile, args.Debug) case profile.TraceReq: - resp, err = profile.Trace(context.TODO(), args.Seconds) + resp, headers, err = profile.Trace(context.TODO(), args.Seconds) default: err = structs.NewErrRPCCoded(404, "Unknown profile request type") } @@ -90,6 +93,7 @@ func (a *Agent) Profile(args *cstructs.AgentPprofRequest, reply *cstructs.AgentP // Copy profile response to reply reply.Payload = resp + reply.HTTPHeaders = headers return nil }