uint32.go 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126
  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. "encoding/binary"
  31. "github.com/gogo/protobuf/proto"
  32. "io"
  33. )
  34. func NewUint32DelimitedWriter(w io.Writer, byteOrder binary.ByteOrder) WriteCloser {
  35. return &uint32Writer{w, byteOrder, nil}
  36. }
  37. func NewSizeUint32DelimitedWriter(w io.Writer, byteOrder binary.ByteOrder, size int) WriteCloser {
  38. return &uint32Writer{w, byteOrder, make([]byte, size)}
  39. }
  40. type uint32Writer struct {
  41. w io.Writer
  42. byteOrder binary.ByteOrder
  43. buffer []byte
  44. }
  45. func (this *uint32Writer) WriteMsg(msg proto.Message) (err error) {
  46. var data []byte
  47. if m, ok := msg.(marshaler); ok {
  48. n, ok := getSize(m)
  49. if !ok {
  50. data, err = proto.Marshal(msg)
  51. if err != nil {
  52. return err
  53. }
  54. }
  55. if n >= len(this.buffer) {
  56. this.buffer = make([]byte, n)
  57. }
  58. _, err = m.MarshalTo(this.buffer)
  59. if err != nil {
  60. return err
  61. }
  62. data = this.buffer[:n]
  63. } else {
  64. data, err = proto.Marshal(msg)
  65. if err != nil {
  66. return err
  67. }
  68. }
  69. length := uint32(len(data))
  70. if err = binary.Write(this.w, this.byteOrder, &length); err != nil {
  71. return err
  72. }
  73. _, err = this.w.Write(data)
  74. return err
  75. }
  76. func (this *uint32Writer) Close() error {
  77. if closer, ok := this.w.(io.Closer); ok {
  78. return closer.Close()
  79. }
  80. return nil
  81. }
  82. type uint32Reader struct {
  83. r io.Reader
  84. byteOrder binary.ByteOrder
  85. lenBuf []byte
  86. buf []byte
  87. maxSize int
  88. }
  89. func NewUint32DelimitedReader(r io.Reader, byteOrder binary.ByteOrder, maxSize int) ReadCloser {
  90. return &uint32Reader{r, byteOrder, make([]byte, 4), nil, maxSize}
  91. }
  92. func (this *uint32Reader) ReadMsg(msg proto.Message) error {
  93. if _, err := io.ReadFull(this.r, this.lenBuf); err != nil {
  94. return err
  95. }
  96. length32 := this.byteOrder.Uint32(this.lenBuf)
  97. length := int(length32)
  98. if length < 0 || length > this.maxSize {
  99. return io.ErrShortBuffer
  100. }
  101. if length >= len(this.buf) {
  102. this.buf = make([]byte, length)
  103. }
  104. _, err := io.ReadFull(this.r, this.buf[:length])
  105. if err != nil {
  106. return err
  107. }
  108. return proto.Unmarshal(this.buf[:length], msg)
  109. }
  110. func (this *uint32Reader) Close() error {
  111. if closer, ok := this.r.(io.Closer); ok {
  112. return closer.Close()
  113. }
  114. return nil
  115. }