buffer.go 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324
  1. // Copyright 2019 The Go Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package proto
  5. import (
  6. "errors"
  7. "fmt"
  8. "google.golang.org/protobuf/encoding/prototext"
  9. "google.golang.org/protobuf/encoding/protowire"
  10. "google.golang.org/protobuf/runtime/protoimpl"
  11. )
  12. const (
  13. WireVarint = 0
  14. WireFixed32 = 5
  15. WireFixed64 = 1
  16. WireBytes = 2
  17. WireStartGroup = 3
  18. WireEndGroup = 4
  19. )
  20. // EncodeVarint returns the varint encoded bytes of v.
  21. func EncodeVarint(v uint64) []byte {
  22. return protowire.AppendVarint(nil, v)
  23. }
  24. // SizeVarint returns the length of the varint encoded bytes of v.
  25. // This is equal to len(EncodeVarint(v)).
  26. func SizeVarint(v uint64) int {
  27. return protowire.SizeVarint(v)
  28. }
  29. // DecodeVarint parses a varint encoded integer from b,
  30. // returning the integer value and the length of the varint.
  31. // It returns (0, 0) if there is a parse error.
  32. func DecodeVarint(b []byte) (uint64, int) {
  33. v, n := protowire.ConsumeVarint(b)
  34. if n < 0 {
  35. return 0, 0
  36. }
  37. return v, n
  38. }
  39. // Buffer is a buffer for encoding and decoding the protobuf wire format.
  40. // It may be reused between invocations to reduce memory usage.
  41. type Buffer struct {
  42. buf []byte
  43. idx int
  44. deterministic bool
  45. }
  46. // NewBuffer allocates a new Buffer initialized with buf,
  47. // where the contents of buf are considered the unread portion of the buffer.
  48. func NewBuffer(buf []byte) *Buffer {
  49. return &Buffer{buf: buf}
  50. }
  51. // SetDeterministic specifies whether to use deterministic serialization.
  52. //
  53. // Deterministic serialization guarantees that for a given binary, equal
  54. // messages will always be serialized to the same bytes. This implies:
  55. //
  56. // - Repeated serialization of a message will return the same bytes.
  57. // - Different processes of the same binary (which may be executing on
  58. // different machines) will serialize equal messages to the same bytes.
  59. //
  60. // Note that the deterministic serialization is NOT canonical across
  61. // languages. It is not guaranteed to remain stable over time. It is unstable
  62. // across different builds with schema changes due to unknown fields.
  63. // Users who need canonical serialization (e.g., persistent storage in a
  64. // canonical form, fingerprinting, etc.) should define their own
  65. // canonicalization specification and implement their own serializer rather
  66. // than relying on this API.
  67. //
  68. // If deterministic serialization is requested, map entries will be sorted
  69. // by keys in lexographical order. This is an implementation detail and
  70. // subject to change.
  71. func (b *Buffer) SetDeterministic(deterministic bool) {
  72. b.deterministic = deterministic
  73. }
  74. // SetBuf sets buf as the internal buffer,
  75. // where the contents of buf are considered the unread portion of the buffer.
  76. func (b *Buffer) SetBuf(buf []byte) {
  77. b.buf = buf
  78. b.idx = 0
  79. }
  80. // Reset clears the internal buffer of all written and unread data.
  81. func (b *Buffer) Reset() {
  82. b.buf = b.buf[:0]
  83. b.idx = 0
  84. }
  85. // Bytes returns the internal buffer.
  86. func (b *Buffer) Bytes() []byte {
  87. return b.buf
  88. }
  89. // Unread returns the unread portion of the buffer.
  90. func (b *Buffer) Unread() []byte {
  91. return b.buf[b.idx:]
  92. }
  93. // Marshal appends the wire-format encoding of m to the buffer.
  94. func (b *Buffer) Marshal(m Message) error {
  95. var err error
  96. b.buf, err = marshalAppend(b.buf, m, b.deterministic)
  97. return err
  98. }
  99. // Unmarshal parses the wire-format message in the buffer and
  100. // places the decoded results in m.
  101. // It does not reset m before unmarshaling.
  102. func (b *Buffer) Unmarshal(m Message) error {
  103. err := UnmarshalMerge(b.Unread(), m)
  104. b.idx = len(b.buf)
  105. return err
  106. }
  107. type unknownFields struct{ XXX_unrecognized protoimpl.UnknownFields }
  108. func (m *unknownFields) String() string { panic("not implemented") }
  109. func (m *unknownFields) Reset() { panic("not implemented") }
  110. func (m *unknownFields) ProtoMessage() { panic("not implemented") }
  111. // DebugPrint dumps the encoded bytes of b with a header and footer including s
  112. // to stdout. This is only intended for debugging.
  113. func (*Buffer) DebugPrint(s string, b []byte) {
  114. m := MessageReflect(new(unknownFields))
  115. m.SetUnknown(b)
  116. b, _ = prototext.MarshalOptions{AllowPartial: true, Indent: "\t"}.Marshal(m.Interface())
  117. fmt.Printf("==== %s ====\n%s==== %s ====\n", s, b, s)
  118. }
  119. // EncodeVarint appends an unsigned varint encoding to the buffer.
  120. func (b *Buffer) EncodeVarint(v uint64) error {
  121. b.buf = protowire.AppendVarint(b.buf, v)
  122. return nil
  123. }
  124. // EncodeZigzag32 appends a 32-bit zig-zag varint encoding to the buffer.
  125. func (b *Buffer) EncodeZigzag32(v uint64) error {
  126. return b.EncodeVarint(uint64((uint32(v) << 1) ^ uint32((int32(v) >> 31))))
  127. }
  128. // EncodeZigzag64 appends a 64-bit zig-zag varint encoding to the buffer.
  129. func (b *Buffer) EncodeZigzag64(v uint64) error {
  130. return b.EncodeVarint(uint64((uint64(v) << 1) ^ uint64((int64(v) >> 63))))
  131. }
  132. // EncodeFixed32 appends a 32-bit little-endian integer to the buffer.
  133. func (b *Buffer) EncodeFixed32(v uint64) error {
  134. b.buf = protowire.AppendFixed32(b.buf, uint32(v))
  135. return nil
  136. }
  137. // EncodeFixed64 appends a 64-bit little-endian integer to the buffer.
  138. func (b *Buffer) EncodeFixed64(v uint64) error {
  139. b.buf = protowire.AppendFixed64(b.buf, uint64(v))
  140. return nil
  141. }
  142. // EncodeRawBytes appends a length-prefixed raw bytes to the buffer.
  143. func (b *Buffer) EncodeRawBytes(v []byte) error {
  144. b.buf = protowire.AppendBytes(b.buf, v)
  145. return nil
  146. }
  147. // EncodeStringBytes appends a length-prefixed raw bytes to the buffer.
  148. // It does not validate whether v contains valid UTF-8.
  149. func (b *Buffer) EncodeStringBytes(v string) error {
  150. b.buf = protowire.AppendString(b.buf, v)
  151. return nil
  152. }
  153. // EncodeMessage appends a length-prefixed encoded message to the buffer.
  154. func (b *Buffer) EncodeMessage(m Message) error {
  155. var err error
  156. b.buf = protowire.AppendVarint(b.buf, uint64(Size(m)))
  157. b.buf, err = marshalAppend(b.buf, m, b.deterministic)
  158. return err
  159. }
  160. // DecodeVarint consumes an encoded unsigned varint from the buffer.
  161. func (b *Buffer) DecodeVarint() (uint64, error) {
  162. v, n := protowire.ConsumeVarint(b.buf[b.idx:])
  163. if n < 0 {
  164. return 0, protowire.ParseError(n)
  165. }
  166. b.idx += n
  167. return uint64(v), nil
  168. }
  169. // DecodeZigzag32 consumes an encoded 32-bit zig-zag varint from the buffer.
  170. func (b *Buffer) DecodeZigzag32() (uint64, error) {
  171. v, err := b.DecodeVarint()
  172. if err != nil {
  173. return 0, err
  174. }
  175. return uint64((uint32(v) >> 1) ^ uint32((int32(v&1)<<31)>>31)), nil
  176. }
  177. // DecodeZigzag64 consumes an encoded 64-bit zig-zag varint from the buffer.
  178. func (b *Buffer) DecodeZigzag64() (uint64, error) {
  179. v, err := b.DecodeVarint()
  180. if err != nil {
  181. return 0, err
  182. }
  183. return uint64((uint64(v) >> 1) ^ uint64((int64(v&1)<<63)>>63)), nil
  184. }
  185. // DecodeFixed32 consumes a 32-bit little-endian integer from the buffer.
  186. func (b *Buffer) DecodeFixed32() (uint64, error) {
  187. v, n := protowire.ConsumeFixed32(b.buf[b.idx:])
  188. if n < 0 {
  189. return 0, protowire.ParseError(n)
  190. }
  191. b.idx += n
  192. return uint64(v), nil
  193. }
  194. // DecodeFixed64 consumes a 64-bit little-endian integer from the buffer.
  195. func (b *Buffer) DecodeFixed64() (uint64, error) {
  196. v, n := protowire.ConsumeFixed64(b.buf[b.idx:])
  197. if n < 0 {
  198. return 0, protowire.ParseError(n)
  199. }
  200. b.idx += n
  201. return uint64(v), nil
  202. }
  203. // DecodeRawBytes consumes a length-prefixed raw bytes from the buffer.
  204. // If alloc is specified, it returns a copy the raw bytes
  205. // rather than a sub-slice of the buffer.
  206. func (b *Buffer) DecodeRawBytes(alloc bool) ([]byte, error) {
  207. v, n := protowire.ConsumeBytes(b.buf[b.idx:])
  208. if n < 0 {
  209. return nil, protowire.ParseError(n)
  210. }
  211. b.idx += n
  212. if alloc {
  213. v = append([]byte(nil), v...)
  214. }
  215. return v, nil
  216. }
  217. // DecodeStringBytes consumes a length-prefixed raw bytes from the buffer.
  218. // It does not validate whether the raw bytes contain valid UTF-8.
  219. func (b *Buffer) DecodeStringBytes() (string, error) {
  220. v, n := protowire.ConsumeString(b.buf[b.idx:])
  221. if n < 0 {
  222. return "", protowire.ParseError(n)
  223. }
  224. b.idx += n
  225. return v, nil
  226. }
  227. // DecodeMessage consumes a length-prefixed message from the buffer.
  228. // It does not reset m before unmarshaling.
  229. func (b *Buffer) DecodeMessage(m Message) error {
  230. v, err := b.DecodeRawBytes(false)
  231. if err != nil {
  232. return err
  233. }
  234. return UnmarshalMerge(v, m)
  235. }
  236. // DecodeGroup consumes a message group from the buffer.
  237. // It assumes that the start group marker has already been consumed and
  238. // consumes all bytes until (and including the end group marker).
  239. // It does not reset m before unmarshaling.
  240. func (b *Buffer) DecodeGroup(m Message) error {
  241. v, n, err := consumeGroup(b.buf[b.idx:])
  242. if err != nil {
  243. return err
  244. }
  245. b.idx += n
  246. return UnmarshalMerge(v, m)
  247. }
  248. // consumeGroup parses b until it finds an end group marker, returning
  249. // the raw bytes of the message (excluding the end group marker) and the
  250. // the total length of the message (including the end group marker).
  251. func consumeGroup(b []byte) ([]byte, int, error) {
  252. b0 := b
  253. depth := 1 // assume this follows a start group marker
  254. for {
  255. _, wtyp, tagLen := protowire.ConsumeTag(b)
  256. if tagLen < 0 {
  257. return nil, 0, protowire.ParseError(tagLen)
  258. }
  259. b = b[tagLen:]
  260. var valLen int
  261. switch wtyp {
  262. case protowire.VarintType:
  263. _, valLen = protowire.ConsumeVarint(b)
  264. case protowire.Fixed32Type:
  265. _, valLen = protowire.ConsumeFixed32(b)
  266. case protowire.Fixed64Type:
  267. _, valLen = protowire.ConsumeFixed64(b)
  268. case protowire.BytesType:
  269. _, valLen = protowire.ConsumeBytes(b)
  270. case protowire.StartGroupType:
  271. depth++
  272. case protowire.EndGroupType:
  273. depth--
  274. default:
  275. return nil, 0, errors.New("proto: cannot parse reserved wire type")
  276. }
  277. if valLen < 0 {
  278. return nil, 0, protowire.ParseError(valLen)
  279. }
  280. b = b[valLen:]
  281. if depth == 0 {
  282. return b0[:len(b0)-len(b)-tagLen], len(b0) - len(b), nil
  283. }
  284. }
  285. }