beater/api/profile: fix pprof Content-Type

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,10 +46,12 @@ 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
@@ -68,7 +70,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 +115,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 +154,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 +230,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 {

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 {