perler-beads-server/internal/server/websocket.go

package server

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"net/http"
	"net/url"
	"perler-beads-server/internal/domain"
	"strings"
	"sync"
	"time"
)

const (
	wsAuthTimeout  = 5 * time.Second
	wsHeartbeat    = 20 * time.Second
	wsReadTimeout  = 45 * time.Second
	wsWriteTimeout = 5 * time.Second
	wsRedisTimeout = 3 * time.Second
	wsMessageLimit = 4096
)

type wsClose struct {
	code   int
	reason string
}
type taskSocket struct {
	conn    *websocket.Conn
	user    domain.User
	token   string
	changed chan struct{}
	pong    chan string
	stop    chan wsClose
}
type taskHub struct {
	mu     sync.Mutex
	stores map[string]map[*taskSocket]struct{}
	closed bool
}

func newTaskHub() *taskHub { return &taskHub{stores: map[string]map[*taskSocket]struct{}{}} }
func (h *taskHub) add(client *taskSocket) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.closed {
		return false
	}
	if h.stores[client.user.StoreID] == nil {
		h.stores[client.user.StoreID] = map[*taskSocket]struct{}{}
	}
	h.stores[client.user.StoreID][client] = struct{}{}
	return true
}
func (h *taskHub) remove(client *taskSocket) {
	h.mu.Lock()
	defer h.mu.Unlock()
	store := h.stores[client.user.StoreID]
	delete(store, client)
	if len(store) == 0 {
		delete(h.stores, client.user.StoreID)
	}
}

// Capacity one merges invalidations; slow browsers never block a committed business request.
func (h *taskHub) notify(storeID string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	for client := range h.stores[storeID] {
		select {
		case client.changed <- struct{}{}:
		default:
		}
	}
}
func (h *taskHub) shutdown() {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.closed = true
	for _, store := range h.stores {
		for client := range store {
			client.close(1001, "server shutting down")
		}
	}
}
func (c *taskSocket) close(code int, reason string) {
	select {
	case c.stop <- wsClose{code, reason}:
	default:
	}
}
func (s *Server) notifyTasks(storeID string) {
	if s.wsHub != nil {
		s.wsHub.notify(storeID)
	}
}
func (s *Server) CloseWebSockets() {
	if s.wsHub != nil {
		s.wsHub.shutdown()
	}
}
func (s *Server) wsOrigin(r *http.Request) bool {
	origin := r.Header.Get("Origin")
	if origin == "" || len(r.Header.Values("Origin")) != 1 {
		return false
	}
	u, e := url.Parse(origin)
	if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil || u.Path != "" || u.RawQuery != "" || u.Fragment != "" {
		return false
	}
	if strings.EqualFold(u.Host, r.Host) {
		return true
	}
	for _, allowed := range s.Origins {
		if allowed != "" && origin == allowed {
			return true
		}
	}
	return false
}
func (s *Server) wsUser(token string) (domain.User, error) {
	var u domain.User
	if token == "" || len(token) > 256 || s.Redis == nil {
		return u, errors.New("invalid session")
	}
	ctx, cancel := context.WithTimeout(context.Background(), wsRedisTimeout)
	defer cancel()
	h := sha256.Sum256([]byte(token))
	b, e := s.Redis.Get(ctx, "session:"+hex.EncodeToString(h[:])).Bytes()
	if e != nil {
		return u, e
	}
	if e = json.Unmarshal(b, &u); e != nil {
		return u, e
	}
	if u.ID == "" || u.StoreID == "" || (u.Role != "admin" && u.Role != "staff") {
		return u, errors.New("staff session required")
	}
	return u, nil
}
func (s *Server) websocket(c *gin.Context) {
	if !s.wsOrigin(c.Request) {
		c.Status(http.StatusForbidden)
		return
	}
	if c.Request.URL.RawQuery != "" {
		c.Status(http.StatusBadRequest)
		return
	}
	upgrader := websocket.Upgrader{HandshakeTimeout: wsAuthTimeout, ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: s.wsOrigin}
	conn, e := upgrader.Upgrade(c.Writer, c.Request, nil)
	if e != nil {
		return
	}
	conn.SetReadLimit(wsMessageLimit)
	conn.SetReadDeadline(time.Now().Add(wsAuthTimeout))
	conn.SetCloseHandler(func(code int, text string) error { return &websocket.CloseError{Code: code, Text: text} })
	var auth struct {
		Type  string `json:"type"`
		Token string `json:"token"`
	}
	kind, payload, e := conn.ReadMessage()
	if e != nil {
		conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(4403, "authentication required"), time.Now().Add(wsWriteTimeout))
		conn.Close()
		return
	}
	if kind != websocket.TextMessage || json.Unmarshal(payload, &auth) != nil || auth.Type != "auth" {
		conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(4403, "authentication required"), time.Now().Add(wsWriteTimeout))
		conn.Close()
		return
	}
	u, e := s.wsUser(auth.Token)
	if e != nil {
		conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(4403, "invalid staff session"), time.Now().Add(wsWriteTimeout))
		conn.Close()
		return
	}
	client := &taskSocket{conn: conn, user: u, token: auth.Token, changed: make(chan struct{}, 1), pong: make(chan string, 4), stop: make(chan wsClose, 1)}
	if !s.wsHub.add(client) {
		conn.Close()
		return
	}
	defer s.wsHub.remove(client)
	readerDone := make(chan struct{})
	writerDone := make(chan struct{})
	conn.SetReadDeadline(time.Now().Add(wsReadTimeout))
	conn.SetPongHandler(func(string) error { return conn.SetReadDeadline(time.Now().Add(wsReadTimeout)) })
	conn.SetPingHandler(func(data string) error {
		select {
		case client.pong <- data:
			return nil
		default:
			return errors.New("too many pings")
		}
	})
	go func() { defer close(writerDone); s.writeSocket(client, readerDone) }()
	defer func() { close(readerDone); <-writerDone }()
	for {
		kind, payload, e = conn.ReadMessage()
		if e != nil {
			return
		}
		var message struct {
			Type string `json:"type"`
		}
		if kind != websocket.TextMessage || json.Unmarshal(payload, &message) != nil || message.Type != "pong" {
			client.close(4400, "unexpected message")
			return
		}
	}
}
func (s *Server) writeSocket(client *taskSocket, readerDone <-chan struct{}) {
	defer client.conn.Close()
	write := func(kind string) error {
		client.conn.SetWriteDeadline(time.Now().Add(wsWriteTimeout))
		return client.conn.WriteJSON(map[string]string{"type": kind})
	}
	closeFrame := func(reason wsClose) {
		client.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(reason.code, reason.reason), time.Now().Add(wsWriteTimeout))
	}
	if e := write("ready"); e != nil {
		return
	}
	ticker := time.NewTicker(wsHeartbeat)
	defer ticker.Stop()
	for {
		select {
		case reason := <-client.stop:
			closeFrame(reason)
			return
		case <-readerDone:
			select {
			case reason := <-client.stop:
				closeFrame(reason)
			default:
				closeFrame(wsClose{1000, "connection closed"})
			}
			return
		case data := <-client.pong:
			if e := client.conn.WriteControl(websocket.PongMessage, []byte(data), time.Now().Add(wsWriteTimeout)); e != nil {
				return
			}
		case <-client.changed:
			if e := write("tasks.changed"); e != nil {
				return
			}
		case <-ticker.C:
			u, e := s.wsUser(client.token)
			if e != nil || u.ID != client.user.ID || u.Role != client.user.Role || u.StoreID != client.user.StoreID {
				closeFrame(wsClose{4403, "session expired"})
				return
			}
			if e := client.conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(wsWriteTimeout)); e != nil {
				return
			}
			if e := write("heartbeat"); e != nil {
				return
			}
		}
	}
}