Support non-message fields in JSONPb marshaler

It allows us to specify non-message fields in google.apis.HttpRule.body
while using JSONPb marshaler

Author: yugui

runtime/marshal_jsonpb.go

@@ -3,8 +3,9 @@ package runtime
 import (
 	"bytes"
 	"encoding/json"
-	"errors"
+	"fmt"
 	"io"
+	"reflect"
 
 	"github.com/golang/protobuf/jsonpb"
 	"github.com/golang/protobuf/proto"
@@ -39,79 +40,159 @@ func (*JSONPb) ContentType() string {
 // Currently it can marshal only proto.Message.
 // TODO(yugui) Support fields of primitive types in a message.
 func (j *JSONPb) Marshal(v interface{}) ([]byte, error) {
+	if _, ok := v.(proto.Message); !ok {
+		return j.marshalNonProtoField(v)
+	}
+
+	var buf bytes.Buffer
+	if err := j.marshalTo(&buf, v); err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+func (j *JSONPb) marshalTo(w io.Writer, v interface{}) error {
+	p, ok := v.(proto.Message)
+	if !ok {
+		buf, err := j.marshalNonProtoField(v)
+		if err != nil {
+			return err
+		}
+		_, err = w.Write(buf)
+		return err
+	}
 	m := &jsonpb.Marshaler{
 		EnumsAsInts:  j.EnumsAsInts,
 		EmitDefaults: j.EmitDefaults,
 		Indent:       j.Indent,
 		OrigName:     j.OrigName,
 	}
-	p, ok := v.(proto.Message)
-	if !ok {
-		return nil, errors.New("interface is not proto.Message")
-	}
+	return m.Marshal(w, p)
+}
 
-	var buf bytes.Buffer
-	if err := m.Marshal(&buf, p); err != nil {
-		return nil, err
+// marshalNonProto marshals a non-message field of a protobuf message.
+// This function does not correctly marshals arbitary data structure into JSON,
+// but 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) {
+	rv := reflect.ValueOf(v)
+	for rv.Kind() == reflect.Ptr {
+		if rv.IsNil() {
+			return []byte("null"), nil
+		}
+		rv = rv.Elem()
 	}
-	return buf.Bytes(), nil
 
+	if rv.Kind() == reflect.Map {
+		m := make(map[string]*json.RawMessage)
+		for _, k := range rv.MapKeys() {
+			buf, err := j.Marshal(rv.MapIndex(k).Interface())
+			if err != nil {
+				return nil, err
+			}
+			m[fmt.Sprintf("%v", k.Interface())] = (*json.RawMessage)(&buf)
+		}
+		if j.Indent != "" {
+			return json.MarshalIndent(m, "", j.Indent)
+		}
+		return json.Marshal(m)
+	}
+	if enum, ok := rv.Interface().(protoEnum); ok && !j.EnumsAsInts {
+		return json.Marshal(enum.String())
+	}
+	return json.Marshal(rv.Interface())
 }
 
 // Unmarshal unmarshals JSON "data" into "v"
 // Currently it can marshal only proto.Message.
 // TODO(yugui) Support fields of primitive types in a message.
 func (j *JSONPb) Unmarshal(data []byte, v interface{}) error {
-	r := bytes.NewReader(data)
-	p, ok := v.(proto.Message)
-	if !ok {
-		return errors.New("interface is not proto.Message")
-	}
-	return jsonpb.Unmarshal(r, p)
+	return unmarshalJSONPb(data, v)
 }
 
 // NewDecoder returns a Decoder which reads JSON stream from "r".
 func (j *JSONPb) NewDecoder(r io.Reader) Decoder {
-	return &jsonPbDecoder{decoder: json.NewDecoder(r)}
+	d := json.NewDecoder(r)
+	return DecoderFunc(func(v interface{}) error { return decodeJSONPb(d, v) })
 }
 
 // NewEncoder returns an Encoder which writes JSON stream into "w".
 func (j *JSONPb) NewEncoder(w io.Writer) Encoder {
-	m := &jsonpb.Marshaler{
-		EnumsAsInts:  j.EnumsAsInts,
-		EmitDefaults: j.EmitDefaults,
-		Indent:       j.Indent,
-		OrigName:     j.OrigName,
-	}
-
-	return &jsonPbEncoder{
-		marshal: m,
-		writer:  w,
-	}
+	return EncoderFunc(func(v interface{}) error { return j.marshalTo(w, v) })
 }
 
-type jsonPbDecoder struct {
-	decoder *json.Decoder
+func unmarshalJSONPb(data []byte, v interface{}) error {
+	d := json.NewDecoder(bytes.NewReader(data))
+	return decodeJSONPb(d, v)
 }
 
-func (j *jsonPbDecoder) Decode(v interface{}) error {
+func decodeJSONPb(d *json.Decoder, v interface{}) error {
 	p, ok := v.(proto.Message)
 	if !ok {
-		return errors.New("interface is not proto.Message")
+		return decodeNonProtoField(d, v)
 	}
-
-	return jsonpb.UnmarshalNext(j.decoder, p)
+	return jsonpb.UnmarshalNext(d, p)
 }
 
-type jsonPbEncoder struct {
-	marshal *jsonpb.Marshaler
-	writer  io.Writer
+func decodeNonProtoField(d *json.Decoder, v interface{}) error {
+	rv := reflect.ValueOf(v)
+	if rv.Kind() != reflect.Ptr {
+		return fmt.Errorf("%T is not a pointer", v)
+	}
+	for rv.Kind() == reflect.Ptr {
+		if rv.IsNil() {
+			rv.Set(reflect.New(rv.Type().Elem()))
+		}
+		rv = rv.Elem()
+	}
+	if rv.Kind() == reflect.Map {
+		if rv.IsNil() {
+			rv.Set(reflect.MakeMap(rv.Type()))
+		}
+		conv, ok := convFromType[rv.Type().Key().Kind()]
+		if !ok {
+			return fmt.Errorf("unsupported type of map field key: %v", rv.Type().Key())
+		}
+
+		m := make(map[string]*json.RawMessage)
+		if err := d.Decode(&m); err != nil {
+			return err
+		}
+		for k, v := range m {
+			result := conv.Call([]reflect.Value{reflect.ValueOf(k)})
+			if err := result[1].Interface(); err != nil {
+				return err.(error)
+			}
+			bk := result[0]
+			bv := reflect.New(rv.Type().Elem())
+			if err := unmarshalJSONPb([]byte(*v), bv.Interface()); err != nil {
+				return err
+			}
+			rv.SetMapIndex(bk, bv.Elem())
+		}
+		return nil
+	}
+	if _, ok := rv.Interface().(protoEnum); ok {
+		var repr interface{}
+		if err := d.Decode(&repr); err != nil {
+			return err
+		}
+		switch repr.(type) {
+		case string:
+			// TODO(yugui) Should use proto.StructProperties?
+			return fmt.Errorf("unmarshaling of symbolic enum %q not supported: %T", repr, rv.Interface())
+		case float64:
+			rv.Set(reflect.ValueOf(int32(repr.(float64))).Convert(rv.Type()))
+			return nil
+		default:
+			return fmt.Errorf("cannot assign %#v into Go type %T", repr, rv.Interface())
+		}
+	}
+	return d.Decode(v)
 }
 
-func (j *jsonPbEncoder) Encode(v interface{}) error {
-	p, ok := v.(proto.Message)
-	if !ok {
-		return errors.New("interface is not proto.Message")
-	}
-	return j.marshal.Marshal(j.writer, p)
+type protoEnum interface {
+	fmt.Stringer
+	EnumDescriptor() ([]byte, []int)
 }