package main import ( "encoding/json" "errors" "fmt" "log" "os" "strconv" "strings" "time" "github.com/glebarez/sqlite" "github.com/gorilla/websocket" "gorm.io/gorm" ) type signalMessage struct { Type string `json:"type"` Payload map[string]any `json:"payload"` } func main() { if len(os.Args) < 2 { log.Fatal("usage: live_smoke ...") } var err error switch os.Args[1] { case "seed": err = seed(os.Args[2:]) case "publisher": err = publisher(os.Args[2:]) case "wscheck": err = wscheck(os.Args[2:]) case "inspect": err = inspect(os.Args[2:]) default: err = fmt.Errorf("unknown command %q", os.Args[1]) } if err != nil { log.Fatal(err) } } func openDB(path string) (*gorm.DB, error) { return gorm.Open(sqlite.Open(path), &gorm.Config{}) } func seed(args []string) error { if len(args) != 2 { return errors.New("seed requires ") } userID, err := strconv.ParseUint(args[1], 10, 64) if err != nil || userID == 0 { return errors.New("invalid user id") } db, err := openDB(args[0]) if err != nil { return err } if result := db.Exec("UPDATE sys_integration_configs SET enabled = ?, updated_at = ? WHERE kind = ? AND provider = ?", true, time.Now(), "websocket", "melody"); result.Error != nil { return result.Error } else if result.RowsAffected != 1 { return fmt.Errorf("websocket config update affected %d rows", result.RowsAffected) } var categoryID uint if err = db.Raw("SELECT id FROM live_categories WHERE enabled = ? ORDER BY sort, id LIMIT 1", true).Scan(&categoryID).Error; err != nil { return err } if categoryID == 0 { return errors.New("no enabled live category") } name := fmt.Sprintf("Smoke Live Room %d", userID) if err = db.Exec("DELETE FROM live_rooms WHERE name = ?", name).Error; err != nil { return err } if err = db.Exec(`INSERT INTO live_rooms (name, category_id, cover_url, intro, sort, anchor_user_id, stream_url, replay_url, enabled, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, name, categoryID, "", "isolated live smoke test", 1, uint(userID), "", "", true, time.Now(), time.Now()).Error; err != nil { return err } var roomID uint if err = db.Raw("SELECT id FROM live_rooms WHERE name = ? ORDER BY id DESC LIMIT 1", name).Scan(&roomID).Error; err != nil { return err } if roomID == 0 { return errors.New("room insert did not return an id") } fmt.Printf("{\"roomId\":%d,\"categoryId\":%d}\n", roomID, categoryID) return nil } func publisher(args []string) error { if len(args) != 3 { return errors.New("publisher requires ") } token, err := cookieToken(args[1]) if err != nil { return err } connection, _, err := websocket.DefaultDialer.Dial(args[0]+"?token="+token, nil) if err != nil { return err } defer connection.Close() if err = writeSignal(connection, "Join", map[string]any{"role": "publisher", "streamId": "room-" + args[2]}); err != nil { return err } if _, err = readUntil(connection, "Joined", 5*time.Second); err != nil { return err } fmt.Println("publisher joined") ticker := time.NewTicker(15 * time.Second) defer ticker.Stop() for range ticker.C { if err = writeSignal(connection, "Heartbeat", map[string]any{"streamId": "room-" + args[2]}); err != nil { return err } } return nil } func wscheck(args []string) error { if len(args) != 3 { return errors.New("wscheck requires ") } streamID := "room-" + args[2] anonymous, _, err := websocket.DefaultDialer.Dial(args[0], nil) if err != nil { return fmt.Errorf("anonymous dial: %w", err) } defer anonymous.Close() if err = writeSignal(anonymous, "Join", map[string]any{"role": "subscriber", "streamId": streamID}); err != nil { return err } if _, err = readUntil(anonymous, "Joined", 5*time.Second); err != nil { return fmt.Errorf("anonymous join: %w", err) } if err = writeSignal(anonymous, "ChatSend", map[string]any{"content": "anonymous must fail"}); err != nil { return err } message, err := readUntil(anonymous, "Error", 5*time.Second) if err != nil { return fmt.Errorf("anonymous chat result: %w", err) } if message.Payload["code"] != "AUTH_REQUIRED" { return fmt.Errorf("anonymous chat code=%v", message.Payload["code"]) } token, err := cookieToken(args[1]) if err != nil { return err } authenticated, _, err := websocket.DefaultDialer.Dial(args[0]+"?token="+token, nil) if err != nil { return fmt.Errorf("authenticated dial: %w", err) } defer authenticated.Close() if err = writeSignal(authenticated, "Join", map[string]any{"role": "subscriber", "streamId": streamID}); err != nil { return err } if _, err = readUntil(authenticated, "Joined", 5*time.Second); err != nil { return fmt.Errorf("authenticated join: %w", err) } content := "smoke chat 2026-09-22" if err = writeSignal(authenticated, "ChatSend", map[string]any{"content": content}); err != nil { return err } message, err = readUntil(authenticated, "ChatMessage", 5*time.Second) if err != nil { return fmt.Errorf("authenticated chat result: %w", err) } if message.Payload["content"] != content { return fmt.Errorf("chat content=%v", message.Payload["content"]) } fmt.Println("anonymous join/auth rejection and authenticated chat broadcast passed") return nil } func inspect(args []string) error { if len(args) != 3 { return errors.New("inspect requires ") } db, err := openDB(args[0]) if err != nil { return err } userID, err := strconv.ParseUint(args[1], 10, 64) if err != nil { return err } roomID, err := strconv.ParseUint(args[2], 10, 64) if err != nil { return err } result := map[string]int64{} queries := map[string]struct { query string args []any }{ "balance": {"SELECT diamond_balance FROM live_user_wallets WHERE user_id = ?", []any{userID}}, "followCount": {"SELECT COUNT(*) FROM live_room_follows WHERE user_id = ? AND room_id = ?", []any{userID, roomID}}, "giftSendCount": {"SELECT COUNT(*) FROM live_gift_sends WHERE user_id = ? AND room_id = ?", []any{userID, roomID}}, "giftDiamonds": {"SELECT COALESCE(SUM(diamond_total), 0) FROM live_gift_sends WHERE user_id = ? AND room_id = ?", []any{userID, roomID}}, "chatRows": {"SELECT COUNT(*) FROM live_chat_messages WHERE room_id = ?", []any{roomID}}, "rechargeCount": {"SELECT COUNT(*) FROM live_recharge_orders WHERE user_id = ?", []any{userID}}, } for key, item := range queries { var value int64 if err = db.Raw(item.query, item.args...).Scan(&value).Error; err != nil { return fmt.Errorf("%s: %w", key, err) } result[key] = value } encoded, _ := json.Marshal(result) fmt.Println(string(encoded)) return nil } func writeSignal(connection *websocket.Conn, kind string, payload map[string]any) error { return connection.WriteJSON(signalMessage{Type: kind, Payload: payload}) } func readUntil(connection *websocket.Conn, kind string, timeout time.Duration) (*signalMessage, error) { if err := connection.SetReadDeadline(time.Now().Add(timeout)); err != nil { return nil, err } for { var message signalMessage if err := connection.ReadJSON(&message); err != nil { return nil, err } if message.Type == kind { return &message, nil } if message.Type == "JoinRejected" { return nil, fmt.Errorf("join rejected: %v", message.Payload) } } } func cookieToken(path string) (string, error) { content, err := os.ReadFile(path) if err != nil { return "", err } for _, line := range strings.Split(string(content), "\n") { fields := strings.Fields(line) if len(fields) >= 7 && fields[5] == "x-token" { return fields[6], nil } } return "", errors.New("x-token cookie not found") }