diff --git a/beater/api/profile/handler.go b/beater/api/profile/handler.go index a307426d798..79db49c07a5 100644 --- a/beater/api/profile/handler.go +++ b/beater/api/profile/handler.go @@ -20,8 +20,8 @@ package profile import ( "fmt" "io" + "mime" "net/http" - "strings" pprof_profile "github.com/google/pprof/profile" "github.com/pkg/errors" @@ -46,15 +46,21 @@ var ( ) const ( - // TODO(axw) include messageType in pprofContentType; needs fix in agent - pprofContentType = "application/x-protobuf" - metadataContentType = "application/json" - requestContentType = "multipart/form-data" + pprofMediaType = "application/x-protobuf" + metadataMediaType = "application/json" + requestMediaType = "multipart/form-data" + + // value for the "messagetype" param + pprofMessageType = "perftools.profiles.Profile" metadataContentLengthLimit = 10 * 1024 profileContentLengthLimit = 10 * 1024 * 1024 ) +var ( + pprofMediaTypeParams = map[string]string{"messagetype": "perftools.profiles.Profile"} +) + // Handler returns a request.Handler for managing profile requests. func Handler( dec decoder.ReqDecoder, @@ -68,7 +74,7 @@ func Handler( err: errors.New("only POST requests are supported"), } } - if err := validateContentType(c.Request.Header, requestContentType); err != nil { + if _, err := validateContentType(c.Request.Header, requestMediaType); err != nil { return nil, requestError{ id: request.IDResponseErrorsValidate, err: err, @@ -113,7 +119,7 @@ func Handler( switch part.FormName() { case "metadata": - if err := validateContentType(http.Header(part.Header), metadataContentType); err != nil { + if _, err := validateContentType(http.Header(part.Header), metadataMediaType); err != nil { return nil, requestError{ id: request.IDResponseErrorsValidate, err: errors.Wrap(err, "invalid metadata"), @@ -152,12 +158,22 @@ func Handler( tctx.Metadata = *metadata case "profile": - if err := validateContentType(http.Header(part.Header), pprofContentType); err != nil { + params, err := validateContentType(http.Header(part.Header), pprofMediaType) + if err != nil { return nil, requestError{ id: request.IDResponseErrorsValidate, err: errors.Wrap(err, "invalid profile"), } } + if v, ok := params["messagetype"]; ok && v != pprofMessageType { + // If messagetype is specified, require that it matches. + // Otherwise we assume that it's pprof, and we'll error + // out below if it doesn't decode. + return nil, requestError{ + id: request.IDResponseErrorsValidate, + err: errors.Wrapf(err, "expected messagetype %q, got %q", pprofMessageType, v), + } + } r := &decoder.LimitedReader{R: part, N: totalLimitRemaining} profile, err := pprof_profile.Parse(r) if err != nil { @@ -218,12 +234,15 @@ func Handler( } } -func validateContentType(header http.Header, contentType string) error { - got := header.Get(headers.ContentType) - if !strings.Contains(got, contentType) { - return fmt.Errorf("invalid content type %q, expected %q", got, contentType) +func validateContentType(header http.Header, expectedMediatype string) (params map[string]string, err error) { + mediatype, params, err := mime.ParseMediaType(header.Get(headers.ContentType)) + if err != nil { + return nil, err + } + if mediatype != expectedMediatype { + return nil, fmt.Errorf("invalid content type %q, expected %q", mediatype, expectedMediatype) } - return nil + return params, nil } type result struct { diff --git a/beater/api/profile/handler_test.go b/beater/api/profile/handler_test.go index ea5bb0620f9..1cef437384e 100644 --- a/beater/api/profile/handler_test.go +++ b/beater/api/profile/handler_test.go @@ -44,6 +44,8 @@ import ( "github.com/elastic/apm-server/publish" ) +const pprofContentType = `application/x-protobuf; messageType="perftools.profiles.Profile"` + func TestHandler(t *testing.T) { var rateLimit, err = ratelimit.NewStore(1, 0, 0) require.NoError(t, err) @@ -127,7 +129,12 @@ func TestHandler(t *testing.T) { id: request.IDResponseValidAccepted, parts: []part{ heapProfilePart(), - heapProfilePart(), + part{ + name: "profile", + // No messageType param specified, so pprof is assumed. + contentType: "application/x-protobuf", + body: heapProfileBody(), + }, part{ name: "metadata", contentType: "application/json", @@ -147,17 +154,28 @@ func TestHandler(t *testing.T) { "ProfileInvalidContentType": { id: request.IDResponseErrorsValidate, parts: []part{{ - name: "metadata", + name: "profile", contentType: "text/plain", body: strings.NewReader(""), }}, body: prettyJSON(map[string]interface{}{"accepted": 0}), }, + "ProfileInvalidMessageType": { + id: request.IDResponseErrorsValidate, + parts: []part{{ + name: "profile", + // Improperly formatted "messageType" param + // in Content-Type from APM Agent Go v1.6.0. + contentType: "application/x-protobuf; messageType=ā€¯perftools.profiles.Profile", + body: strings.NewReader(""), + }}, + body: prettyJSON(map[string]interface{}{"accepted": 0}), + }, "ProfileInvalid": { id: request.IDResponseErrorsDecode, parts: []part{{ name: "profile", - contentType: "application/x-protobuf", + contentType: pprofContentType, body: strings.NewReader("foo"), }}, body: prettyJSON(map[string]interface{}{"accepted": 0}), @@ -168,7 +186,7 @@ func TestHandler(t *testing.T) { heapProfilePart(), part{ name: "profile", - contentType: "application/x-protobuf", + contentType: pprofContentType, body: strings.NewReader(strings.Repeat("*", 10*1024*1024)), }, }, @@ -261,15 +279,15 @@ func emptyDec(_ *http.Request) (map[string]interface{}, error) { } func heapProfilePart() part { + return part{name: "profile", contentType: pprofContentType, body: heapProfileBody()} +} + +func heapProfileBody() io.Reader { var buf bytes.Buffer if err := pprof.WriteHeapProfile(&buf); err != nil { panic(err) } - return part{ - name: "profile", - contentType: "application/x-protobuf", - body: &buf, - } + return &buf } type part struct {