github.com/whiteCcinn/protobuf-go@v1.0.9/encoding/protodelim/protodelim.go (about)

     1  // Copyright 2022 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package protodelim marshals and unmarshals varint size-delimited messages.
     6  package protodelim
     7  
     8  import (
     9  	"encoding/binary"
    10  	"fmt"
    11  	"io"
    12  
    13  	"github.com/whiteCcinn/protobuf-go/encoding/protowire"
    14  	"github.com/whiteCcinn/protobuf-go/internal/errors"
    15  	"github.com/whiteCcinn/protobuf-go/proto"
    16  )
    17  
    18  // MarshalOptions is a configurable varint size-delimited marshaler.
    19  type MarshalOptions struct{ proto.MarshalOptions }
    20  
    21  // MarshalTo writes a varint size-delimited wire-format message to w.
    22  // If w returns an error, MarshalTo returns it unchanged.
    23  func (o MarshalOptions) MarshalTo(w io.Writer, m proto.Message) (int, error) {
    24  	msgBytes, err := o.MarshalOptions.Marshal(m)
    25  	if err != nil {
    26  		return 0, err
    27  	}
    28  
    29  	sizeBytes := protowire.AppendVarint(nil, uint64(len(msgBytes)))
    30  	sizeWritten, err := w.Write(sizeBytes)
    31  	if err != nil {
    32  		return sizeWritten, err
    33  	}
    34  	msgWritten, err := w.Write(msgBytes)
    35  	if err != nil {
    36  		return sizeWritten + msgWritten, err
    37  	}
    38  	return sizeWritten + msgWritten, nil
    39  }
    40  
    41  // MarshalTo writes a varint size-delimited wire-format message to w
    42  // with the default options.
    43  //
    44  // See the documentation for MarshalOptions.MarshalTo.
    45  func MarshalTo(w io.Writer, m proto.Message) (int, error) {
    46  	return MarshalOptions{}.MarshalTo(w, m)
    47  }
    48  
    49  // UnmarshalOptions is a configurable varint size-delimited unmarshaler.
    50  type UnmarshalOptions struct {
    51  	proto.UnmarshalOptions
    52  
    53  	// MaxSize is the maximum size in wire-format bytes of a single message.
    54  	// Unmarshaling a message larger than MaxSize will return an error.
    55  	// A zero MaxSize will default to 4 MiB.
    56  	// Setting MaxSize to -1 disables the limit.
    57  	MaxSize int64
    58  }
    59  
    60  const defaultMaxSize = 4 << 20 // 4 MiB, corresponds to the default gRPC max request/response size
    61  
    62  // SizeTooLargeError is an error that is returned when the unmarshaler encounters a message size
    63  // that is larger than its configured MaxSize.
    64  type SizeTooLargeError struct {
    65  	// Size is the varint size of the message encountered
    66  	// that was larger than the provided MaxSize.
    67  	Size uint64
    68  
    69  	// MaxSize is the MaxSize limit configured in UnmarshalOptions, which Size exceeded.
    70  	MaxSize uint64
    71  }
    72  
    73  func (e *SizeTooLargeError) Error() string {
    74  	return fmt.Sprintf("message size %d exceeded unmarshaler's maximum configured size %d", e.Size, e.MaxSize)
    75  }
    76  
    77  // Reader is the interface expected by UnmarshalFrom.
    78  // It is implemented by *bufio.Reader.
    79  type Reader interface {
    80  	io.Reader
    81  	io.ByteReader
    82  }
    83  
    84  // UnmarshalFrom parses and consumes a varint size-delimited wire-format message
    85  // from r.
    86  // The provided message must be mutable (e.g., a non-nil pointer to a message).
    87  //
    88  // The error is io.EOF error only if no bytes are read.
    89  // If an EOF happens after reading some but not all the bytes,
    90  // UnmarshalFrom returns a non-io.EOF error.
    91  // In particular if r returns a non-io.EOF error, UnmarshalFrom returns it unchanged,
    92  // and if only a size is read with no subsequent message, io.ErrUnexpectedEOF is returned.
    93  func (o UnmarshalOptions) UnmarshalFrom(r Reader, m proto.Message) error {
    94  	var sizeArr [binary.MaxVarintLen64]byte
    95  	sizeBuf := sizeArr[:0]
    96  	for i := range sizeArr {
    97  		b, err := r.ReadByte()
    98  		if err != nil && (err != io.EOF || i == 0) {
    99  			return err
   100  		}
   101  		sizeBuf = append(sizeBuf, b)
   102  		if b < 0x80 {
   103  			break
   104  		}
   105  	}
   106  	size, n := protowire.ConsumeVarint(sizeBuf)
   107  	if n < 0 {
   108  		return protowire.ParseError(n)
   109  	}
   110  
   111  	maxSize := o.MaxSize
   112  	if maxSize == 0 {
   113  		maxSize = defaultMaxSize
   114  	}
   115  	if maxSize != -1 && size > uint64(maxSize) {
   116  		return errors.Wrap(&SizeTooLargeError{Size: size, MaxSize: uint64(maxSize)}, "")
   117  	}
   118  
   119  	b := make([]byte, size)
   120  	_, err := io.ReadFull(r, b)
   121  	if err == io.EOF {
   122  		return io.ErrUnexpectedEOF
   123  	}
   124  	if err != nil {
   125  		return err
   126  	}
   127  	if err := o.Unmarshal(b, m); err != nil {
   128  		return err
   129  	}
   130  	return nil
   131  }
   132  
   133  // UnmarshalFrom parses and consumes a varint size-delimited wire-format message
   134  // from r with the default options.
   135  // The provided message must be mutable (e.g., a non-nil pointer to a message).
   136  //
   137  // See the documentation for UnmarshalOptions.UnmarshalFrom.
   138  func UnmarshalFrom(r Reader, m proto.Message) error {
   139  	return UnmarshalOptions{}.UnmarshalFrom(r, m)
   140  }