huya-go/internal/taf/writer.go

package taf

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

// Writer 负责按 TAF/JCE 格式编码字段。
type Writer struct {
	buf bytes.Buffer
}

func NewWriter() *Writer {
	return &Writer{}
}

func (w *Writer) Bytes() []byte {
	return w.buf.Bytes()
}

func (w *Writer) WriteHead(tag int, typ Type) error {
	if tag < 0 || tag > 255 {
		return fmt.Errorf("taf: tag 超出范围: %d", tag)
	}
	if tag < 15 {
		return w.buf.WriteByte(byte(tag<<4) | byte(typ))
	}
	if err := w.buf.WriteByte(0xf0 | byte(typ)); err != nil {
		return err
	}
	return w.buf.WriteByte(byte(tag))
}

func (w *Writer) WriteInt8(tag int, value int64) error {
	if value == 0 {
		return w.WriteHead(tag, TypeZero)
	}
	if value < math.MinInt8 || value > math.MaxInt8 {
		return fmt.Errorf("taf: int8 溢出: %d", value)
	}
	if err := w.WriteHead(tag, TypeInt8); err != nil {
		return err
	}
	return w.buf.WriteByte(byte(int8(value)))
}

func (w *Writer) WriteInt16(tag int, value int64) error {
	if value >= math.MinInt8 && value <= math.MaxInt8 {
		return w.WriteInt8(tag, value)
	}
	if value < math.MinInt16 || value > math.MaxInt16 {
		return fmt.Errorf("taf: int16 溢出: %d", value)
	}
	if err := w.WriteHead(tag, TypeInt16); err != nil {
		return err
	}
	var b [2]byte
	binary.BigEndian.PutUint16(b[:], uint16(int16(value)))
	_, err := w.buf.Write(b[:])
	return err
}

func (w *Writer) WriteInt32(tag int, value int64) error {
	if value >= math.MinInt16 && value <= math.MaxInt16 {
		return w.WriteInt16(tag, value)
	}
	if value < math.MinInt32 || value > math.MaxInt32 {
		return fmt.Errorf("taf: int32 溢出: %d", value)
	}
	if err := w.WriteHead(tag, TypeInt32); err != nil {
		return err
	}
	var b [4]byte
	binary.BigEndian.PutUint32(b[:], uint32(int32(value)))
	_, err := w.buf.Write(b[:])
	return err
}

func (w *Writer) WriteInt64(tag int, value int64) error {
	if value >= math.MinInt32 && value <= math.MaxInt32 {
		return w.WriteInt32(tag, value)
	}
	if err := w.WriteHead(tag, TypeInt64); err != nil {
		return err
	}
	var b [8]byte
	binary.BigEndian.PutUint64(b[:], uint64(value))
	_, err := w.buf.Write(b[:])
	return err
}

func (w *Writer) WriteBool(tag int, value bool) error {
	if value {
		return w.WriteInt8(tag, 1)
	}
	return w.WriteInt8(tag, 0)
}

func (w *Writer) WriteFloat(tag int, value float32) error {
	if err := w.WriteHead(tag, TypeFloat); err != nil {
		return err
	}
	var b [4]byte
	binary.BigEndian.PutUint32(b[:], math.Float32bits(value))
	_, err := w.buf.Write(b[:])
	return err
}

func (w *Writer) WriteDouble(tag int, value float64) error {
	if err := w.WriteHead(tag, TypeDouble); err != nil {
		return err
	}
	var b [8]byte
	binary.BigEndian.PutUint64(b[:], math.Float64bits(value))
	_, err := w.buf.Write(b[:])
	return err
}

func (w *Writer) WriteString(tag int, value string) error {
	raw := []byte(value)
	if len(raw) > math.MaxUint8 {
		if err := w.WriteHead(tag, TypeString4); err != nil {
			return err
		}
		var b [4]byte
		binary.BigEndian.PutUint32(b[:], uint32(len(raw)))
		if _, err := w.buf.Write(b[:]); err != nil {
			return err
		}
		_, err := w.buf.Write(raw)
		return err
	}
	if err := w.WriteHead(tag, TypeString1); err != nil {
		return err
	}
	if err := w.buf.WriteByte(byte(len(raw))); err != nil {
		return err
	}
	_, err := w.buf.Write(raw)
	return err
}

func (w *Writer) WriteBytes(tag int, value []byte) error {
	if err := w.WriteHead(tag, TypeSimpleList); err != nil {
		return err
	}
	if err := w.WriteHead(0, TypeInt8); err != nil {
		return err
	}
	if err := w.WriteInt32(0, int64(len(value))); err != nil {
		return err
	}
	_, err := w.buf.Write(value)
	return err
}

func (w *Writer) WriteStructBegin(tag int) error {
	return w.WriteHead(tag, TypeStructBegin)
}

func (w *Writer) WriteStructEnd() error {
	return w.WriteHead(0, TypeStructEnd)
}

func (w *Writer) WriteStruct(tag int, value Marshaler) error {
	if err := w.WriteStructBegin(tag); err != nil {
		return err
	}
	if err := value.WriteTo(w); err != nil {
		return err
	}
	return w.WriteStructEnd()
}

func (w *Writer) WriteMapStringString(tag int, value map[string]string) error {
	if value == nil {
		value = map[string]string{}
	}
	if err := w.WriteHead(tag, TypeMap); err != nil {
		return err
	}
	if err := w.WriteInt32(0, int64(len(value))); err != nil {
		return err
	}
	for k, v := range value {
		if err := w.WriteString(0, k); err != nil {
			return err
		}
		if err := w.WriteString(1, v); err != nil {
			return err
		}
	}
	return nil
}

func (w *Writer) WriteMapStringBytes(tag int, value map[string][]byte) error {
	if value == nil {
		value = map[string][]byte{}
	}
	if err := w.WriteHead(tag, TypeMap); err != nil {
		return err
	}
	if err := w.WriteInt32(0, int64(len(value))); err != nil {
		return err
	}
	for k, v := range value {
		if err := w.WriteString(0, k); err != nil {
			return err
		}
		if err := w.WriteBytes(1, v); err != nil {
			return err
		}
	}
	return nil
}

func (w *Writer) WriteListStruct(tag int, values []Marshaler) error {
	if err := w.WriteHead(tag, TypeList); err != nil {
		return err
	}
	if err := w.WriteInt32(0, int64(len(values))); err != nil {
		return err
	}
	for _, value := range values {
		if err := w.WriteStruct(0, value); err != nil {
			return err
		}
	}
	return nil
}