extensions_gogo.go 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294
  1. // Protocol Buffers for Go with Gadgets
  2. //
  3. // Copyright (c) 2013, The GoGo Authors. All rights reserved.
  4. // http://github.com/gogo/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. //
  17. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  18. // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  19. // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  20. // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  21. // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  22. // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  23. // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  24. // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  25. // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  26. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  27. // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  28. package proto
  29. import (
  30. "bytes"
  31. "errors"
  32. "fmt"
  33. "reflect"
  34. "sort"
  35. "strings"
  36. "sync"
  37. )
  38. func GetBoolExtension(pb Message, extension *ExtensionDesc, ifnotset bool) bool {
  39. if reflect.ValueOf(pb).IsNil() {
  40. return ifnotset
  41. }
  42. value, err := GetExtension(pb, extension)
  43. if err != nil {
  44. return ifnotset
  45. }
  46. if value == nil {
  47. return ifnotset
  48. }
  49. if value.(*bool) == nil {
  50. return ifnotset
  51. }
  52. return *(value.(*bool))
  53. }
  54. func (this *Extension) Equal(that *Extension) bool {
  55. return bytes.Equal(this.enc, that.enc)
  56. }
  57. func (this *Extension) Compare(that *Extension) int {
  58. return bytes.Compare(this.enc, that.enc)
  59. }
  60. func SizeOfInternalExtension(m extendableProto) (n int) {
  61. return SizeOfExtensionMap(m.extensionsWrite())
  62. }
  63. func SizeOfExtensionMap(m map[int32]Extension) (n int) {
  64. return extensionsMapSize(m)
  65. }
  66. type sortableMapElem struct {
  67. field int32
  68. ext Extension
  69. }
  70. func newSortableExtensionsFromMap(m map[int32]Extension) sortableExtensions {
  71. s := make(sortableExtensions, 0, len(m))
  72. for k, v := range m {
  73. s = append(s, &sortableMapElem{field: k, ext: v})
  74. }
  75. return s
  76. }
  77. type sortableExtensions []*sortableMapElem
  78. func (this sortableExtensions) Len() int { return len(this) }
  79. func (this sortableExtensions) Swap(i, j int) { this[i], this[j] = this[j], this[i] }
  80. func (this sortableExtensions) Less(i, j int) bool { return this[i].field < this[j].field }
  81. func (this sortableExtensions) String() string {
  82. sort.Sort(this)
  83. ss := make([]string, len(this))
  84. for i := range this {
  85. ss[i] = fmt.Sprintf("%d: %v", this[i].field, this[i].ext)
  86. }
  87. return "map[" + strings.Join(ss, ",") + "]"
  88. }
  89. func StringFromInternalExtension(m extendableProto) string {
  90. return StringFromExtensionsMap(m.extensionsWrite())
  91. }
  92. func StringFromExtensionsMap(m map[int32]Extension) string {
  93. return newSortableExtensionsFromMap(m).String()
  94. }
  95. func StringFromExtensionsBytes(ext []byte) string {
  96. m, err := BytesToExtensionsMap(ext)
  97. if err != nil {
  98. panic(err)
  99. }
  100. return StringFromExtensionsMap(m)
  101. }
  102. func EncodeInternalExtension(m extendableProto, data []byte) (n int, err error) {
  103. return EncodeExtensionMap(m.extensionsWrite(), data)
  104. }
  105. func EncodeExtensionMap(m map[int32]Extension, data []byte) (n int, err error) {
  106. if err := encodeExtensionsMap(m); err != nil {
  107. return 0, err
  108. }
  109. keys := make([]int, 0, len(m))
  110. for k := range m {
  111. keys = append(keys, int(k))
  112. }
  113. sort.Ints(keys)
  114. for _, k := range keys {
  115. n += copy(data[n:], m[int32(k)].enc)
  116. }
  117. return n, nil
  118. }
  119. func GetRawExtension(m map[int32]Extension, id int32) ([]byte, error) {
  120. if m[id].value == nil || m[id].desc == nil {
  121. return m[id].enc, nil
  122. }
  123. if err := encodeExtensionsMap(m); err != nil {
  124. return nil, err
  125. }
  126. return m[id].enc, nil
  127. }
  128. func size(buf []byte, wire int) (int, error) {
  129. switch wire {
  130. case WireVarint:
  131. _, n := DecodeVarint(buf)
  132. return n, nil
  133. case WireFixed64:
  134. return 8, nil
  135. case WireBytes:
  136. v, n := DecodeVarint(buf)
  137. return int(v) + n, nil
  138. case WireFixed32:
  139. return 4, nil
  140. case WireStartGroup:
  141. offset := 0
  142. for {
  143. u, n := DecodeVarint(buf[offset:])
  144. fwire := int(u & 0x7)
  145. offset += n
  146. if fwire == WireEndGroup {
  147. return offset, nil
  148. }
  149. s, err := size(buf[offset:], wire)
  150. if err != nil {
  151. return 0, err
  152. }
  153. offset += s
  154. }
  155. }
  156. return 0, fmt.Errorf("proto: can't get size for unknown wire type %d", wire)
  157. }
  158. func BytesToExtensionsMap(buf []byte) (map[int32]Extension, error) {
  159. m := make(map[int32]Extension)
  160. i := 0
  161. for i < len(buf) {
  162. tag, n := DecodeVarint(buf[i:])
  163. if n <= 0 {
  164. return nil, fmt.Errorf("unable to decode varint")
  165. }
  166. fieldNum := int32(tag >> 3)
  167. wireType := int(tag & 0x7)
  168. l, err := size(buf[i+n:], wireType)
  169. if err != nil {
  170. return nil, err
  171. }
  172. end := i + int(l) + n
  173. m[int32(fieldNum)] = Extension{enc: buf[i:end]}
  174. i = end
  175. }
  176. return m, nil
  177. }
  178. func NewExtension(e []byte) Extension {
  179. ee := Extension{enc: make([]byte, len(e))}
  180. copy(ee.enc, e)
  181. return ee
  182. }
  183. func AppendExtension(e Message, tag int32, buf []byte) {
  184. if ee, eok := e.(extensionsBytes); eok {
  185. ext := ee.GetExtensions()
  186. *ext = append(*ext, buf...)
  187. return
  188. }
  189. if ee, eok := e.(extendableProto); eok {
  190. m := ee.extensionsWrite()
  191. ext := m[int32(tag)] // may be missing
  192. ext.enc = append(ext.enc, buf...)
  193. m[int32(tag)] = ext
  194. }
  195. }
  196. func encodeExtension(e *Extension) error {
  197. if e.value == nil || e.desc == nil {
  198. // Extension is only in its encoded form.
  199. return nil
  200. }
  201. // We don't skip extensions that have an encoded form set,
  202. // because the extension value may have been mutated after
  203. // the last time this function was called.
  204. et := reflect.TypeOf(e.desc.ExtensionType)
  205. props := extensionProperties(e.desc)
  206. p := NewBuffer(nil)
  207. // If e.value has type T, the encoder expects a *struct{ X T }.
  208. // Pass a *T with a zero field and hope it all works out.
  209. x := reflect.New(et)
  210. x.Elem().Set(reflect.ValueOf(e.value))
  211. if err := props.enc(p, props, toStructPointer(x)); err != nil {
  212. return err
  213. }
  214. e.enc = p.buf
  215. return nil
  216. }
  217. func (this Extension) GoString() string {
  218. if this.enc == nil {
  219. if err := encodeExtension(&this); err != nil {
  220. panic(err)
  221. }
  222. }
  223. return fmt.Sprintf("proto.NewExtension(%#v)", this.enc)
  224. }
  225. func SetUnsafeExtension(pb Message, fieldNum int32, value interface{}) error {
  226. typ := reflect.TypeOf(pb).Elem()
  227. ext, ok := extensionMaps[typ]
  228. if !ok {
  229. return fmt.Errorf("proto: bad extended type; %s is not extendable", typ.String())
  230. }
  231. desc, ok := ext[fieldNum]
  232. if !ok {
  233. return errors.New("proto: bad extension number; not in declared ranges")
  234. }
  235. return SetExtension(pb, desc, value)
  236. }
  237. func GetUnsafeExtension(pb Message, fieldNum int32) (interface{}, error) {
  238. typ := reflect.TypeOf(pb).Elem()
  239. ext, ok := extensionMaps[typ]
  240. if !ok {
  241. return nil, fmt.Errorf("proto: bad extended type; %s is not extendable", typ.String())
  242. }
  243. desc, ok := ext[fieldNum]
  244. if !ok {
  245. return nil, fmt.Errorf("unregistered field number %d", fieldNum)
  246. }
  247. return GetExtension(pb, desc)
  248. }
  249. func NewUnsafeXXX_InternalExtensions(m map[int32]Extension) XXX_InternalExtensions {
  250. x := &XXX_InternalExtensions{
  251. p: new(struct {
  252. mu sync.Mutex
  253. extensionMap map[int32]Extension
  254. }),
  255. }
  256. x.p.extensionMap = m
  257. return *x
  258. }
  259. func GetUnsafeExtensionsMap(extendable Message) map[int32]Extension {
  260. pb := extendable.(extendableProto)
  261. return pb.extensionsWrite()
  262. }