huya-go/internal/wup/wup.go

package wup

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

	"huya-go/internal/taf"
)

// Request 是 WUP v3 请求封装。
type Request struct {
	Version    int16
	PacketType int8
	Message    int32
	RequestID  int32
	Servant    string
	Func       string
	Timeout    int32
	Context    map[string]string
	Status     map[string]string
	NewData    map[string][]byte
}

func NewRequest(servant, fn string, requestID int32) *Request {
	return &Request{
		Version:   3,
		Servant:   servant,
		Func:      fn,
		RequestID: requestID,
		Timeout:   3000,
		Context:   map[string]string{},
		Status:    map[string]string{},
		NewData:   map[string][]byte{},
	}
}

func (r *Request) WriteStruct(key string, value taf.Marshaler) error {
	w := taf.NewWriter()
	if err := w.WriteStruct(0, value); err != nil {
		return err
	}
	r.NewData[key] = w.Bytes()
	return nil
}

func (r *Request) Encode() ([]byte, error) {
	data := taf.NewWriter()
	if err := data.WriteMapStringBytes(0, r.NewData); err != nil {
		return nil, err
	}

	body := taf.NewWriter()
	if err := body.WriteInt16(1, int64(r.Version)); err != nil {
		return nil, err
	}
	if err := body.WriteInt8(2, int64(r.PacketType)); err != nil {
		return nil, err
	}
	if err := body.WriteInt32(3, int64(r.Message)); err != nil {
		return nil, err
	}
	if err := body.WriteInt32(4, int64(r.RequestID)); err != nil {
		return nil, err
	}
	if err := body.WriteString(5, r.Servant); err != nil {
		return nil, err
	}
	if err := body.WriteString(6, r.Func); err != nil {
		return nil, err
	}
	if err := body.WriteBytes(7, data.Bytes()); err != nil {
		return nil, err
	}
	if err := body.WriteInt32(8, int64(r.Timeout)); err != nil {
		return nil, err
	}
	if err := body.WriteMapStringString(9, r.Context); err != nil {
		return nil, err
	}
	if err := body.WriteMapStringString(10, r.Status); err != nil {
		return nil, err
	}

	rawBody := body.Bytes()
	out := bytes.NewBuffer(make([]byte, 0, len(rawBody)+4))
	var lenBuf [4]byte
	binary.BigEndian.PutUint32(lenBuf[:], uint32(len(rawBody)+4))
	out.Write(lenBuf[:])
	out.Write(rawBody)
	return out.Bytes(), nil
}

// Response 是 WUP 响应封装。
type Response struct {
	Version    int16
	PacketType int8
	Message    int32
	RequestID  int32
	Servant    string
	Func       string
	Buffer     []byte
	Timeout    int32
	Context    map[string]string
	Status     map[string]string
	NewData    map[string][]byte
}

func (r *Response) Decode(data []byte) error {
	body := NormalizePayload(data)
	reader := taf.NewReader(body)
	r.Context = map[string]string{}
	r.Status = map[string]string{}
	r.NewData = map[string][]byte{}

	for {
		tag, typ, err := reader.ReadHead()
		if err != nil {
			if errors.Is(err, io.EOF) {
				return nil
			}
			return err
		}
		if tag > 10 {
			return nil
		}
		switch tag {
		case 1:
			v, err := reader.ReadIntValue(typ)
			r.Version = int16(v)
			if err != nil {
				return err
			}
		case 2:
			v, err := reader.ReadIntValue(typ)
			r.PacketType = int8(v)
			if err != nil {
				return err
			}
		case 3:
			v, err := reader.ReadIntValue(typ)
			r.Message = int32(v)
			if err != nil {
				return err
			}
		case 4:
			v, err := reader.ReadIntValue(typ)
			r.RequestID = int32(v)
			if err != nil {
				return err
			}
		case 5:
			v, err := reader.ReadStringValue(typ)
			r.Servant = v
			if err != nil {
				return err
			}
		case 6:
			v, err := reader.ReadStringValue(typ)
			r.Func = v
			if err != nil {
				return err
			}
		case 7:
			v, err := reader.ReadBytesValue(typ)
			if err != nil {
				return err
			}
			r.Buffer = v
			if err := r.decodeBuffer(); err != nil {
				return err
			}
		case 8:
			v, err := reader.ReadIntValue(typ)
			r.Timeout = int32(v)
			if err != nil {
				return err
			}
		case 9:
			if typ != taf.TypeMap {
				if err := reader.Skip(typ); err != nil {
					return err
				}
				continue
			}
			_ = reader.Seek(reader.Position() - 1)
			v, err := reader.ReadMapStringString(9, false)
			r.Context = v
			if err != nil {
				return err
			}
		case 10:
			if typ != taf.TypeMap {
				if err := reader.Skip(typ); err != nil {
					return err
				}
				return nil
			}
			_ = reader.Seek(reader.Position() - 1)
			v, err := reader.ReadMapStringString(10, false)
			r.Status = v
			if err != nil {
				return err
			}
			return nil
		default:
			if err := reader.Skip(typ); err != nil {
				return err
			}
		}
	}
}

func (r *Response) decodeBuffer() error {
	if len(r.Buffer) == 0 {
		return nil
	}
	reader := taf.NewReader(r.Buffer)
	m, err := reader.ReadMapStringBytes(0, false)
	if err != nil {
		return err
	}
	r.NewData = m
	return nil
}

func (r *Response) ReadStruct(key string, obj taf.Unmarshaler) (bool, error) {
	data := r.NewData[key]
	if len(data) == 0 && key == "tRsp" {
		data = r.NewData["tResp"]
	}
	if len(data) == 0 && key == "tResp" {
		data = r.NewData["tRsp"]
	}
	if len(data) == 0 {
		return false, nil
	}
	reader := taf.NewReader(data)
	_, typ, err := reader.ReadHead()
	if err != nil {
		return false, err
	}
	if typ != taf.TypeStructBegin {
		return false, fmt.Errorf("wup: %s 不是结构体响应: 0x%02x", key, byte(typ))
	}
	if err := obj.ReadFrom(reader); err != nil {
		return false, err
	}
	return true, reader.SkipToStructEnd()
}

func NormalizePayload(data []byte) []byte {
	if len(data) < 4 {
		return data
	}
	n := int(binary.BigEndian.Uint32(data[:4]))
	if n == len(data) || n+4 == len(data) {
		return data[4:]
	}
	return data
}