kra-new/tmp/live_smoke/main.go

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 <seed|publisher|wscheck|inspect> ...")
	}
	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 <db-dsn> <user-id>")
	}
	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 <ws-url> <cookie-file> <room-id>")
	}
	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 <ws-url> <cookie-file> <room-id>")
	}
	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-dsn> <user-id> <room-id>")
	}
	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")
}