extensions_gogo.go 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232
  1. // Copyright (c) 2013, Vastech SA (PTY) LTD. All rights reserved.
  2. // http://github.com/gogo/protobuf/gogoproto
  3. //
  4. // Redistribution and use in source and binary forms, with or without
  5. // modification, are permitted provided that the following conditions are
  6. // met:
  7. //
  8. // * Redistributions of source code must retain the above copyright
  9. // notice, this list of conditions and the following disclaimer.
  10. // * Redistributions in binary form must reproduce the above
  11. // copyright notice, this list of conditions and the following disclaimer
  12. // in the documentation and/or other materials provided with the
  13. // distribution.
  14. //
  15. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  16. // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  17. // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  18. // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  19. // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  20. // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  21. // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  22. // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  23. // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  24. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  25. // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  26. package proto
  27. import (
  28. "bytes"
  29. "errors"
  30. "fmt"
  31. "reflect"
  32. "sort"
  33. "strings"
  34. )
  35. func GetBoolExtension(pb extendableProto, extension *ExtensionDesc, ifnotset bool) bool {
  36. if reflect.ValueOf(pb).IsNil() {
  37. return ifnotset
  38. }
  39. value, err := GetExtension(pb, extension)
  40. if err != nil {
  41. return ifnotset
  42. }
  43. if value == nil {
  44. return ifnotset
  45. }
  46. if value.(*bool) == nil {
  47. return ifnotset
  48. }
  49. return *(value.(*bool))
  50. }
  51. func (this *Extension) Equal(that *Extension) bool {
  52. return bytes.Equal(this.enc, that.enc)
  53. }
  54. func SizeOfExtensionMap(m map[int32]Extension) (n int) {
  55. return sizeExtensionMap(m)
  56. }
  57. type sortableMapElem struct {
  58. field int32
  59. ext Extension
  60. }
  61. func newSortableExtensionsFromMap(m map[int32]Extension) sortableExtensions {
  62. s := make(sortableExtensions, 0, len(m))
  63. for k, v := range m {
  64. s = append(s, &sortableMapElem{field: k, ext: v})
  65. }
  66. return s
  67. }
  68. type sortableExtensions []*sortableMapElem
  69. func (this sortableExtensions) Len() int { return len(this) }
  70. func (this sortableExtensions) Swap(i, j int) { this[i], this[j] = this[j], this[i] }
  71. func (this sortableExtensions) Less(i, j int) bool { return this[i].field < this[j].field }
  72. func (this sortableExtensions) String() string {
  73. sort.Sort(this)
  74. ss := make([]string, len(this))
  75. for i := range this {
  76. ss[i] = fmt.Sprintf("%d: %v", this[i].field, this[i].ext)
  77. }
  78. return "map[" + strings.Join(ss, ",") + "]"
  79. }
  80. func StringFromExtensionsMap(m map[int32]Extension) string {
  81. return newSortableExtensionsFromMap(m).String()
  82. }
  83. func StringFromExtensionsBytes(ext []byte) string {
  84. m, err := BytesToExtensionsMap(ext)
  85. if err != nil {
  86. panic(err)
  87. }
  88. return StringFromExtensionsMap(m)
  89. }
  90. func EncodeExtensionMap(m map[int32]Extension, data []byte) (n int, err error) {
  91. if err := encodeExtensionMap(m); err != nil {
  92. return 0, err
  93. }
  94. keys := make([]int, 0, len(m))
  95. for k := range m {
  96. keys = append(keys, int(k))
  97. }
  98. sort.Ints(keys)
  99. for _, k := range keys {
  100. n += copy(data[n:], m[int32(k)].enc)
  101. }
  102. return n, nil
  103. }
  104. func GetRawExtension(m map[int32]Extension, id int32) ([]byte, error) {
  105. if m[id].value == nil || m[id].desc == nil {
  106. return m[id].enc, nil
  107. }
  108. if err := encodeExtensionMap(m); err != nil {
  109. return nil, err
  110. }
  111. return m[id].enc, nil
  112. }
  113. func size(buf []byte, wire int) (int, error) {
  114. switch wire {
  115. case WireVarint:
  116. _, n := DecodeVarint(buf)
  117. return n, nil
  118. case WireFixed64:
  119. return 8, nil
  120. case WireBytes:
  121. v, n := DecodeVarint(buf)
  122. return int(v) + n, nil
  123. case WireFixed32:
  124. return 4, nil
  125. case WireStartGroup:
  126. offset := 0
  127. for {
  128. u, n := DecodeVarint(buf[offset:])
  129. fwire := int(u & 0x7)
  130. offset += n
  131. if fwire == WireEndGroup {
  132. return offset, nil
  133. }
  134. s, err := size(buf[offset:], wire)
  135. if err != nil {
  136. return 0, err
  137. }
  138. offset += s
  139. }
  140. }
  141. return 0, fmt.Errorf("proto: can't get size for unknown wire type %d", wire)
  142. }
  143. func BytesToExtensionsMap(buf []byte) (map[int32]Extension, error) {
  144. m := make(map[int32]Extension)
  145. i := 0
  146. for i < len(buf) {
  147. tag, n := DecodeVarint(buf[i:])
  148. if n <= 0 {
  149. return nil, fmt.Errorf("unable to decode varint")
  150. }
  151. fieldNum := int32(tag >> 3)
  152. wireType := int(tag & 0x7)
  153. l, err := size(buf[i+n:], wireType)
  154. if err != nil {
  155. return nil, err
  156. }
  157. end := i + int(l) + n
  158. m[int32(fieldNum)] = Extension{enc: buf[i:end]}
  159. i = end
  160. }
  161. return m, nil
  162. }
  163. func NewExtension(e []byte) Extension {
  164. ee := Extension{enc: make([]byte, len(e))}
  165. copy(ee.enc, e)
  166. return ee
  167. }
  168. func AppendExtension(e extendableProto, tag int32, buf []byte) {
  169. if ee, eok := e.(extensionsMap); eok {
  170. ext := ee.ExtensionMap()[int32(tag)] // may be missing
  171. ext.enc = append(ext.enc, buf...)
  172. ee.ExtensionMap()[int32(tag)] = ext
  173. } else if ee, eok := e.(extensionsBytes); eok {
  174. ext := ee.GetExtensions()
  175. *ext = append(*ext, buf...)
  176. }
  177. }
  178. func (this Extension) GoString() string {
  179. if this.enc == nil {
  180. if err := encodeExtension(&this); err != nil {
  181. panic(err)
  182. }
  183. }
  184. return fmt.Sprintf("proto.NewExtension(%#v)", this.enc)
  185. }
  186. func SetUnsafeExtension(pb extendableProto, fieldNum int32, value interface{}) error {
  187. typ := reflect.TypeOf(pb).Elem()
  188. ext, ok := extensionMaps[typ]
  189. if !ok {
  190. return fmt.Errorf("proto: bad extended type; %s is not extendable", typ.String())
  191. }
  192. desc, ok := ext[fieldNum]
  193. if !ok {
  194. return errors.New("proto: bad extension number; not in declared ranges")
  195. }
  196. return setExtension(pb, desc, value)
  197. }
  198. func GetUnsafeExtension(pb extendableProto, fieldNum int32) (interface{}, error) {
  199. typ := reflect.TypeOf(pb).Elem()
  200. ext, ok := extensionMaps[typ]
  201. if !ok {
  202. return nil, fmt.Errorf("proto: bad extended type; %s is not extendable", typ.String())
  203. }
  204. desc, ok := ext[fieldNum]
  205. if !ok {
  206. return nil, fmt.Errorf("unregistered field number %d", fieldNum)
  207. }
  208. return GetExtension(pb, desc)
  209. }