github.com/kamalshkeir/kencoding@v0.0.2-0.20230409043843-44b609a0475a/thrift/binary.go

package thrift

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"math"
)

// BinaryProtocol is a Protocol implementation for the binary thrift protocol.
//
// https://github.com/apache/thrift/blob/master/doc/specs/thrift-binary-protocol.md
type BinaryProtocol struct {
	NonStrict bool
}

func (p *BinaryProtocol) NewReader(r io.Reader) Reader {
	return &binaryReader{p: p, r: r}
}

func (p *BinaryProtocol) NewWriter(w io.Writer) Writer {
	return &binaryWriter{p: p, w: w}
}

func (p *BinaryProtocol) Features() Features {
	return 0
}

type binaryReader struct {
	p *BinaryProtocol
	r io.Reader
	b [8]byte
}

func (r *binaryReader) Protocol() Protocol {
	return r.p
}

func (r *binaryReader) Reader() io.Reader {
	return r.r
}

func (r *binaryReader) ReadBool() (bool, error) {
	v, err := r.ReadByte()
	return v != 0, err
}

func (r *binaryReader) ReadInt8() (int8, error) {
	b, err := r.ReadByte()
	return int8(b), err
}

func (r *binaryReader) ReadInt16() (int16, error) {
	b, err := r.read(2)
	if len(b) < 2 {
		return 0, err
	}
	return int16(binary.BigEndian.Uint16(b)), nil
}

func (r *binaryReader) ReadInt32() (int32, error) {
	b, err := r.read(4)
	if len(b) < 4 {
		return 0, err
	}
	return int32(binary.BigEndian.Uint32(b)), nil
}

func (r *binaryReader) ReadInt64() (int64, error) {
	b, err := r.read(8)
	if len(b) < 8 {
		return 0, err
	}
	return int64(binary.BigEndian.Uint64(b)), nil
}

func (r *binaryReader) ReadFloat64() (float64, error) {
	b, err := r.read(8)
	if len(b) < 8 {
		return 0, err
	}
	return math.Float64frombits(binary.BigEndian.Uint64(b)), nil
}

func (r *binaryReader) ReadBytes() ([]byte, error) {
	n, err := r.ReadLength()
	if err != nil {
		return nil, err
	}
	b := make([]byte, n)
	_, err = io.ReadFull(r.r, b)
	return b, err
}

func (r *binaryReader) ReadString() (string, error) {
	b, err := r.ReadBytes()
	return unsafeBytesToString(b), err
}

func (r *binaryReader) ReadLength() (int, error) {
	b, err := r.read(4)
	if len(b) < 4 {
		return 0, err
	}
	n := binary.BigEndian.Uint32(b)
	if n > math.MaxInt32 {
		return 0, fmt.Errorf("length out of range: %d", n)
	}
	return int(n), nil
}

func (r *binaryReader) ReadMessage() (Message, error) {
	m := Message{}

	b, err := r.read(4)
	if len(b) < 4 {
		return m, err
	}

	if (b[0] >> 7) == 0 { // non-strict
		n := int(binary.BigEndian.Uint32(b))
		s := make([]byte, n)
		_, err := io.ReadFull(r.r, s)
		if err != nil {
			return m, dontExpectEOF(err)
		}
		m.Name = unsafeBytesToString(s)

		t, err := r.ReadInt8()
		if err != nil {
			return m, dontExpectEOF(err)
		}

		m.Type = MessageType(t & 0x7)
	} else {
		m.Type = MessageType(b[3] & 0x7)

		if m.Name, err = r.ReadString(); err != nil {
			return m, dontExpectEOF(err)
		}
	}

	m.SeqID, err = r.ReadInt32()
	return m, err
}

func (r *binaryReader) ReadField() (Field, error) {
	t, err := r.ReadInt8()
	if err != nil {
		return Field{}, err
	}
	i, err := r.ReadInt16()
	if err != nil {
		return Field{}, err
	}
	return Field{ID: i, Type: Type(t)}, nil
}

func (r *binaryReader) ReadList() (List, error) {
	t, err := r.ReadInt8()
	if err != nil {
		return List{}, err
	}
	n, err := r.ReadInt32()
	if err != nil {
		return List{}, dontExpectEOF(err)
	}
	return List{Size: n, Type: Type(t)}, nil
}

func (r *binaryReader) ReadSet() (Set, error) {
	l, err := r.ReadList()
	return Set(l), err
}

func (r *binaryReader) ReadMap() (Map, error) {
	k, err := r.ReadByte()
	if err != nil {
		return Map{}, err
	}
	v, err := r.ReadByte()
	if err != nil {
		return Map{}, dontExpectEOF(err)
	}
	n, err := r.ReadInt32()
	if err != nil {
		return Map{}, dontExpectEOF(err)
	}
	return Map{Size: n, Key: Type(k), Value: Type(v)}, nil
}

func (r *binaryReader) ReadByte() (byte, error) {
	switch x := r.r.(type) {
	case *bytes.Buffer:
		return x.ReadByte()
	case *bytes.Reader:
		return x.ReadByte()
	case *bufio.Reader:
		return x.ReadByte()
	case io.ByteReader:
		return x.ReadByte()
	default:
		b, err := r.read(1)
		if err != nil {
			return 0, err
		}
		return b[0], nil
	}
}

func (r *binaryReader) read(n int) ([]byte, error) {
	_, err := io.ReadFull(r.r, r.b[:n])
	return r.b[:n], err
}

type binaryWriter struct {
	p *BinaryProtocol
	b [8]byte
	w io.Writer
}

func (w *binaryWriter) Protocol() Protocol {
	return w.p
}

func (w *binaryWriter) Writer() io.Writer {
	return w.w
}

func (w *binaryWriter) WriteBool(v bool) error {
	var b byte
	if v {
		b = 1
	}
	return w.writeByte(b)
}

func (w *binaryWriter) WriteInt8(v int8) error {
	return w.writeByte(byte(v))
}

func (w *binaryWriter) WriteInt16(v int16) error {
	binary.BigEndian.PutUint16(w.b[:2], uint16(v))
	return w.write(w.b[:2])
}

func (w *binaryWriter) WriteInt32(v int32) error {
	binary.BigEndian.PutUint32(w.b[:4], uint32(v))
	return w.write(w.b[:4])
}

func (w *binaryWriter) WriteInt64(v int64) error {
	binary.BigEndian.PutUint64(w.b[:8], uint64(v))
	return w.write(w.b[:8])
}

func (w *binaryWriter) WriteFloat64(v float64) error {
	binary.BigEndian.PutUint64(w.b[:8], math.Float64bits(v))
	return w.write(w.b[:8])
}

func (w *binaryWriter) WriteBytes(v []byte) error {
	if err := w.WriteLength(len(v)); err != nil {
		return err
	}
	return w.write(v)
}

func (w *binaryWriter) WriteString(v string) error {
	if err := w.WriteLength(len(v)); err != nil {
		return err
	}
	return w.writeString(v)
}

func (w *binaryWriter) WriteLength(n int) error {
	if n < 0 {
		return fmt.Errorf("negative length cannot be encoded in thrift: %d", n)
	}
	if n > math.MaxInt32 {
		return fmt.Errorf("length is too large to be encoded in thrift: %d", n)
	}
	return w.WriteInt32(int32(n))
}

func (w *binaryWriter) WriteMessage(m Message) error {
	if w.p.NonStrict {
		if err := w.WriteString(m.Name); err != nil {
			return err
		}
		if err := w.writeByte(byte(m.Type)); err != nil {
			return err
		}
	} else {
		w.b[0] = 1 << 7
		w.b[1] = 0
		w.b[2] = 0
		w.b[3] = byte(m.Type) & 0x7
		binary.BigEndian.PutUint32(w.b[4:], uint32(len(m.Name)))

		if err := w.write(w.b[:8]); err != nil {
			return err
		}
		if err := w.writeString(m.Name); err != nil {
			return err
		}
	}
	return w.WriteInt32(m.SeqID)
}

func (w *binaryWriter) WriteField(f Field) error {
	if err := w.writeByte(byte(f.Type)); err != nil {
		return err
	}
	return w.WriteInt16(f.ID)
}

func (w *binaryWriter) WriteList(l List) error {
	if err := w.writeByte(byte(l.Type)); err != nil {
		return err
	}
	return w.WriteInt32(l.Size)
}

func (w *binaryWriter) WriteSet(s Set) error {
	return w.WriteList(List(s))
}

func (w *binaryWriter) WriteMap(m Map) error {
	if err := w.writeByte(byte(m.Key)); err != nil {
		return err
	}
	if err := w.writeByte(byte(m.Value)); err != nil {
		return err
	}
	return w.WriteInt32(m.Size)
}

func (w *binaryWriter) write(b []byte) error {
	_, err := w.w.Write(b)
	return err
}

func (w *binaryWriter) writeString(s string) error {
	_, err := io.WriteString(w.w, s)
	return err
}

func (w *binaryWriter) writeByte(b byte) error {
	// The special cases are intended to reduce the runtime overheadof testing
	// for the io.ByteWriter interface for common types. Type assertions on a
	// concrete type is just a pointer comparison, instead of requiring a
	// complex lookup in the type metadata.
	switch x := w.w.(type) {
	case *bytes.Buffer:
		return x.WriteByte(b)
	case *bufio.Writer:
		return x.WriteByte(b)
	case io.ByteWriter:
		return x.WriteByte(b)
	default:
		w.b[0] = b
		return w.write(w.b[:1])
	}
}