gochat/backend/internal/handler/ws/handler.go

package ws

import (
	"encoding/json"
	"net/http"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"

	wspkg "github.com/gochat/gochat/internal/ws"
	"github.com/gochat/gochat/pkg/logger"
)

// uintToStr converts a uint to its string representation.
// Kept for potential use in identifier formatting.
func uintToStr(u uint) string {
	return strconv.FormatUint(uint64(u), 10)
}

// Handler manages WebSocket connections, upgrade, JWT auth, and message dispatch.
// Reference: Chatwoot ActionCable — WebSocket upgrade at /cable, JWT-based auth,
// room-based subscription model (AccountChannel, ConversationChannel).
//
// Supports two authentication paths:
//   1. Agent/User auth: JWT token (from query param or Authorization header)
//   2. Contact auth: pubsub_token + user_id (Chatwoot RoomChannel pattern)
type Handler struct {
	hub         *Hub
	authenticator *wspkg.WSAuthenticator
	upgrader    websocket.Upgrader
}

// NewHandler creates a WebSocket handler with the given hub and authenticator.
// The authenticator provides both JWT (agent) and pubsub_token (contact) auth paths.
func NewHandler(hub *Hub, authenticator *wspkg.WSAuthenticator) *Handler {
	return &Handler{
		hub:         hub,
		authenticator: authenticator,
		upgrader: websocket.Upgrader{
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
			// Allow all origins — CORS is handled at the Gin middleware layer
			CheckOrigin: func(r *http.Request) bool { return true },
		},
	}
}

// ServeWS handles the WebSocket upgrade request at /ws.
// URL: /ws?token=<JWT_ACCESS_TOKEN>  OR  /ws?pubsub_token=<TOKEN>&user_id=<ID>
// On successful upgrade, the handler:
//  1. Authenticates the request via WSAuthenticator (JWT or pubsub_token)
//  2. Authorizes the user/contact for the requested account
//  3. Upgrades HTTP connection to WebSocket
//  4. Creates a Client and registers it with the Hub
//  5. Starts readPump (incoming commands) and writePump (outgoing events) goroutines
func (h *Handler) ServeWS(c *gin.Context) {
	// Step 1: Authenticate (JWT or pubsub_token)
	claims, err := h.authenticator.Authenticate(c)
	if err != nil {
		logger.L().Errorf("ws: authentication failed: %v", err)
		c.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
		return
	}

	// Step 2: Authorize (verify account access)
	if err := h.authenticator.Authorize(claims, c); err != nil {
		logger.L().Errorf("ws: authorization failed: %v", err)
		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
		return
	}

	// Step 3: Upgrade HTTP connection to WebSocket
	conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logger.L().Errorf("ws: upgrade failed: %v", err)
		return
	}

	// Step 4: Create client and register with hub
	client := NewClient(claims.UserID, claims.AccountID, conn, h.hub)
	// Populate extended claims fields on the client
	client.Role = claims.Role
	client.IsContact = claims.IsContact
	client.PubsubToken = claims.PubsubToken
	client.ContactID = claims.ContactID
	client.InboxID = claims.InboxID

	h.hub.Register(client)

	logger.L().Infof("ws: connection established (user=%d, account=%d, is_contact=%v)", claims.UserID, claims.AccountID, claims.IsContact)

	// Start pumps in separate goroutines
	go h.writePump(client)
	go h.readPump(client)
}

// ServeCable handles the WebSocket upgrade request at /cable.
// This is the ActionCable-compatible endpoint (Chatwoot uses /cable for WS).
// It is functionally identical to ServeWS but uses the ActionCable naming convention.
// Some Chatwoot clients specifically connect to /cable.
func (h *Handler) ServeCable(c *gin.Context) {
	h.ServeWS(c)
}

// readPump reads messages from the WebSocket connection and dispatches commands.
// Reference: Chatwoot ActionCable consumer — processes subscribe/unsubscribe/ping commands.
// One readPump runs per client connection.
func (h *Handler) readPump(client *Client) {
	defer func() {
		h.hub.Unregister(client)
		client.Conn.Close()
	}()

	client.Conn.SetReadLimit(MaxMessageSize)
	client.Conn.SetReadDeadline(time.Now().Add(PongWait))
	client.Conn.SetPongHandler(func(string) error {
		client.Conn.SetReadDeadline(time.Now().Add(PongWait))
		return nil
	})

	for {
		_, message, err := client.Conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				logger.L().Errorf("ws: unexpected close for user=%d: %v", client.UserID, err)
			}
			break
		}

		// Decode the command frame
		var cmd CommandFrame
		if err := json.Unmarshal(message, &cmd); err != nil {
			logger.L().Warnf("ws: invalid command from user=%d: %v", client.UserID, err)
			continue
		}

		// Dispatch command
		switch cmd.Command {
		case CommandSubscribe:
			h.handleSubscribe(client, cmd)
		case CommandUnsubscribe:
			h.handleUnsubscribe(client, cmd)
		case CommandPing:
			h.handlePing(client)
		default:
			logger.L().Warnf("ws: unknown command '%s' from user=%d", cmd.Command, client.UserID)
		}
	}
}

// writePump sends messages from the client's Send channel to the WebSocket connection.
// It also sends periodic ping frames for heartbeat detection.
// One writePump runs per client connection.
func (h *Handler) writePump(client *Client) {
	ticker := time.NewTicker(time.Duration(PingInterval) * time.Second)
	defer func() {
		ticker.Stop()
		client.Conn.Close()
	}()

	for {
		select {
		case message, ok := <-client.Send:
			client.Conn.SetWriteDeadline(time.Now().Add(WriteWait))
			if !ok {
				// Hub closed the channel — send close frame
				client.Conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			// Write text message (all our frames are JSON text)
			if err := client.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
				logger.L().Errorf("ws: write error for user=%d: %v", client.UserID, err)
				return
			}

		case <-ticker.C:
			// Send ping frame for heartbeat
			client.Conn.SetWriteDeadline(time.Now().Add(WriteWait))
			if err := client.Conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				logger.L().Errorf("ws: ping write failed for user=%d: %v", client.UserID, err)
				return
			}
		}
	}
}

// handleSubscribe processes a subscribe command.
// Validates the ChannelIdentifier and adds the client to the appropriate room.
func (h *Handler) handleSubscribe(client *Client, cmd CommandFrame) {
	var identifier ChannelIdentifier
	if err := json.Unmarshal([]byte(cmd.Identifier), &identifier); err != nil {
		logger.L().Warnf("ws: invalid identifier from user=%d: %v", client.UserID, err)
		// Send reject frame
		rejectData, _ := json.Marshal(RejectFrame{
			Type:       ServerRejectSubscribe,
			Identifier: cmd.Identifier,
			Reason:     "invalid identifier format",
		})
		client.Send <- rejectData
		return
	}

	// Validate account_id matches the client's authenticated account
	if identifier.AccountID != client.AccountID {
		rejectData, _ := json.Marshal(RejectFrame{
			Type:       ServerRejectSubscribe,
			Identifier: cmd.Identifier,
			Reason:     "account_id mismatch",
		})
		client.Send <- rejectData
		return
	}

	// Determine room name based on channel type (uses Hub's canonical naming)
	room := ""
	switch identifier.Channel {
	case ChannelAccount:
		room = accountRoomName(identifier.AccountID)
	case ChannelConversation:
		if identifier.ConversationID == 0 {
			rejectData, _ := json.Marshal(RejectFrame{
				Type:       ServerRejectSubscribe,
				Identifier: cmd.Identifier,
				Reason:     "conversation_id required for ConversationChannel",
			})
			client.Send <- rejectData
			return
		}
		room = conversationRoomName(identifier.AccountID, identifier.ConversationID)
	default:
		rejectData, _ := json.Marshal(RejectFrame{
			Type:       ServerRejectSubscribe,
			Identifier: cmd.Identifier,
			Reason:     "unknown channel type: " + identifier.Channel,
		})
		client.Send <- rejectData
		return
	}

	// Subscribe the client to the room
	client.Subscribe(room)

	// Send confirmation frame
	confirmData, _ := json.Marshal(ConfirmFrame{
		Type:       ServerConfirmSubscribe,
		Identifier: cmd.Identifier,
	})
	client.Send <- confirmData
}

// handleUnsubscribe processes an unsubscribe command.
func (h *Handler) handleUnsubscribe(client *Client, cmd CommandFrame) {
	var identifier ChannelIdentifier
	if err := json.Unmarshal([]byte(cmd.Identifier), &identifier); err != nil {
		logger.L().Warnf("ws: invalid identifier in unsubscribe from user=%d: %v", client.UserID, err)
		return
	}

	// Determine room name (uses Hub's canonical naming)
	room := ""
	switch identifier.Channel {
	case ChannelAccount:
		room = accountRoomName(identifier.AccountID)
	case ChannelConversation:
		room = conversationRoomName(identifier.AccountID, identifier.ConversationID)
	default:
		return
	}

	client.Unsubscribe(room)

	confirmData, _ := json.Marshal(ConfirmFrame{
		Type:       ServerConfirmUnsubscribe,
		Identifier: cmd.Identifier,
	})
	client.Send <- confirmData
}

// handlePing responds to a client-initiated ping command with a pong frame.
func (h *Handler) handlePing(client *Client) {
	pingData, _ := json.Marshal(PingFrame{
		Type:    ServerPing,
		Message: time.Now().UTC().Format(time.RFC3339),
	})
	client.Send <- pingData
}