Files
kra-new/tmp/live_acceptance/websocket.go
T
2026-09-23 17:42:55 +08:00

155 lines
3.7 KiB
Go

package main
import (
"fmt"
"log"
"net/url"
"sync"
"time"
"github.com/gorilla/websocket"
)
type signalMessage struct {
Type string `json:"type"`
Payload map[string]any `json:"payload"`
}
type signalSocket struct {
conn *websocket.Conn
messages chan signalMessage
done chan struct{}
writeMu sync.Mutex
}
func connectSocket(token string) (*signalSocket, error) {
endpoint := wsURL
if token != "" {
endpoint += "?token=" + url.QueryEscape(token)
}
conn, _, err := websocket.DefaultDialer.Dial(endpoint, nil)
if err != nil {
return nil, err
}
s := &signalSocket{conn: conn, messages: make(chan signalMessage, 512), done: make(chan struct{})}
// The read loop processes WebSocket ping frames, allowing the server to
// receive pongs while HTTP assertions and application heartbeats run.
go func() {
defer close(s.done)
defer close(s.messages)
for {
var message signalMessage
if err := conn.ReadJSON(&message); err != nil {
return
}
select {
case s.messages <- message:
case <-time.After(5 * time.Second):
_ = conn.Close()
return
}
}
}()
go func() {
ticker := time.NewTicker(10 * time.Second)
defer ticker.Stop()
for {
select {
case <-s.done:
return
case <-ticker.C:
if s.write("Heartbeat", map[string]any{}) != nil {
_ = s.conn.Close()
return
}
}
}
}()
return s, nil
}
func (s *signalSocket) write(kind string, payload map[string]any) error {
s.writeMu.Lock()
defer s.writeMu.Unlock()
_ = s.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
return s.conn.WriteJSON(signalMessage{Type: kind, Payload: payload})
}
func (s *signalSocket) wait(kind string, match func(map[string]any) bool, timeout time.Duration) (signalMessage, error) {
timer := time.NewTimer(timeout)
defer timer.Stop()
for {
select {
case message, ok := <-s.messages:
if !ok {
return signalMessage{}, fmt.Errorf("WebSocket closed waiting for %s", kind)
}
if message.Type == kind && (match == nil || match(message.Payload)) {
return message, nil
}
if message.Type == "JoinRejected" {
return message, fmt.Errorf("join rejected: %v", message.Payload)
}
case <-timer.C:
return signalMessage{}, fmt.Errorf("timed out waiting for %s", kind)
}
}
}
func (s *signalSocket) quietChat(duration time.Duration) error {
timer := time.NewTimer(duration)
defer timer.Stop()
for {
select {
case message, ok := <-s.messages:
if !ok {
return fmt.Errorf("WebSocket closed while checking no chat history")
}
if message.Type == "ChatMessage" || message.Type == "ChatHistory" {
return fmt.Errorf("unexpected chat replay: %s %v", message.Type, message.Payload)
}
case <-timer.C:
return nil
}
}
}
func joinSocket(token, role string, room uint) (*signalSocket, error) {
s, err := connectSocket(token)
if err != nil {
return nil, err
}
if err = s.write("Join", map[string]any{"role": role, "streamId": fmt.Sprintf("room-%d", room)}); err != nil {
_ = s.conn.Close()
return nil, err
}
if _, err = s.wait("Joined", nil, 5*time.Second); err != nil {
_ = s.conn.Close()
return nil, err
}
return s, nil
}
func holdPublisher() error {
f, err := readFixtures()
if err != nil {
return err
}
client, err := login(f.Anchor, false)
if err != nil {
return err
}
socket, err := joinSocket(client.token, "publisher", f.RoomID)
if err != nil {
return err
}
defer socket.conn.Close()
log.Printf("Local fixture publisher joined room=%d with continuous reads and heartbeats", f.RoomID)
for message := range socket.messages {
if message.Type == "Error" || message.Type == "JoinRejected" {
return fmt.Errorf("publisher error: %v", message.Payload)
}
}
return fmt.Errorf("publisher disconnected")
}