marshal_jsonpb.go 8.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348
  1. package runtime
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "reflect"
  8. "strconv"
  9. "google.golang.org/protobuf/encoding/protojson"
  10. "google.golang.org/protobuf/proto"
  11. )
  12. // JSONPb is a Marshaler which marshals/unmarshals into/from JSON
  13. // with the "google.golang.org/protobuf/encoding/protojson" marshaler.
  14. // It supports the full functionality of protobuf unlike JSONBuiltin.
  15. //
  16. // The NewDecoder method returns a DecoderWrapper, so the underlying
  17. // *json.Decoder methods can be used.
  18. type JSONPb struct {
  19. protojson.MarshalOptions
  20. protojson.UnmarshalOptions
  21. }
  22. // ContentType always returns "application/json".
  23. func (*JSONPb) ContentType(_ interface{}) string {
  24. return "application/json"
  25. }
  26. // Marshal marshals "v" into JSON.
  27. func (j *JSONPb) Marshal(v interface{}) ([]byte, error) {
  28. if _, ok := v.(proto.Message); !ok {
  29. return j.marshalNonProtoField(v)
  30. }
  31. var buf bytes.Buffer
  32. if err := j.marshalTo(&buf, v); err != nil {
  33. return nil, err
  34. }
  35. return buf.Bytes(), nil
  36. }
  37. func (j *JSONPb) marshalTo(w io.Writer, v interface{}) error {
  38. p, ok := v.(proto.Message)
  39. if !ok {
  40. buf, err := j.marshalNonProtoField(v)
  41. if err != nil {
  42. return err
  43. }
  44. _, err = w.Write(buf)
  45. return err
  46. }
  47. b, err := j.MarshalOptions.Marshal(p)
  48. if err != nil {
  49. return err
  50. }
  51. _, err = w.Write(b)
  52. return err
  53. }
  54. var (
  55. // protoMessageType is stored to prevent constant lookup of the same type at runtime.
  56. protoMessageType = reflect.TypeOf((*proto.Message)(nil)).Elem()
  57. )
  58. // marshalNonProto marshals a non-message field of a protobuf message.
  59. // This function does not correctly marshal arbitrary data structures into JSON,
  60. // it is only capable of marshaling non-message field values of protobuf,
  61. // i.e. primitive types, enums; pointers to primitives or enums; maps from
  62. // integer/string types to primitives/enums/pointers to messages.
  63. func (j *JSONPb) marshalNonProtoField(v interface{}) ([]byte, error) {
  64. if v == nil {
  65. return []byte("null"), nil
  66. }
  67. rv := reflect.ValueOf(v)
  68. for rv.Kind() == reflect.Ptr {
  69. if rv.IsNil() {
  70. return []byte("null"), nil
  71. }
  72. rv = rv.Elem()
  73. }
  74. if rv.Kind() == reflect.Slice {
  75. if rv.IsNil() {
  76. if j.EmitUnpopulated {
  77. return []byte("[]"), nil
  78. }
  79. return []byte("null"), nil
  80. }
  81. if rv.Type().Elem().Implements(protoMessageType) {
  82. var buf bytes.Buffer
  83. if err := buf.WriteByte('['); err != nil {
  84. return nil, err
  85. }
  86. for i := 0; i < rv.Len(); i++ {
  87. if i != 0 {
  88. if err := buf.WriteByte(','); err != nil {
  89. return nil, err
  90. }
  91. }
  92. if err := j.marshalTo(&buf, rv.Index(i).Interface().(proto.Message)); err != nil {
  93. return nil, err
  94. }
  95. }
  96. if err := buf.WriteByte(']'); err != nil {
  97. return nil, err
  98. }
  99. return buf.Bytes(), nil
  100. }
  101. if rv.Type().Elem().Implements(typeProtoEnum) {
  102. var buf bytes.Buffer
  103. if err := buf.WriteByte('['); err != nil {
  104. return nil, err
  105. }
  106. for i := 0; i < rv.Len(); i++ {
  107. if i != 0 {
  108. if err := buf.WriteByte(','); err != nil {
  109. return nil, err
  110. }
  111. }
  112. var err error
  113. if j.UseEnumNumbers {
  114. _, err = buf.WriteString(strconv.FormatInt(rv.Index(i).Int(), 10))
  115. } else {
  116. _, err = buf.WriteString("\"" + rv.Index(i).Interface().(protoEnum).String() + "\"")
  117. }
  118. if err != nil {
  119. return nil, err
  120. }
  121. }
  122. if err := buf.WriteByte(']'); err != nil {
  123. return nil, err
  124. }
  125. return buf.Bytes(), nil
  126. }
  127. }
  128. if rv.Kind() == reflect.Map {
  129. m := make(map[string]*json.RawMessage)
  130. for _, k := range rv.MapKeys() {
  131. buf, err := j.Marshal(rv.MapIndex(k).Interface())
  132. if err != nil {
  133. return nil, err
  134. }
  135. m[fmt.Sprintf("%v", k.Interface())] = (*json.RawMessage)(&buf)
  136. }
  137. if j.Indent != "" {
  138. return json.MarshalIndent(m, "", j.Indent)
  139. }
  140. return json.Marshal(m)
  141. }
  142. if enum, ok := rv.Interface().(protoEnum); ok && !j.UseEnumNumbers {
  143. return json.Marshal(enum.String())
  144. }
  145. return json.Marshal(rv.Interface())
  146. }
  147. // Unmarshal unmarshals JSON "data" into "v"
  148. func (j *JSONPb) Unmarshal(data []byte, v interface{}) error {
  149. return unmarshalJSONPb(data, j.UnmarshalOptions, v)
  150. }
  151. // NewDecoder returns a Decoder which reads JSON stream from "r".
  152. func (j *JSONPb) NewDecoder(r io.Reader) Decoder {
  153. d := json.NewDecoder(r)
  154. return DecoderWrapper{
  155. Decoder: d,
  156. UnmarshalOptions: j.UnmarshalOptions,
  157. }
  158. }
  159. // DecoderWrapper is a wrapper around a *json.Decoder that adds
  160. // support for protos to the Decode method.
  161. type DecoderWrapper struct {
  162. *json.Decoder
  163. protojson.UnmarshalOptions
  164. }
  165. // Decode wraps the embedded decoder's Decode method to support
  166. // protos using a jsonpb.Unmarshaler.
  167. func (d DecoderWrapper) Decode(v interface{}) error {
  168. return decodeJSONPb(d.Decoder, d.UnmarshalOptions, v)
  169. }
  170. // NewEncoder returns an Encoder which writes JSON stream into "w".
  171. func (j *JSONPb) NewEncoder(w io.Writer) Encoder {
  172. return EncoderFunc(func(v interface{}) error {
  173. if err := j.marshalTo(w, v); err != nil {
  174. return err
  175. }
  176. // mimic json.Encoder by adding a newline (makes output
  177. // easier to read when it contains multiple encoded items)
  178. _, err := w.Write(j.Delimiter())
  179. return err
  180. })
  181. }
  182. func unmarshalJSONPb(data []byte, unmarshaler protojson.UnmarshalOptions, v interface{}) error {
  183. d := json.NewDecoder(bytes.NewReader(data))
  184. return decodeJSONPb(d, unmarshaler, v)
  185. }
  186. func decodeJSONPb(d *json.Decoder, unmarshaler protojson.UnmarshalOptions, v interface{}) error {
  187. p, ok := v.(proto.Message)
  188. if !ok {
  189. return decodeNonProtoField(d, unmarshaler, v)
  190. }
  191. // Decode into bytes for marshalling
  192. var b json.RawMessage
  193. if err := d.Decode(&b); err != nil {
  194. return err
  195. }
  196. return unmarshaler.Unmarshal([]byte(b), p)
  197. }
  198. func decodeNonProtoField(d *json.Decoder, unmarshaler protojson.UnmarshalOptions, v interface{}) error {
  199. rv := reflect.ValueOf(v)
  200. if rv.Kind() != reflect.Ptr {
  201. return fmt.Errorf("%T is not a pointer", v)
  202. }
  203. for rv.Kind() == reflect.Ptr {
  204. if rv.IsNil() {
  205. rv.Set(reflect.New(rv.Type().Elem()))
  206. }
  207. if rv.Type().ConvertibleTo(typeProtoMessage) {
  208. // Decode into bytes for marshalling
  209. var b json.RawMessage
  210. if err := d.Decode(&b); err != nil {
  211. return err
  212. }
  213. return unmarshaler.Unmarshal([]byte(b), rv.Interface().(proto.Message))
  214. }
  215. rv = rv.Elem()
  216. }
  217. if rv.Kind() == reflect.Map {
  218. if rv.IsNil() {
  219. rv.Set(reflect.MakeMap(rv.Type()))
  220. }
  221. conv, ok := convFromType[rv.Type().Key().Kind()]
  222. if !ok {
  223. return fmt.Errorf("unsupported type of map field key: %v", rv.Type().Key())
  224. }
  225. m := make(map[string]*json.RawMessage)
  226. if err := d.Decode(&m); err != nil {
  227. return err
  228. }
  229. for k, v := range m {
  230. result := conv.Call([]reflect.Value{reflect.ValueOf(k)})
  231. if err := result[1].Interface(); err != nil {
  232. return err.(error)
  233. }
  234. bk := result[0]
  235. bv := reflect.New(rv.Type().Elem())
  236. if v == nil {
  237. null := json.RawMessage("null")
  238. v = &null
  239. }
  240. if err := unmarshalJSONPb([]byte(*v), unmarshaler, bv.Interface()); err != nil {
  241. return err
  242. }
  243. rv.SetMapIndex(bk, bv.Elem())
  244. }
  245. return nil
  246. }
  247. if rv.Kind() == reflect.Slice {
  248. if rv.Type().Elem().Kind() == reflect.Uint8 {
  249. var sl []byte
  250. if err := d.Decode(&sl); err != nil {
  251. return err
  252. }
  253. if sl != nil {
  254. rv.SetBytes(sl)
  255. }
  256. return nil
  257. }
  258. var sl []json.RawMessage
  259. if err := d.Decode(&sl); err != nil {
  260. return err
  261. }
  262. if sl != nil {
  263. rv.Set(reflect.MakeSlice(rv.Type(), 0, 0))
  264. }
  265. for _, item := range sl {
  266. bv := reflect.New(rv.Type().Elem())
  267. if err := unmarshalJSONPb([]byte(item), unmarshaler, bv.Interface()); err != nil {
  268. return err
  269. }
  270. rv.Set(reflect.Append(rv, bv.Elem()))
  271. }
  272. return nil
  273. }
  274. if _, ok := rv.Interface().(protoEnum); ok {
  275. var repr interface{}
  276. if err := d.Decode(&repr); err != nil {
  277. return err
  278. }
  279. switch v := repr.(type) {
  280. case string:
  281. // TODO(yugui) Should use proto.StructProperties?
  282. return fmt.Errorf("unmarshaling of symbolic enum %q not supported: %T", repr, rv.Interface())
  283. case float64:
  284. rv.Set(reflect.ValueOf(int32(v)).Convert(rv.Type()))
  285. return nil
  286. default:
  287. return fmt.Errorf("cannot assign %#v into Go type %T", repr, rv.Interface())
  288. }
  289. }
  290. return d.Decode(v)
  291. }
  292. type protoEnum interface {
  293. fmt.Stringer
  294. EnumDescriptor() ([]byte, []int)
  295. }
  296. var typeProtoEnum = reflect.TypeOf((*protoEnum)(nil)).Elem()
  297. var typeProtoMessage = reflect.TypeOf((*proto.Message)(nil)).Elem()
  298. // Delimiter for newline encoded JSON streams.
  299. func (j *JSONPb) Delimiter() []byte {
  300. return []byte("\n")
  301. }
  302. var (
  303. convFromType = map[reflect.Kind]reflect.Value{
  304. reflect.String: reflect.ValueOf(String),
  305. reflect.Bool: reflect.ValueOf(Bool),
  306. reflect.Float64: reflect.ValueOf(Float64),
  307. reflect.Float32: reflect.ValueOf(Float32),
  308. reflect.Int64: reflect.ValueOf(Int64),
  309. reflect.Int32: reflect.ValueOf(Int32),
  310. reflect.Uint64: reflect.ValueOf(Uint64),
  311. reflect.Uint32: reflect.ValueOf(Uint32),
  312. reflect.Slice: reflect.ValueOf(Bytes),
  313. }
  314. )