65 lines
1.5 KiB
Go
65 lines
1.5 KiB
Go
|
package channel
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bufio"
|
||
|
|
"bytes"
|
||
|
|
"encoding/binary"
|
||
|
|
"io"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Varint is a framing that transmits and receives messages on r and wc, with
|
||
|
|
// each message prefixed by its length encoded in a varint as defined by the
|
||
|
|
// encoding/binary package.
|
||
|
|
func Varint(r io.Reader, wc io.WriteCloser) Channel {
|
||
|
|
return &varint{
|
||
|
|
wc: wc,
|
||
|
|
rd: bufio.NewReader(r),
|
||
|
|
buf: bytes.NewBuffer(nil),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// A varint implements Channel for varint-prefixed messages.
|
||
|
|
type varint struct {
|
||
|
|
wc io.WriteCloser
|
||
|
|
rd *bufio.Reader
|
||
|
|
buf *bytes.Buffer
|
||
|
|
}
|
||
|
|
|
||
|
|
// Send implements part of the Channel interface. It encodes len(msg) as a
|
||
|
|
// varint, concatenates it with the message body, and writes the framed message
|
||
|
|
// to the underlying writer.
|
||
|
|
func (v *varint) Send(msg []byte) error {
|
||
|
|
var ln [binary.MaxVarintLen64]byte
|
||
|
|
nb := binary.PutUvarint(ln[:], uint64(len(msg)))
|
||
|
|
|
||
|
|
v.buf.Reset()
|
||
|
|
v.buf.Grow(nb + len(msg))
|
||
|
|
v.buf.Write(ln[:nb])
|
||
|
|
v.buf.Write(msg)
|
||
|
|
_, err := v.wc.Write(v.buf.Next(v.buf.Len()))
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
// Recv implements part of the Channel interface. It decodes a varint message
|
||
|
|
// length then reads that many bytes from the underlying reader.
|
||
|
|
func (v *varint) Recv() ([]byte, error) {
|
||
|
|
ln, err := v.decode()
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
out := make([]byte, ln)
|
||
|
|
nr, err := io.ReadFull(v.rd, out)
|
||
|
|
return out[:nr], err
|
||
|
|
}
|
||
|
|
|
||
|
|
// Close implements part of the Channel interface.
|
||
|
|
func (v *varint) Close() error { return v.wc.Close() }
|
||
|
|
|
||
|
|
func (v *varint) decode() (int, error) {
|
||
|
|
ln, err := binary.ReadUvarint(v.rd)
|
||
|
|
if err != nil {
|
||
|
|
return 0, err
|
||
|
|
}
|
||
|
|
return int(ln), nil
|
||
|
|
}
|