github.com/pachyderm/pachyderm@v1.13.4/src/client/pkg/pbutil/pbutil.go (about)

     1  package pbutil
     2  
     3  import (
     4  	"encoding/binary"
     5  	"io"
     6  	"unsafe"
     7  
     8  	"github.com/gogo/protobuf/proto"
     9  
    10  	"github.com/pachyderm/pachyderm/src/client/pkg/errors"
    11  )
    12  
    13  // Reader is io.Reader for proto.Message instead of []byte.
    14  type Reader interface {
    15  	Read(val proto.Message) error
    16  	ReadBytes() ([]byte, error)
    17  }
    18  
    19  // Writer is io.Writer for proto.Message instead of []byte.
    20  type Writer interface {
    21  	Write(val proto.Message) (int64, error)
    22  	WriteBytes([]byte) (int64, error)
    23  }
    24  
    25  // ReadWriter is io.ReadWriter for proto.Message instead of []byte.
    26  type ReadWriter interface {
    27  	Reader
    28  	Writer
    29  }
    30  
    31  type readWriter struct {
    32  	w   io.Writer
    33  	r   io.Reader
    34  	buf []byte
    35  }
    36  
    37  func (r *readWriter) ReadBytes() ([]byte, error) {
    38  	var l int64
    39  	if err := binary.Read(r.r, binary.LittleEndian, &l); err != nil {
    40  		return nil, errors.EnsureStack(err)
    41  	}
    42  	if r.buf == nil || len(r.buf) < int(l) {
    43  		r.buf = make([]byte, l)
    44  	}
    45  	buf := r.buf[0:l]
    46  	if _, err := io.ReadFull(r.r, buf); err != nil {
    47  		if errors.Is(err, io.EOF) {
    48  			return nil, io.ErrUnexpectedEOF
    49  		}
    50  		return nil, errors.EnsureStack(err)
    51  	}
    52  	return buf, nil
    53  }
    54  
    55  // Read reads val from r.
    56  func (r *readWriter) Read(val proto.Message) error {
    57  	buf, err := r.ReadBytes()
    58  	if err != nil {
    59  		return errors.EnsureStack(err)
    60  	}
    61  	return proto.Unmarshal(buf, val)
    62  }
    63  
    64  func (r *readWriter) WriteBytes(bytes []byte) (int64, error) {
    65  	if err := binary.Write(r.w, binary.LittleEndian, int64(len(bytes))); err != nil {
    66  		return 0, errors.EnsureStack(err)
    67  	}
    68  	lenByteSize := unsafe.Sizeof(int64(len(bytes)))
    69  	n, err := r.w.Write(bytes)
    70  	return int64(lenByteSize) + int64(n), errors.EnsureStack(err)
    71  }
    72  
    73  // Write writes val to r.
    74  func (r *readWriter) Write(val proto.Message) (int64, error) {
    75  	bytes, err := proto.Marshal(val)
    76  	if err != nil {
    77  		return 0, errors.EnsureStack(err)
    78  	}
    79  	return r.WriteBytes(bytes)
    80  }
    81  
    82  // NewReader returns a new Reader with r as its source.
    83  func NewReader(r io.Reader) Reader {
    84  	return &readWriter{r: r}
    85  }
    86  
    87  // NewWriter returns a new Writer with w as its sink.
    88  func NewWriter(w io.Writer) Writer {
    89  	return &readWriter{w: w}
    90  }
    91  
    92  // NewReadWriter returns a new ReadWriter with rw as both its source and its sink.
    93  func NewReadWriter(rw io.ReadWriter) ReadWriter {
    94  	return &readWriter{r: rw, w: rw}
    95  }