payload.go

package parser

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"strconv"
	"sync"
)

// payloadEncoder is the encoder to encode packets as payload. It can be used in multi-thread.
type PayloadEncoder struct {
	buffers  [][]byte
	locker   sync.Mutex
	isString bool
}

// NewStringPayloadEncoder returns the encoder which encode as string.
func NewStringPayloadEncoder() *PayloadEncoder {
	return &PayloadEncoder{
		isString: true,
	}
}

// NewStringPayloadEncoder returns the encoder which encode as binary.
func NewBinaryPayloadEncoder() *PayloadEncoder {
	return &PayloadEncoder{
		isString: false,
	}
}

type encoder struct {
	*PacketEncoder
	buf          *bytes.Buffer
	binaryPrefix string
	payload      *PayloadEncoder
}

func (e encoder) Close() error {
	if err := e.PacketEncoder.Close(); err != nil {
		return err
	}
	var buffer []byte
	if e.payload.isString {
		buffer = []byte(fmt.Sprintf("%d:%s", e.buf.Len(), e.buf.String()))
	} else {
		buffer = []byte(fmt.Sprintf("%s%d", e.binaryPrefix, e.buf.Len()))
		for i, n := 0, len(buffer); i < n; i++ {
			buffer[i] = buffer[i] - '0'
		}
		buffer = append(buffer, 0xff)
		buffer = append(buffer, e.buf.Bytes()...)
	}

	e.payload.locker.Lock()
	e.payload.buffers = append(e.payload.buffers, buffer)
	e.payload.locker.Unlock()

	return nil
}

// NextString returns the encoder with packet type t and encode as string.
func (e *PayloadEncoder) NextString(t PacketType) (io.WriteCloser, error) {
	buf := bytes.NewBuffer(nil)
	pEncoder, err := NewStringEncoder(buf, t)
	if err != nil {
		return nil, err
	}
	return encoder{
		PacketEncoder: pEncoder,
		buf:           buf,
		binaryPrefix:  "0",
		payload:       e,
	}, nil
}

// NextBinary returns the encoder with packet type t and encode as binary.
func (e *PayloadEncoder) NextBinary(t PacketType) (io.WriteCloser, error) {
	buf := bytes.NewBuffer(nil)
	var pEncoder *PacketEncoder
	var err error
	if e.isString {
		pEncoder, err = NewB64Encoder(buf, t)
	} else {
		pEncoder, err = NewBinaryEncoder(buf, t)
	}
	if err != nil {
		return nil, err
	}
	return encoder{
		PacketEncoder: pEncoder,
		buf:           buf,
		binaryPrefix:  "1",
		payload:       e,
	}, nil
}

// EncodeTo writes encoded payload to writer w. It will clear the buffer of encoder.
func (e *PayloadEncoder) EncodeTo(w io.Writer) error {
	e.locker.Lock()
	buffers := e.buffers
	e.buffers = nil
	e.locker.Unlock()

	for _, b := range buffers {
		for len(b) > 0 {
			n, err := w.Write(b)
			if err != nil {
				return err
			}
			b = b[n:]
		}
	}
	return nil
}

//IsString returns true if payload encode to string, otherwise returns false.
func (e *PayloadEncoder) IsString() bool {
	return e.isString
}

// payloadDecoder is the decoder to decode payload.
type PayloadDecoder struct {
	r *bufio.Reader
}

// NewPaylaodDecoder returns the payload decoder which read from reader r.
func NewPayloadDecoder(r io.Reader) *PayloadDecoder {
	br, ok := r.(*bufio.Reader)
	if !ok {
		br = bufio.NewReader(r)
	}
	return &PayloadDecoder{
		r: br,
	}
}

// Next returns the packet decoder. Make sure it will be closed after used.
func (d *PayloadDecoder) Next() (*PacketDecoder, error) {
	firstByte, err := d.r.Peek(1)
	if err != nil {
		return nil, err
	}
	isBinary := firstByte[0] < '0'
	delim := byte(':')
	if isBinary {
		d.r.ReadByte()
		delim = 0xff
	}
	line, err := d.r.ReadBytes(delim)
	if err != nil {
		return nil, err
	}
	l := len(line)
	if l < 1 {
		return nil, fmt.Errorf("invalid input")
	}
	lenByte := line[:l-1]
	if isBinary {
		for i, n := 0, l; i < n; i++ {
			line[i] = line[i] + '0'
		}
	}
	packetLen, err := strconv.ParseInt(string(lenByte), 10, 64)
	if err != nil {
		return nil, fmt.Errorf("invalid input")
	}
	return NewDecoder(newLimitReader(d.r, int(packetLen)))
}