|
|
@@ -5,30 +5,50 @@ import (
|
|
|
"log"
|
|
|
"net/http"
|
|
|
"sync"
|
|
|
+ "time"
|
|
|
|
|
|
"github.com/gorilla/websocket"
|
|
|
)
|
|
|
|
|
|
-var upgrader = websocket.Upgrader{
|
|
|
- CheckOrigin: func(r *http.Request) bool { return true },
|
|
|
-}
|
|
|
-
|
|
|
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
|
|
|
+ mu sync.RWMutex
|
|
|
+ clients map[*websocket.Conn]VisitorDot
|
|
|
+ allowedOrigins []string // разрешённые Origin для CORS WebSocket
|
|
|
}
|
|
|
|
|
|
-func NewWSHub() *WSHub {
|
|
|
- return &WSHub{clients: make(map[*websocket.Conn]VisitorDot)}
|
|
|
+// 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)
|
|
|
@@ -41,13 +61,7 @@ func (h *WSHub) HandleWS(w http.ResponseWriter, r *http.Request) {
|
|
|
h.clients[conn] = dot
|
|
|
h.mu.Unlock()
|
|
|
|
|
|
- defer func() {
|
|
|
- h.mu.Lock()
|
|
|
- delete(h.clients, conn)
|
|
|
- h.mu.Unlock()
|
|
|
- conn.Close()
|
|
|
- }()
|
|
|
-
|
|
|
+ // Цикл чтения сообщений от клиента
|
|
|
for {
|
|
|
var msg VisitorDot
|
|
|
if err := conn.ReadJSON(&msg); err != nil {
|
|
|
@@ -61,12 +75,30 @@ func (h *WSHub) HandleWS(w http.ResponseWriter, r *http.Request) {
|
|
|
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 {
|
|
|
- c.Close()
|
|
|
- delete(h.clients, c)
|
|
|
+ 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()
|
|
|
}
|