ruby-sea-battle/game-server/internal/ws/handler.go

// Package ws implements the websocket endpoint of the game server:
// connection upgrade, JOIN handling and the per-connection read loop.
package ws

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

	"github.com/coder/websocket"

	"github.com/ruby-sea-battle/game-server/internal/auth"
	"github.com/ruby-sea-battle/game-server/internal/game"
	"github.com/ruby-sea-battle/game-server/internal/protocol"
)

const (
	readLimit       = 16 * 1024 // game messages are tiny
	joinReadTimeout = 10 * time.Second
	writeTimeout    = 5 * time.Second
	pingInterval    = 25 * time.Second
	pingTimeout     = 10 * time.Second
)

// Conn adapts a websocket connection to the game.Sender interface.
type Conn struct {
	ws  *websocket.Conn
	wmu sync.Mutex // serializes writes
}

func (c *Conn) Send(_ context.Context, data []byte) error {
	ctx, cancel := context.WithTimeout(context.Background(), writeTimeout)
	defer cancel()
	c.wmu.Lock()
	defer c.wmu.Unlock()
	return c.ws.Write(ctx, websocket.MessageText, data)
}

func (c *Conn) Close(code int, reason string) error {
	return c.ws.Close(websocket.StatusCode(code), reason)
}

// Handler is the http.Handler of the websocket endpoint.
type Handler struct {
	Manager    *game.Manager
	JWTSecret  []byte
	PublicAddr string // must match the "game_server" claim strictly

	acceptOpts *websocket.AcceptOptions
	log        *slog.Logger
}

func NewHandler(manager *game.Manager, jwtSecret []byte, publicAddr string, allowedOrigins []string, log *slog.Logger) *Handler {
	opts := &websocket.AcceptOptions{}
	if len(allowedOrigins) == 1 && allowedOrigins[0] == "*" {
		opts.InsecureSkipVerify = true
	} else {
		opts.OriginPatterns = allowedOrigins
	}
	return &Handler{
		Manager:    manager,
		JWTSecret:  jwtSecret,
		PublicAddr: publicAddr,
		acceptOpts: opts,
		log:        log,
	}
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	wsc, err := websocket.Accept(w, r, h.acceptOpts)
	if err != nil {
		h.log.Debug("websocket accept failed", "err", err)
		return
	}
	wsc.SetReadLimit(readLimit)

	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()
	defer wsc.Close(websocket.StatusInternalError, "")

	// The first message on a fresh connection must be JOIN.
	readCtx, cancelRead := context.WithTimeout(ctx, joinReadTimeout)
	_, data, err := wsc.Read(readCtx)
	cancelRead()
	if err != nil {
		wsc.Close(protocol.CloseCodeExpectedJoin, "join timeout")
		return
	}

	opcode, args, err := protocol.ParseMessage(data)
	if err != nil {
		wsc.Close(protocol.CloseCodeBadMessage, "malformed message")
		return
	}
	if opcode != protocol.OpJoin {
		wsc.Close(protocol.CloseCodeExpectedJoin, "expected JOIN")
		return
	}
	gameToken, userToken, err := protocol.ParseJoin(args)
	if err != nil {
		wsc.Close(protocol.CloseCodeBadMessage, "malformed JOIN")
		return
	}

	// Validate the game token and make sure the token was issued for this
	// particular server instance (strict comparison with the public address).
	gc, err := auth.ParseGameToken(h.JWTSecret, gameToken)
	if err != nil {
		h.log.Debug("bad game token", "err", err)
		wsc.Close(protocol.CloseCodeBadGameToken, "invalid game token")
		return
	}
	if gc.GameServer != h.PublicAddr {
		h.log.Debug("game token issued for another server", "want", h.PublicAddr, "got", gc.GameServer)
		wsc.Close(protocol.CloseCodeBadGameToken, "game server mismatch")
		return
	}
	if gc.Player1ID == gc.Player2ID {
		wsc.Close(protocol.CloseCodeBadGameToken, "invalid match")
		return
	}

	// Validate the user token and make sure the user is a participant.
	uc, err := auth.ParseUserToken(h.JWTSecret, userToken)
	if err != nil {
		h.log.Debug("bad user token", "err", err)
		wsc.Close(protocol.CloseCodeBadUserToken, "invalid user token")
		return
	}
	if uc.UserID != gc.Player1ID && uc.UserID != gc.Player2ID {
		wsc.Close(protocol.CloseCodeBadUserToken, "not a participant of this match")
		return
	}

	conn := &Conn{ws: wsc}
	match, slotIdx, err := h.Manager.Join(ctx, conn, gc, uc.UserID)
	if err != nil {
		// ErrMatchFinished: the END message has already been replayed and
		// the connection closed by the manager.
		if err != game.ErrMatchFinished {
			h.log.Debug("join failed", "err", err)
		}
		return
	}
	h.log.Info("player joined", "match", match.ID, "user", uc.UserID, "slot", slotIdx)

	go h.pingLoop(ctx, wsc)

	// Read loop: process messages until the connection dies.
	for {
		_, data, err := wsc.Read(ctx)
		if err != nil {
			break
		}
		opcode, args, err := protocol.ParseMessage(data)
		if err != nil {
			h.log.Debug("malformed message", "match", match.ID, "user", uc.UserID, "err", err)
			continue
		}
		switch opcode {
		case protocol.OpAck:
			seq, err := protocol.ParseAck(args)
			if err != nil {
				h.log.Debug("malformed ACK", "match", match.ID, "user", uc.UserID, "err", err)
				continue
			}
			h.Manager.Ack(match, slotIdx, seq)
		default:
			// Unknown opcodes are ignored for forward compatibility.
			h.log.Debug("unknown opcode", "match", match.ID, "user", uc.UserID, "opcode", opcode)
		}
	}

	h.Manager.Leave(conn, match, slotIdx)
	h.log.Info("player disconnected", "match", match.ID, "user", uc.UserID, "slot", slotIdx)
}

// pingLoop keeps the connection alive and detects dead peers. When the ping
// fails the connection is closed, which terminates the read loop.
func (h *Handler) pingLoop(ctx context.Context, wsc *websocket.Conn) {
	ticker := time.NewTicker(pingInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			pingCtx, cancel := context.WithTimeout(ctx, pingTimeout)
			err := wsc.Ping(pingCtx)
			cancel()
			if err != nil {
				wsc.Close(websocket.StatusGoingAway, "ping failed")
				return
			}
		}
	}
}