MMGame/backend/internal/handlers/ws.go

package handlers

import (
	"encoding/json"
	"net/http"
	"sync"

	"github.com/go-chi/chi/v5"
	"github.com/gorilla/websocket"
)

type Client struct {
	Conn   *websocket.Conn
	RoleID string
}

var (
	connsMu     sync.RWMutex
	activeConns = map[string][]*Client{}
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

func removeConn(gameID string, c *Client) {
	connsMu.Lock()
	defer connsMu.Unlock()
	clients := activeConns[gameID]
	for i, x := range clients {
		if x == c {
			activeConns[gameID] = append(clients[:i], clients[i+1:]...)
			break
		}
	}
	if len(activeConns[gameID]) == 0 {
		delete(activeConns, gameID)
	}
}

func removeGameConns(gameID string) {
	connsMu.Lock()
	defer connsMu.Unlock()
	if clients, ok := activeConns[gameID]; ok {
		for _, c := range clients {
			_ = c.Conn.Close()
		}
	}
	delete(activeConns, gameID)
}

func broadcast(gameID string, message interface{}) {
	data, err := json.Marshal(message)
	if err != nil {
		return
	}
	connsMu.RLock()
	clients := append([]*Client{}, activeConns[gameID]...)
	connsMu.RUnlock()
	for _, c := range clients {
		if err := c.Conn.WriteMessage(websocket.TextMessage, data); err != nil {
			removeConn(gameID, c)
		}
	}
}

func broadcastToRole(gameID, targetRoleID string, message interface{}) {
	data, err := json.Marshal(message)
	if err != nil {
		return
	}
	connsMu.RLock()
	clients := append([]*Client{}, activeConns[gameID]...)
	connsMu.RUnlock()
	for _, c := range clients {
		if c.RoleID != "" && c.RoleID != targetRoleID {
			continue
		}
		if err := c.Conn.WriteMessage(websocket.TextMessage, data); err != nil {
			removeConn(gameID, c)
		}
	}
}

func GameWebSocket(w http.ResponseWriter, r *http.Request) {
	gameID := chi.URLParam(r, "game_id")
	roleID := r.URL.Query().Get("role_id")

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	client := &Client{Conn: conn, RoleID: roleID}

	connsMu.Lock()
	activeConns[gameID] = append(activeConns[gameID], client)
	connsMu.Unlock()

	defer func() {
		removeConn(gameID, client)
		_ = conn.Close()
	}()

	for {
		_, data, err := conn.ReadMessage()
		if err != nil {
			return
		}
		var msg struct {
			Type string `json:"type"`
		}
		if err := json.Unmarshal(data, &msg); err != nil {
			continue
		}
		if msg.Type == "ping" {
			if err := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"pong"}`)); err != nil {
				return
			}
		}
	}
}