huya-go/internal/taf/reader.go

package taf

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

// Reader 负责从 TAF/JCE 字节流中按 tag 顺序读取字段。
type Reader struct {
	buf *bytes.Reader
}

func NewReader(data []byte) *Reader {
	return &Reader{buf: bytes.NewReader(data)}
}

func (r *Reader) Position() int64 {
	pos, _ := r.buf.Seek(0, io.SeekCurrent)
	return pos
}

func (r *Reader) Seek(pos int64) error {
	_, err := r.buf.Seek(pos, io.SeekStart)
	return err
}

func (r *Reader) ReadHead() (int, Type, error) {
	b, err := r.buf.ReadByte()
	if err != nil {
		return 0, 0, err
	}
	tag := int((b >> 4) & 0x0f)
	typ := Type(b & 0x0f)
	if tag == 15 {
		ext, err := r.buf.ReadByte()
		if err != nil {
			return 0, 0, err
		}
		tag = int(ext)
	}
	return tag, typ, nil
}

func (r *Reader) PeekHead() (int, Type, error) {
	pos := r.Position()
	tag, typ, err := r.ReadHead()
	_ = r.Seek(pos)
	return tag, typ, err
}

func (r *Reader) ReadIntValue(typ Type) (int64, error) {
	switch typ {
	case TypeZero:
		return 0, nil
	case TypeInt8:
		b, err := r.buf.ReadByte()
		return int64(int8(b)), err
	case TypeInt16:
		var b [2]byte
		if _, err := io.ReadFull(r.buf, b[:]); err != nil {
			return 0, err
		}
		return int64(int16(binary.BigEndian.Uint16(b[:]))), nil
	case TypeInt32:
		var b [4]byte
		if _, err := io.ReadFull(r.buf, b[:]); err != nil {
			return 0, err
		}
		return int64(int32(binary.BigEndian.Uint32(b[:]))), nil
	case TypeInt64:
		var b [8]byte
		if _, err := io.ReadFull(r.buf, b[:]); err != nil {
			return 0, err
		}
		return int64(binary.BigEndian.Uint64(b[:])), nil
	default:
		return 0, fmt.Errorf("taf: 期望整数类型，实际 0x%02x", byte(typ))
	}
}

func (r *Reader) ReadStringValue(typ Type) (string, error) {
	var n uint32
	switch typ {
	case TypeString1:
		b, err := r.buf.ReadByte()
		if err != nil {
			return "", err
		}
		n = uint32(b)
	case TypeString4:
		var b [4]byte
		if _, err := io.ReadFull(r.buf, b[:]); err != nil {
			return "", err
		}
		n = binary.BigEndian.Uint32(b[:])
	default:
		return "", fmt.Errorf("taf: 期望字符串类型，实际 0x%02x", byte(typ))
	}
	raw := make([]byte, n)
	if _, err := io.ReadFull(r.buf, raw); err != nil {
		return "", err
	}
	return string(raw), nil
}

func (r *Reader) ReadBytesValue(typ Type) ([]byte, error) {
	if typ != TypeSimpleList {
		return nil, fmt.Errorf("taf: 期望 bytes，实际 0x%02x", byte(typ))
	}
	_, itemType, err := r.ReadHead()
	if err != nil {
		return nil, err
	}
	if itemType != TypeInt8 {
		return nil, fmt.Errorf("taf: bytes 元素类型不是 int8: 0x%02x", byte(itemType))
	}
	n, err := r.readLength()
	if err != nil {
		return nil, err
	}
	if n < 0 || n > math.MaxInt32 {
		return nil, fmt.Errorf("taf: bytes 长度异常: %d", n)
	}
	raw := make([]byte, int(n))
	_, err = io.ReadFull(r.buf, raw)
	return raw, err
}

func (r *Reader) Skip(typ Type) error {
	switch typ {
	case TypeZero:
		return nil
	case TypeInt8:
		_, err := r.buf.ReadByte()
		return err
	case TypeInt16:
		_, err := r.buf.Seek(2, io.SeekCurrent)
		return err
	case TypeInt32, TypeFloat:
		_, err := r.buf.Seek(4, io.SeekCurrent)
		return err
	case TypeInt64, TypeDouble:
		_, err := r.buf.Seek(8, io.SeekCurrent)
		return err
	case TypeString1:
		b, err := r.buf.ReadByte()
		if err != nil {
			return err
		}
		_, err = r.buf.Seek(int64(b), io.SeekCurrent)
		return err
	case TypeString4:
		var b [4]byte
		if _, err := io.ReadFull(r.buf, b[:]); err != nil {
			return err
		}
		_, err := r.buf.Seek(int64(binary.BigEndian.Uint32(b[:])), io.SeekCurrent)
		return err
	case TypeMap:
		n, err := r.readLength()
		if err != nil {
			return err
		}
		for i := int64(0); i < n; i++ {
			if _, kt, err := r.ReadHead(); err != nil {
				return err
			} else if err := r.Skip(kt); err != nil {
				return err
			}
			if _, vt, err := r.ReadHead(); err != nil {
				return err
			} else if err := r.Skip(vt); err != nil {
				return err
			}
		}
		return nil
	case TypeList:
		n, err := r.readLength()
		if err != nil {
			return err
		}
		for i := int64(0); i < n; i++ {
			if _, it, err := r.ReadHead(); err != nil {
				return err
			} else if err := r.Skip(it); err != nil {
				return err
			}
		}
		return nil
	case TypeStructBegin:
		return r.SkipToStructEnd()
	case TypeStructEnd:
		return nil
	case TypeSimpleList:
		_, itemType, err := r.ReadHead()
		if err != nil {
			return err
		}
		if itemType != TypeInt8 {
			return fmt.Errorf("taf: simple list 元素类型异常: 0x%02x", byte(itemType))
		}
		n, err := r.readLength()
		if err != nil {
			return err
		}
		_, err = r.buf.Seek(n, io.SeekCurrent)
		return err
	default:
		return fmt.Errorf("taf: 未知字段类型 0x%02x", byte(typ))
	}
}

func (r *Reader) readLength() (int64, error) {
	_, typ, err := r.ReadHead()
	if err != nil {
		return 0, err
	}
	return r.ReadIntValue(typ)
}

func (r *Reader) findTag(target int, required bool) (Type, bool, error) {
	for {
		pos := r.Position()
		tag, typ, err := r.ReadHead()
		if err != nil {
			if errors.Is(err, io.EOF) && !required {
				return 0, false, nil
			}
			return 0, false, err
		}
		if typ == TypeStructEnd {
			_ = r.Seek(pos)
			if required {
				return 0, false, fmt.Errorf("taf: 未找到 tag=%d，已到结构尾", target)
			}
			return 0, false, nil
		}
		if tag == target {
			return typ, true, nil
		}
		if tag > target {
			_ = r.Seek(pos)
			if required {
				return 0, false, fmt.Errorf("taf: 未找到 tag=%d，遇到 tag=%d", target, tag)
			}
			return 0, false, nil
		}
		if err := r.Skip(typ); err != nil {
			return 0, false, err
		}
	}
}

func (r *Reader) ReadInt64(tag int, required bool, def int64) (int64, error) {
	typ, ok, err := r.findTag(tag, required)
	if err != nil || !ok {
		return def, err
	}
	return r.ReadIntValue(typ)
}

func (r *Reader) ReadInt32(tag int, required bool, def int32) (int32, error) {
	value, err := r.ReadInt64(tag, required, int64(def))
	return int32(value), err
}

func (r *Reader) ReadString(tag int, required bool, def string) (string, error) {
	typ, ok, err := r.findTag(tag, required)
	if err != nil || !ok {
		return def, err
	}
	return r.ReadStringValue(typ)
}

func (r *Reader) ReadBytes(tag int, required bool, def []byte) ([]byte, error) {
	typ, ok, err := r.findTag(tag, required)
	if err != nil || !ok {
		return def, err
	}
	return r.ReadBytesValue(typ)
}

func (r *Reader) ReadMapStringString(tag int, required bool) (map[string]string, error) {
	typ, ok, err := r.findTag(tag, required)
	if err != nil || !ok {
		return map[string]string{}, err
	}
	if typ != TypeMap {
		return nil, fmt.Errorf("taf: tag=%d 不是 map: 0x%02x", tag, byte(typ))
	}
	n, err := r.readLength()
	if err != nil {
		return nil, err
	}
	out := make(map[string]string, n)
	for i := int64(0); i < n; i++ {
		_, kt, err := r.ReadHead()
		if err != nil {
			return nil, err
		}
		k, err := r.ReadStringValue(kt)
		if err != nil {
			return nil, err
		}
		_, vt, err := r.ReadHead()
		if err != nil {
			return nil, err
		}
		v, err := r.ReadStringValue(vt)
		if err != nil {
			return nil, err
		}
		out[k] = v
	}
	return out, nil
}

func (r *Reader) ReadMapStringBytes(tag int, required bool) (map[string][]byte, error) {
	typ, ok, err := r.findTag(tag, required)
	if err != nil || !ok {
		return map[string][]byte{}, err
	}
	if typ != TypeMap {
		return nil, fmt.Errorf("taf: tag=%d 不是 map: 0x%02x", tag, byte(typ))
	}
	n, err := r.readLength()
	if err != nil {
		return nil, err
	}
	out := make(map[string][]byte, n)
	for i := int64(0); i < n; i++ {
		_, kt, err := r.ReadHead()
		if err != nil {
			return nil, err
		}
		k, err := r.ReadStringValue(kt)
		if err != nil {
			return nil, err
		}
		_, vt, err := r.ReadHead()
		if err != nil {
			return nil, err
		}
		v, err := r.ReadBytesValue(vt)
		if err != nil {
			return nil, err
		}
		out[k] = v
	}
	return out, nil
}

func (r *Reader) ReadStruct(tag int, required bool, obj Unmarshaler) (bool, error) {
	typ, ok, err := r.findTag(tag, required)
	if err != nil || !ok {
		return false, err
	}
	if typ != TypeStructBegin {
		return false, fmt.Errorf("taf: tag=%d 不是结构体: 0x%02x", tag, byte(typ))
	}
	if err := obj.ReadFrom(r); err != nil {
		return false, err
	}
	return true, r.SkipToStructEnd()
}

func (r *Reader) ReadListStruct(tag int, required bool, factory func() Unmarshaler) ([]Unmarshaler, error) {
	typ, ok, err := r.findTag(tag, required)
	if err != nil || !ok {
		return nil, err
	}
	if typ != TypeList {
		return nil, fmt.Errorf("taf: tag=%d 不是 list: 0x%02x", tag, byte(typ))
	}
	n, err := r.readLength()
	if err != nil {
		return nil, err
	}
	out := make([]Unmarshaler, 0, n)
	for i := int64(0); i < n; i++ {
		_, it, err := r.ReadHead()
		if err != nil {
			return nil, err
		}
		if it != TypeStructBegin {
			if err := r.Skip(it); err != nil {
				return nil, err
			}
			continue
		}
		obj := factory()
		if err := obj.ReadFrom(r); err != nil {
			return nil, err
		}
		if err := r.SkipToStructEnd(); err != nil {
			return nil, err
		}
		out = append(out, obj)
	}
	return out, nil
}

func (r *Reader) SkipToStructEnd() error {
	for {
		_, typ, err := r.ReadHead()
		if err != nil {
			return err
		}
		if typ == TypeStructEnd {
			return nil
		}
		if err := r.Skip(typ); err != nil {
			return err
		}
	}
}