live-gateway/douyu/stt/stream.go

package stt

import (
	"io"
)

// Stream is an io.Writer like object, with STT specific write functions.
// Error is not returned as return value, but stored as Error member on this stream instance.
type Stream struct {
	cfg        *frozenConfig
	out        io.Writer
	depth      int
	buf        []byte
	Error      error
	Attachment interface{} // open for customized encoder
}

// NewStream create new stream instance.
// cfg can be stt.ConfigDefault.
// out can be nil if write to internal buffer.
// bufSize is the initial size for the internal buffer in bytes.
func NewStream(cfg API, out io.Writer, bufSize int) *Stream {
	return &Stream{
		cfg:   cfg.(*frozenConfig),
		out:   out,
		depth: 0,
		buf:   make([]byte, 0, bufSize),
		Error: nil,
	}
}

// Pool returns a pool can provide more stream with same configuration
func (stream *Stream) Pool() StreamPool {
	return stream.cfg
}

// Reset reuse this stream instance by assign a new writer
func (stream *Stream) Reset(out io.Writer) {
	stream.depth = 0
	stream.out = out
	stream.buf = stream.buf[:0]
}

// Available returns how many bytes are unused in the buffer.
func (stream *Stream) Available() int {
	return cap(stream.buf) - len(stream.buf)
}

// Buffered returns the number of bytes that have been written into the current buffer.
func (stream *Stream) Buffered() int {
	return len(stream.buf)
}

// Buffer if writer is nil, use this method to take the result
func (stream *Stream) Buffer() []byte {
	return stream.buf
}

// SetBuffer allows to append to the internal buffer directly
func (stream *Stream) SetBuffer(buf []byte) {
	stream.buf = buf
}

// Write writes the contents of p into the buffer.
// It returns the number of bytes written.
// If nn < len(p), it also returns an error explaining
// why write is short.
func (stream *Stream) Write(p []byte) (nn int, err error) {
	stream.buf = append(stream.buf, p...)
	if stream.out != nil {
		nn, err = stream.out.Write(stream.buf)
		stream.buf = stream.buf[nn:]
		return
	}
	return len(p), nil
}

// WriteByte writes a single byte.
func (stream *Stream) writeBytes(c ...byte) {
	stream.buf = append(stream.buf, c...)
}

// Flush writes any buffered data to the underlying io.Writer.
func (stream *Stream) Flush() error {
	if stream.out == nil {
		return nil
	}
	if stream.Error != nil {
		return stream.Error
	}
	_, err := stream.out.Write(stream.buf)
	if err != nil {
		if stream.Error == nil {
			stream.Error = err
		}
		return err
	}
	stream.buf = stream.buf[:0]
	return nil
}

// WriteRaw write string out without quotes, just like []byte
func (stream *Stream) WriteRaw(s string) {
	stream.buf = append(stream.buf, s...)
}

// WriteNil write null to stream
func (stream *Stream) WriteNil() {
	stream.writeBytes('n', 'u', 'l', 'l')
}

// WriteTrue write true to stream
func (stream *Stream) WriteTrue() {
	stream.writeBytes('t', 'r', 'u', 'e')
}

// WriteFalse write false to stream
func (stream *Stream) WriteFalse() {
	stream.writeBytes('f', 'a', 'l', 's', 'e')
}

// WriteBool write true or false into stream
func (stream *Stream) WriteBool(val bool) {
	if val {
		stream.WriteTrue()
	} else {
		stream.WriteFalse()
	}
}

// WriteObjectField write field@= with possible indention
func (stream *Stream) WriteObjectField(field string) {
	stream.WriteString(field)
	// TODO 尚不清楚斗鱼STT的转义规则，大体来看是2层深度转义2次，1层深度作为operator是不转义的?
	bb := []byte("@=")
	if stream.depth > 1 {
		bb = encryptBytes(bb, stream.depth)
	}
	stream.writeBytes(bb...)
}

// WriteObjectEnd write / with possible indention
func (stream *Stream) WriteObjectEnd() {
	if stream.depth > 1 {
		stream.writeBytes(encryptBytes([]byte{'/'}, 1)...)
	} else {
		stream.writeBytes('/')
	}
}

// WriteEmptyObject write /
func (stream *Stream) WriteEmptyObject() {
	//stream.writeBytes('/')
}

// WriteMore write / with possible indention
func (stream *Stream) WriteMore() {
	if stream.depth > 1 {
		stream.writeBytes(encryptBytes([]byte{'/'}, 1)...)
	} else {
		stream.writeBytes('/')
	}
}

func (stream *Stream) incrementDepth() (success bool) {
	stream.depth++
	if stream.depth <= maxDepth {
		return true
	}
	return false
}

func (stream *Stream) decrementDepth() (success bool) {
	stream.depth--
	if stream.depth >= 0 {
		return true
	}
	return false
}