ceremonyclient/go-libp2p/p2p/protocol/circuitv2/util/io.go

67 lines
1.7 KiB
Go
Raw Normal View History

2023-08-21 03:50:38 +00:00
package util
import (
"errors"
"io"
pool "github.com/libp2p/go-buffer-pool"
"github.com/libp2p/go-msgio/pbio"
"github.com/multiformats/go-varint"
"google.golang.org/protobuf/proto"
)
type DelimitedReader struct {
r io.Reader
buf []byte
}
// The gogo protobuf NewDelimitedReader is buffered, which may eat up stream data.
// So we need to implement a compatible delimited reader that reads unbuffered.
// There is a slowdown from unbuffered reading: when reading the message
// it can take multiple single byte Reads to read the length and another Read
// to read the message payload.
// However, this is not critical performance degradation as
// - the reader is utilized to read one (dialer, stop) or two messages (hop) during
// the handshake, so it's a drop in the water for the connection lifetime.
// - messages are small (max 4k) and the length fits in a couple of bytes,
// so overall we have at most three reads per message.
func NewDelimitedReader(r io.Reader, maxSize int) *DelimitedReader {
return &DelimitedReader{r: r, buf: pool.Get(maxSize)}
}
func (d *DelimitedReader) Close() {
if d.buf != nil {
pool.Put(d.buf)
d.buf = nil
}
}
func (d *DelimitedReader) ReadByte() (byte, error) {
buf := d.buf[:1]
_, err := d.r.Read(buf)
return buf[0], err
}
func (d *DelimitedReader) ReadMsg(msg proto.Message) error {
mlen, err := varint.ReadUvarint(d)
if err != nil {
return err
}
if uint64(len(d.buf)) < mlen {
return errors.New("message too large")
}
buf := d.buf[:mlen]
_, err = io.ReadFull(d.r, buf)
if err != nil {
return err
}
return proto.Unmarshal(buf, msg)
}
func NewDelimitedWriter(w io.Writer) pbio.WriteCloser {
return pbio.NewDelimitedWriter(w)
}