photoplaces/backend/internal/handlers/websocket.go

package handlers

import (
	"encoding/json"
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

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

type visitorsPayload struct {
	Visitors []VisitorDot `json:"visitors"`
}

type WSHub struct {
	clients        map[*Client]bool
	register       chan *Client
	unregister     chan *Client
	broadcast      chan []byte
	allowedOrigins []string
	isProd         bool
}

func NewWSHub(allowedOrigins []string, isProd bool) *WSHub {
	return &WSHub{
		clients:        make(map[*Client]bool),
		register:       make(chan *Client),
		unregister:     make(chan *Client),
		broadcast:      make(chan []byte, 256),
		allowedOrigins: allowedOrigins,
		isProd:         isProd,
	}
}

func (h *WSHub) Run() {
	for {
		select {
		case client := <-h.register:
			h.clients[client] = true
			h.broadcastVisitors()

		case client := <-h.unregister:
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				close(client.send)
				h.broadcastVisitors()
			}

		case message := <-h.broadcast:
			for client := range h.clients {
				select {
				case client.send <- message:
				default:
					close(client.send)
					delete(h.clients, client)
				}
			}
		}
	}
}

func (h *WSHub) broadcastVisitors() {
	visitors := make([]VisitorDot, 0, len(h.clients))
	for c := range h.clients {
		visitors = append(visitors, VisitorDot{
			UserID: c.userID,
			Lat:    c.lat,
			Lng:    c.lng,
		})
	}

	data, err := json.Marshal(visitorsPayload{Visitors: visitors})
	if err != nil {
		return
	}
	msg, err := json.Marshal(wsMessage{Type: "visitors", Data: data})
	if err != nil {
		return
	}

	for client := range h.clients {
		select {
		case client.send <- msg:
		default:
			close(client.send)
			delete(h.clients, client)
		}
	}
}

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
	}

	client := &Client{
		hub:  h,
		conn: conn,
		send: make(chan []byte, 256),
	}

	h.register <- client

	go client.writePump()
	go client.readPump()
}