decode.go 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081
  1. // Copyright 2013 Matt T. Proud
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package pbutil
  15. import (
  16. "encoding/binary"
  17. "errors"
  18. "io"
  19. "google.golang.org/protobuf/proto"
  20. )
  21. // TODO: Give error package name prefix in next minor release.
  22. var errInvalidVarint = errors.New("invalid varint32 encountered")
  23. // ReadDelimited decodes a message from the provided length-delimited stream,
  24. // where the length is encoded as 32-bit varint prefix to the message body.
  25. // It returns the total number of bytes read and any applicable error. This is
  26. // roughly equivalent to the companion Java API's
  27. // MessageLite#parseDelimitedFrom. As per the reader contract, this function
  28. // calls r.Read repeatedly as required until exactly one message including its
  29. // prefix is read and decoded (or an error has occurred). The function never
  30. // reads more bytes from the stream than required. The function never returns
  31. // an error if a message has been read and decoded correctly, even if the end
  32. // of the stream has been reached in doing so. In that case, any subsequent
  33. // calls return (0, io.EOF).
  34. func ReadDelimited(r io.Reader, m proto.Message) (n int, err error) {
  35. // TODO: Consider allowing the caller to specify a decode buffer in the
  36. // next major version.
  37. // TODO: Consider using error wrapping to annotate error state in pass-
  38. // through cases in the next minor version.
  39. // Per AbstractParser#parsePartialDelimitedFrom with
  40. // CodedInputStream#readRawVarint32.
  41. var headerBuf [binary.MaxVarintLen32]byte
  42. var bytesRead, varIntBytes int
  43. var messageLength uint64
  44. for varIntBytes == 0 { // i.e. no varint has been decoded yet.
  45. if bytesRead >= len(headerBuf) {
  46. return bytesRead, errInvalidVarint
  47. }
  48. // We have to read byte by byte here to avoid reading more bytes
  49. // than required. Each read byte is appended to what we have
  50. // read before.
  51. newBytesRead, err := r.Read(headerBuf[bytesRead : bytesRead+1])
  52. if newBytesRead == 0 {
  53. if err != nil {
  54. return bytesRead, err
  55. }
  56. // A Reader should not return (0, nil); but if it does, it should
  57. // be treated as no-op according to the Reader contract.
  58. continue
  59. }
  60. bytesRead += newBytesRead
  61. // Now present everything read so far to the varint decoder and
  62. // see if a varint can be decoded already.
  63. messageLength, varIntBytes = binary.Uvarint(headerBuf[:bytesRead])
  64. }
  65. messageBuf := make([]byte, messageLength)
  66. newBytesRead, err := io.ReadFull(r, messageBuf)
  67. bytesRead += newBytesRead
  68. if err != nil {
  69. return bytesRead, err
  70. }
  71. return bytesRead, proto.Unmarshal(messageBuf, m)
  72. }