gochat/internal/ws/auth.go

package ws

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/gin-gonic/gin"

	"github.com/gochat/gochat/internal/auth"
	"github.com/gochat/gochat/internal/model"
	"github.com/gochat/gochat/internal/repository"
	"github.com/gochat/gochat/pkg/logger"
)

// WSClaims represents authenticated WebSocket connection claims.
// Extends auth.Claims with PubsubToken for contact-based auth
// (mirrors Chatwoot's RoomChannel where contacts connect via pubsub_token).
type WSClaims struct {
	UserID       uint   `json:"user_id"`
	AccountID    uint   `json:"account_id"`
	Role         string `json:"role"`
	Provider     string `json:"provider"`
	PubsubToken  string `json:"pubsub_token,omitempty"` // contact auth token (Chatwoot RoomChannel)
	IsContact    bool   `json:"is_contact"`              // true when authenticated via pubsub_token
	ContactID    uint   `json:"contact_id,omitempty"`    // resolved contact ID for contact auth
	InboxID      uint   `json:"inbox_id,omitempty"`      // resolved inbox ID for contact auth
}

// WSAuthenticator handles WebSocket authentication and authorization.
// Reference: Chatwoot ActionCable RoomChannel — authenticates both
// agent users (via JWT) and contacts (via pubsub_token).
type WSAuthenticator struct {
	jwtService         *auth.JWTService
	contactInboxRepo   *repository.ContactInboxRepo
}

// NewWSAuthenticator creates a new WebSocket authenticator.
func NewWSAuthenticator(jwtService *auth.JWTService, contactInboxRepo *repository.ContactInboxRepo) *WSAuthenticator {
	return &WSAuthenticator{
		jwtService:       jwtService,
		contactInboxRepo: contactInboxRepo,
	}
}

// Authenticate validates WebSocket upgrade request parameters and returns WSClaims.
// It supports two authentication paths (mirrors Chatwoot's RoomChannel):
//
//  1. Agent/User auth (primary): JWT token from 'token' query param or Authorization header.
//     Validates via jwtService.ValidateAccessToken, populates WSClaims from auth.Claims.
//
//  2. Contact auth (secondary): pubsub_token + user_id query params.
//     Looks up ContactInbox by pubsub_token, verifies the contact belongs to the account,
//     populates WSClaims with contact identity.
//
// Returns WSClaims on success, or an error suitable for HTTP 401 rejection.
func (a *WSAuthenticator) Authenticate(c *gin.Context) (*WSClaims, error) {
	// --- Path 1: Agent/User authentication via JWT ---
	token := extractWSToken(c)
	if token != "" {
		claims, err := a.jwtService.ValidateAccessToken(token)
		if err != nil {
			logger.L().Debugf("ws auth: JWT validation failed: %v", err)
			return nil, fmt.Errorf("invalid JWT token: %w", err)
		}

		wsClaims := &WSClaims{
			UserID:    claims.UserID,
			AccountID: claims.AccountID,
			Role:      claims.Role,
			Provider:  claims.Provider,
			IsContact: false,
		}

		// Also extract pubsub_token and user_id if present (for dual auth context)
		pubsubToken := c.Query("pubsub_token")
		if pubsubToken != "" {
			wsClaims.PubsubToken = pubsubToken
		}

		logger.L().Infof("ws auth: agent authenticated (user_id=%d, account_id=%d, role=%s)",
			wsClaims.UserID, wsClaims.AccountID, wsClaims.Role)
		return wsClaims, nil
	}

	// --- Path 2: Contact authentication via pubsub_token ---
	pubsubToken := c.Query("pubsub_token")
	if pubsubToken == "" {
		return nil, errors.New("authentication required: provide 'token' (JWT) or 'pubsub_token' + 'user_id' params")
	}

	// Contact auth requires user_id param (Chatwoot RoomChannel: contact_id from params)
	userIDStr := c.Query("user_id")
	if userIDStr == "" {
		return nil, errors.New("contact auth requires 'user_id' parameter alongside 'pubsub_token'")
	}

	contactID, err := strconv.ParseUint(userIDStr, 10, 32)
	if err != nil {
		return nil, fmt.Errorf("invalid user_id parameter: %w", err)
	}

	// Lookup ContactInbox by pubsub_token
	contactInbox, err := a.findContactInboxByPubsubToken(c.Request.Context(), pubsubToken)
	if err != nil {
		logger.L().Debugf("ws auth: contact inbox lookup failed for pubsub_token: %v", err)
		return nil, fmt.Errorf("invalid pubsub_token: %w", err)
	}

	// Verify the contact ID matches
	if contactInbox.ContactID != uint(contactID) {
		logger.L().Debugf("ws auth: contact mismatch (expected=%d, found=%d)", uint(contactID), contactInbox.ContactID)
		return nil, errors.New("pubsub_token does not belong to the specified contact")
	}

	wsClaims := &WSClaims{
		UserID:      uint(contactID), // for contacts, UserID maps to contact_id (Chatwoot convention)
		AccountID:   contactInbox.Contact.AccountID,
		Role:        "contact",
		Provider:    "pubsub_token",
		PubsubToken: pubsubToken,
		IsContact:   true,
		ContactID:   contactInbox.ContactID,
		InboxID:     contactInbox.InboxID,
	}

	logger.L().Infof("ws auth: contact authenticated (contact_id=%d, account_id=%d, inbox_id=%d)",
		wsClaims.ContactID, wsClaims.AccountID, wsClaims.InboxID)
	return wsClaims, nil
}

// Authorize verifies the authenticated user/contact has access to the requested account.
// For agent auth: verifies the account_id param matches the JWT claims' AccountID.
// For contact auth: verifies the contact's account matches the requested account_id.
//
// Call this after Authenticate succeeds, before upgrading the WebSocket connection.
// Returns nil on success, or an error suitable for HTTP 403 rejection.
func (a *WSAuthenticator) Authorize(claims *WSClaims, c *gin.Context) error {
	// Extract requested account_id from query params
	accountIDStr := c.Query("account_id")
	if accountIDStr == "" {
		// If no explicit account_id requested, use the one from claims
		return nil
	}

	requestedAccountID, err := strconv.ParseUint(accountIDStr, 10, 32)
	if err != nil {
		return fmt.Errorf("invalid account_id parameter: %w", err)
	}

	if claims.AccountID != uint(requestedAccountID) {
		logger.L().Debugf("ws auth: account mismatch (claims=%d, requested=%d)",
			claims.AccountID, uint(requestedAccountID))
		return fmt.Errorf("user does not have access to account %d", requestedAccountID)
	}

	return nil
}

// AuthenticateAndServeWS is a combined authentication + upgrade handler.
// It authenticates the request, authorizes access, then upgrades to WebSocket.
// Rejects with HTTP 401 if authentication fails, 403 if authorization fails.
//
// Usage: register as a Gin handler for the WebSocket upgrade route.
//   router.GET("/ws", wsAuth.AuthenticateAndServeWS(hub, upgrader, onConnect))
func (a *WSAuthenticator) AuthenticateAndServeWS(c *gin.Context) {
	// Step 1: Authenticate
	claims, err := a.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
	if err := a.Authorize(claims, c); err != nil {
		logger.L().Errorf("ws: authorization failed: %v", err)
		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
		return
	}

	// Authentication and authorization succeeded — set claims in context
	// for downstream handler (e.g. the actual upgrade handler) to use.
	c.Set("ws_claims", claims)
	c.Set("ws_authenticated", true)
	c.Next()
}

// findContactInboxByPubsubToken looks up a ContactInbox by its PubsubToken field.
// Preloads the associated Contact to resolve the AccountID for authorization.
func (a *WSAuthenticator) findContactInboxByPubsubToken(ctx context.Context, pubsubToken string) (*model.ContactInbox, error) {
	return a.contactInboxRepo.FindByPubsubToken(ctx, pubsubToken)
}

// extractWSToken pulls the JWT token from websocket upgrade request.
// Order of precedence:
//   1. 'token' query parameter (browser WebSocket API can't set custom headers)
//   2. Authorization header (Bearer token, for non-browser clients)
//   3. Sec-WebSocket-Protocol header (some ActionCable-compatible clients)
func extractWSToken(c *gin.Context) string {
	// Primary: 'token' query param
	token := c.Query("token")
	if token != "" {
		return token
	}

	// Fallback: Authorization header
	authHeader := c.GetHeader("Authorization")
	if authHeader != "" {
		parts := strings.SplitN(authHeader, " ", 2)
		if len(parts) == 2 && strings.ToLower(parts[0]) == "bearer" {
			return parts[1]
		}
	}

	// Fallback: Sec-WebSocket-Protocol header (ActionCable convention)
	proto := c.GetHeader("Sec-WebSocket-Protocol")
	if proto != "" {
		token = strings.TrimSpace(proto)
		if token != "" {
			return token
		}
	}

	return ""
}

// ParseWSQueryParams extracts all WebSocket-relevant query parameters
// from the upgrade request URL. Useful for subscription authorization
// after initial authentication.
// Reference: Chatwoot RoomChannel subscribe params — account_id, conversation_id, etc.
func ParseWSQueryParams(query url.Values) map[string]string {
	params := make(map[string]string)
	for _, key := range []string{
		"account_id",
		"conversation_id",
		"inbox_id",
		"pubsub_token",
		"user_id",
		"token",
	} {
		if v := query.Get(key); v != "" {
			params[key] = v
		}
	}
	return params
}