extensions.go 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457
  1. // Go support for Protocol Buffers - Google's data interchange format
  2. //
  3. // Copyright 2010 The Go Authors. All rights reserved.
  4. // http://code.google.com/p/goprotobuf/
  5. //
  6. // Redistribution and use in source and binary forms, with or without
  7. // modification, are permitted provided that the following conditions are
  8. // met:
  9. //
  10. // * Redistributions of source code must retain the above copyright
  11. // notice, this list of conditions and the following disclaimer.
  12. // * Redistributions in binary form must reproduce the above
  13. // copyright notice, this list of conditions and the following disclaimer
  14. // in the documentation and/or other materials provided with the
  15. // distribution.
  16. // * Neither the name of Google Inc. nor the names of its
  17. // contributors may be used to endorse or promote products derived from
  18. // this software without specific prior written permission.
  19. //
  20. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  21. // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  22. // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  23. // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  24. // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  25. // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  26. // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  27. // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  28. // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  29. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  30. // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  31. package proto
  32. /*
  33. * Types and routines for supporting protocol buffer extensions.
  34. */
  35. import (
  36. "errors"
  37. "reflect"
  38. "strconv"
  39. "sync"
  40. )
  41. // ErrMissingExtension is the error returned by GetExtension if the named extension is not in the message.
  42. var ErrMissingExtension = errors.New("proto: missing extension")
  43. // ExtensionRange represents a range of message extensions for a protocol buffer.
  44. // Used in code generated by the protocol compiler.
  45. type ExtensionRange struct {
  46. Start, End int32 // both inclusive
  47. }
  48. // extendableProto is an interface implemented by any protocol buffer that may be extended.
  49. type extendableProto interface {
  50. Message
  51. ExtensionRangeArray() []ExtensionRange
  52. }
  53. type extensionsMap interface {
  54. extendableProto
  55. ExtensionMap() map[int32]Extension
  56. }
  57. type extensionsBytes interface {
  58. extendableProto
  59. GetExtensions() *[]byte
  60. }
  61. var extendableProtoType = reflect.TypeOf((*extendableProto)(nil)).Elem()
  62. // ExtensionDesc represents an extension specification.
  63. // Used in generated code from the protocol compiler.
  64. type ExtensionDesc struct {
  65. ExtendedType Message // nil pointer to the type that is being extended
  66. ExtensionType interface{} // nil pointer to the extension type
  67. Field int32 // field number
  68. Name string // fully-qualified name of extension, for text formatting
  69. Tag string // protobuf tag style
  70. }
  71. func (ed *ExtensionDesc) repeated() bool {
  72. t := reflect.TypeOf(ed.ExtensionType)
  73. return t.Kind() == reflect.Slice && t.Elem().Kind() != reflect.Uint8
  74. }
  75. // Extension represents an extension in a message.
  76. type Extension struct {
  77. // When an extension is stored in a message using SetExtension
  78. // only desc and value are set. When the message is marshaled
  79. // enc will be set to the encoded form of the message.
  80. //
  81. // When a message is unmarshaled and contains extensions, each
  82. // extension will have only enc set. When such an extension is
  83. // accessed using GetExtension (or GetExtensions) desc and value
  84. // will be set.
  85. desc *ExtensionDesc
  86. value interface{}
  87. enc []byte
  88. }
  89. // SetRawExtension is for testing only.
  90. func SetRawExtension(base extendableProto, id int32, b []byte) {
  91. if ebase, ok := base.(extensionsMap); ok {
  92. ebase.ExtensionMap()[id] = Extension{enc: b}
  93. } else if ebase, ok := base.(extensionsBytes); ok {
  94. clearExtension(base, id)
  95. ext := ebase.GetExtensions()
  96. *ext = append(*ext, b...)
  97. } else {
  98. panic("unreachable")
  99. }
  100. }
  101. // isExtensionField returns true iff the given field number is in an extension range.
  102. func isExtensionField(pb extendableProto, field int32) bool {
  103. for _, er := range pb.ExtensionRangeArray() {
  104. if er.Start <= field && field <= er.End {
  105. return true
  106. }
  107. }
  108. return false
  109. }
  110. // checkExtensionTypes checks that the given extension is valid for pb.
  111. func checkExtensionTypes(pb extendableProto, extension *ExtensionDesc) error {
  112. // Check the extended type.
  113. if a, b := reflect.TypeOf(pb), reflect.TypeOf(extension.ExtendedType); a != b {
  114. return errors.New("proto: bad extended type; " + b.String() + " does not extend " + a.String())
  115. }
  116. // Check the range.
  117. if !isExtensionField(pb, extension.Field) {
  118. return errors.New("proto: bad extension number; not in declared ranges")
  119. }
  120. return nil
  121. }
  122. // extPropKey is sufficient to uniquely identify an extension.
  123. type extPropKey struct {
  124. base reflect.Type
  125. field int32
  126. }
  127. var extProp = struct {
  128. sync.RWMutex
  129. m map[extPropKey]*Properties
  130. }{
  131. m: make(map[extPropKey]*Properties),
  132. }
  133. func extensionProperties(ed *ExtensionDesc) *Properties {
  134. key := extPropKey{base: reflect.TypeOf(ed.ExtendedType), field: ed.Field}
  135. extProp.RLock()
  136. if prop, ok := extProp.m[key]; ok {
  137. extProp.RUnlock()
  138. return prop
  139. }
  140. extProp.RUnlock()
  141. extProp.Lock()
  142. defer extProp.Unlock()
  143. // Check again.
  144. if prop, ok := extProp.m[key]; ok {
  145. return prop
  146. }
  147. prop := new(Properties)
  148. prop.Init(reflect.TypeOf(ed.ExtensionType), "unknown_name", ed.Tag, nil)
  149. extProp.m[key] = prop
  150. return prop
  151. }
  152. // encodeExtensionMap encodes any unmarshaled (unencoded) extensions in m.
  153. func encodeExtensionMap(m map[int32]Extension) error {
  154. for k, e := range m {
  155. if e.value == nil || e.desc == nil {
  156. // Extension is only in its encoded form.
  157. continue
  158. }
  159. // We don't skip extensions that have an encoded form set,
  160. // because the extension value may have been mutated after
  161. // the last time this function was called.
  162. et := reflect.TypeOf(e.desc.ExtensionType)
  163. props := extensionProperties(e.desc)
  164. p := NewBuffer(nil)
  165. // If e.value has type T, the encoder expects a *struct{ X T }.
  166. // Pass a *T with a zero field and hope it all works out.
  167. x := reflect.New(et)
  168. x.Elem().Set(reflect.ValueOf(e.value))
  169. if err := props.enc(p, props, toStructPointer(x)); err != nil {
  170. return err
  171. }
  172. e.enc = p.buf
  173. m[k] = e
  174. }
  175. return nil
  176. }
  177. func sizeExtensionMap(m map[int32]Extension) (n int) {
  178. for _, e := range m {
  179. if e.value == nil || e.desc == nil {
  180. // Extension is only in its encoded form.
  181. n += len(e.enc)
  182. continue
  183. }
  184. // We don't skip extensions that have an encoded form set,
  185. // because the extension value may have been mutated after
  186. // the last time this function was called.
  187. et := reflect.TypeOf(e.desc.ExtensionType)
  188. props := extensionProperties(e.desc)
  189. // If e.value has type T, the encoder expects a *struct{ X T }.
  190. // Pass a *T with a zero field and hope it all works out.
  191. x := reflect.New(et)
  192. x.Elem().Set(reflect.ValueOf(e.value))
  193. n += props.size(props, toStructPointer(x))
  194. }
  195. return
  196. }
  197. // HasExtension returns whether the given extension is present in pb.
  198. func HasExtension(pb extendableProto, extension *ExtensionDesc) bool {
  199. // TODO: Check types, field numbers, etc.?
  200. if epb, doki := pb.(extensionsMap); doki {
  201. _, ok := epb.ExtensionMap()[extension.Field]
  202. return ok
  203. } else if epb, doki := pb.(extensionsBytes); doki {
  204. ext := epb.GetExtensions()
  205. buf := *ext
  206. o := 0
  207. for o < len(buf) {
  208. tag, n := DecodeVarint(buf[o:])
  209. fieldNum := int32(tag >> 3)
  210. if int32(fieldNum) == extension.Field {
  211. return true
  212. }
  213. wireType := int(tag & 0x7)
  214. o += n
  215. l, err := size(buf[o:], wireType)
  216. if err != nil {
  217. return false
  218. }
  219. o += l
  220. }
  221. return false
  222. }
  223. panic("unreachable")
  224. }
  225. func deleteExtension(pb extensionsBytes, theFieldNum int32, offset int) int {
  226. ext := pb.GetExtensions()
  227. for offset < len(*ext) {
  228. tag, n1 := DecodeVarint((*ext)[offset:])
  229. fieldNum := int32(tag >> 3)
  230. wireType := int(tag & 0x7)
  231. n2, err := size((*ext)[offset+n1:], wireType)
  232. if err != nil {
  233. panic(err)
  234. }
  235. newOffset := offset + n1 + n2
  236. if fieldNum == theFieldNum {
  237. *ext = append((*ext)[:offset], (*ext)[newOffset:]...)
  238. return offset
  239. }
  240. offset = newOffset
  241. }
  242. return -1
  243. }
  244. func clearExtension(pb extendableProto, fieldNum int32) {
  245. if epb, doki := pb.(extensionsMap); doki {
  246. delete(epb.ExtensionMap(), fieldNum)
  247. } else if epb, doki := pb.(extensionsBytes); doki {
  248. offset := 0
  249. for offset != -1 {
  250. offset = deleteExtension(epb, fieldNum, offset)
  251. }
  252. } else {
  253. panic("unreachable")
  254. }
  255. }
  256. // ClearExtension removes the given extension from pb.
  257. func ClearExtension(pb extendableProto, extension *ExtensionDesc) {
  258. // TODO: Check types, field numbers, etc.?
  259. clearExtension(pb, extension.Field)
  260. }
  261. // GetExtension parses and returns the given extension of pb.
  262. // If the extension is not present it returns ErrMissingExtension.
  263. func GetExtension(pb extendableProto, extension *ExtensionDesc) (interface{}, error) {
  264. if err := checkExtensionTypes(pb, extension); err != nil {
  265. return nil, err
  266. }
  267. if epb, doki := pb.(extensionsMap); doki {
  268. e, ok := epb.ExtensionMap()[extension.Field]
  269. if !ok {
  270. return nil, ErrMissingExtension
  271. }
  272. if e.value != nil {
  273. // Already decoded. Check the descriptor, though.
  274. if e.desc != extension {
  275. // This shouldn't happen. If it does, it means that
  276. // GetExtension was called twice with two different
  277. // descriptors with the same field number.
  278. return nil, errors.New("proto: descriptor conflict")
  279. }
  280. return e.value, nil
  281. }
  282. v, err := decodeExtension(e.enc, extension)
  283. if err != nil {
  284. return nil, err
  285. }
  286. // Remember the decoded version and drop the encoded version.
  287. // That way it is safe to mutate what we return.
  288. e.value = v
  289. e.desc = extension
  290. e.enc = nil
  291. return e.value, nil
  292. } else if epb, doki := pb.(extensionsBytes); doki {
  293. ext := epb.GetExtensions()
  294. o := 0
  295. for o < len(*ext) {
  296. tag, n := DecodeVarint((*ext)[o:])
  297. fieldNum := int32(tag >> 3)
  298. wireType := int(tag & 0x7)
  299. l, err := size((*ext)[o+n:], wireType)
  300. if err != nil {
  301. return nil, err
  302. }
  303. if int32(fieldNum) == extension.Field {
  304. v, err := decodeExtension((*ext)[o:o+n+l], extension)
  305. if err != nil {
  306. return nil, err
  307. }
  308. return v, nil
  309. }
  310. o += n + l
  311. }
  312. }
  313. panic("unreachable")
  314. }
  315. // decodeExtension decodes an extension encoded in b.
  316. func decodeExtension(b []byte, extension *ExtensionDesc) (interface{}, error) {
  317. o := NewBuffer(b)
  318. t := reflect.TypeOf(extension.ExtensionType)
  319. rep := extension.repeated()
  320. props := extensionProperties(extension)
  321. // t is a pointer to a struct, pointer to basic type or a slice.
  322. // Allocate a "field" to store the pointer/slice itself; the
  323. // pointer/slice will be stored here. We pass
  324. // the address of this field to props.dec.
  325. // This passes a zero field and a *t and lets props.dec
  326. // interpret it as a *struct{ x t }.
  327. value := reflect.New(t).Elem()
  328. for {
  329. // Discard wire type and field number varint. It isn't needed.
  330. if _, err := o.DecodeVarint(); err != nil {
  331. return nil, err
  332. }
  333. if err := props.dec(o, props, toStructPointer(value.Addr())); err != nil {
  334. return nil, err
  335. }
  336. if !rep || o.index >= len(o.buf) {
  337. break
  338. }
  339. }
  340. return value.Interface(), nil
  341. }
  342. // GetExtensions returns a slice of the extensions present in pb that are also listed in es.
  343. // The returned slice has the same length as es; missing extensions will appear as nil elements.
  344. func GetExtensions(pb Message, es []*ExtensionDesc) (extensions []interface{}, err error) {
  345. epb, ok := pb.(extendableProto)
  346. if !ok {
  347. err = errors.New("proto: not an extendable proto")
  348. return
  349. }
  350. extensions = make([]interface{}, len(es))
  351. for i, e := range es {
  352. extensions[i], err = GetExtension(epb, e)
  353. if err != nil {
  354. return
  355. }
  356. }
  357. return
  358. }
  359. // SetExtension sets the specified extension of pb to the specified value.
  360. func SetExtension(pb extendableProto, extension *ExtensionDesc, value interface{}) error {
  361. if err := checkExtensionTypes(pb, extension); err != nil {
  362. return err
  363. }
  364. typ := reflect.TypeOf(extension.ExtensionType)
  365. if typ != reflect.TypeOf(value) {
  366. return errors.New("proto: bad extension value type")
  367. }
  368. if epb, doki := pb.(extensionsMap); doki {
  369. epb.ExtensionMap()[extension.Field] = Extension{desc: extension, value: value}
  370. } else if epb, doki := pb.(extensionsBytes); doki {
  371. ClearExtension(pb, extension)
  372. ext := epb.GetExtensions()
  373. et := reflect.TypeOf(extension.ExtensionType)
  374. props := extensionProperties(extension)
  375. p := NewBuffer(nil)
  376. x := reflect.New(et)
  377. x.Elem().Set(reflect.ValueOf(value))
  378. if err := props.enc(p, props, toStructPointer(x)); err != nil {
  379. return err
  380. }
  381. *ext = append(*ext, p.buf...)
  382. }
  383. return nil
  384. }
  385. // A global registry of extensions.
  386. // The generated code will register the generated descriptors by calling RegisterExtension.
  387. var extensionMaps = make(map[reflect.Type]map[int32]*ExtensionDesc)
  388. // RegisterExtension is called from the generated code.
  389. func RegisterExtension(desc *ExtensionDesc) {
  390. st := reflect.TypeOf(desc.ExtendedType).Elem()
  391. m := extensionMaps[st]
  392. if m == nil {
  393. m = make(map[int32]*ExtensionDesc)
  394. extensionMaps[st] = m
  395. }
  396. if _, ok := m[desc.Field]; ok {
  397. panic("proto: duplicate extension registered: " + st.String() + " " + strconv.Itoa(int(desc.Field)))
  398. }
  399. m[desc.Field] = desc
  400. }
  401. // RegisteredExtensions returns a map of the registered extensions of a
  402. // protocol buffer struct, indexed by the extension number.
  403. // The argument pb should be a nil pointer to the struct type.
  404. func RegisteredExtensions(pb Message) map[int32]*ExtensionDesc {
  405. return extensionMaps[reflect.TypeOf(pb).Elem()]
  406. }