extensions.go 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518
  1. // Go support for Protocol Buffers - Google's data interchange format
  2. //
  3. // Copyright 2010 The Go Authors. All rights reserved.
  4. // https://github.com/golang/protobuf
  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. "fmt"
  38. "reflect"
  39. "strconv"
  40. "sync"
  41. )
  42. // ErrMissingExtension is the error returned by GetExtension if the named extension is not in the message.
  43. var ErrMissingExtension = errors.New("proto: missing extension")
  44. // ExtensionRange represents a range of message extensions for a protocol buffer.
  45. // Used in code generated by the protocol compiler.
  46. type ExtensionRange struct {
  47. Start, End int32 // both inclusive
  48. }
  49. // extendableProto is an interface implemented by any protocol buffer that may be extended.
  50. type extendableProto interface {
  51. Message
  52. ExtensionRangeArray() []ExtensionRange
  53. }
  54. type extensionsMap interface {
  55. extendableProto
  56. ExtensionMap() map[int32]Extension
  57. }
  58. type extensionsBytes interface {
  59. extendableProto
  60. GetExtensions() *[]byte
  61. }
  62. var extendableProtoType = reflect.TypeOf((*extendableProto)(nil)).Elem()
  63. // ExtensionDesc represents an extension specification.
  64. // Used in generated code from the protocol compiler.
  65. type ExtensionDesc struct {
  66. ExtendedType Message // nil pointer to the type that is being extended
  67. ExtensionType interface{} // nil pointer to the extension type
  68. Field int32 // field number
  69. Name string // fully-qualified name of extension, for text formatting
  70. Tag string // protobuf tag style
  71. }
  72. func (ed *ExtensionDesc) repeated() bool {
  73. t := reflect.TypeOf(ed.ExtensionType)
  74. return t.Kind() == reflect.Slice && t.Elem().Kind() != reflect.Uint8
  75. }
  76. // Extension represents an extension in a message.
  77. type Extension struct {
  78. // When an extension is stored in a message using SetExtension
  79. // only desc and value are set. When the message is marshaled
  80. // enc will be set to the encoded form of the message.
  81. //
  82. // When a message is unmarshaled and contains extensions, each
  83. // extension will have only enc set. When such an extension is
  84. // accessed using GetExtension (or GetExtensions) desc and value
  85. // will be set.
  86. desc *ExtensionDesc
  87. value interface{}
  88. enc []byte
  89. }
  90. // SetRawExtension is for testing only.
  91. func SetRawExtension(base extendableProto, id int32, b []byte) {
  92. if ebase, ok := base.(extensionsMap); ok {
  93. ebase.ExtensionMap()[id] = Extension{enc: b}
  94. } else if ebase, ok := base.(extensionsBytes); ok {
  95. clearExtension(base, id)
  96. ext := ebase.GetExtensions()
  97. *ext = append(*ext, b...)
  98. } else {
  99. panic("unreachable")
  100. }
  101. }
  102. // isExtensionField returns true iff the given field number is in an extension range.
  103. func isExtensionField(pb extendableProto, field int32) bool {
  104. for _, er := range pb.ExtensionRangeArray() {
  105. if er.Start <= field && field <= er.End {
  106. return true
  107. }
  108. }
  109. return false
  110. }
  111. // checkExtensionTypes checks that the given extension is valid for pb.
  112. func checkExtensionTypes(pb extendableProto, extension *ExtensionDesc) error {
  113. // Check the extended type.
  114. if a, b := reflect.TypeOf(pb), reflect.TypeOf(extension.ExtendedType); a != b {
  115. return errors.New("proto: bad extended type; " + b.String() + " does not extend " + a.String())
  116. }
  117. // Check the range.
  118. if !isExtensionField(pb, extension.Field) {
  119. return errors.New("proto: bad extension number; not in declared ranges")
  120. }
  121. return nil
  122. }
  123. // extPropKey is sufficient to uniquely identify an extension.
  124. type extPropKey struct {
  125. base reflect.Type
  126. field int32
  127. }
  128. var extProp = struct {
  129. sync.RWMutex
  130. m map[extPropKey]*Properties
  131. }{
  132. m: make(map[extPropKey]*Properties),
  133. }
  134. func extensionProperties(ed *ExtensionDesc) *Properties {
  135. key := extPropKey{base: reflect.TypeOf(ed.ExtendedType), field: ed.Field}
  136. extProp.RLock()
  137. if prop, ok := extProp.m[key]; ok {
  138. extProp.RUnlock()
  139. return prop
  140. }
  141. extProp.RUnlock()
  142. extProp.Lock()
  143. defer extProp.Unlock()
  144. // Check again.
  145. if prop, ok := extProp.m[key]; ok {
  146. return prop
  147. }
  148. prop := new(Properties)
  149. prop.Init(reflect.TypeOf(ed.ExtensionType), "unknown_name", ed.Tag, nil)
  150. extProp.m[key] = prop
  151. return prop
  152. }
  153. // encodeExtensionMap encodes any unmarshaled (unencoded) extensions in m.
  154. func encodeExtensionMap(m map[int32]Extension) error {
  155. for k, e := range m {
  156. err := encodeExtension(&e)
  157. if err != nil {
  158. return err
  159. }
  160. m[k] = e
  161. }
  162. return nil
  163. }
  164. func encodeExtension(e *Extension) error {
  165. if e.value == nil || e.desc == nil {
  166. // Extension is only in its encoded form.
  167. return nil
  168. }
  169. // We don't skip extensions that have an encoded form set,
  170. // because the extension value may have been mutated after
  171. // the last time this function was called.
  172. et := reflect.TypeOf(e.desc.ExtensionType)
  173. props := extensionProperties(e.desc)
  174. p := NewBuffer(nil)
  175. // If e.value has type T, the encoder expects a *struct{ X T }.
  176. // Pass a *T with a zero field and hope it all works out.
  177. x := reflect.New(et)
  178. x.Elem().Set(reflect.ValueOf(e.value))
  179. if err := props.enc(p, props, toStructPointer(x)); err != nil {
  180. return err
  181. }
  182. e.enc = p.buf
  183. return nil
  184. }
  185. func sizeExtensionMap(m map[int32]Extension) (n int) {
  186. for _, e := range m {
  187. if e.value == nil || e.desc == nil {
  188. // Extension is only in its encoded form.
  189. n += len(e.enc)
  190. continue
  191. }
  192. // We don't skip extensions that have an encoded form set,
  193. // because the extension value may have been mutated after
  194. // the last time this function was called.
  195. et := reflect.TypeOf(e.desc.ExtensionType)
  196. props := extensionProperties(e.desc)
  197. // If e.value has type T, the encoder expects a *struct{ X T }.
  198. // Pass a *T with a zero field and hope it all works out.
  199. x := reflect.New(et)
  200. x.Elem().Set(reflect.ValueOf(e.value))
  201. n += props.size(props, toStructPointer(x))
  202. }
  203. return
  204. }
  205. // HasExtension returns whether the given extension is present in pb.
  206. func HasExtension(pb extendableProto, extension *ExtensionDesc) bool {
  207. // TODO: Check types, field numbers, etc.?
  208. if epb, doki := pb.(extensionsMap); doki {
  209. _, ok := epb.ExtensionMap()[extension.Field]
  210. return ok
  211. } else if epb, doki := pb.(extensionsBytes); doki {
  212. ext := epb.GetExtensions()
  213. buf := *ext
  214. o := 0
  215. for o < len(buf) {
  216. tag, n := DecodeVarint(buf[o:])
  217. fieldNum := int32(tag >> 3)
  218. if int32(fieldNum) == extension.Field {
  219. return true
  220. }
  221. wireType := int(tag & 0x7)
  222. o += n
  223. l, err := size(buf[o:], wireType)
  224. if err != nil {
  225. return false
  226. }
  227. o += l
  228. }
  229. return false
  230. }
  231. panic("unreachable")
  232. }
  233. func deleteExtension(pb extensionsBytes, theFieldNum int32, offset int) int {
  234. ext := pb.GetExtensions()
  235. for offset < len(*ext) {
  236. tag, n1 := DecodeVarint((*ext)[offset:])
  237. fieldNum := int32(tag >> 3)
  238. wireType := int(tag & 0x7)
  239. n2, err := size((*ext)[offset+n1:], wireType)
  240. if err != nil {
  241. panic(err)
  242. }
  243. newOffset := offset + n1 + n2
  244. if fieldNum == theFieldNum {
  245. *ext = append((*ext)[:offset], (*ext)[newOffset:]...)
  246. return offset
  247. }
  248. offset = newOffset
  249. }
  250. return -1
  251. }
  252. func clearExtension(pb extendableProto, fieldNum int32) {
  253. if epb, doki := pb.(extensionsMap); doki {
  254. delete(epb.ExtensionMap(), fieldNum)
  255. } else if epb, doki := pb.(extensionsBytes); doki {
  256. offset := 0
  257. for offset != -1 {
  258. offset = deleteExtension(epb, fieldNum, offset)
  259. }
  260. } else {
  261. panic("unreachable")
  262. }
  263. }
  264. // ClearExtension removes the given extension from pb.
  265. func ClearExtension(pb extendableProto, extension *ExtensionDesc) {
  266. // TODO: Check types, field numbers, etc.?
  267. clearExtension(pb, extension.Field)
  268. }
  269. // GetExtension parses and returns the given extension of pb.
  270. // If the extension is not present it returns ErrMissingExtension.
  271. func GetExtension(pb extendableProto, extension *ExtensionDesc) (interface{}, error) {
  272. if err := checkExtensionTypes(pb, extension); err != nil {
  273. return nil, err
  274. }
  275. if epb, doki := pb.(extensionsMap); doki {
  276. emap := epb.ExtensionMap()
  277. e, ok := emap[extension.Field]
  278. if !ok {
  279. // defaultExtensionValue returns the default value or
  280. // ErrMissingExtension if there is no default.
  281. return defaultExtensionValue(extension)
  282. }
  283. if e.value != nil {
  284. // Already decoded. Check the descriptor, though.
  285. if e.desc != extension {
  286. // This shouldn't happen. If it does, it means that
  287. // GetExtension was called twice with two different
  288. // descriptors with the same field number.
  289. return nil, errors.New("proto: descriptor conflict")
  290. }
  291. return e.value, nil
  292. }
  293. v, err := decodeExtension(e.enc, extension)
  294. if err != nil {
  295. return nil, err
  296. }
  297. // Remember the decoded version and drop the encoded version.
  298. // That way it is safe to mutate what we return.
  299. e.value = v
  300. e.desc = extension
  301. e.enc = nil
  302. emap[extension.Field] = e
  303. return e.value, nil
  304. } else if epb, doki := pb.(extensionsBytes); doki {
  305. ext := epb.GetExtensions()
  306. o := 0
  307. for o < len(*ext) {
  308. tag, n := DecodeVarint((*ext)[o:])
  309. fieldNum := int32(tag >> 3)
  310. wireType := int(tag & 0x7)
  311. l, err := size((*ext)[o+n:], wireType)
  312. if err != nil {
  313. return nil, err
  314. }
  315. if int32(fieldNum) == extension.Field {
  316. v, err := decodeExtension((*ext)[o:o+n+l], extension)
  317. if err != nil {
  318. return nil, err
  319. }
  320. return v, nil
  321. }
  322. o += n + l
  323. }
  324. return defaultExtensionValue(extension)
  325. }
  326. panic("unreachable")
  327. }
  328. // defaultExtensionValue returns the default value for extension.
  329. // If no default for an extension is defined ErrMissingExtension is returned.
  330. func defaultExtensionValue(extension *ExtensionDesc) (interface{}, error) {
  331. t := reflect.TypeOf(extension.ExtensionType)
  332. props := extensionProperties(extension)
  333. sf, _, err := fieldDefault(t, props)
  334. if err != nil {
  335. return nil, err
  336. }
  337. if sf == nil || sf.value == nil {
  338. // There is no default value.
  339. return nil, ErrMissingExtension
  340. }
  341. if t.Kind() != reflect.Ptr {
  342. // We do not need to return a Ptr, we can directly return sf.value.
  343. return sf.value, nil
  344. }
  345. // We need to return an interface{} that is a pointer to sf.value.
  346. value := reflect.New(t).Elem()
  347. value.Set(reflect.New(value.Type().Elem()))
  348. if sf.kind == reflect.Int32 {
  349. // We may have an int32 or an enum, but the underlying data is int32.
  350. // Since we can't set an int32 into a non int32 reflect.value directly
  351. // set it as a int32.
  352. value.Elem().SetInt(int64(sf.value.(int32)))
  353. } else {
  354. value.Elem().Set(reflect.ValueOf(sf.value))
  355. }
  356. return value.Interface(), nil
  357. }
  358. // decodeExtension decodes an extension encoded in b.
  359. func decodeExtension(b []byte, extension *ExtensionDesc) (interface{}, error) {
  360. o := NewBuffer(b)
  361. t := reflect.TypeOf(extension.ExtensionType)
  362. props := extensionProperties(extension)
  363. // t is a pointer to a struct, pointer to basic type or a slice.
  364. // Allocate a "field" to store the pointer/slice itself; the
  365. // pointer/slice will be stored here. We pass
  366. // the address of this field to props.dec.
  367. // This passes a zero field and a *t and lets props.dec
  368. // interpret it as a *struct{ x t }.
  369. value := reflect.New(t).Elem()
  370. for {
  371. // Discard wire type and field number varint. It isn't needed.
  372. if _, err := o.DecodeVarint(); err != nil {
  373. return nil, err
  374. }
  375. if err := props.dec(o, props, toStructPointer(value.Addr())); err != nil {
  376. return nil, err
  377. }
  378. if o.index >= len(o.buf) {
  379. break
  380. }
  381. }
  382. return value.Interface(), nil
  383. }
  384. // GetExtensions returns a slice of the extensions present in pb that are also listed in es.
  385. // The returned slice has the same length as es; missing extensions will appear as nil elements.
  386. func GetExtensions(pb Message, es []*ExtensionDesc) (extensions []interface{}, err error) {
  387. epb, ok := pb.(extendableProto)
  388. if !ok {
  389. err = errors.New("proto: not an extendable proto")
  390. return
  391. }
  392. extensions = make([]interface{}, len(es))
  393. for i, e := range es {
  394. extensions[i], err = GetExtension(epb, e)
  395. if err == ErrMissingExtension {
  396. err = nil
  397. }
  398. if err != nil {
  399. return
  400. }
  401. }
  402. return
  403. }
  404. // SetExtension sets the specified extension of pb to the specified value.
  405. func SetExtension(pb extendableProto, extension *ExtensionDesc, value interface{}) error {
  406. if err := checkExtensionTypes(pb, extension); err != nil {
  407. return err
  408. }
  409. typ := reflect.TypeOf(extension.ExtensionType)
  410. if typ != reflect.TypeOf(value) {
  411. return errors.New("proto: bad extension value type")
  412. }
  413. // nil extension values need to be caught early, because the
  414. // encoder can't distinguish an ErrNil due to a nil extension
  415. // from an ErrNil due to a missing field. Extensions are
  416. // always optional, so the encoder would just swallow the error
  417. // and drop all the extensions from the encoded message.
  418. if reflect.ValueOf(value).IsNil() {
  419. return fmt.Errorf("proto: SetExtension called with nil value of type %T", value)
  420. }
  421. return setExtension(pb, extension, value)
  422. }
  423. func setExtension(pb extendableProto, extension *ExtensionDesc, value interface{}) error {
  424. if epb, doki := pb.(extensionsMap); doki {
  425. epb.ExtensionMap()[extension.Field] = Extension{desc: extension, value: value}
  426. } else if epb, doki := pb.(extensionsBytes); doki {
  427. ClearExtension(pb, extension)
  428. ext := epb.GetExtensions()
  429. et := reflect.TypeOf(extension.ExtensionType)
  430. props := extensionProperties(extension)
  431. p := NewBuffer(nil)
  432. x := reflect.New(et)
  433. x.Elem().Set(reflect.ValueOf(value))
  434. if err := props.enc(p, props, toStructPointer(x)); err != nil {
  435. return err
  436. }
  437. *ext = append(*ext, p.buf...)
  438. }
  439. return nil
  440. }
  441. // A global registry of extensions.
  442. // The generated code will register the generated descriptors by calling RegisterExtension.
  443. var extensionMaps = make(map[reflect.Type]map[int32]*ExtensionDesc)
  444. // RegisterExtension is called from the generated code.
  445. func RegisterExtension(desc *ExtensionDesc) {
  446. st := reflect.TypeOf(desc.ExtendedType).Elem()
  447. m := extensionMaps[st]
  448. if m == nil {
  449. m = make(map[int32]*ExtensionDesc)
  450. extensionMaps[st] = m
  451. }
  452. if _, ok := m[desc.Field]; ok {
  453. panic("proto: duplicate extension registered: " + st.String() + " " + strconv.Itoa(int(desc.Field)))
  454. }
  455. m[desc.Field] = desc
  456. }
  457. // RegisteredExtensions returns a map of the registered extensions of a
  458. // protocol buffer struct, indexed by the extension number.
  459. // The argument pb should be a nil pointer to the struct type.
  460. func RegisteredExtensions(pb Message) map[int32]*ExtensionDesc {
  461. return extensionMaps[reflect.TypeOf(pb).Elem()]
  462. }