protodelim.go

Documentation: google.golang.org/protobuf/encoding/protodelim

     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  	"bufio"
    10  	"encoding/binary"
    11  	"fmt"
    12  	"io"
    13  
    14  	"google.golang.org/protobuf/encoding/protowire"
    15  	"google.golang.org/protobuf/internal/errors"
    16  	"google.golang.org/protobuf/proto"
    17  )
    18  
    19  // MarshalOptions is a configurable varint size-delimited marshaler.
    20  type MarshalOptions struct{ proto.MarshalOptions }
    21  
    22  // MarshalTo writes a varint size-delimited wire-format message to w.
    23  // If w returns an error, MarshalTo returns it unchanged.
    24  func (o MarshalOptions) MarshalTo(w io.Writer, m proto.Message) (int, error) {
    25  	msgBytes, err := o.MarshalOptions.Marshal(m)
    26  	if err != nil {
    27  		return 0, err
    28  	}
    29  
    30  	sizeBytes := protowire.AppendVarint(nil, uint64(len(msgBytes)))
    31  	sizeWritten, err := w.Write(sizeBytes)
    32  	if err != nil {
    33  		return sizeWritten, err
    34  	}
    35  	msgWritten, err := w.Write(msgBytes)
    36  	if err != nil {
    37  		return sizeWritten + msgWritten, err
    38  	}
    39  	return sizeWritten + msgWritten, nil
    40  }
    41  
    42  // MarshalTo writes a varint size-delimited wire-format message to w
    43  // with the default options.
    44  //
    45  // See the documentation for [MarshalOptions.MarshalTo].
    46  func MarshalTo(w io.Writer, m proto.Message) (int, error) {
    47  	return MarshalOptions{}.MarshalTo(w, m)
    48  }
    49  
    50  // UnmarshalOptions is a configurable varint size-delimited unmarshaler.
    51  type UnmarshalOptions struct {
    52  	proto.UnmarshalOptions
    53  
    54  	// MaxSize is the maximum size in wire-format bytes of a single message.
    55  	// Unmarshaling a message larger than MaxSize will return an error.
    56  	// A zero MaxSize will default to 4 MiB.
    57  	// Setting MaxSize to -1 disables the limit.
    58  	MaxSize int64
    59  }
    60  
    61  const defaultMaxSize = 4 << 20 // 4 MiB, corresponds to the default gRPC max request/response size
    62  
    63  // SizeTooLargeError is an error that is returned when the unmarshaler encounters a message size
    64  // that is larger than its configured [UnmarshalOptions.MaxSize].
    65  type SizeTooLargeError struct {
    66  	// Size is the varint size of the message encountered
    67  	// that was larger than the provided MaxSize.
    68  	Size uint64
    69  
    70  	// MaxSize is the MaxSize limit configured in UnmarshalOptions, which Size exceeded.
    71  	MaxSize uint64
    72  }
    73  
    74  func (e *SizeTooLargeError) Error() string {
    75  	return fmt.Sprintf("message size %d exceeded unmarshaler's maximum configured size %d", e.Size, e.MaxSize)
    76  }
    77  
    78  // Reader is the interface expected by [UnmarshalFrom].
    79  // It is implemented by *[bufio.Reader].
    80  type Reader interface {
    81  	io.Reader
    82  	io.ByteReader
    83  }
    84  
    85  // UnmarshalFrom parses and consumes a varint size-delimited wire-format message
    86  // from r.
    87  // The provided message must be mutable (e.g., a non-nil pointer to a message).
    88  //
    89  // The error is [io.EOF] error only if no bytes are read.
    90  // If an EOF happens after reading some but not all the bytes,
    91  // UnmarshalFrom returns a non-io.EOF error.
    92  // In particular if r returns a non-io.EOF error, UnmarshalFrom returns it unchanged,
    93  // and if only a size is read with no subsequent message, [io.ErrUnexpectedEOF] is returned.
    94  func (o UnmarshalOptions) UnmarshalFrom(r Reader, m proto.Message) error {
    95  	var sizeArr [binary.MaxVarintLen64]byte
    96  	sizeBuf := sizeArr[:0]
    97  	for i := range sizeArr {
    98  		b, err := r.ReadByte()
    99  		if err != nil {
   100  			// Immediate EOF is unexpected.
   101  			if err == io.EOF && i != 0 {
   102  				break
   103  			}
   104  			return err
   105  		}
   106  		sizeBuf = append(sizeBuf, b)
   107  		if b < 0x80 {
   108  			break
   109  		}
   110  	}
   111  	size, n := protowire.ConsumeVarint(sizeBuf)
   112  	if n < 0 {
   113  		return protowire.ParseError(n)
   114  	}
   115  
   116  	maxSize := o.MaxSize
   117  	if maxSize == 0 {
   118  		maxSize = defaultMaxSize
   119  	}
   120  	if maxSize != -1 && size > uint64(maxSize) {
   121  		return errors.Wrap(&SizeTooLargeError{Size: size, MaxSize: uint64(maxSize)}, "")
   122  	}
   123  
   124  	var b []byte
   125  	var err error
   126  	if br, ok := r.(*bufio.Reader); ok {
   127  		// Use the []byte from the bufio.Reader instead of having to allocate one.
   128  		// This reduces CPU usage and allocated bytes.
   129  		b, err = br.Peek(int(size))
   130  		if err == nil {
   131  			defer br.Discard(int(size))
   132  		} else {
   133  			b = nil
   134  		}
   135  	}
   136  	if b == nil {
   137  		b = make([]byte, size)
   138  		_, err = io.ReadFull(r, b)
   139  	}
   140  
   141  	if err == io.EOF {
   142  		return io.ErrUnexpectedEOF
   143  	}
   144  	if err != nil {
   145  		return err
   146  	}
   147  	if err := o.Unmarshal(b, m); err != nil {
   148  		return err
   149  	}
   150  	return nil
   151  }
   152  
   153  // UnmarshalFrom parses and consumes a varint size-delimited wire-format message
   154  // from r with the default options.
   155  // The provided message must be mutable (e.g., a non-nil pointer to a message).
   156  //
   157  // See the documentation for [UnmarshalOptions.UnmarshalFrom].
   158  func UnmarshalFrom(r Reader, m proto.Message) error {
   159  	return UnmarshalOptions{}.UnmarshalFrom(r, m)
   160  }
   161
View as plain text