marshal_jsonpb.go 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262
  1. package runtime
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "reflect"
  8. "github.com/golang/protobuf/jsonpb"
  9. "github.com/golang/protobuf/proto"
  10. )
  11. // JSONPb is a Marshaler which marshals/unmarshals into/from JSON
  12. // with the "github.com/golang/protobuf/jsonpb".
  13. // It supports fully functionality of protobuf unlike JSONBuiltin.
  14. //
  15. // The NewDecoder method returns a DecoderWrapper, so the underlying
  16. // *json.Decoder methods can be used.
  17. type JSONPb jsonpb.Marshaler
  18. // ContentType always returns "application/json".
  19. func (*JSONPb) ContentType() string {
  20. return "application/json"
  21. }
  22. // Marshal marshals "v" into JSON.
  23. func (j *JSONPb) Marshal(v interface{}) ([]byte, error) {
  24. if _, ok := v.(proto.Message); !ok {
  25. return j.marshalNonProtoField(v)
  26. }
  27. var buf bytes.Buffer
  28. if err := j.marshalTo(&buf, v); err != nil {
  29. return nil, err
  30. }
  31. return buf.Bytes(), nil
  32. }
  33. func (j *JSONPb) marshalTo(w io.Writer, v interface{}) error {
  34. p, ok := v.(proto.Message)
  35. if !ok {
  36. buf, err := j.marshalNonProtoField(v)
  37. if err != nil {
  38. return err
  39. }
  40. _, err = w.Write(buf)
  41. return err
  42. }
  43. return (*jsonpb.Marshaler)(j).Marshal(w, p)
  44. }
  45. var (
  46. // protoMessageType is stored to prevent constant lookup of the same type at runtime.
  47. protoMessageType = reflect.TypeOf((*proto.Message)(nil)).Elem()
  48. )
  49. // marshalNonProto marshals a non-message field of a protobuf message.
  50. // This function does not correctly marshals arbitrary data structure into JSON,
  51. // but it is only capable of marshaling non-message field values of protobuf,
  52. // i.e. primitive types, enums; pointers to primitives or enums; maps from
  53. // integer/string types to primitives/enums/pointers to messages.
  54. func (j *JSONPb) marshalNonProtoField(v interface{}) ([]byte, error) {
  55. if v == nil {
  56. return []byte("null"), nil
  57. }
  58. rv := reflect.ValueOf(v)
  59. for rv.Kind() == reflect.Ptr {
  60. if rv.IsNil() {
  61. return []byte("null"), nil
  62. }
  63. rv = rv.Elem()
  64. }
  65. if rv.Kind() == reflect.Slice {
  66. if rv.IsNil() {
  67. if j.EmitDefaults {
  68. return []byte("[]"), nil
  69. }
  70. return []byte("null"), nil
  71. }
  72. if rv.Type().Elem().Implements(protoMessageType) {
  73. var buf bytes.Buffer
  74. err := buf.WriteByte('[')
  75. if err != nil {
  76. return nil, err
  77. }
  78. for i := 0; i < rv.Len(); i++ {
  79. if i != 0 {
  80. err = buf.WriteByte(',')
  81. if err != nil {
  82. return nil, err
  83. }
  84. }
  85. if err = (*jsonpb.Marshaler)(j).Marshal(&buf, rv.Index(i).Interface().(proto.Message)); err != nil {
  86. return nil, err
  87. }
  88. }
  89. err = buf.WriteByte(']')
  90. if err != nil {
  91. return nil, err
  92. }
  93. return buf.Bytes(), nil
  94. }
  95. }
  96. if rv.Kind() == reflect.Map {
  97. m := make(map[string]*json.RawMessage)
  98. for _, k := range rv.MapKeys() {
  99. buf, err := j.Marshal(rv.MapIndex(k).Interface())
  100. if err != nil {
  101. return nil, err
  102. }
  103. m[fmt.Sprintf("%v", k.Interface())] = (*json.RawMessage)(&buf)
  104. }
  105. if j.Indent != "" {
  106. return json.MarshalIndent(m, "", j.Indent)
  107. }
  108. return json.Marshal(m)
  109. }
  110. if enum, ok := rv.Interface().(protoEnum); ok && !j.EnumsAsInts {
  111. return json.Marshal(enum.String())
  112. }
  113. return json.Marshal(rv.Interface())
  114. }
  115. // Unmarshal unmarshals JSON "data" into "v"
  116. func (j *JSONPb) Unmarshal(data []byte, v interface{}) error {
  117. return unmarshalJSONPb(data, v)
  118. }
  119. // NewDecoder returns a Decoder which reads JSON stream from "r".
  120. func (j *JSONPb) NewDecoder(r io.Reader) Decoder {
  121. d := json.NewDecoder(r)
  122. return DecoderWrapper{Decoder: d}
  123. }
  124. // DecoderWrapper is a wrapper around a *json.Decoder that adds
  125. // support for protos to the Decode method.
  126. type DecoderWrapper struct {
  127. *json.Decoder
  128. }
  129. // Decode wraps the embedded decoder's Decode method to support
  130. // protos using a jsonpb.Unmarshaler.
  131. func (d DecoderWrapper) Decode(v interface{}) error {
  132. return decodeJSONPb(d.Decoder, v)
  133. }
  134. // NewEncoder returns an Encoder which writes JSON stream into "w".
  135. func (j *JSONPb) NewEncoder(w io.Writer) Encoder {
  136. return EncoderFunc(func(v interface{}) error {
  137. if err := j.marshalTo(w, v); err != nil {
  138. return err
  139. }
  140. // mimic json.Encoder by adding a newline (makes output
  141. // easier to read when it contains multiple encoded items)
  142. _, err := w.Write(j.Delimiter())
  143. return err
  144. })
  145. }
  146. func unmarshalJSONPb(data []byte, v interface{}) error {
  147. d := json.NewDecoder(bytes.NewReader(data))
  148. return decodeJSONPb(d, v)
  149. }
  150. func decodeJSONPb(d *json.Decoder, v interface{}) error {
  151. p, ok := v.(proto.Message)
  152. if !ok {
  153. return decodeNonProtoField(d, v)
  154. }
  155. unmarshaler := &jsonpb.Unmarshaler{AllowUnknownFields: allowUnknownFields}
  156. return unmarshaler.UnmarshalNext(d, p)
  157. }
  158. func decodeNonProtoField(d *json.Decoder, v interface{}) error {
  159. rv := reflect.ValueOf(v)
  160. if rv.Kind() != reflect.Ptr {
  161. return fmt.Errorf("%T is not a pointer", v)
  162. }
  163. for rv.Kind() == reflect.Ptr {
  164. if rv.IsNil() {
  165. rv.Set(reflect.New(rv.Type().Elem()))
  166. }
  167. if rv.Type().ConvertibleTo(typeProtoMessage) {
  168. unmarshaler := &jsonpb.Unmarshaler{AllowUnknownFields: allowUnknownFields}
  169. return unmarshaler.UnmarshalNext(d, rv.Interface().(proto.Message))
  170. }
  171. rv = rv.Elem()
  172. }
  173. if rv.Kind() == reflect.Map {
  174. if rv.IsNil() {
  175. rv.Set(reflect.MakeMap(rv.Type()))
  176. }
  177. conv, ok := convFromType[rv.Type().Key().Kind()]
  178. if !ok {
  179. return fmt.Errorf("unsupported type of map field key: %v", rv.Type().Key())
  180. }
  181. m := make(map[string]*json.RawMessage)
  182. if err := d.Decode(&m); err != nil {
  183. return err
  184. }
  185. for k, v := range m {
  186. result := conv.Call([]reflect.Value{reflect.ValueOf(k)})
  187. if err := result[1].Interface(); err != nil {
  188. return err.(error)
  189. }
  190. bk := result[0]
  191. bv := reflect.New(rv.Type().Elem())
  192. if err := unmarshalJSONPb([]byte(*v), bv.Interface()); err != nil {
  193. return err
  194. }
  195. rv.SetMapIndex(bk, bv.Elem())
  196. }
  197. return nil
  198. }
  199. if _, ok := rv.Interface().(protoEnum); ok {
  200. var repr interface{}
  201. if err := d.Decode(&repr); err != nil {
  202. return err
  203. }
  204. switch repr.(type) {
  205. case string:
  206. // TODO(yugui) Should use proto.StructProperties?
  207. return fmt.Errorf("unmarshaling of symbolic enum %q not supported: %T", repr, rv.Interface())
  208. case float64:
  209. rv.Set(reflect.ValueOf(int32(repr.(float64))).Convert(rv.Type()))
  210. return nil
  211. default:
  212. return fmt.Errorf("cannot assign %#v into Go type %T", repr, rv.Interface())
  213. }
  214. }
  215. return d.Decode(v)
  216. }
  217. type protoEnum interface {
  218. fmt.Stringer
  219. EnumDescriptor() ([]byte, []int)
  220. }
  221. var typeProtoMessage = reflect.TypeOf((*proto.Message)(nil)).Elem()
  222. // Delimiter for newline encoded JSON streams.
  223. func (j *JSONPb) Delimiter() []byte {
  224. return []byte("\n")
  225. }
  226. // allowUnknownFields helps not to return an error when the destination
  227. // is a struct and the input contains object keys which do not match any
  228. // non-ignored, exported fields in the destination.
  229. var allowUnknownFields = true
  230. // DisallowUnknownFields enables option in decoder (unmarshaller) to
  231. // return an error when it finds an unknown field. This function must be
  232. // called before using the JSON marshaller.
  233. func DisallowUnknownFields() {
  234. allowUnknownFields = false
  235. }