photoplaces/backend/internal/handlers/websocket.go

// Package handlers
package handlers

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

	"github.com/gorilla/websocket"
)

type VisitorDot struct {
	UserID string  `json:"user_id,omitempty"`
	Lat    float64 `json:"lat"`
	Lng    float64 `json:"lng"`
}

// WSHub управляет WebSocket-соединениями для отображения посетителей на карте
type WSHub struct {
	mu             sync.RWMutex
	clients        map[*websocket.Conn]VisitorDot
	allowedOrigins []string // разрешённые Origin для CORS WebSocket
}

// NewWSHub создаёт хаб с переданным списком разрешённых Origin
func NewWSHub(allowedOrigins []string) *WSHub {
	return &WSHub{
		clients:        make(map[*websocket.Conn]VisitorDot),
		allowedOrigins: allowedOrigins,
	}
}

// HandleWS обрабатывает WebSocket-подключение: апгрейд, чтение, рассылка, очистка
func (h *WSHub) HandleWS(w http.ResponseWriter, r *http.Request) {
	// Создаём upgrader локально с проверкой Origin из конфига
	upgrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			if len(h.allowedOrigins) == 0 {
				return true // если origins не настроены — пропускаем все (обратная совместимость)
			}
			origin := r.Header.Get("Origin")
			for _, o := range h.allowedOrigins {
				if o == origin {
					return true
				}
			}
			return false
		},
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("ws upgrade: %v", err)
		return
	}

	dot := VisitorDot{}

	h.mu.Lock()
	h.clients[conn] = dot
	h.mu.Unlock()

	// Цикл чтения сообщений от клиента
	for {
		var msg VisitorDot
		if err := conn.ReadJSON(&msg); err != nil {
			break
		}

		h.mu.Lock()
		h.clients[conn] = msg
		visitors := make([]VisitorDot, 0, len(h.clients))
		for _, v := range h.clients {
			visitors = append(visitors, v)
		}

		// Рассылаем всем клиентам, собираем отвалившиеся соединения
		var failed []*websocket.Conn
		for c := range h.clients {
			// Таймаут записи для предотвращения slow-client DoS
			if err := c.SetWriteDeadline(time.Now().Add(10 * time.Second)); err != nil {
				failed = append(failed, c)
				continue
			}
			if err := c.WriteJSON(map[string]interface{}{"visitors": visitors}); err != nil {
				failed = append(failed, c)
			}
		}
		// Закрываем и удаляем упавшие соединения (вне итерации по map — безопасно)
		for _, c := range failed {
			c.WriteMessage(websocket.CloseMessage, []byte{})
			c.Close()
			delete(h.clients, c)
		}
		h.mu.Unlock()
	}

	// Очистка после завершения цикла чтения
	h.mu.Lock()
	delete(h.clients, conn)
	h.mu.Unlock()
	conn.Close()
}