huya-go/internal/huya/wss_client_test.go

package huya

import (
	"context"
	"encoding/base64"
	"encoding/hex"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"huya-go/internal/taf"
)

func TestWSRegisterGroupReqEncoding(t *testing.T) {
	w := taf.NewWriter()
	req := &wsRegisterGroupReq{GroupIDs: []string{WSSGroupID}}
	if err := req.WriteTo(w); err != nil {
		t.Fatal(err)
	}
	got := hex.EncodeToString(w.Bytes())
	want := "09000106066c6976653a301600"
	if got != want {
		t.Fatalf("register group bytes = %s, want %s", got, want)
	}
}

func TestWSUpdateUserExpsReqEncoding(t *testing.T) {
	w := taf.NewWriter()
	if err := (&wsUpdateUserExpsReq{}).WriteTo(w); err != nil {
		t.Fatal(err)
	}
	got := hex.EncodeToString(w.Bytes())
	want := "080c"
	if got != want {
		t.Fatalf("update user exps bytes = %s, want %s", got, want)
	}
}

func TestGenerateWSSBaseinfoMatchesBrowserShape(t *testing.T) {
	encoded, err := GenerateWSSBaseinfo(123, "guid-value", "58664.107263")
	if err != nil {
		t.Fatal(err)
	}
	unescaped, err := url.QueryUnescape(encoded)
	if err != nil {
		t.Fatal(err)
	}
	raw, err := base64.StdEncoding.DecodeString(unescaped)
	if err != nil {
		t.Fatal(err)
	}

	r := taf.NewReader(raw)
	uid, err := r.ReadInt64(0, true, 0)
	if err != nil {
		t.Fatal(err)
	}
	guid, err := r.ReadString(1, true, "")
	if err != nil {
		t.Fatal(err)
	}
	ua, err := r.ReadString(2, true, "")
	if err != nil {
		t.Fatal(err)
	}
	appSrc, err := r.ReadString(3, true, "")
	if err != nil {
		t.Fatal(err)
	}
	exp, err := r.ReadString(5, true, "")
	if err != nil {
		t.Fatal(err)
	}
	cookie, err := r.ReadString(8, false, "not-empty")
	if err != nil {
		t.Fatal(err)
	}
	traceID, err := r.ReadString(9, false, "not-empty")
	if err != nil {
		t.Fatal(err)
	}

	if uid != 123 || guid != "guid-value" || ua != WSSCookieUA || appSrc != "HUYA&ZH&2052" {
		t.Fatalf("unexpected baseinfo: uid=%d guid=%q ua=%q appSrc=%q", uid, guid, ua, appSrc)
	}
	if exp != "58664.107263" {
		t.Fatalf("baseinfo exp=%q, want 58664.107263", exp)
	}
	if cookie != "" || traceID != "" {
		t.Fatalf("wss baseinfo should not carry cookie/trace: cookie=%q trace=%q", cookie, traceID)
	}
}

func TestNewWSSCommandTraceMatchesBrowserShape(t *testing.T) {
	trace := NewWSSCommandTrace()
	parts := strings.Split(trace, ":")
	if len(parts) != 4 {
		t.Fatalf("trace parts = %v", parts)
	}
	if parts[0] == "" || parts[1] == "" || parts[0] != parts[1] {
		t.Fatalf("unexpected trace prefix: %q", trace)
	}
	if parts[2] != "0" || parts[3] != "0" {
		t.Fatalf("trace suffix = %q, want :0:0", trace)
	}
}

func TestSessionSharesSingleDialPerUID(t *testing.T) {
	client := NewWSSClient(nil)
	dialStarted := make(chan struct{}, 1)
	release := make(chan struct{})
	var dials atomic.Int32
	client.dialFunc = func(ctx context.Context, uid int64, cookie string) (*WSSSession, error) {
		dials.Add(1)
		select {
		case dialStarted <- struct{}{}:
		default:
		}
		<-release
		return &WSSSession{
			client:  client,
			uid:     uid,
			cookie:  cookie,
			closed:  make(chan struct{}),
			pending: map[int32]chan rpcResult{},
		}, nil
	}

	const callers = 4
	sessions := make([]*WSSSession, callers)
	var wg sync.WaitGroup
	for i := 0; i < callers; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			session, err := client.session(context.Background(), 42, "yyuid=42")
			if err != nil {
				t.Errorf("session: %v", err)
				return
			}
			sessions[i] = session
		}(i)
	}

	select {
	case <-dialStarted:
	case <-time.After(time.Second):
		t.Fatal("dial was not started")
	}
	// 留出时间让其余调用进入等待分支。
	time.Sleep(50 * time.Millisecond)
	close(release)
	wg.Wait()

	if got := dials.Load(); got != 1 {
		t.Fatalf("dials=%d, want 1", got)
	}
	for i, session := range sessions {
		if session == nil {
			t.Fatalf("caller %d got nil session", i)
		}
		if session != sessions[0] {
			t.Fatalf("caller %d got a different session", i)
		}
	}
}