serialization.go

package proto

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"time"
)

const (
	maxParseArrayLen = 256
)

var ErrNotEnoughData = errors.New("not enough data")
var ErrInvalidArrayLen = errors.New("invalid array length")

type decoder struct {
	buf []byte
	r   io.Reader
	err error
}

func NewDecoder(r io.Reader) *decoder {
	return &decoder{
		r:   r,
		buf: make([]byte, 1024),
	}
}

func (d *decoder) SetReader(r io.Reader) {
	d.r = r
}

func (d *decoder) DecodeInt8() int8 {
	if d.err != nil {
		return 0
	}
	b := d.buf[:1]
	n, err := io.ReadFull(d.r, b)
	if err != nil {
		d.err = err
		return 0
	}
	if n != 1 {
		d.err = ErrNotEnoughData
		return 0
	}
	return int8(b[0])
}

func (d *decoder) DecodeInt16() int16 {
	if d.err != nil {
		return 0
	}
	b := d.buf[:2]
	n, err := io.ReadFull(d.r, b)
	if err != nil {
		d.err = err
		return 0
	}
	if n != 2 {
		d.err = ErrNotEnoughData
		return 0
	}
	return int16(binary.BigEndian.Uint16(b))
}

func (d *decoder) DecodeInt32() int32 {
	if d.err != nil {
		return 0
	}
	b := d.buf[:4]
	n, err := io.ReadFull(d.r, b)
	if err != nil {
		d.err = err
		return 0
	}
	if n != 4 {
		d.err = ErrNotEnoughData
		return 0
	}
	return int32(binary.BigEndian.Uint32(b))
}

func (d *decoder) DecodeUint32() uint32 {
	if d.err != nil {
		return 0
	}
	b := d.buf[:4]
	n, err := io.ReadFull(d.r, b)
	if err != nil {
		d.err = err
		return 0
	}
	if n != 4 {
		d.err = ErrNotEnoughData
		return 0
	}
	return binary.BigEndian.Uint32(b)
}

func (d *decoder) DecodeInt64() int64 {
	if d.err != nil {
		return 0
	}
	b := d.buf[:8]
	n, err := io.ReadFull(d.r, b)
	if err != nil {
		d.err = err
		return 0
	}
	if n != 8 {
		d.err = ErrNotEnoughData
		return 0
	}
	return int64(binary.BigEndian.Uint64(b))
}

func (d *decoder) DecodeDuration32() time.Duration {
	return time.Duration(d.DecodeInt32()) * time.Millisecond
}

func (d *decoder) DecodeString() string {
	if d.err != nil {
		return ""
	}
	slen := d.DecodeInt16()
	if d.err != nil {
		return ""
	}
	if slen < 1 {
		return ""
	}

	var b []byte
	if int(slen) > len(d.buf) {
		var err error
		b, err = allocParseBuf(int(slen))
		if err != nil {
			d.err = err
			return ""
		}
	} else {
		b = d.buf[:int(slen)]
	}
	n, err := io.ReadFull(d.r, b)
	if err != nil {
		d.err = err
		return ""
	}
	if n != int(slen) {
		d.err = ErrNotEnoughData
		return ""
	}
	return string(b)
}

func (d *decoder) DecodeArrayLen() (int, error) {
	len := int(d.DecodeInt32())

	// Sometime kafka may send -1 as size of array.
	if len == -1 {
		return 0, nil
	}

	if len > maxParseBufSize {
		return 0, ErrInvalidArrayLen
	}

	return len, nil
}

func (d *decoder) DecodeBytes() []byte {
	if d.err != nil {
		return nil
	}
	slen := d.DecodeInt32()
	if d.err != nil {
		return nil
	}
	if slen < 1 {
		return nil
	}

	b, err := allocParseBuf(int(slen))
	if err != nil {
		d.err = err
		return nil
	}
	n, err := io.ReadFull(d.r, b)
	if err != nil {
		d.err = err
		return nil
	}
	if n != int(slen) {
		d.err = ErrNotEnoughData
		return nil
	}
	return b
}

func (d *decoder) DecodeVarInt() int64 {
	// err already stored by ReadByte():
	res, _ := binary.ReadVarint(d)
	return res
}

// ReadByte implements ByteReader
func (d *decoder) ReadByte() (byte, error) {
	_, err := io.ReadFull(d.r, d.buf[:1])
	if err != nil {
		d.err = err
	}
	return d.buf[0], err
}

func (d *decoder) DecodeVarBytes() []byte {
	slen := d.DecodeVarInt()

	if slen < 1 {
		return nil
	}

	b, err := allocParseBuf(int(slen))
	if err != nil {
		d.err = err
		return nil
	}
	n, err := io.ReadFull(d.r, b)
	if err != nil {
		d.err = err
		return nil
	}
	if n != int(slen) {
		d.err = ErrNotEnoughData
		return nil
	}
	return b
}

func (d *decoder) DecodeVarString() string {
	if d.err != nil {
		return ""
	}
	slen := d.DecodeVarInt()
	if d.err != nil {
		return ""
	}
	if slen < 1 {
		return ""
	}

	var b []byte
	if int(slen) > len(d.buf) {
		var err error
		b, err = allocParseBuf(int(slen))
		if err != nil {
			d.err = err
			return ""
		}
	} else {
		b = d.buf[:int(slen)]
	}
	n, err := io.ReadFull(d.r, b)
	if err != nil {
		d.err = err
		return ""
	}
	if n != int(slen) {
		d.err = ErrNotEnoughData
		return ""
	}
	return string(b)
}

func (d *decoder) Err() error {
	return d.err
}

type encoder struct {
	w   io.Writer
	err error
	buf [8]byte
}

func NewEncoder(w io.Writer) *encoder {
	return &encoder{w: w}
}

func (e *encoder) EncodeDuration(val time.Duration) {
	if e.err != nil {
		return
	}

	intVal := uint32(val / time.Millisecond)
	b := e.buf[:4]
	binary.BigEndian.PutUint32(b, intVal)
	_, e.err = e.w.Write(b)
}

func (e *encoder) EncodeInt8(val int8) {
	if e.err != nil {
		return
	}

	b := e.buf[:1]
	b[0] = byte(val)
	_, e.err = e.w.Write(b)
}

func (e *encoder) EncodeInt16(val int16) {
	if e.err != nil {
		return
	}

	b := e.buf[:2]
	binary.BigEndian.PutUint16(b, uint16(val))
	e.err = writeAll(e.w, b)
}

func (e *encoder) EncodeInt32(val int32) {
	if e.err != nil {
		return
	}

	b := e.buf[:4]
	binary.BigEndian.PutUint32(b, uint32(val))
	e.err = writeAll(e.w, b)
}

func (e *encoder) EncodeInt32s(val []int32) {
	if e.err != nil {
		return
	}

	e.EncodeArrayLen(len(val))
	for _, v := range val {
		e.EncodeInt32(v)
	}
}

func (e *encoder) EncodeInt64(val int64) {
	if e.err != nil {
		return
	}

	b := e.buf[:8]
	binary.BigEndian.PutUint64(b, uint64(val))
	e.err = writeAll(e.w, b)
}

func (e *encoder) EncodeInt64s(val []int64) {
	if e.err != nil {
		return
	}

	e.EncodeArrayLen(len(val))
	for _, v := range val {
		e.EncodeInt64(v)
	}
}

func (e *encoder) EncodeUint32(val uint32) {
	if e.err != nil {
		return
	}

	b := e.buf[:4]
	binary.BigEndian.PutUint32(b, val)
	e.err = writeAll(e.w, b)
}

func (e *encoder) EncodeBytes(val []byte) {
	if e.err != nil {
		return
	}

	buf := e.buf[:4]

	if val == nil {
		no := int32(-1)
		binary.BigEndian.PutUint32(buf, uint32(no))
		e.err = writeAll(e.w, buf)
		return
	}

	binary.BigEndian.PutUint32(buf, uint32(len(val)))
	e.err = writeAll(e.w, buf)
	if e.err == nil {
		e.err = writeAll(e.w, val)
	}
}

func (e *encoder) EncodeString(val string) {
	if e.err != nil {
		return
	}

	buf := e.buf[:2]

	binary.BigEndian.PutUint16(buf, uint16(len(val)))
	e.err = writeAll(e.w, buf)
	if e.err == nil {
		e.err = writeAll(e.w, []byte(val))
	}
}

func (e *encoder) EncodeError(err error) {
	b := e.buf[:2]

	if err == nil {
		binary.BigEndian.PutUint16(b, uint16(0))
		e.err = writeAll(e.w, b)
		return
	}
	kerr, ok := err.(*KafkaError)
	if !ok {
		e.err = fmt.Errorf("cannot encode error of type %T", err)
	}

	binary.BigEndian.PutUint16(b, uint16(kerr.errno))
	e.err = writeAll(e.w, b)
}

func (e *encoder) EncodeArrayLen(length int) {
	e.EncodeInt32(int32(length))
}

func (e *encoder) Err() error {
	return e.err
}

func writeAll(w io.Writer, b []byte) error {
	n, err := w.Write(b)
	if err != nil {
		return err
	}
	if n != len(b) {
		return fmt.Errorf("cannot write %d: %d written", len(b), n)
	}
	return nil
}