package protocol import ( "bytes" "encoding/binary" "errors" "io" "sync" lz4 "github.com/bkaradzic/go-lz4" ) const lz4Magic = 0x5e63b278 type lz4Writer struct { wr io.Writer mut sync.Mutex buf []byte } func newLZ4Writer(w io.Writer) *lz4Writer { return &lz4Writer{wr: w} } func (w *lz4Writer) Write(bs []byte) (int, error) { w.mut.Lock() defer w.mut.Unlock() var err error w.buf, err = lz4.Encode(w.buf[:cap(w.buf)], bs) if err != nil { return 0, err } var hdr [8]byte binary.BigEndian.PutUint32(hdr[0:], lz4Magic) binary.BigEndian.PutUint32(hdr[4:], uint32(len(w.buf))) _, err = w.wr.Write(hdr[:]) if err != nil { return 0, err } _, err = w.wr.Write(w.buf) if err != nil { return 0, err } if debug { l.Debugf("lz4 write; %d / %d bytes", len(bs), 8+len(w.buf)) } return len(bs), nil } type lz4Reader struct { rd io.Reader mut sync.Mutex buf []byte ebuf []byte obuf *bytes.Buffer ibytes uint64 obytes uint64 } func newLZ4Reader(r io.Reader) *lz4Reader { return &lz4Reader{rd: r} } func (r *lz4Reader) Read(bs []byte) (int, error) { r.mut.Lock() defer r.mut.Unlock() if r.obuf == nil { r.obuf = bytes.NewBuffer(nil) } if r.obuf.Len() == 0 { if err := r.moreBits(); err != nil { return 0, err } } n, err := r.obuf.Read(bs) if debug { l.Debugf("lz4 read; %d bytes", n) } return n, err } func (r *lz4Reader) moreBits() error { var hdr [8]byte _, err := io.ReadFull(r.rd, hdr[:]) if binary.BigEndian.Uint32(hdr[0:]) != lz4Magic { return errors.New("bad magic") } ln := int(binary.BigEndian.Uint32(hdr[4:])) if len(r.buf) < ln { r.buf = make([]byte, int(ln)) } else { r.buf = r.buf[:ln] } _, err = io.ReadFull(r.rd, r.buf) if err != nil { return err } r.ebuf, err = lz4.Decode(r.ebuf[:cap(r.ebuf)], r.buf) if err != nil { return err } if debug { l.Debugf("lz4 moreBits: %d / %d bytes", ln+8, len(r.ebuf)) } _, err = r.obuf.Write(r.ebuf) return err }