codec_extension.go 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223
  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. "sync"
  7. "sync/atomic"
  8. "google.golang.org/protobuf/encoding/protowire"
  9. "google.golang.org/protobuf/internal/errors"
  10. pref "google.golang.org/protobuf/reflect/protoreflect"
  11. )
  12. type extensionFieldInfo struct {
  13. wiretag uint64
  14. tagsize int
  15. unmarshalNeedsValue bool
  16. funcs valueCoderFuncs
  17. validation validationInfo
  18. }
  19. var legacyExtensionFieldInfoCache sync.Map // map[protoreflect.ExtensionType]*extensionFieldInfo
  20. func getExtensionFieldInfo(xt pref.ExtensionType) *extensionFieldInfo {
  21. if xi, ok := xt.(*ExtensionInfo); ok {
  22. xi.lazyInit()
  23. return xi.info
  24. }
  25. return legacyLoadExtensionFieldInfo(xt)
  26. }
  27. // legacyLoadExtensionFieldInfo dynamically loads a *ExtensionInfo for xt.
  28. func legacyLoadExtensionFieldInfo(xt pref.ExtensionType) *extensionFieldInfo {
  29. if xi, ok := legacyExtensionFieldInfoCache.Load(xt); ok {
  30. return xi.(*extensionFieldInfo)
  31. }
  32. e := makeExtensionFieldInfo(xt.TypeDescriptor())
  33. if e, ok := legacyMessageTypeCache.LoadOrStore(xt, e); ok {
  34. return e.(*extensionFieldInfo)
  35. }
  36. return e
  37. }
  38. func makeExtensionFieldInfo(xd pref.ExtensionDescriptor) *extensionFieldInfo {
  39. var wiretag uint64
  40. if !xd.IsPacked() {
  41. wiretag = protowire.EncodeTag(xd.Number(), wireTypes[xd.Kind()])
  42. } else {
  43. wiretag = protowire.EncodeTag(xd.Number(), protowire.BytesType)
  44. }
  45. e := &extensionFieldInfo{
  46. wiretag: wiretag,
  47. tagsize: protowire.SizeVarint(wiretag),
  48. funcs: encoderFuncsForValue(xd),
  49. }
  50. // Does the unmarshal function need a value passed to it?
  51. // This is true for composite types, where we pass in a message, list, or map to fill in,
  52. // and for enums, where we pass in a prototype value to specify the concrete enum type.
  53. switch xd.Kind() {
  54. case pref.MessageKind, pref.GroupKind, pref.EnumKind:
  55. e.unmarshalNeedsValue = true
  56. default:
  57. if xd.Cardinality() == pref.Repeated {
  58. e.unmarshalNeedsValue = true
  59. }
  60. }
  61. return e
  62. }
  63. type lazyExtensionValue struct {
  64. atomicOnce uint32 // atomically set if value is valid
  65. mu sync.Mutex
  66. xi *extensionFieldInfo
  67. value pref.Value
  68. b []byte
  69. fn func() pref.Value
  70. }
  71. type ExtensionField struct {
  72. typ pref.ExtensionType
  73. // value is either the value of GetValue,
  74. // or a *lazyExtensionValue that then returns the value of GetValue.
  75. value pref.Value
  76. lazy *lazyExtensionValue
  77. }
  78. func (f *ExtensionField) appendLazyBytes(xt pref.ExtensionType, xi *extensionFieldInfo, num protowire.Number, wtyp protowire.Type, b []byte) {
  79. if f.lazy == nil {
  80. f.lazy = &lazyExtensionValue{xi: xi}
  81. }
  82. f.typ = xt
  83. f.lazy.xi = xi
  84. f.lazy.b = protowire.AppendTag(f.lazy.b, num, wtyp)
  85. f.lazy.b = append(f.lazy.b, b...)
  86. }
  87. func (f *ExtensionField) canLazy(xt pref.ExtensionType) bool {
  88. if f.typ == nil {
  89. return true
  90. }
  91. if f.typ == xt && f.lazy != nil && atomic.LoadUint32(&f.lazy.atomicOnce) == 0 {
  92. return true
  93. }
  94. return false
  95. }
  96. func (f *ExtensionField) lazyInit() {
  97. f.lazy.mu.Lock()
  98. defer f.lazy.mu.Unlock()
  99. if atomic.LoadUint32(&f.lazy.atomicOnce) == 1 {
  100. return
  101. }
  102. if f.lazy.xi != nil {
  103. b := f.lazy.b
  104. val := f.typ.New()
  105. for len(b) > 0 {
  106. var tag uint64
  107. if b[0] < 0x80 {
  108. tag = uint64(b[0])
  109. b = b[1:]
  110. } else if len(b) >= 2 && b[1] < 128 {
  111. tag = uint64(b[0]&0x7f) + uint64(b[1])<<7
  112. b = b[2:]
  113. } else {
  114. var n int
  115. tag, n = protowire.ConsumeVarint(b)
  116. if n < 0 {
  117. panic(errors.New("bad tag in lazy extension decoding"))
  118. }
  119. b = b[n:]
  120. }
  121. num := protowire.Number(tag >> 3)
  122. wtyp := protowire.Type(tag & 7)
  123. var out unmarshalOutput
  124. var err error
  125. val, out, err = f.lazy.xi.funcs.unmarshal(b, val, num, wtyp, lazyUnmarshalOptions)
  126. if err != nil {
  127. panic(errors.New("decode failure in lazy extension decoding: %v", err))
  128. }
  129. b = b[out.n:]
  130. }
  131. f.lazy.value = val
  132. } else {
  133. f.lazy.value = f.lazy.fn()
  134. }
  135. f.lazy.xi = nil
  136. f.lazy.fn = nil
  137. f.lazy.b = nil
  138. atomic.StoreUint32(&f.lazy.atomicOnce, 1)
  139. }
  140. // Set sets the type and value of the extension field.
  141. // This must not be called concurrently.
  142. func (f *ExtensionField) Set(t pref.ExtensionType, v pref.Value) {
  143. f.typ = t
  144. f.value = v
  145. f.lazy = nil
  146. }
  147. // SetLazy sets the type and a value that is to be lazily evaluated upon first use.
  148. // This must not be called concurrently.
  149. func (f *ExtensionField) SetLazy(t pref.ExtensionType, fn func() pref.Value) {
  150. f.typ = t
  151. f.lazy = &lazyExtensionValue{fn: fn}
  152. }
  153. // Value returns the value of the extension field.
  154. // This may be called concurrently.
  155. func (f *ExtensionField) Value() pref.Value {
  156. if f.lazy != nil {
  157. if atomic.LoadUint32(&f.lazy.atomicOnce) == 0 {
  158. f.lazyInit()
  159. }
  160. return f.lazy.value
  161. }
  162. return f.value
  163. }
  164. // Type returns the type of the extension field.
  165. // This may be called concurrently.
  166. func (f ExtensionField) Type() pref.ExtensionType {
  167. return f.typ
  168. }
  169. // IsSet returns whether the extension field is set.
  170. // This may be called concurrently.
  171. func (f ExtensionField) IsSet() bool {
  172. return f.typ != nil
  173. }
  174. // IsLazy reports whether a field is lazily encoded.
  175. // It is exported for testing.
  176. func IsLazy(m pref.Message, fd pref.FieldDescriptor) bool {
  177. var mi *MessageInfo
  178. var p pointer
  179. switch m := m.(type) {
  180. case *messageState:
  181. mi = m.messageInfo()
  182. p = m.pointer()
  183. case *messageReflectWrapper:
  184. mi = m.messageInfo()
  185. p = m.pointer()
  186. default:
  187. return false
  188. }
  189. xd, ok := fd.(pref.ExtensionTypeDescriptor)
  190. if !ok {
  191. return false
  192. }
  193. xt := xd.Type()
  194. ext := mi.extensionMap(p)
  195. if ext == nil {
  196. return false
  197. }
  198. f, ok := (*ext)[int32(fd.Number())]
  199. if !ok {
  200. return false
  201. }
  202. return f.typ == xt && f.lazy != nil && atomic.LoadUint32(&f.lazy.atomicOnce) == 0
  203. }