photoplaces/backend/internal/handlers/websocket.go

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"`
}

type WSHub struct {
	mu             sync.RWMutex
	clients        map[*websocket.Conn]VisitorDot
	allowedOrigins []string
	isProd         bool
}

func NewWSHub(allowedOrigins []string, isProd bool) *WSHub {
	if isProd && len(allowedOrigins) == 0 {
		log.Println("WARN: WebSocket AllowedOrigins is empty in production — connections will be rejected")
	}
	return &WSHub{
		clients:        make(map[*websocket.Conn]VisitorDot),
		allowedOrigins: allowedOrigins,
		isProd:         isProd,
	}
}

func (h *WSHub) HandleWS(w http.ResponseWriter, r *http.Request) {
	upgrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			if len(h.allowedOrigins) == 0 {
				return !h.isProd
			}
			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 {
			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)
			}
		}
		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()
}