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