1 files changed, 140 insertions, 0 deletions
diff --git a/encoding/protodelim/protodelim.go b/encoding/protodelim/protodelim.go
new file mode 100644
index 00000000..e2b6cd44
--- /dev/null
+++ b/encoding/protodelim/protodelim.go
@@ -0,0 +1,140 @@
+// Copyright 2022 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package protodelim marshals and unmarshals varint size-delimited messages.
+package protodelim
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+
+	"google.golang.org/protobuf/encoding/protowire"
+	"google.golang.org/protobuf/internal/errors"
+	"google.golang.org/protobuf/proto"
+)
+
+// MarshalOptions is a configurable varint size-delimited marshaler.
+type MarshalOptions struct{ proto.MarshalOptions }
+
+// MarshalTo writes a varint size-delimited wire-format message to w.
+// If w returns an error, MarshalTo returns it unchanged.
+func (o MarshalOptions) MarshalTo(w io.Writer, m proto.Message) (int, error) {
+	msgBytes, err := o.MarshalOptions.Marshal(m)
+	if err != nil {
+		return 0, err
+	}
+
+	sizeBytes := protowire.AppendVarint(nil, uint64(len(msgBytes)))
+	sizeWritten, err := w.Write(sizeBytes)
+	if err != nil {
+		return sizeWritten, err
+	}
+	msgWritten, err := w.Write(msgBytes)
+	if err != nil {
+		return sizeWritten + msgWritten, err
+	}
+	return sizeWritten + msgWritten, nil
+}
+
+// MarshalTo writes a varint size-delimited wire-format message to w
+// with the default options.
+//
+// See the documentation for MarshalOptions.MarshalTo.
+func MarshalTo(w io.Writer, m proto.Message) (int, error) {
+	return MarshalOptions{}.MarshalTo(w, m)
+}
+
+// UnmarshalOptions is a configurable varint size-delimited unmarshaler.
+type UnmarshalOptions struct {
+	proto.UnmarshalOptions
+
+	// MaxSize is the maximum size in wire-format bytes of a single message.
+	// Unmarshaling a message larger than MaxSize will return an error.
+	// A zero MaxSize will default to 4 MiB.
+	// Setting MaxSize to -1 disables the limit.
+	MaxSize int64
+}
+
+const defaultMaxSize = 4 << 20 // 4 MiB, corresponds to the default gRPC max request/response size
+
+// SizeTooLargeError is an error that is returned when the unmarshaler encounters a message size
+// that is larger than its configured MaxSize.
+type SizeTooLargeError struct {
+	// Size is the varint size of the message encountered
+	// that was larger than the provided MaxSize.
+	Size uint64
+
+	// MaxSize is the MaxSize limit configured in UnmarshalOptions, which Size exceeded.
+	MaxSize uint64
+}
+
+func (e *SizeTooLargeError) Error() string {
+	return fmt.Sprintf("message size %d exceeded unmarshaler's maximum configured size %d", e.Size, e.MaxSize)
+}
+
+// Reader is the interface expected by UnmarshalFrom.
+// It is implemented by *bufio.Reader.
+type Reader interface {
+	io.Reader
+	io.ByteReader
+}
+
+// UnmarshalFrom parses and consumes a varint size-delimited wire-format message
+// from r.
+// The provided message must be mutable (e.g., a non-nil pointer to a message).
+//
+// The error is io.EOF error only if no bytes are read.
+// If an EOF happens after reading some but not all the bytes,
+// UnmarshalFrom returns a non-io.EOF error.
+// In particular if r returns a non-io.EOF error, UnmarshalFrom returns it unchanged,
+// and if only a size is read with no subsequent message, io.ErrUnexpectedEOF is returned.
+func (o UnmarshalOptions) UnmarshalFrom(r Reader, m proto.Message) error {
+	var sizeArr [binary.MaxVarintLen64]byte
+	sizeBuf := sizeArr[:0]
+	for i := range sizeArr {
+		b, err := r.ReadByte()
+		if err != nil && (err != io.EOF || i == 0) {
+			return err
+		}
+		sizeBuf = append(sizeBuf, b)
+		if b < 0x80 {
+			break
+		}
+	}
+	size, n := protowire.ConsumeVarint(sizeBuf)
+	if n < 0 {
+		return protowire.ParseError(n)
+	}
+
+	maxSize := o.MaxSize
+	if maxSize == 0 {
+		maxSize = defaultMaxSize
+	}
+	if maxSize != -1 && size > uint64(maxSize) {
+		return errors.Wrap(&SizeTooLargeError{Size: size, MaxSize: uint64(maxSize)}, "")
+	}
+
+	b := make([]byte, size)
+	_, err := io.ReadFull(r, b)
+	if err == io.EOF {
+		return io.ErrUnexpectedEOF
+	}
+	if err != nil {
+		return err
+	}
+	if err := o.Unmarshal(b, m); err != nil {
+		return err
+	}
+	return nil
+}
+
+// UnmarshalFrom parses and consumes a varint size-delimited wire-format message
+// from r with the default options.
+// The provided message must be mutable (e.g., a non-nil pointer to a message).
+//
+// See the documentation for UnmarshalOptions.UnmarshalFrom.
+func UnmarshalFrom(r Reader, m proto.Message) error {
+	return UnmarshalOptions{}.UnmarshalFrom(r, m)
+}