pulse-lets-go/internal/api/ws.go

package api

import (
	"fmt"
	"log"
	"net/http"
	"sync"

	"github.com/golang-jwt/jwt/v5"
	"github.com/gorilla/websocket"

	"github.com/pulse-lets-go/internal/models"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true // Разрешаем все origins для dev
	},
}

// wsClient — соединение WebSocket одного клиента.
type wsClient struct {
	conn *websocket.Conn
	mu   sync.Mutex
}

// wsHub управляет подключёнными WebSocket-клиентами и рассылает метрики.
type wsHub struct {
	mu      sync.RWMutex
	clients map[*wsClient]bool
}

// newWSHub создаёт хаб WebSocket.
func newWSHub() *wsHub {
	return &wsHub{
		clients: make(map[*wsClient]bool),
	}
}

// add добавляет клиента в хаб.
func (h *wsHub) add(c *wsClient) {
	h.mu.Lock()
	h.clients[c] = true
	h.mu.Unlock()
}

// remove удаляет клиента из хаба.
func (h *wsHub) remove(c *wsClient) {
	h.mu.Lock()
	if h.clients[c] {
		c.conn.Close()
		delete(h.clients, c)
	}
	h.mu.Unlock()
}

// broadcast рассылает метрики всем подключённым клиентам.
func (h *wsHub) broadcast(msg *models.WsMetricsMessage) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	for client := range h.clients {
		client.mu.Lock()
		err := client.conn.WriteJSON(msg)
		client.mu.Unlock()
		if err != nil {
			log.Printf("[ws] ошибка отправки: %v", err)
			go h.remove(client)
		}
	}
}

// handleWS — WebSocket /ws/metrics.
func (a *API) handleWS(w http.ResponseWriter, r *http.Request) {
	token := r.URL.Query().Get("token")
	if token == "" {
		// пробуем заголовок
		token = extractBearerToken(r)
	}
	if token == "" {
		writeError(w, http.StatusUnauthorized, "требуется токен")
		return
	}

	// Валидация JWT
	userInfo, err := a.validateToken(token)
	if err != nil {
		writeError(w, http.StatusUnauthorized, "недействительный токен")
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("[ws] ошибка upgrade: %v", err)
		return
	}

	client := &wsClient{conn: conn}
	a.wsHub.add(client)

	log.Printf("[ws] клиент подключён: %s (%s)", userInfo.Username, r.RemoteAddr)

	// Отправляем текущее состояние при подключении
	go func() {
		a.BroadcastAll()
	}()

	// Читаем из вебсокета (ping/pong + закрытие)
	go func() {
		defer a.wsHub.remove(client)
		for {
			_, _, err := conn.ReadMessage()
			if err != nil {
				log.Printf("[ws] клиент отключён: %s", r.RemoteAddr)
				return
			}
		}
	}()
}

// BroadcastMetrics рассылает метрики нод всем WS-клиентам.
func (a *API) BroadcastMetrics() {
	nodes := a.engine.GetAllNodes()
	msg := &models.WsMetricsMessage{
		Type:    "nodes_update",
		Payload: nodes,
	}
	a.wsHub.broadcast(msg)
}

// BroadcastAll рассылает ноды и состояние балансировщика.
func (a *API) BroadcastAll() {
	a.BroadcastMetrics()
	a.broadcastBalancerUpdate()
}

func (a *API) broadcastBalancerUpdate() {
	stats := a.engine.GetHealthStats()
	gwTotal, gwUp, gwDown := a.gatewayStats()

	info := models.BalancerInfo{
		NatsConnected:  a.natsConnected != nil && a.natsConnected(),
		GatewaysTotal:  gwTotal,
		GatewaysUp:     gwUp,
		GatewaysDown:   gwDown,
		RouteTotal:     stats.RouteRequests,
		RouteFallbacks: stats.RouteFallbacks,
		UptimeSec:      stats.UptimeSeconds,
	}

	if a.eslClient != nil {
		eslStats := a.eslClient.GetStats()
		info.EslConnected = eslStats.Status == "connected"
		info.EslUptimeSec = eslStats.UptimeSec
		info.EslReconnects = eslStats.Reconnects

		a.fsStatsMu.RLock()
		info.FsActiveCalls = a.fsStats.ActiveCalls
		info.FsMaxSessions = a.fsStats.MaxSessions
		info.FsUptime = a.fsStats.FsUptime
		a.fsStatsMu.RUnlock()
	}

	msg := &models.WsMetricsMessage{
		Type:    "balancer_update",
		Payload: info,
	}
	a.wsHub.broadcast(msg)
}

// validateToken проверяет JWT и возвращает UserInfo.
func (a *API) validateToken(tokenStr string) (*UserInfo, error) {
	token, err := parseJWT(tokenStr, a.jwtSecret)
	if err != nil {
		return nil, err
	}
	return token, nil
}

// parseJWT парсит JWT токен.
func parseJWT(tokenStr string, secret string) (*UserInfo, error) {
	c := &claims{}
	token, err := jwt.ParseWithClaims(tokenStr, c, func(token *jwt.Token) (interface{}, error) {
		return []byte(secret), nil
	})
	if err != nil || !token.Valid {
		return nil, fmt.Errorf("недействительный токен: %w", err)
	}
	return &UserInfo{
		UserID:   c.UserID,
		Username: c.Username,
		Role:     c.Role,
	}, nil
}