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. 6package protodelim 7 8import ( 9 "encoding/binary" 10 "fmt" 11 "io" 12 13 "google.golang.org/protobuf/encoding/protowire" 14 "google.golang.org/protobuf/internal/errors" 15 "google.golang.org/protobuf/proto" 16) 17 18// MarshalOptions is a configurable varint size-delimited marshaler. 19type 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. 23func (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. 45func 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. 50type 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 60const 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. 64type 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 73func (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. 79type 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. 93func (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. 138func UnmarshalFrom(r Reader, m proto.Message) error { 139 return UnmarshalOptions{}.UnmarshalFrom(r, m) 140} 141