runtime: make HTTPBodyMarshaler the default

runtime/errors.go

@@ -80,13 +80,7 @@ func defaultHTTPErrorHandler(ctx context.Context, mux *ServeMux, marshaler Marsh
 
 	w.Header().Del("Trailer")
 
-	contentType := marshaler.ContentType()
-	// Check marshaler at runtime in order to keep backwards compatibility.
-	// An interface param needs to be added to the ContentType() function on
-	// the Marshal interface to be able to remove this check
-	if typeMarshaler, ok := marshaler.(contentTypeMarshaler); ok {
-		contentType = typeMarshaler.ContentTypeFromMessage(pb)
-	}
+	contentType := marshaler.ContentType(pb)
 	w.Header().Set("Content-Type", contentType)
 
 	buf, merr := marshaler.Marshal(pb)

runtime/handler.go

@@ -32,7 +32,6 @@ func ForwardResponseStream(ctx context.Context, mux *ServeMux, marshaler Marshal
 	handleForwardResponseServerMetadata(w, mux, md)
 
 	w.Header().Set("Transfer-Encoding", "chunked")
-	w.Header().Set("Content-Type", marshaler.ContentType())
 	if err := handleForwardResponseOptions(ctx, w, nil, opts); err != nil {
 		HTTPError(ctx, mux, marshaler, w, req, err)
 		return
@@ -60,14 +59,17 @@ func ForwardResponseStream(ctx context.Context, mux *ServeMux, marshaler Marshal
 			return
 		}
 
+		if !wroteHeader {
+			w.Header().Set("Content-Type", marshaler.ContentType(resp))
+		}
+
 		var buf []byte
 		httpBody, isHTTPBody := resp.(*httpbody.HttpBody)
 		switch {
 		case resp == nil:
 			buf, err = marshaler.Marshal(errorChunk(status.New(codes.Internal, "empty response")))
 		case isHTTPBody:
 			buf = httpBody.GetData()
-			w.Header().Set("Content-Type", httpBody.GetContentType())
 		default:
 			result := map[string]interface{}{"result": resp}
 			if rb, ok := resp.(responseBody); ok {
@@ -137,13 +139,7 @@ func ForwardResponseMessage(ctx context.Context, mux *ServeMux, marshaler Marsha
 	handleForwardResponseServerMetadata(w, mux, md)
 	handleForwardResponseTrailerHeader(w, md)
 
-	contentType := marshaler.ContentType()
-	// Check marshaler on run time in order to keep backwards compatibility
-	// An interface param needs to be added to the ContentType() function on
-	// the Marshal interface to be able to remove this check
-	if typeMarshaler, ok := marshaler.(contentTypeMarshaler); ok {
-		contentType = typeMarshaler.ContentTypeFromMessage(resp)
-	}
+	contentType := marshaler.ContentType(resp)
 	w.Header().Set("Content-Type", contentType)
 
 	if err := handleForwardResponseOptions(ctx, w, resp, opts); err != nil {

runtime/handler_test.go

@@ -166,12 +166,11 @@ type CustomMarshaler struct {
 	m *runtime.JSONPb
 }
 
-func (c *CustomMarshaler) Marshal(v interface{}) ([]byte, error)       { return c.m.Marshal(v) }
-func (c *CustomMarshaler) Unmarshal(data []byte, v interface{}) error  { return c.m.Unmarshal(data, v) }
-func (c *CustomMarshaler) NewDecoder(r io.Reader) runtime.Decoder      { return c.m.NewDecoder(r) }
-func (c *CustomMarshaler) NewEncoder(w io.Writer) runtime.Encoder      { return c.m.NewEncoder(w) }
-func (c *CustomMarshaler) ContentType() string                         { return c.m.ContentType() }
-func (c *CustomMarshaler) ContentTypeFromMessage(v interface{}) string { return "Custom-Content-Type" }
+func (c *CustomMarshaler) Marshal(v interface{}) ([]byte, error)      { return c.m.Marshal(v) }
+func (c *CustomMarshaler) Unmarshal(data []byte, v interface{}) error { return c.m.Unmarshal(data, v) }
+func (c *CustomMarshaler) NewDecoder(r io.Reader) runtime.Decoder     { return c.m.NewDecoder(r) }
+func (c *CustomMarshaler) NewEncoder(w io.Writer) runtime.Encoder     { return c.m.NewEncoder(w) }
+func (c *CustomMarshaler) ContentType(v interface{}) string           { return "Custom-Content-Type" }
 
 func TestForwardResponseStreamCustomMarshaler(t *testing.T) {
 	type msg struct {

runtime/marshal_httpbodyproto.go

@@ -4,13 +4,6 @@ import (
 	"google.golang.org/genproto/googleapis/api/httpbody"
 )
 
-// SetHTTPBodyMarshaler overwrites the default marshaler with the HTTPBodyMarshaler
-func SetHTTPBodyMarshaler(serveMux *ServeMux) {
-	serveMux.marshalers.mimeMap[MIMEWildcard] = &HTTPBodyMarshaler{
-		Marshaler: defaultMarshaler,
-	}
-}
-
 // HTTPBodyMarshaler is a Marshaler which supports marshaling of a
 // google.api.HttpBody message as the full response body if it is
 // the actual message used as the response. If not, then this will
@@ -19,18 +12,14 @@ type HTTPBodyMarshaler struct {
 	Marshaler
 }
 
-// ContentType implementation to keep backwards compatibility with marshal interface
-func (h *HTTPBodyMarshaler) ContentType() string {
-	return h.ContentTypeFromMessage(nil)
-}
-
-// ContentTypeFromMessage in case v is a google.api.HttpBody message it returns
-// its specified content type otherwise fall back to the default Marshaler.
-func (h *HTTPBodyMarshaler) ContentTypeFromMessage(v interface{}) string {
+// ContentType returns its specified content type in case v is a
+// google.api.HttpBody message, otherwise it will fall back to the default Marshalers
+// content type.
+func (h *HTTPBodyMarshaler) ContentType(v interface{}) string {
 	if httpBody, ok := v.(*httpbody.HttpBody); ok {
 		return httpBody.GetContentType()
 	}
-	return h.Marshaler.ContentType()
+	return h.Marshaler.ContentType(v)
 }
 
 // Marshal marshals "v" by returning the body bytes if v is a

runtime/marshal_httpbodyproto_test.go

@@ -21,11 +21,11 @@ func TestHTTPBodyContentType(t *testing.T) {
 	message := &httpbody.HttpBody{
 		ContentType: expected,
 	}
-	res := m.ContentType()
+	res := m.ContentType(nil)
 	if res != "application/json" {
 		t.Errorf("content type not equal (%q, %q)", res, expected)
 	}
-	res = m.ContentTypeFromMessage(message)
+	res = m.ContentType(message)
 	if res != expected {
 		t.Errorf("content type not equal (%q, %q)", res, expected)
 	}

runtime/marshal_json.go

@@ -15,7 +15,7 @@ import (
 type JSONBuiltin struct{}
 
 // ContentType always Returns "application/json".
-func (*JSONBuiltin) ContentType() string {
+func (*JSONBuiltin) ContentType(_ interface{}) string {
 	return "application/json"
 }
 

runtime/marshal_jsonpb.go

@@ -23,7 +23,7 @@ type JSONPb struct {
 }
 
 // ContentType always returns "application/json".
-func (*JSONPb) ContentType() string {
+func (*JSONPb) ContentType(_ interface{}) string {
 	return "application/json"
 }
 
@@ -65,8 +65,8 @@ var (
 )
 
 // marshalNonProto marshals a non-message field of a protobuf message.
-// This function does not correctly marshals arbitrary data structure into JSON,
-// but it is only capable of marshaling non-message field values of protobuf,
+// This function does not correctly marshal arbitrary data structures into JSON,
+// it is only capable of marshaling non-message field values of protobuf,
 // i.e. primitive types, enums; pointers to primitives or enums; maps from
 // integer/string types to primitives/enums/pointers to messages.
 func (j *JSONPb) marshalNonProtoField(v interface{}) ([]byte, error) {

runtime/marshal_proto.go

@@ -13,7 +13,7 @@ import (
 type ProtoMarshaller struct{}
 
 // ContentType always returns "application/octet-stream".
-func (*ProtoMarshaller) ContentType() string {
+func (*ProtoMarshaller) ContentType(_ interface{}) string {
 	return "application/octet-stream"
 }
 

runtime/marshaler.go

@@ -16,14 +16,9 @@ type Marshaler interface {
 	// NewEncoder returns an Encoder which writes bytes sequence into "w".
 	NewEncoder(w io.Writer) Encoder
 	// ContentType returns the Content-Type which this marshaler is responsible for.
-	ContentType() string
-}
-
-// Marshalers that implement contentTypeMarshaler will have their ContentTypeFromMessage method called
-// to set the Content-Type header on the response
-type contentTypeMarshaler interface {
-	// ContentTypeFromMessage returns the Content-Type this marshaler produces from the provided message
-	ContentTypeFromMessage(v interface{}) string
+	// The parameter describes the type which is being marshalled, which can sometimes
+	// affect the content type returned.
+	ContentType(v interface{}) string
 }
 
 // Decoder decodes a byte sequence

runtime/marshaler_registry.go

@@ -15,9 +15,11 @@ var (
 	acceptHeader      = http.CanonicalHeaderKey("Accept")
 	contentTypeHeader = http.CanonicalHeaderKey("Content-Type")
 
-	defaultMarshaler = &JSONPb{
-		MarshalOptions: protojson.MarshalOptions{
-			EmitUnpopulated: true,
+	defaultMarshaler = &HTTPBodyMarshaler{
+		Marshaler: &JSONPb{
+			MarshalOptions: protojson.MarshalOptions{
+				EmitUnpopulated: true,
+			},
 		},
 	}
 )

runtime/marshaler_registry_test.go

@@ -20,11 +20,11 @@ func TestMarshalerForRequest(t *testing.T) {
 	mux := runtime.NewServeMux()
 
 	in, out := runtime.MarshalerForRequest(mux, r)
-	if _, ok := in.(*runtime.JSONPb); !ok {
-		t.Errorf("in = %#v; want a runtime.JSONPb", in)
+	if _, ok := in.(*runtime.HTTPBodyMarshaler); !ok {
+		t.Errorf("in = %#v; want a runtime.HTTPBodyMarshaler", in)
 	}
-	if _, ok := out.(*runtime.JSONPb); !ok {
-		t.Errorf("out = %#v; want a runtime.JSONPb", in)
+	if _, ok := out.(*runtime.HTTPBodyMarshaler); !ok {
+		t.Errorf("out = %#v; want a runtime.HTTPBodyMarshaler", in)
 	}
 
 	var marshalers [3]dummyMarshaler
@@ -78,7 +78,7 @@ func TestMarshalerForRequest(t *testing.T) {
 
 type dummyMarshaler struct{}
 
-func (dummyMarshaler) ContentType() string { return "" }
+func (dummyMarshaler) ContentType(_ interface{}) string { return "" }
 func (dummyMarshaler) Marshal(interface{}) ([]byte, error) {
 	return nil, errors.New("not implemented")
 }