162 lines
3.8 KiB
Go
162 lines
3.8 KiB
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() {
|
||
nodes := a.engine.GetAllNodes()
|
||
msg := &models.WsMetricsMessage{
|
||
Type: "nodes_update",
|
||
Payload: nodes,
|
||
}
|
||
a.wsHub.broadcast(msg)
|
||
}()
|
||
|
||
// Читаем из вебсокета (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-клиентам (вызывается engine при обновлении).
|
||
func (a *API) BroadcastMetrics() {
|
||
nodes := a.engine.GetAllNodes()
|
||
msg := &models.WsMetricsMessage{
|
||
Type: "nodes_update",
|
||
Payload: nodes,
|
||
}
|
||
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
|
||
}
|