veza/veza-backend-api/internal/handlers/chat_websocket_handler.go

package handlers

import (
	"time"

	"github.com/coder/websocket"
	"github.com/gin-gonic/gin"
	"go.uber.org/zap"

	apperrors "veza-backend-api/internal/errors"
	"veza-backend-api/internal/services"
	chatws "veza-backend-api/internal/websocket/chat"
)

type ChatWebSocketHandler struct {
	chatService *services.ChatService
	hub         *chatws.Hub
	handler     *chatws.MessageHandler
	logger      *zap.Logger
}

func NewChatWebSocketHandler(chatService *services.ChatService, hub *chatws.Hub, handler *chatws.MessageHandler, logger *zap.Logger) *ChatWebSocketHandler {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &ChatWebSocketHandler{
		chatService: chatService,
		hub:         hub,
		handler:     handler,
		logger:      logger,
	}
}

func (h *ChatWebSocketHandler) HandleWebSocket(c *gin.Context) {
	tokenString := c.Query("token")
	if tokenString == "" {
		RespondWithAppError(c, apperrors.NewUnauthorizedError("missing token"))
		return
	}

	claims, err := h.chatService.ValidateChatToken(tokenString)
	if err != nil {
		h.logger.Warn("Invalid chat token",
			zap.Error(err))
		RespondWithAppError(c, apperrors.NewUnauthorizedError("invalid or expired token"))
		return
	}

	conn, err := websocket.Accept(c.Writer, c.Request, &websocket.AcceptOptions{
		OriginPatterns: GetAllowedWebSocketOrigins(),
	})
	if err != nil {
		h.logger.Error("Failed to accept WebSocket",
			zap.Error(err),
			zap.String("user_id", claims.UserID.String()))
		return
	}

	client := chatws.NewClient(h.hub, conn, claims.UserID, claims.Username, h.handler, h.logger)

	h.hub.Register(client)

	client.SendJSON(chatws.NewActionConfirmedResponse("connected", true))

	ctx := c.Request.Context()

	// SECURITY(MEDIUM-010): Periodic token re-validation.
	// Re-validate every 60s; close connection if token is expired or revoked.
	go func() {
		ticker := time.NewTicker(60 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				if _, err := h.chatService.ValidateChatToken(tokenString); err != nil {
					h.logger.Info("WebSocket token expired, closing connection",
						zap.String("user_id", claims.UserID.String()))
					conn.Close(websocket.StatusPolicyViolation, "token expired")
					return
				}
			}
		}
	}()

	go client.WritePump(ctx)
	go client.ReadPump(ctx)
}