varint.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134
  1. // Protocol Buffers for Go with Gadgets
  2. //
  3. // Copyright (c) 2013, The GoGo Authors. All rights reserved.
  4. // http://github.com/gogo/protobuf
  5. //
  6. // Redistribution and use in source and binary forms, with or without
  7. // modification, are permitted provided that the following conditions are
  8. // met:
  9. //
  10. // * Redistributions of source code must retain the above copyright
  11. // notice, this list of conditions and the following disclaimer.
  12. // * Redistributions in binary form must reproduce the above
  13. // copyright notice, this list of conditions and the following disclaimer
  14. // in the documentation and/or other materials provided with the
  15. // distribution.
  16. //
  17. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  18. // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  19. // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  20. // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  21. // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  22. // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  23. // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  24. // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  25. // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  26. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  27. // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  28. package io
  29. import (
  30. "bufio"
  31. "encoding/binary"
  32. "errors"
  33. "github.com/gogo/protobuf/proto"
  34. "io"
  35. )
  36. var (
  37. errSmallBuffer = errors.New("Buffer Too Small")
  38. errLargeValue = errors.New("Value is Larger than 64 bits")
  39. )
  40. func NewDelimitedWriter(w io.Writer) WriteCloser {
  41. return &varintWriter{w, make([]byte, 10), nil}
  42. }
  43. type varintWriter struct {
  44. w io.Writer
  45. lenBuf []byte
  46. buffer []byte
  47. }
  48. func (this *varintWriter) WriteMsg(msg proto.Message) (err error) {
  49. var data []byte
  50. if m, ok := msg.(marshaler); ok {
  51. n, ok := getSize(m)
  52. if !ok {
  53. data, err = proto.Marshal(msg)
  54. if err != nil {
  55. return err
  56. }
  57. }
  58. if n >= len(this.buffer) {
  59. this.buffer = make([]byte, n)
  60. }
  61. _, err = m.MarshalTo(this.buffer)
  62. if err != nil {
  63. return err
  64. }
  65. data = this.buffer[:n]
  66. } else {
  67. data, err = proto.Marshal(msg)
  68. if err != nil {
  69. return err
  70. }
  71. }
  72. length := uint64(len(data))
  73. n := binary.PutUvarint(this.lenBuf, length)
  74. _, err = this.w.Write(this.lenBuf[:n])
  75. if err != nil {
  76. return err
  77. }
  78. _, err = this.w.Write(data)
  79. return err
  80. }
  81. func (this *varintWriter) Close() error {
  82. if closer, ok := this.w.(io.Closer); ok {
  83. return closer.Close()
  84. }
  85. return nil
  86. }
  87. func NewDelimitedReader(r io.Reader, maxSize int) ReadCloser {
  88. var closer io.Closer
  89. if c, ok := r.(io.Closer); ok {
  90. closer = c
  91. }
  92. return &varintReader{bufio.NewReader(r), nil, maxSize, closer}
  93. }
  94. type varintReader struct {
  95. r *bufio.Reader
  96. buf []byte
  97. maxSize int
  98. closer io.Closer
  99. }
  100. func (this *varintReader) ReadMsg(msg proto.Message) error {
  101. length64, err := binary.ReadUvarint(this.r)
  102. if err != nil {
  103. return err
  104. }
  105. length := int(length64)
  106. if length < 0 || length > this.maxSize {
  107. return io.ErrShortBuffer
  108. }
  109. if len(this.buf) < length {
  110. this.buf = make([]byte, length)
  111. }
  112. buf := this.buf[:length]
  113. if _, err := io.ReadFull(this.r, buf); err != nil {
  114. return err
  115. }
  116. return proto.Unmarshal(buf, msg)
  117. }
  118. func (this *varintReader) Close() error {
  119. if this.closer != nil {
  120. return this.closer.Close()
  121. }
  122. return nil
  123. }