huya-go/internal/huya/wss_client.go

package huya

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"huya-go/internal/taf"
	"huya-go/internal/wup"
)

const (
	DefaultWSSHost   = "77bc035c-ws.va.huya.com"
	DefaultWSSOrigin = "https://zt.huya.com"

	WSSCookieUA = "webh5&0.0.1&websocket&&diypc_52775"
	WSSBizUA    = WSSCookieUA
	WSSInitUA   = WSSCookieUA
	WSSGroupID  = "live:0"
)

type WSSClient struct {
	Host        string
	Origin      string
	UserAgent   string
	Dialer      *websocket.Dialer
	Logger      func(string)
	CallTimeout time.Duration

	// dialFunc 仅用于测试注入，为 nil 时走 dialSession。
	dialFunc func(ctx context.Context, uid int64, cookie string) (*WSSSession, error)

	mu       sync.Mutex
	sessions map[string]*WSSSession
	// dialing 记录正在建连的 uid，避免并发调用重复握手。
	dialing map[string]chan struct{}
}

type WSSSession struct {
	client *WSSClient
	uid    int64
	cookie string
	guid   string

	conn *websocket.Conn

	writeMu   sync.Mutex
	pendingMu sync.Mutex
	pending   map[int32]chan rpcResult
	requestMu sync.Mutex
	closed    chan struct{}
	closeOnce sync.Once
	closeErr  error

	nextWUPRequestID int32
	launchGUID       string
	launchIP         string
}

type rpcResult struct {
	resp *wup.Response
	err  error
}

func NewWSSClient(logger func(string)) *WSSClient {
	return &WSSClient{
		Host:        DefaultWSSHost,
		Origin:      DefaultWSSOrigin,
		UserAgent:   PCUserAgent,
		Logger:      logger,
		CallTimeout: 10 * time.Second,
		sessions:    map[string]*WSSSession{},
		dialing:     map[string]chan struct{}{},
	}
}

func (c *WSSClient) logf(format string, args ...any) {
	if c.Logger != nil {
		c.Logger(fmt.Sprintf(format, args...))
	}
}

func (c *WSSClient) Close() error {
	c.mu.Lock()
	sessions := make([]*WSSSession, 0, len(c.sessions))
	for _, session := range c.sessions {
		sessions = append(sessions, session)
	}
	c.sessions = map[string]*WSSSession{}
	c.mu.Unlock()

	var firstErr error
	for _, session := range sessions {
		if err := session.Close(); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}

func (c *WSSClient) PrewarmActivity(ctx context.Context, uid int64, cookie string, sid int32) error {
	_, err := c.session(ctx, uid, cookie)
	return err
}

// session 返回 uid 对应的连接，必要时新建。
//
// 握手与初始化在锁外完成，同一 uid 的并发调用共享同一次建连（single-flight），
// 因此多账号预热/抢购时不会因为持有全局锁而串行化。
func (c *WSSClient) session(ctx context.Context, uid int64, cookie string) (*WSSSession, error) {
	key := strconv.FormatInt(uid, 10)
	for {
		c.mu.Lock()
		if session := c.sessions[key]; session != nil && !session.isClosed() {
			c.mu.Unlock()
			return session, nil
		}
		dialing := c.dialing[key]
		if dialing == nil {
			dialing = make(chan struct{})
			c.dialing[key] = dialing
			c.mu.Unlock()

			session, err := c.dial(ctx, uid, cookie)

			c.mu.Lock()
			delete(c.dialing, key)
			if err == nil {
				c.sessions[key] = session
			}
			c.mu.Unlock()
			close(dialing)
			if err != nil {
				return nil, err
			}
			return session, nil
		}
		c.mu.Unlock()

		// 已有同 uid 的建连在途，等它结束后复用结果或自行重试。
		select {
		case <-dialing:
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
}

func (c *WSSClient) dial(ctx context.Context, uid int64, cookie string) (*WSSSession, error) {
	if c.dialFunc != nil {
		return c.dialFunc(ctx, uid, cookie)
	}
	return c.dialSession(ctx, uid, cookie)
}

func (c *WSSClient) dialSession(ctx context.Context, uid int64, cookie string) (*WSSSession, error) {
	host := c.Host
	if host == "" {
		host = DefaultWSSHost
	}
	origin := c.Origin
	if origin == "" {
		origin = DefaultWSSOrigin
	}
	userAgent := c.UserAgent
	if userAgent == "" {
		userAgent = PCUserAgent
	}

	guid := CookieValue(cookie, "guid")
	baseinfo, err := GenerateWSSBaseinfo(uid, guid)
	if err != nil {
		return nil, err
	}
	wsURL := url.URL{Scheme: "wss", Host: host, Path: "/"}
	wsURL.RawQuery = "baseinfo=" + baseinfo

	dialer := c.Dialer
	if dialer == nil {
		dialer = &websocket.Dialer{
			HandshakeTimeout:  10 * time.Second,
			EnableCompression: true,
			Proxy:             http.ProxyFromEnvironment,
		}
	}
	headers := http.Header{}
	headers.Set("Origin", origin)
	headers.Set("User-Agent", userAgent)
	headers.Set("Accept-Language", "zh-CN,zh;q=0.9")

	c.logf("[WSS] 连接 %s uid=%d", host, uid)
	conn, resp, err := dialer.DialContext(ctx, wsURL.String(), headers)
	if err != nil {
		if resp != nil {
			return nil, fmt.Errorf("wss dial: %w (status=%s)", err, resp.Status)
		}
		return nil, fmt.Errorf("wss dial: %w", err)
	}
	conn.EnableWriteCompression(true)
	conn.SetReadLimit(4 << 20)

	session := &WSSSession{
		client:           c,
		uid:              uid,
		cookie:           cookie,
		guid:             guid,
		conn:             conn,
		closed:           make(chan struct{}),
		pending:          map[int32]chan rpcResult{},
		nextWUPRequestID: 1,
	}
	go session.readLoop()
	go session.heartbeatLoop()

	if err := session.initialize(ctx); err != nil {
		_ = session.Close()
		return nil, err
	}
	return session, nil
}

func (s *WSSSession) isClosed() bool {
	select {
	case <-s.closed:
		return true
	default:
		return false
	}
}

func (s *WSSSession) Close() error {
	var err error
	s.closeOnce.Do(func() {
		close(s.closed)
		err = s.conn.Close()
		s.failPending(errors.New("wss closed"))
	})
	return err
}

func (s *WSSSession) readLoop() {
	for {
		messageType, raw, err := s.conn.ReadMessage()
		if err != nil {
			s.closeWithError(err)
			return
		}
		if messageType != websocket.BinaryMessage {
			continue
		}
		cmd, err := DecodeWSSCommand(raw)
		if err != nil {
			s.client.logf("[WSS] 解码消息失败: %v", err)
			continue
		}
		switch cmd.CmdType {
		case WSSCmdWupRsp:
			var resp wup.Response
			if err := resp.Decode(cmd.Data); err != nil {
				s.client.logf("[WSS] 解码 WUP 响应失败: %v", err)
				continue
			}
			s.resolveRequest(resp.RequestID, rpcResult{resp: &resp})
		case WSSCmdHeartBeatAck, WSSCmdHeartBeatRsp:
			// 心跳响应无需处理。
		case WSSCmdRegisterGroupRsp, WSSCmdVerifyCookieRsp, WSSCmdUpdateUserInfoRsp, WSSCmdUpdateUserExpsRsp:
			// 控制类响应频率高且没有业务字段，避免刷屏。
		default:
			s.client.logf("[WSS] 收到 cmd=%d data=%dB requestID=%d", cmd.CmdType, len(cmd.Data), cmd.RequestID)
		}
	}
}

func (s *WSSSession) heartbeatLoop() {
	ticker := time.NewTicker(8 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
			_ = s.sendHeartbeat(ctx)
			cancel()
		case <-s.closed:
			return
		}
	}
}

func (s *WSSSession) closeWithError(err error) {
	s.closeOnce.Do(func() {
		s.closeErr = err
		close(s.closed)
		_ = s.conn.Close()
		s.failPending(err)
	})
}

func (s *WSSSession) failPending(err error) {
	s.pendingMu.Lock()
	pending := s.pending
	s.pending = map[int32]chan rpcResult{}
	s.pendingMu.Unlock()
	for _, ch := range pending {
		ch <- rpcResult{err: err}
	}
}

func (s *WSSSession) resolveRequest(requestID int32, result rpcResult) {
	s.pendingMu.Lock()
	ch := s.pending[requestID]
	if ch == nil {
		s.pendingMu.Unlock()
		s.client.logf("[WSS] 未匹配响应 requestID=%d", requestID)
		return
	}
	delete(s.pending, requestID)
	s.pendingMu.Unlock()
	ch <- result
}

func (s *WSSSession) removePending(requestID int32) {
	s.pendingMu.Lock()
	defer s.pendingMu.Unlock()
	delete(s.pending, requestID)
}

func (s *WSSSession) writePayload(ctx context.Context, payload []byte) error {
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	if deadline, ok := ctx.Deadline(); ok {
		_ = s.conn.SetWriteDeadline(deadline)
		defer s.conn.SetWriteDeadline(time.Time{})
	}
	return s.conn.WriteMessage(websocket.BinaryMessage, payload)
}

func (s *WSSSession) initialize(ctx context.Context) error {
	if err := s.sendRegisterGroup(ctx); err != nil {
		return err
	}
	if err := s.callWSLaunch(ctx); err != nil {
		return err
	}
	if err := s.sendVerifyCookie(ctx); err != nil {
		return err
	}
	if err := s.callGetConfig(ctx); err != nil {
		s.client.logf("[WSS] getConfig 失败但继续: %v", err)
	}
	if err := s.sendUpdateUserExps(ctx); err != nil {
		return err
	}
	if err := s.sendRegisterGroup(ctx); err != nil {
		return err
	}
	if err := s.sendHeartbeat(ctx); err != nil {
		return err
	}
	return nil
}

func (s *WSSSession) withTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
	if _, ok := ctx.Deadline(); ok || s.client.CallTimeout <= 0 {
		return ctx, func() {}
	}
	return context.WithTimeout(ctx, s.client.CallTimeout)
}

func (s *WSSSession) sendHeartbeat(ctx context.Context) error {
	payload, err := EncodeWSSCommand(WSSCmdHeartBeatReq, nil, "")
	if err != nil {
		return err
	}
	return s.writePayload(ctx, payload)
}

func (s *WSSSession) sendRegisterGroup(ctx context.Context) error {
	req := &wsRegisterGroupReq{GroupIDs: []string{WSSGroupID}}
	payload, err := encodeWSSStructCommand(WSSCmdRegisterGroupReq, req)
	if err != nil {
		return err
	}
	return s.writePayload(ctx, payload)
}

func (s *WSSSession) callWSLaunch(ctx context.Context) error {
	req := &wsLaunchReq{UID: s.uid, GUID: s.guid, UA: WSSInitUA}
	resp, err := s.callWUPWithRequestIDAndTrace(ctx, "launch", "wsLaunch", req, -1, "")
	if err != nil {
		return err
	}
	var launch wsLaunchResp
	ok, err := resp.ReadStruct("tRsp", &launch)
	if err != nil {
		return err
	}
	if !ok {
		return errors.New("wsLaunch 响应缺少 tRsp")
	}
	s.launchGUID = launch.GUID
	s.launchIP = launch.ClientIP
	s.client.logf("[WSS] wsLaunch guid=%s ip=%s", launch.GUID, launch.ClientIP)
	return nil
}

func (s *WSSSession) callGetConfig(ctx context.Context) error {
	req := &getConfigReq{UID: s.uid, GUID: s.launchGUID, Cookie: s.cookie}
	_, err := s.callWUP(ctx, "mobileui", "getConfig", req)
	return err
}

func (s *WSSSession) sendUpdateUserExps(ctx context.Context) error {
	payload, err := encodeWSSStructCommand(WSSCmdUpdateUserExpsReq, &wsUpdateUserExpsReq{})
	if err != nil {
		return err
	}
	return s.writePayload(ctx, payload)
}

func (s *WSSSession) sendVerifyCookie(ctx context.Context) error {
	guid := s.launchGUID
	if guid == "" {
		guid = s.guid
	}
	req := &wsVerifyCookieReq{
		UID:             s.uid,
		UA:              WSSInitUA,
		Cookie:          NormalizeWSSCookie(s.cookie),
		GUID:            guid,
		AutoRegisterUID: 1,
		AppSrc:          "HUYA&ZH&2052",
	}
	payload, err := encodeWSSStructCommand(WSSCmdVerifyCookieReq, req)
	if err != nil {
		return err
	}
	return s.writePayload(ctx, payload)
}

func (s *WSSSession) nextRequestID() int32 {
	s.requestMu.Lock()
	defer s.requestMu.Unlock()
	id := s.nextWUPRequestID
	s.nextWUPRequestID++
	return id
}

func (s *WSSSession) callWUP(ctx context.Context, service, method string, reqStruct taf.Marshaler) (*wup.Response, error) {
	return s.callWUPWithRequestID(ctx, service, method, reqStruct, s.nextRequestID())
}

func (s *WSSSession) callWUPWithRequestID(ctx context.Context, service, method string, reqStruct taf.Marshaler, requestID int32) (*wup.Response, error) {
	return s.callWUPWithRequestIDAndTrace(ctx, service, method, reqStruct, requestID, NewWSSCommandTrace())
}

func (s *WSSSession) callWUPWithRequestIDAndTrace(ctx context.Context, service, method string, reqStruct taf.Marshaler, requestID int32, traceID string) (*wup.Response, error) {
	wupReq := wup.NewRequest(service, method, requestID)
	wupReq.Timeout = 0
	if err := wupReq.WriteStruct("tReq", reqStruct); err != nil {
		return nil, err
	}
	data, err := wupReq.Encode()
	if err != nil {
		return nil, err
	}
	ctx, cancel := s.withTimeout(ctx)
	defer cancel()
	payload, err := EncodeWSSCommand(WSSCmdWupReq, data, traceID)
	if err != nil {
		return nil, err
	}
	ch := make(chan rpcResult, 1)
	s.pendingMu.Lock()
	s.pending[requestID] = ch
	s.pendingMu.Unlock()
	s.client.logf("[WSS] 请求 %s.%s requestID=%d", service, method, requestID)
	if err := s.writePayload(ctx, payload); err != nil {
		s.removePending(requestID)
		return nil, err
	}
	select {
	case result := <-ch:
		return result.resp, result.err
	case <-ctx.Done():
		s.removePending(requestID)
		return nil, ctx.Err()
	case <-s.closed:
		s.removePending(requestID)
		if s.closeErr != nil {
			return nil, s.closeErr
		}
		return nil, errors.New("wss closed")
	}
}

func (s *WSSSession) callRPC(ctx context.Context, service, method string, reqStruct taf.Marshaler, respStruct taf.Unmarshaler) (*wup.Response, error) {
	resp, err := s.callWUP(ctx, service, method, reqStruct)
	if err != nil {
		return nil, err
	}
	if respStruct != nil {
		ok, err := resp.ReadStruct("tRsp", respStruct)
		if err != nil {
			return nil, err
		}
		if !ok {
			return nil, fmt.Errorf("%s.%s 响应缺少 tRsp/tResp", service, method)
		}
	}
	return resp, nil
}

func (c *WSSClient) buildUser(uid int64, cookie string) ActivityUserID {
	user := NewActivityUserID(uid, NormalizeCookie(cookie), WSSBizUA)
	user.GUID = ""
	user.Token = ""
	user.DeviceInfo = ""
	return user
}

func (c *WSSClient) QueryUserScore(ctx context.Context, uid int64, cookie string, sid int32) (*GetUserScoreResp, error) {
	session, err := c.session(ctx, uid, cookie)
	if err != nil {
		return nil, err
	}
	req := &GetUserScoreReq{User: c.buildUser(uid, cookie), SID: sid}
	var resp GetUserScoreResp
	_, err = session.callRPC(ctx, ServantWebActUI, FuncGetUserScore, req, &resp)
	if err == nil {
		c.logf("[WSS] 积分已解码 uid=%d sid=%d status=%d 可用=%d msg=%s", uid, sid, resp.Status, resp.AvailableScore(), resp.Msg)
	}
	return &resp, err
}

func (c *WSSClient) GetActPrizeList(ctx context.Context, uid int64, cookie string, sid int32) (*GetActPrizeListResp, error) {
	session, err := c.session(ctx, uid, cookie)
	if err != nil {
		return nil, err
	}
	req := &GetActPrizeListReq{User: c.buildUser(uid, cookie), SID: sid}
	var resp GetActPrizeListResp
	_, err = session.callRPC(ctx, ServantWebActUI, FuncPrizeList, req, &resp)
	if err == nil {
		left, total, leftKnown, totalKnown := prizeListInventorySummary(resp.Prizes)
		c.logf("[WSS] 商品列表已解码 uid=%d sid=%d status=%d count=%d 剩余=%d/%d项 总量=%d/%d项 msg=%s", uid, sid, resp.Status, len(resp.Prizes), left, leftKnown, total, totalKnown, resp.Msg)
	}
	return &resp, err
}

func (c *WSSClient) GetUserPrizeRecords(ctx context.Context, uid int64, cookie string, sid int32) (*GetUserPrizeRecordsResp, error) {
	session, err := c.session(ctx, uid, cookie)
	if err != nil {
		return nil, err
	}
	req := &GetUserPrizeRecordsReq{User: c.buildUser(uid, cookie), SID: sid}
	var resp GetUserPrizeRecordsResp
	_, err = session.callRPC(ctx, ServantWebActUI, FuncPrizeRecords, req, &resp)
	if err == nil {
		c.logf("[WSS] 兑换记录已解码 uid=%d sid=%d status=%d count=%d msg=%s", uid, sid, resp.Status, len(resp.Records), resp.Msg)
	}
	return &resp, err
}

func (c *WSSClient) ScoreExchangePrize(ctx context.Context, uid int64, cookie string, sid, pid int32) (*ScoreExchangePrizeResp, error) {
	session, err := c.session(ctx, uid, cookie)
	if err != nil {
		return nil, err
	}
	req := &ScoreExchangePrizeReq{
		User:      c.buildUser(uid, cookie),
		SID:       sid,
		PID:       pid,
		ClientEnv: map[string]string{},
		IsRole:    1,
	}
	var resp ScoreExchangePrizeResp
	_, err = session.callRPC(ctx, ServantWebActUI, FuncExchangePrize, req, &resp)
	if err == nil {
		c.logf("[WSS] 兑换结果已解码 uid=%d sid=%d pid=%d status=%d order=%s msg=%s", uid, sid, pid, resp.Status, resp.OrderID, exchangeResponseMessage(&resp))
	}
	return &resp, err
}

func NormalizeWSSCookie(cookie string) string {
	cookie = NormalizeCookie(cookie)
	pairs := CookiePairs(cookie)
	out := []string{"huya_ua=" + WSSCookieUA}
	for _, pair := range pairs {
		if strings.EqualFold(pair.Key, "huya_ua") {
			continue
		}
		out = append(out, pair.Key+"="+pair.Value)
	}
	return strings.Join(out, "; ")
}

type wsRegisterGroupReq struct {
	GroupIDs []string
	Token    string
}

func (r *wsRegisterGroupReq) WriteTo(w *taf.Writer) error {
	if err := w.WriteHead(0, taf.TypeList); err != nil {
		return err
	}
	if err := w.WriteInt32(0, int64(len(r.GroupIDs))); err != nil {
		return err
	}
	for _, groupID := range r.GroupIDs {
		if err := w.WriteString(0, groupID); err != nil {
			return err
		}
	}
	return w.WriteString(1, r.Token)
}

type wsLaunchReq struct {
	UID  int64
	GUID string
	UA   string
}

func (r *wsLaunchReq) WriteTo(w *taf.Writer) error {
	if err := w.WriteInt64(0, r.UID); err != nil {
		return err
	}
	if err := w.WriteString(1, r.GUID); err != nil {
		return err
	}
	ua := r.UA
	if ua == "" {
		ua = WSSInitUA
	}
	if err := w.WriteString(2, ua); err != nil {
		return err
	}
	if err := w.WriteString(3, "HUYA&ZH&2052"); err != nil {
		return err
	}
	if err := w.WriteStructBegin(4); err != nil {
		return err
	}
	for i := 0; i < 5; i++ {
		if err := w.WriteString(i, ""); err != nil {
			return err
		}
	}
	return w.WriteStructEnd()
}

type wsUpdateUserExpsReq struct{}

func (r *wsUpdateUserExpsReq) WriteTo(w *taf.Writer) error {
	return w.WriteMapStringString(0, nil)
}

type wsVerifyCookieReq struct {
	UID             int64
	UA              string
	Cookie          string
	GUID            string
	AutoRegisterUID int32
	AppSrc          string
}

func (r *wsVerifyCookieReq) WriteTo(w *taf.Writer) error {
	if err := w.WriteInt64(0, r.UID); err != nil {
		return err
	}
	if err := w.WriteString(1, r.UA); err != nil {
		return err
	}
	if err := w.WriteString(2, r.Cookie); err != nil {
		return err
	}
	if err := w.WriteString(3, r.GUID); err != nil {
		return err
	}
	if err := w.WriteInt32(4, int64(r.AutoRegisterUID)); err != nil {
		return err
	}
	return w.WriteString(5, r.AppSrc)
}

type wsLaunchResp struct {
	GUID     string
	ClientIP string
}

func (r *wsLaunchResp) ReadFrom(rd *taf.Reader) error {
	var err error
	if r.GUID, err = rd.ReadString(0, false, r.GUID); err != nil {
		return err
	}
	r.ClientIP, err = rd.ReadString(1, false, r.ClientIP)
	return err
}

type getConfigReq struct {
	UID    int64
	GUID   string
	Cookie string
}

func (r *getConfigReq) WriteTo(w *taf.Writer) error {
	user := NewActivityUserID(r.UID, NormalizeWSSCookie(r.Cookie), WSSInitUA)
	user.GUID = r.GUID
	if err := w.WriteStruct(0, &user); err != nil {
		return err
	}
	if err := w.WriteMapStringString(1, nil); err != nil {
		return err
	}
	if err := w.WriteString(2, "huya"); err != nil {
		return err
	}
	if err := w.WriteHead(3, taf.TypeList); err != nil {
		return err
	}
	return w.WriteInt32(0, 0)
}