lazy.go 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359
  1. // Copyright 2024 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 protolazy contains internal data structures for lazy message decoding.
  5. package protolazy
  6. import (
  7. "fmt"
  8. "sort"
  9. "google.golang.org/protobuf/encoding/protowire"
  10. piface "google.golang.org/protobuf/runtime/protoiface"
  11. )
  12. // IndexEntry is the structure for an index of the fields in a message of a
  13. // proto (not descending to sub-messages)
  14. type IndexEntry struct {
  15. FieldNum uint32
  16. // first byte of this tag/field
  17. Start uint32
  18. // first byte after a contiguous sequence of bytes for this tag/field, which could
  19. // include a single encoding of the field, or multiple encodings for the field
  20. End uint32
  21. // True if this protobuf segment includes multiple encodings of the field
  22. MultipleContiguous bool
  23. }
  24. // XXX_lazyUnmarshalInfo has information about a particular lazily decoded message
  25. //
  26. // Deprecated: Do not use. This will be deleted in the near future.
  27. type XXX_lazyUnmarshalInfo struct {
  28. // Index of fields and their positions in the protobuf for this
  29. // message. Make index be a pointer to a slice so it can be updated
  30. // atomically. The index pointer is only set once (lazily when/if
  31. // the index is first needed), and must always be SET and LOADED
  32. // ATOMICALLY.
  33. index *[]IndexEntry
  34. // The protobuf associated with this lazily decoded message. It is
  35. // only set during proto.Unmarshal(). It doesn't need to be set and
  36. // loaded atomically, since any simultaneous set (Unmarshal) and read
  37. // (during a get) would already be a race in the app code.
  38. Protobuf []byte
  39. // The flags present when Unmarshal was originally called for this particular message
  40. unmarshalFlags piface.UnmarshalInputFlags
  41. }
  42. // The Buffer and SetBuffer methods let v2/internal/impl interact with
  43. // XXX_lazyUnmarshalInfo via an interface, to avoid an import cycle.
  44. // Buffer returns the lazy unmarshal buffer.
  45. //
  46. // Deprecated: Do not use. This will be deleted in the near future.
  47. func (lazy *XXX_lazyUnmarshalInfo) Buffer() []byte {
  48. return lazy.Protobuf
  49. }
  50. // SetBuffer sets the lazy unmarshal buffer.
  51. //
  52. // Deprecated: Do not use. This will be deleted in the near future.
  53. func (lazy *XXX_lazyUnmarshalInfo) SetBuffer(b []byte) {
  54. lazy.Protobuf = b
  55. }
  56. // SetUnmarshalFlags is called to set a copy of the original unmarshalInputFlags.
  57. // The flags should reflect how Unmarshal was called.
  58. func (lazy *XXX_lazyUnmarshalInfo) SetUnmarshalFlags(f piface.UnmarshalInputFlags) {
  59. lazy.unmarshalFlags = f
  60. }
  61. // UnmarshalFlags returns the original unmarshalInputFlags.
  62. func (lazy *XXX_lazyUnmarshalInfo) UnmarshalFlags() piface.UnmarshalInputFlags {
  63. return lazy.unmarshalFlags
  64. }
  65. // AllowedPartial returns true if the user originally unmarshalled this message with
  66. // AllowPartial set to true
  67. func (lazy *XXX_lazyUnmarshalInfo) AllowedPartial() bool {
  68. return (lazy.unmarshalFlags & piface.UnmarshalCheckRequired) == 0
  69. }
  70. func protoFieldNumber(tag uint32) uint32 {
  71. return tag >> 3
  72. }
  73. // buildIndex builds an index of the specified protobuf, return the index
  74. // array and an error.
  75. func buildIndex(buf []byte) ([]IndexEntry, error) {
  76. index := make([]IndexEntry, 0, 16)
  77. var lastProtoFieldNum uint32
  78. var outOfOrder bool
  79. var r BufferReader = NewBufferReader(buf)
  80. for !r.Done() {
  81. var tag uint32
  82. var err error
  83. var curPos = r.Pos
  84. // INLINED: tag, err = r.DecodeVarint32()
  85. {
  86. i := r.Pos
  87. buf := r.Buf
  88. if i >= len(buf) {
  89. return nil, errOutOfBounds
  90. } else if buf[i] < 0x80 {
  91. r.Pos++
  92. tag = uint32(buf[i])
  93. } else if r.Remaining() < 5 {
  94. var v uint64
  95. v, err = r.DecodeVarintSlow()
  96. tag = uint32(v)
  97. } else {
  98. var v uint32
  99. // we already checked the first byte
  100. tag = uint32(buf[i]) & 127
  101. i++
  102. v = uint32(buf[i])
  103. i++
  104. tag |= (v & 127) << 7
  105. if v < 128 {
  106. goto done
  107. }
  108. v = uint32(buf[i])
  109. i++
  110. tag |= (v & 127) << 14
  111. if v < 128 {
  112. goto done
  113. }
  114. v = uint32(buf[i])
  115. i++
  116. tag |= (v & 127) << 21
  117. if v < 128 {
  118. goto done
  119. }
  120. v = uint32(buf[i])
  121. i++
  122. tag |= (v & 127) << 28
  123. if v < 128 {
  124. goto done
  125. }
  126. return nil, errOutOfBounds
  127. done:
  128. r.Pos = i
  129. }
  130. }
  131. // DONE: tag, err = r.DecodeVarint32()
  132. fieldNum := protoFieldNumber(tag)
  133. if fieldNum < lastProtoFieldNum {
  134. outOfOrder = true
  135. }
  136. // Skip the current value -- will skip over an entire group as well.
  137. // INLINED: err = r.SkipValue(tag)
  138. wireType := tag & 0x7
  139. switch protowire.Type(wireType) {
  140. case protowire.VarintType:
  141. // INLINED: err = r.SkipVarint()
  142. i := r.Pos
  143. if len(r.Buf)-i < 10 {
  144. // Use DecodeVarintSlow() to skip while
  145. // checking for buffer overflow, but ignore result
  146. _, err = r.DecodeVarintSlow()
  147. goto out2
  148. }
  149. if r.Buf[i] < 0x80 {
  150. goto out
  151. }
  152. i++
  153. if r.Buf[i] < 0x80 {
  154. goto out
  155. }
  156. i++
  157. if r.Buf[i] < 0x80 {
  158. goto out
  159. }
  160. i++
  161. if r.Buf[i] < 0x80 {
  162. goto out
  163. }
  164. i++
  165. if r.Buf[i] < 0x80 {
  166. goto out
  167. }
  168. i++
  169. if r.Buf[i] < 0x80 {
  170. goto out
  171. }
  172. i++
  173. if r.Buf[i] < 0x80 {
  174. goto out
  175. }
  176. i++
  177. if r.Buf[i] < 0x80 {
  178. goto out
  179. }
  180. i++
  181. if r.Buf[i] < 0x80 {
  182. goto out
  183. }
  184. i++
  185. if r.Buf[i] < 0x80 {
  186. goto out
  187. }
  188. return nil, errOverflow
  189. out:
  190. r.Pos = i + 1
  191. // DONE: err = r.SkipVarint()
  192. case protowire.Fixed64Type:
  193. err = r.SkipFixed64()
  194. case protowire.BytesType:
  195. var n uint32
  196. n, err = r.DecodeVarint32()
  197. if err == nil {
  198. err = r.Skip(int(n))
  199. }
  200. case protowire.StartGroupType:
  201. err = r.SkipGroup(tag)
  202. case protowire.Fixed32Type:
  203. err = r.SkipFixed32()
  204. default:
  205. err = fmt.Errorf("Unexpected wire type (%d)", wireType)
  206. }
  207. // DONE: err = r.SkipValue(tag)
  208. out2:
  209. if err != nil {
  210. return nil, err
  211. }
  212. if fieldNum != lastProtoFieldNum {
  213. index = append(index, IndexEntry{FieldNum: fieldNum,
  214. Start: uint32(curPos),
  215. End: uint32(r.Pos)},
  216. )
  217. } else {
  218. index[len(index)-1].End = uint32(r.Pos)
  219. index[len(index)-1].MultipleContiguous = true
  220. }
  221. lastProtoFieldNum = fieldNum
  222. }
  223. if outOfOrder {
  224. sort.Slice(index, func(i, j int) bool {
  225. return index[i].FieldNum < index[j].FieldNum ||
  226. (index[i].FieldNum == index[j].FieldNum &&
  227. index[i].Start < index[j].Start)
  228. })
  229. }
  230. return index, nil
  231. }
  232. func (lazy *XXX_lazyUnmarshalInfo) SizeField(num uint32) (size int) {
  233. start, end, found, _, multipleEntries := lazy.FindFieldInProto(num)
  234. if multipleEntries != nil {
  235. for _, entry := range multipleEntries {
  236. size += int(entry.End - entry.Start)
  237. }
  238. return size
  239. }
  240. if !found {
  241. return 0
  242. }
  243. return int(end - start)
  244. }
  245. func (lazy *XXX_lazyUnmarshalInfo) AppendField(b []byte, num uint32) ([]byte, bool) {
  246. start, end, found, _, multipleEntries := lazy.FindFieldInProto(num)
  247. if multipleEntries != nil {
  248. for _, entry := range multipleEntries {
  249. b = append(b, lazy.Protobuf[entry.Start:entry.End]...)
  250. }
  251. return b, true
  252. }
  253. if !found {
  254. return nil, false
  255. }
  256. b = append(b, lazy.Protobuf[start:end]...)
  257. return b, true
  258. }
  259. func (lazy *XXX_lazyUnmarshalInfo) SetIndex(index []IndexEntry) {
  260. atomicStoreIndex(&lazy.index, &index)
  261. }
  262. // FindFieldInProto looks for field fieldNum in lazyUnmarshalInfo information
  263. // (including protobuf), returns startOffset/endOffset/found.
  264. func (lazy *XXX_lazyUnmarshalInfo) FindFieldInProto(fieldNum uint32) (start, end uint32, found, multipleContiguous bool, multipleEntries []IndexEntry) {
  265. if lazy.Protobuf == nil {
  266. // There is no backing protobuf for this message -- it was made from a builder
  267. return 0, 0, false, false, nil
  268. }
  269. index := atomicLoadIndex(&lazy.index)
  270. if index == nil {
  271. r, err := buildIndex(lazy.Protobuf)
  272. if err != nil {
  273. panic(fmt.Sprintf("findFieldInfo: error building index when looking for field %d: %v", fieldNum, err))
  274. }
  275. // lazy.index is a pointer to the slice returned by BuildIndex
  276. index = &r
  277. atomicStoreIndex(&lazy.index, index)
  278. }
  279. return lookupField(index, fieldNum)
  280. }
  281. // lookupField returns the offset at which the indicated field starts using
  282. // the index, offset immediately after field ends (including all instances of
  283. // a repeated field), and bools indicating if field was found and if there
  284. // are multiple encodings of the field in the byte range.
  285. //
  286. // To hande the uncommon case where there are repeated encodings for the same
  287. // field which are not consecutive in the protobuf (so we need to returns
  288. // multiple start/end offsets), we also return a slice multipleEntries. If
  289. // multipleEntries is non-nil, then multiple entries were found, and the
  290. // values in the slice should be used, rather than start/end/found.
  291. func lookupField(indexp *[]IndexEntry, fieldNum uint32) (start, end uint32, found bool, multipleContiguous bool, multipleEntries []IndexEntry) {
  292. // The pointer indexp to the index was already loaded atomically.
  293. // The slice is uniquely associated with the pointer, so it doesn't
  294. // need to be loaded atomically.
  295. index := *indexp
  296. for i, entry := range index {
  297. if fieldNum == entry.FieldNum {
  298. if i < len(index)-1 && entry.FieldNum == index[i+1].FieldNum {
  299. // Handle the uncommon case where there are
  300. // repeated entries for the same field which
  301. // are not contiguous in the protobuf.
  302. multiple := make([]IndexEntry, 1, 2)
  303. multiple[0] = IndexEntry{fieldNum, entry.Start, entry.End, entry.MultipleContiguous}
  304. i++
  305. for i < len(index) && index[i].FieldNum == fieldNum {
  306. multiple = append(multiple, IndexEntry{fieldNum, index[i].Start, index[i].End, index[i].MultipleContiguous})
  307. i++
  308. }
  309. return 0, 0, false, false, multiple
  310. }
  311. return entry.Start, entry.End, true, entry.MultipleContiguous, nil
  312. }
  313. if fieldNum < entry.FieldNum {
  314. return 0, 0, false, false, nil
  315. }
  316. }
  317. return 0, 0, false, false, nil
  318. }