Harden media type negotiation and preserve routed params

errors.go

@@ -125,8 +125,8 @@ func ErrorFilter(req Request, svc Service) Response {
 				var err error
 				tp := &terrorsproto.Error{}
 
-				switch rsp.Header.Get("Content-Type") {
-				case "application/octet-stream", "application/x-protobuf", "application/protobuf":
+				switch {
+				case isProtobufMediaType(rsp.Header.Get("Content-Type")):
 					err = legacyproto.Unmarshal(b, tp)
 				default:
 					err = json.Unmarshal(b, tp)

mediatype.go

@@ -0,0 +1,59 @@
+package typhon
+
+import (
+	"mime"
+	"net/http"
+	"strings"
+)
+
+var protobufMediaTypes = map[string]struct{}{
+	"application/octet-stream":      {},
+	"application/x-google-protobuf": {},
+	"application/protobuf":          {},
+	"application/x-protobuf":        {},
+}
+
+func canonicalMediaType(value string) string {
+	if value == "" {
+		return ""
+	}
+
+	mediaType, _, err := mime.ParseMediaType(value)
+	if err == nil {
+		return strings.ToLower(mediaType)
+	}
+
+	return strings.ToLower(strings.TrimSpace(strings.SplitN(value, ";", 2)[0]))
+}
+
+func isProtobufMediaType(value string) bool {
+	_, ok := protobufMediaTypes[canonicalMediaType(value)]
+	return ok
+}
+
+func acceptsProtobuf(header http.Header) bool {
+	if header == nil {
+		return false
+	}
+
+	for _, acceptValue := range header.Values("Accept") {
+		for _, part := range strings.Split(acceptValue, ",") {
+			mediaType, params, err := mime.ParseMediaType(strings.TrimSpace(part))
+			if err != nil {
+				mediaType = canonicalMediaType(part)
+				params = nil
+			}
+
+			if params != nil && params["q"] == "0" {
+				continue
+			}
+
+			mediaType = strings.ToLower(mediaType)
+			if mediaType == "*/*" || mediaType == "application/*" || isProtobufMediaType(mediaType) {
+				return true
+			}
+		}
+	}
+
+	return false
+}

request.go

@@ -93,11 +93,11 @@ func (r Request) Decode(v interface{}) error {
 		return terrors.WrapWithCode(err, nil, terrors.ErrBadRequest)
 	}
 
-	switch r.Header.Get("Content-Type") {
+	switch {
 	// application/x-protobuf is the "canonical" use, application/protobuf is defined in an expired IETF draft.
 	// See: https://datatracker.ietf.org/doc/html/draft-rfernando-protocol-buffers-00#section-3.2
 	// See: https://github.com/google/protorpc/blob/eb03145/python/protorpc/protobuf.py#L49-L51
-	case "application/octet-stream", "application/x-google-protobuf", "application/protobuf", "application/x-protobuf":
+	case isProtobufMediaType(r.Header.Get("Content-Type")):
 		switch m := v.(type) {
 		case proto.Message:
 			err = proto.Unmarshal(b, m)

request_test.go

@@ -77,6 +77,18 @@ func TestRequestDecodeProto(t *testing.T) {
 	assert.Equal(t, "Hello world!", g2.Message)
 }
 
+func TestRequestDecodeProtoWithMediaTypeParameters(t *testing.T) {
+	req := NewRequest(nil, "GET", "/", nil)
+	b, _ := proto.Marshal(&prototest.Greeting{Message: "Hello world!"})
+	req.Header.Set("Content-Type", "application/protobuf; charset=binary")
+	req.Body = newDoneReader(ioutil.NopCloser(bytes.NewReader(b)), -1)
+
+	g := &prototest.Greeting{}
+	err := req.Decode(g)
+	assert.NoError(t, err)
+	assert.Equal(t, "Hello world!", g.Message)
+}
+
 func TestRequestDecodeProtoMaskingAsJSON(t *testing.T) {
 	req := NewRequest(nil, "GET", "/", nil)
 	b := []byte("{\"message\":\"Hello world!\"}\n")

response.go

@@ -7,7 +7,6 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
-	"strings"
 
 	legacyproto "github.com/golang/protobuf/proto"
 	"github.com/monzo/terrors"
@@ -50,7 +49,7 @@ func (r *Response) Encode(v interface{}) {
 	switch m := v.(type) {
 	case proto.Message:
 		// if we didn't ask for protobuf, send JSON
-		if !strings.Contains(r.Request.Header.Get("Accept"), "application/protobuf") {
+		if !acceptsProtobuf(requestHeader(r.Request)) {
 			r.EncodeAsProtobufJSON(m)
 			return
 		}
@@ -59,7 +58,7 @@ func (r *Response) Encode(v interface{}) {
 		return
 	case legacyproto.Message:
 		// if we asked for protobuf, send it using the legacy encoder for the error filter.
-		if strings.Contains(r.Request.Header.Get("Accept"), "application/protobuf") {
+		if acceptsProtobuf(requestHeader(r.Request)) {
 			r.EncodeAsLegacyProtobuf(m)
 			return
 		}
@@ -152,7 +151,7 @@ func (r *Response) Decode(v interface{}) error {
 		return r.Error
 	}
 
-	contentType := r.Header.Get("Content-Type")
+	contentType := canonicalMediaType(r.Header.Get("Content-Type"))
 
 	params := map[string]string{
 		"response_content_type": contentType,
@@ -165,15 +164,10 @@ func (r *Response) Decode(v interface{}) error {
 	case proto.Message:
 		params["response_object_type"] = "protobuf"
 
-		switch contentType {
-		case "application/octet-stream",
-			"application/x-google-protobuf",
-			"application/protobuf",
-			"application/x-protobuf":
-
+		switch {
+		case isProtobufMediaType(contentType):
 			err = proto.Unmarshal(b, m)
 		default:
-
 			err = protojson.Unmarshal(b, m)
 		}
 
@@ -182,12 +176,8 @@ func (r *Response) Decode(v interface{}) error {
 	// Upgrade to google.golang.org/protobuf/proto.Message as soon as possible.
 	case legacyproto.Message:
 		params["response_object_type"] = "legacyproto"
-		switch contentType {
-		case "application/octet-stream",
-			"application/x-google-protobuf",
-			"application/protobuf",
-			"application/x-protobuf":
-
+		switch {
+		case isProtobufMediaType(contentType):
 			err = legacyproto.Unmarshal(b, m)
 		default:
 			err = json.Unmarshal(b, m)
@@ -302,6 +292,13 @@ func (r Response) String() string {
 	return b.String()
 }
 
+func requestHeader(req *Request) http.Header {
+	if req == nil {
+		return nil
+	}
+	return req.Header
+}
+
 func newHTTPResponse(req Request, statusCode int) *http.Response {
 	return &http.Response{
 		StatusCode:    statusCode,

response_test.go

@@ -222,6 +222,24 @@ func TestResponseDecodeProtobufWithAltType(t *testing.T) {
 	assert.EqualValues(t, 1, gout.Priority)
 }
 
+func TestResponseDecodeProtobufWithMediaTypeParameters(t *testing.T) {
+	t.Parallel()
+
+	g := &prototest.Greeting{
+		Message:  "Hello world!",
+		Priority: 1,
+	}
+	b, _ := proto.Marshal(g)
+	rsp := NewResponse(Request{})
+	rsp.Body = ioutil.NopCloser(bytes.NewReader(b))
+	rsp.Header.Set("Content-Type", "application/protobuf; charset=binary")
+
+	gout := &prototest.Greeting{}
+	assert.NoError(t, rsp.Decode(gout))
+	assert.Equal(t, "Hello world!", gout.Message)
+	assert.EqualValues(t, 1, gout.Priority)
+}
+
 // TestResponseDecodeLegacyProtobuf verifies decoding of a legacy protobuf message
 func TestResponseDecodeLegacyProtobuf(t *testing.T) {
 	t.Parallel()
@@ -409,3 +427,31 @@ func TestResponseEncodeErrorGivesTerror(t *testing.T) {
 	assert.True(t, terrors.Is(rsp.Error, "internal_service"))
 	assert.True(t, terrors.Matches(rsp.Error, "unsupported value"))
 }
+
+func TestResponseEncodeProtobufWithoutRequestFallsBackToJSON(t *testing.T) {
+	t.Parallel()
+
+	rsp := Response{}
+	rsp.Encode(&prototest.Greeting{Message: "hello"})
+
+	assert.NoError(t, rsp.Error)
+	assert.Equal(t, "application/json", rsp.Header.Get("Content-Type"))
+	body, err := ioutil.ReadAll(rsp.Body)
+	require.NoError(t, err)
+	assert.Contains(t, string(body), "\"message\":\"hello\"")
+}
+
+func TestResponseEncodeProtobufWithComplexAcceptHeader(t *testing.T) {
+	t.Parallel()
+
+	req := NewRequest(nil, "GET", "/", nil)
+	req.Header.Set("Accept", "application/json;q=0.5, application/protobuf; charset=binary")
+	rsp := Response{Request: &req}
+	rsp.Encode(&prototest.Greeting{Message: "hello"})
+
+	assert.NoError(t, rsp.Error)
+	assert.Equal(t, "application/protobuf", rsp.Header.Get("Content-Type"))
+	body, err := ioutil.ReadAll(rsp.Body)
+	require.NoError(t, err)
+	assert.Subset(t, body, []byte("hello"), "'hello' should appear in the wire format")
+}

router.go

@@ -15,11 +15,13 @@ import (
 type routerContextKeyType struct{}
 type routerRequestPatternContextKeyType struct{}
 type routerRequestMethodContextKeyType struct{}
+type routerRequestParamsContextKeyType struct{}
 
 var (
 	routerContextKey               = routerContextKeyType{}
 	routerRequestPatternContextKey = routerRequestPatternContextKeyType{}
 	routerRequestMethodContextKey  = routerRequestMethodContextKeyType{}
+	routerRequestParamsContextKey  = routerRequestParamsContextKeyType{}
 	routerComponentsRe             = regexp.MustCompile(`(?:^|/)(\*\w*|:\w+)`)
 )
 
@@ -71,6 +73,27 @@ func RequestMethodFromContext(ctx context.Context) (string, bool) {
 	return "", false
 }
 
+// RequestParamsFromContext returns the route params captured for the request, if available.
+func RequestParamsFromContext(ctx context.Context) (map[string]string, bool) {
+	if v := ctx.Value(routerRequestParamsContextKey); v != nil {
+		params := v.(map[string]string)
+		return cloneRouteParams(params), true
+	}
+	return nil, false
+}
+
+func cloneRouteParams(params map[string]string) map[string]string {
+	if len(params) == 0 {
+		return map[string]string{}
+	}
+
+	cloned := make(map[string]string, len(params))
+	for key, value := range params {
+		cloned[key] = value
+	}
+	return cloned
+}
+
 func (r *Router) compile(pattern string) *regexp.Regexp {
 	re, pos := ``, 0
 	for _, m := range routerComponentsRe.FindAllStringSubmatchIndex(pattern, -1) {
@@ -143,7 +166,8 @@ func (r Router) Lookup(method, path string) (Service, string, map[string]string,
 // Serve returns a Service which will route inbound requests to the enclosed routes.
 func (r Router) Serve() Service {
 	return func(req Request) Response {
-		svc, pathPattern, ok := r.lookup(req.Method, req.URL.Path, nil)
+		params := map[string]string{}
+		svc, pathPattern, ok := r.lookup(req.Method, req.URL.Path, params)
 		if !ok {
 			txt := fmt.Sprintf("No handler for %s %s", req.Method, req.URL.Path)
 			rsp := NewResponse(req)
@@ -153,6 +177,7 @@ func (r Router) Serve() Service {
 		req.Context = context.WithValue(req.Context, routerContextKey, &r)
 		req.Context = context.WithValue(req.Context, routerRequestPatternContextKey, pathPattern)
 		req.Context = context.WithValue(req.Context, routerRequestMethodContextKey, req.Method)
+		req.Context = context.WithValue(req.Context, routerRequestParamsContextKey, cloneRouteParams(params))
 		rsp := svc(req)
 		if rsp.Request == nil {
 			rsp.Request = &req
@@ -169,6 +194,9 @@ func (r Router) Pattern(req Request) string {
 
 // Params returns extracted path parameters, assuming the request has been routed and has captured parameters.
 func (r Router) Params(req Request) map[string]string {
+	if params, ok := RequestParamsFromContext(req.Context); ok {
+		return params
+	}
 	_, _, params, _ := r.Lookup(req.Method, req.URL.Path)
 	return params
 }

router_test.go

@@ -148,3 +148,48 @@ func TestRouterSetsContextValues(t *testing.T) {
 	assert.True(t, ok)
 	assert.Equal(t, "GET", ctxMethod)
 }
+
+func TestRouterSetsContextParams(t *testing.T) {
+	t.Parallel()
+
+	router := Router{}
+	router.GET("/users/:userID/orders/:orderID", func(req Request) Response {
+		return Response{}
+	})
+
+	req := NewRequest(context.Background(), "GET", "/users/u-123/orders/o-456", nil)
+	rsp := router.Serve()(req)
+	require.NotNil(t, rsp.Request)
+
+	params, ok := RequestParamsFromContext(rsp.Request.Context)
+	require.True(t, ok)
+	assert.Equal(t, map[string]string{
+		"userID":  "u-123",
+		"orderID": "o-456",
+	}, params)
+}
+
+func TestRouterParamsUsesCapturedParamsWhenPathChanges(t *testing.T) {
+	t.Parallel()
+
+	router := Router{}
+	router.GET("/users/:userID/orders/:orderID", func(req Request) Response {
+		req.URL.Path = "/mutated/path"
+
+		params := router.Params(req)
+		assert.Equal(t, map[string]string{
+			"userID":  "u-123",
+			"orderID": "o-456",
+		}, params)
+
+		// Ensure callers don't get a shared mutable map back.
+		params["userID"] = "changed"
+		again := router.Params(req)
+		assert.Equal(t, "u-123", again["userID"])
+		return Response{}
+	})
+
+	req := NewRequest(context.Background(), "GET", "/users/u-123/orders/o-456", nil)
+	rsp := router.Serve()(req)
+	require.NotNil(t, rsp.Request)
+}