extensions_gogo.go 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236
  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 (this *Extension) Compare(that *Extension) int {
  55. return bytes.Compare(this.enc, that.enc)
  56. }
  57. func SizeOfExtensionMap(m map[int32]Extension) (n int) {
  58. return sizeExtensionMap(m)
  59. }
  60. type sortableMapElem struct {
  61. field int32
  62. ext Extension
  63. }
  64. func newSortableExtensionsFromMap(m map[int32]Extension) sortableExtensions {
  65. s := make(sortableExtensions, 0, len(m))
  66. for k, v := range m {
  67. s = append(s, &sortableMapElem{field: k, ext: v})
  68. }
  69. return s
  70. }
  71. type sortableExtensions []*sortableMapElem
  72. func (this sortableExtensions) Len() int { return len(this) }
  73. func (this sortableExtensions) Swap(i, j int) { this[i], this[j] = this[j], this[i] }
  74. func (this sortableExtensions) Less(i, j int) bool { return this[i].field < this[j].field }
  75. func (this sortableExtensions) String() string {
  76. sort.Sort(this)
  77. ss := make([]string, len(this))
  78. for i := range this {
  79. ss[i] = fmt.Sprintf("%d: %v", this[i].field, this[i].ext)
  80. }
  81. return "map[" + strings.Join(ss, ",") + "]"
  82. }
  83. func StringFromExtensionsMap(m map[int32]Extension) string {
  84. return newSortableExtensionsFromMap(m).String()
  85. }
  86. func StringFromExtensionsBytes(ext []byte) string {
  87. m, err := BytesToExtensionsMap(ext)
  88. if err != nil {
  89. panic(err)
  90. }
  91. return StringFromExtensionsMap(m)
  92. }
  93. func EncodeExtensionMap(m map[int32]Extension, data []byte) (n int, err error) {
  94. if err := encodeExtensionMap(m); err != nil {
  95. return 0, err
  96. }
  97. keys := make([]int, 0, len(m))
  98. for k := range m {
  99. keys = append(keys, int(k))
  100. }
  101. sort.Ints(keys)
  102. for _, k := range keys {
  103. n += copy(data[n:], m[int32(k)].enc)
  104. }
  105. return n, nil
  106. }
  107. func GetRawExtension(m map[int32]Extension, id int32) ([]byte, error) {
  108. if m[id].value == nil || m[id].desc == nil {
  109. return m[id].enc, nil
  110. }
  111. if err := encodeExtensionMap(m); err != nil {
  112. return nil, err
  113. }
  114. return m[id].enc, nil
  115. }
  116. func size(buf []byte, wire int) (int, error) {
  117. switch wire {
  118. case WireVarint:
  119. _, n := DecodeVarint(buf)
  120. return n, nil
  121. case WireFixed64:
  122. return 8, nil
  123. case WireBytes:
  124. v, n := DecodeVarint(buf)
  125. return int(v) + n, nil
  126. case WireFixed32:
  127. return 4, nil
  128. case WireStartGroup:
  129. offset := 0
  130. for {
  131. u, n := DecodeVarint(buf[offset:])
  132. fwire := int(u & 0x7)
  133. offset += n
  134. if fwire == WireEndGroup {
  135. return offset, nil
  136. }
  137. s, err := size(buf[offset:], wire)
  138. if err != nil {
  139. return 0, err
  140. }
  141. offset += s
  142. }
  143. }
  144. return 0, fmt.Errorf("proto: can't get size for unknown wire type %d", wire)
  145. }
  146. func BytesToExtensionsMap(buf []byte) (map[int32]Extension, error) {
  147. m := make(map[int32]Extension)
  148. i := 0
  149. for i < len(buf) {
  150. tag, n := DecodeVarint(buf[i:])
  151. if n <= 0 {
  152. return nil, fmt.Errorf("unable to decode varint")
  153. }
  154. fieldNum := int32(tag >> 3)
  155. wireType := int(tag & 0x7)
  156. l, err := size(buf[i+n:], wireType)
  157. if err != nil {
  158. return nil, err
  159. }
  160. end := i + int(l) + n
  161. m[int32(fieldNum)] = Extension{enc: buf[i:end]}
  162. i = end
  163. }
  164. return m, nil
  165. }
  166. func NewExtension(e []byte) Extension {
  167. ee := Extension{enc: make([]byte, len(e))}
  168. copy(ee.enc, e)
  169. return ee
  170. }
  171. func AppendExtension(e extendableProto, tag int32, buf []byte) {
  172. if ee, eok := e.(extensionsMap); eok {
  173. ext := ee.ExtensionMap()[int32(tag)] // may be missing
  174. ext.enc = append(ext.enc, buf...)
  175. ee.ExtensionMap()[int32(tag)] = ext
  176. } else if ee, eok := e.(extensionsBytes); eok {
  177. ext := ee.GetExtensions()
  178. *ext = append(*ext, buf...)
  179. }
  180. }
  181. func (this Extension) GoString() string {
  182. if this.enc == nil {
  183. if err := encodeExtension(&this); err != nil {
  184. panic(err)
  185. }
  186. }
  187. return fmt.Sprintf("proto.NewExtension(%#v)", this.enc)
  188. }
  189. func SetUnsafeExtension(pb extendableProto, fieldNum int32, value interface{}) error {
  190. typ := reflect.TypeOf(pb).Elem()
  191. ext, ok := extensionMaps[typ]
  192. if !ok {
  193. return fmt.Errorf("proto: bad extended type; %s is not extendable", typ.String())
  194. }
  195. desc, ok := ext[fieldNum]
  196. if !ok {
  197. return errors.New("proto: bad extension number; not in declared ranges")
  198. }
  199. return setExtension(pb, desc, value)
  200. }
  201. func GetUnsafeExtension(pb extendableProto, fieldNum int32) (interface{}, error) {
  202. typ := reflect.TypeOf(pb).Elem()
  203. ext, ok := extensionMaps[typ]
  204. if !ok {
  205. return nil, fmt.Errorf("proto: bad extended type; %s is not extendable", typ.String())
  206. }
  207. desc, ok := ext[fieldNum]
  208. if !ok {
  209. return nil, fmt.Errorf("unregistered field number %d", fieldNum)
  210. }
  211. return GetExtension(pb, desc)
  212. }