github.com/lazyledger/lazyledger-core@v0.35.0-dev.0.20210613111200-4c651f053571/libs/protoio/reader.go (about)

     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  //
    29  // Modified from original GoGo Protobuf to not buffer the reader.
    30  
    31  package protoio
    32  
    33  import (
    34  	"bytes"
    35  	"encoding/binary"
    36  	"fmt"
    37  	"io"
    38  
    39  	"github.com/gogo/protobuf/proto"
    40  )
    41  
    42  // NewDelimitedReader reads varint-delimited Protobuf messages from a reader. Unlike the gogoproto
    43  // NewDelimitedReader, this does not buffer the reader, which may cause poor performance but is
    44  // necessary when only reading single messages (e.g. in the p2p package).
    45  func NewDelimitedReader(r io.Reader, maxSize int) ReadCloser {
    46  	var closer io.Closer
    47  	if c, ok := r.(io.Closer); ok {
    48  		closer = c
    49  	}
    50  	return &varintReader{newByteReader(r), nil, maxSize, closer}
    51  }
    52  
    53  type varintReader struct {
    54  	r       *byteReader
    55  	buf     []byte
    56  	maxSize int
    57  	closer  io.Closer
    58  }
    59  
    60  func (r *varintReader) ReadMsg(msg proto.Message) error {
    61  	length64, err := binary.ReadUvarint(newByteReader(r.r))
    62  	if err != nil {
    63  		return err
    64  	}
    65  	length := int(length64)
    66  	if length < 0 || length > r.maxSize {
    67  		return fmt.Errorf("message exceeds max size (%v > %v)", length, r.maxSize)
    68  	}
    69  	if len(r.buf) < length {
    70  		r.buf = make([]byte, length)
    71  	}
    72  	buf := r.buf[:length]
    73  	if _, err := io.ReadFull(r.r, buf); err != nil {
    74  		return err
    75  	}
    76  	return proto.Unmarshal(buf, msg)
    77  }
    78  
    79  func (r *varintReader) Close() error {
    80  	if r.closer != nil {
    81  		return r.closer.Close()
    82  	}
    83  	return nil
    84  }
    85  
    86  func UnmarshalDelimited(data []byte, msg proto.Message) error {
    87  	return NewDelimitedReader(bytes.NewReader(data), len(data)).ReadMsg(msg)
    88  }