add support for the google.api.HttpBody proto as a response

runtime/errors.go

@@ -84,9 +84,6 @@ func (*errorBody) ProtoMessage()    {}
 func DefaultHTTPError(ctx context.Context, mux *ServeMux, marshaler Marshaler, w http.ResponseWriter, _ *http.Request, err error) {
 	const fallback = `{"error": "failed to marshal error message"}`
 
-	w.Header().Del("Trailer")
-	w.Header().Set("Content-Type", marshaler.ContentType())
-
 	s, ok := status.FromError(err)
 	if !ok {
 		s = status.New(codes.Unknown, err.Error())
@@ -108,6 +105,9 @@ func DefaultHTTPError(ctx context.Context, mux *ServeMux, marshaler Marshaler, w
 		}
 	}
 
+	w.Header().Del("Trailer")
+	w.Header().Set("Content-Type", marshaler.ContentType(body))
+
 	buf, merr := marshaler.Marshal(body)
 	if merr != nil {
 		grpclog.Printf("Failed to marshal error message %q: %v", body, merr)

runtime/handler.go

@@ -33,7 +33,7 @@ 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())
+	w.Header().Set("Content-Type", marshaler.ContentType(nil))
 	if err := handleForwardResponseOptions(ctx, w, nil, opts); err != nil {
 		HTTPError(ctx, mux, marshaler, w, req, err)
 		return
@@ -47,6 +47,7 @@ func ForwardResponseStream(ctx context.Context, mux *ServeMux, marshaler Marshal
 	}
 
 	var wroteHeader bool
+	contentTypeSet := false
 	for {
 		resp, err := recv()
 		if err == io.EOF {
@@ -61,7 +62,12 @@ func ForwardResponseStream(ctx context.Context, mux *ServeMux, marshaler Marshal
 			return
 		}
 
-		buf, err := marshaler.Marshal(streamChunk(resp, nil))
+		chunk := streamChunk(resp, nil)
+		if !contentTypeSet {
+			w.Header().Set("Content-Type", marshaler.ContentType(chunk))
+			contentTypeSet = true
+		}
+		buf, err := marshaler.Marshal(chunk)
 		if err != nil {
 			grpclog.Printf("Failed to marshal response chunk: %v", err)
 			handleForwardResponseStreamError(wroteHeader, marshaler, w, err)
@@ -115,7 +121,7 @@ func ForwardResponseMessage(ctx context.Context, mux *ServeMux, marshaler Marsha
 
 	handleForwardResponseServerMetadata(w, mux, md)
 	handleForwardResponseTrailerHeader(w, md)
-	w.Header().Set("Content-Type", marshaler.ContentType())
+	w.Header().Set("Content-Type", marshaler.ContentType(resp))
 	if err := handleForwardResponseOptions(ctx, w, resp, opts); err != nil {
 		HTTPError(ctx, mux, marshaler, w, req, err)
 		return

runtime/handler_test.go

@@ -137,7 +137,7 @@ func (c *CustomMarshaler) Marshal(v interface{}) ([]byte, error)      { return c
 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) ContentType(i interface{}) string           { return c.m.ContentType(i) }
 
 func TestForwardResponseStreamCustomMarshaler(t *testing.T) {
 	type msg struct {

runtime/marshal_httpbody.go

@@ -0,0 +1,74 @@
+package runtime
+
+import (
+	"io"
+	"reflect"
+
+	"github.com/golang/protobuf/proto"
+	hb "google.golang.org/genproto/googleapis/api/httpbody"
+)
+
+var (
+	backupMarshaler = &JSONPb{OrigName: true}
+)
+
+// 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
+// simply fallback to the JSONPb marshaler.
+type HTTPBodyMarshaler struct{}
+
+// ContentType returns the type specified in the google.api.HttpBody
+// proto if "v" is a google.api.HttpBody proto, otherwise returns
+// "application/json".
+func (*HTTPBodyMarshaler) ContentType(v interface{}) string {
+	if h := tryHTTPBody(v); h != nil {
+		return h.GetContentType()
+	}
+	return "application/json"
+}
+
+// Marshal marshals "v" by returning the body bytes if v is a
+// google.api.HttpBody message, or it marshals to JSON.
+func (*HTTPBodyMarshaler) Marshal(v interface{}) ([]byte, error) {
+	if h := tryHTTPBody(v); h != nil {
+		return h.GetData(), nil
+	}
+	return backupMarshaler.Marshal(v)
+}
+
+// Unmarshal unmarshals JSON data into "v".
+// google.api.HttpBody messages are not supported on the request.
+func (*HTTPBodyMarshaler) Unmarshal(data []byte, v interface{}) error {
+	return backupMarshaler.Unmarshal(data, v)
+}
+
+// NewDecoder returns a Decoder which reads JSON stream from "r".
+func (*HTTPBodyMarshaler) NewDecoder(r io.Reader) Decoder {
+	return backupMarshaler.NewDecoder(r)
+}
+
+// NewEncoder returns an Encoder which writes JSON stream into "w".
+func (*HTTPBodyMarshaler) NewEncoder(w io.Writer) Encoder {
+	return backupMarshaler.NewEncoder(w)
+}
+
+func tryHTTPBody(v interface{}) *hb.HttpBody {
+	rv := reflect.ValueOf(v)
+	if rv.Kind() != reflect.Ptr {
+		return nil
+	}
+	for rv.Kind() == reflect.Ptr {
+		if rv.IsNil() {
+			rv.Set(reflect.New(rv.Type().Elem()))
+		}
+		if rv.Type().ConvertibleTo(typeProtoMessage) {
+			pb := rv.Interface().(proto.Message)
+			if proto.MessageName(pb) == "google.api.HttpBody" {
+				return v.(*hb.HttpBody)
+			}
+		}
+		rv = rv.Elem()
+	}
+	return nil
+}

runtime/marshal_json.go

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

runtime/marshal_jsonpb.go

@@ -17,7 +17,7 @@ import (
 type JSONPb jsonpb.Marshaler
 
 // ContentType always returns "application/json".
-func (*JSONPb) ContentType() string {
+func (*JSONPb) ContentType(v interface{}) string {
 	return "application/json"
 }
 

runtime/marshaler.go

@@ -15,8 +15,8 @@ type Marshaler interface {
 	NewDecoder(r io.Reader) Decoder
 	// 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
+	// ContentType returns the response Content-Type for "v".
+	ContentType(v interface{}) string
 }
 
 // Decoder decodes a byte sequence

runtime/marshaler_registry.go

@@ -13,7 +13,7 @@ var (
 	acceptHeader      = http.CanonicalHeaderKey("Accept")
 	contentTypeHeader = http.CanonicalHeaderKey("Content-Type")
 
-	defaultMarshaler = &JSONPb{OrigName: true}
+	defaultMarshaler = &HTTPBodyMarshaler{}
 )
 
 // MarshalerForRequest returns the inbound/outbound marshalers for this request.

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(v interface{}) string { return "" }
 func (dummyMarshaler) Marshal(interface{}) ([]byte, error) {
 	return nil, errors.New("not implemented")
 }

runtime/proto_errors.go

@@ -26,15 +26,16 @@ func DefaultHTTPProtoErrorHandler(ctx context.Context, mux *ServeMux, marshaler
 	// return Internal when Marshal failed
 	const fallback = `{"code": 13, "message": "failed to marshal error message"}`
 
-	w.Header().Del("Trailer")
-	w.Header().Set("Content-Type", marshaler.ContentType())
-
 	s, ok := status.FromError(err)
 	if !ok {
 		s = status.New(codes.Unknown, err.Error())
 	}
 
-	buf, merr := marshaler.Marshal(s.Proto())
+	pb := s.Proto()
+	w.Header().Del("Trailer")
+	w.Header().Set("Content-Type", marshaler.ContentType(pb))
+
+	buf, merr := marshaler.Marshal(pb)
 	if merr != nil {
 		grpclog.Printf("Failed to marshal error message %q: %v", s.Proto(), merr)
 		w.WriteHeader(http.StatusInternalServerError)