encode.go 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199
  1. // Copyright 2019 The Go Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package impl
  5. import (
  6. "math"
  7. "sort"
  8. "sync/atomic"
  9. "google.golang.org/protobuf/internal/flags"
  10. proto "google.golang.org/protobuf/proto"
  11. piface "google.golang.org/protobuf/runtime/protoiface"
  12. )
  13. type marshalOptions struct {
  14. flags piface.MarshalInputFlags
  15. }
  16. func (o marshalOptions) Options() proto.MarshalOptions {
  17. return proto.MarshalOptions{
  18. AllowPartial: true,
  19. Deterministic: o.Deterministic(),
  20. UseCachedSize: o.UseCachedSize(),
  21. }
  22. }
  23. func (o marshalOptions) Deterministic() bool { return o.flags&piface.MarshalDeterministic != 0 }
  24. func (o marshalOptions) UseCachedSize() bool { return o.flags&piface.MarshalUseCachedSize != 0 }
  25. // size is protoreflect.Methods.Size.
  26. func (mi *MessageInfo) size(in piface.SizeInput) piface.SizeOutput {
  27. var p pointer
  28. if ms, ok := in.Message.(*messageState); ok {
  29. p = ms.pointer()
  30. } else {
  31. p = in.Message.(*messageReflectWrapper).pointer()
  32. }
  33. size := mi.sizePointer(p, marshalOptions{
  34. flags: in.Flags,
  35. })
  36. return piface.SizeOutput{Size: size}
  37. }
  38. func (mi *MessageInfo) sizePointer(p pointer, opts marshalOptions) (size int) {
  39. mi.init()
  40. if p.IsNil() {
  41. return 0
  42. }
  43. if opts.UseCachedSize() && mi.sizecacheOffset.IsValid() {
  44. if size := atomic.LoadInt32(p.Apply(mi.sizecacheOffset).Int32()); size >= 0 {
  45. return int(size)
  46. }
  47. }
  48. return mi.sizePointerSlow(p, opts)
  49. }
  50. func (mi *MessageInfo) sizePointerSlow(p pointer, opts marshalOptions) (size int) {
  51. if flags.ProtoLegacy && mi.isMessageSet {
  52. size = sizeMessageSet(mi, p, opts)
  53. if mi.sizecacheOffset.IsValid() {
  54. atomic.StoreInt32(p.Apply(mi.sizecacheOffset).Int32(), int32(size))
  55. }
  56. return size
  57. }
  58. if mi.extensionOffset.IsValid() {
  59. e := p.Apply(mi.extensionOffset).Extensions()
  60. size += mi.sizeExtensions(e, opts)
  61. }
  62. for _, f := range mi.orderedCoderFields {
  63. if f.funcs.size == nil {
  64. continue
  65. }
  66. fptr := p.Apply(f.offset)
  67. if f.isPointer && fptr.Elem().IsNil() {
  68. continue
  69. }
  70. size += f.funcs.size(fptr, f, opts)
  71. }
  72. if mi.unknownOffset.IsValid() {
  73. u := *p.Apply(mi.unknownOffset).Bytes()
  74. size += len(u)
  75. }
  76. if mi.sizecacheOffset.IsValid() {
  77. if size > math.MaxInt32 {
  78. // The size is too large for the int32 sizecache field.
  79. // We will need to recompute the size when encoding;
  80. // unfortunately expensive, but better than invalid output.
  81. atomic.StoreInt32(p.Apply(mi.sizecacheOffset).Int32(), -1)
  82. } else {
  83. atomic.StoreInt32(p.Apply(mi.sizecacheOffset).Int32(), int32(size))
  84. }
  85. }
  86. return size
  87. }
  88. // marshal is protoreflect.Methods.Marshal.
  89. func (mi *MessageInfo) marshal(in piface.MarshalInput) (out piface.MarshalOutput, err error) {
  90. var p pointer
  91. if ms, ok := in.Message.(*messageState); ok {
  92. p = ms.pointer()
  93. } else {
  94. p = in.Message.(*messageReflectWrapper).pointer()
  95. }
  96. b, err := mi.marshalAppendPointer(in.Buf, p, marshalOptions{
  97. flags: in.Flags,
  98. })
  99. return piface.MarshalOutput{Buf: b}, err
  100. }
  101. func (mi *MessageInfo) marshalAppendPointer(b []byte, p pointer, opts marshalOptions) ([]byte, error) {
  102. mi.init()
  103. if p.IsNil() {
  104. return b, nil
  105. }
  106. if flags.ProtoLegacy && mi.isMessageSet {
  107. return marshalMessageSet(mi, b, p, opts)
  108. }
  109. var err error
  110. // The old marshaler encodes extensions at beginning.
  111. if mi.extensionOffset.IsValid() {
  112. e := p.Apply(mi.extensionOffset).Extensions()
  113. // TODO: Special handling for MessageSet?
  114. b, err = mi.appendExtensions(b, e, opts)
  115. if err != nil {
  116. return b, err
  117. }
  118. }
  119. for _, f := range mi.orderedCoderFields {
  120. if f.funcs.marshal == nil {
  121. continue
  122. }
  123. fptr := p.Apply(f.offset)
  124. if f.isPointer && fptr.Elem().IsNil() {
  125. continue
  126. }
  127. b, err = f.funcs.marshal(b, fptr, f, opts)
  128. if err != nil {
  129. return b, err
  130. }
  131. }
  132. if mi.unknownOffset.IsValid() && !mi.isMessageSet {
  133. u := *p.Apply(mi.unknownOffset).Bytes()
  134. b = append(b, u...)
  135. }
  136. return b, nil
  137. }
  138. func (mi *MessageInfo) sizeExtensions(ext *map[int32]ExtensionField, opts marshalOptions) (n int) {
  139. if ext == nil {
  140. return 0
  141. }
  142. for _, x := range *ext {
  143. xi := getExtensionFieldInfo(x.Type())
  144. if xi.funcs.size == nil {
  145. continue
  146. }
  147. n += xi.funcs.size(x.Value(), xi.tagsize, opts)
  148. }
  149. return n
  150. }
  151. func (mi *MessageInfo) appendExtensions(b []byte, ext *map[int32]ExtensionField, opts marshalOptions) ([]byte, error) {
  152. if ext == nil {
  153. return b, nil
  154. }
  155. switch len(*ext) {
  156. case 0:
  157. return b, nil
  158. case 1:
  159. // Fast-path for one extension: Don't bother sorting the keys.
  160. var err error
  161. for _, x := range *ext {
  162. xi := getExtensionFieldInfo(x.Type())
  163. b, err = xi.funcs.marshal(b, x.Value(), xi.wiretag, opts)
  164. }
  165. return b, err
  166. default:
  167. // Sort the keys to provide a deterministic encoding.
  168. // Not sure this is required, but the old code does it.
  169. keys := make([]int, 0, len(*ext))
  170. for k := range *ext {
  171. keys = append(keys, int(k))
  172. }
  173. sort.Ints(keys)
  174. var err error
  175. for _, k := range keys {
  176. x := (*ext)[int32(k)]
  177. xi := getExtensionFieldInfo(x.Type())
  178. b, err = xi.funcs.marshal(b, x.Value(), xi.wiretag, opts)
  179. if err != nil {
  180. return b, err
  181. }
  182. }
  183. return b, nil
  184. }
  185. }