gochat/backend/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"
	"gorm.io/gorm"
)

// 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"`
	ClientID    string `json:"client_id,omitempty"`
	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
	db               *gorm.DB
}

// NewWSAuthenticator creates a new WebSocket authenticator.
func NewWSAuthenticator(jwtService *auth.JWTService, contactInboxRepo *repository.ContactInboxRepo, db ...*gorm.DB) *WSAuthenticator {
	authenticator := &WSAuthenticator{
		jwtService:       jwtService,
		contactInboxRepo: contactInboxRepo,
	}
	if len(db) > 0 {
		authenticator.db = db[0]
	}
	return authenticator
}

// 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 := auth.ValidateUserAccessToken(c.Request.Context(), a.jwtService, a.db, 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,
			ClientID:  claims.ClientID,
			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
}

// ValidateAgentAccess rechecks mutable access state for a live agent socket.
func (a *WSAuthenticator) ValidateAgentAccess(ctx context.Context, userID uint, clientID string) error {
	_, err := auth.ValidateUserAccess(ctx, a.db, userID, clientID)
	return err
}

// 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.
// Matches the HTTP auth middleware behaviour: accepts both 'token' and
// 'access-token' query params (browser WebSocket API can't set custom headers),
// plus the Authorization header (Bearer token, for non-browser clients).
//
// NOTE: Sec-WebSocket-Protocol header is NOT used as a JWT source.
//
//	ActionCable sets this to "actioncable-v1-json" for sub-protocol
//	negotiation, not for authentication.
func extractWSToken(c *gin.Context) string {
	// Query params (browser WebSocket API compatible)
	if token := c.Query("token"); token != "" {
		return token
	}
	if token := c.Query("access-token"); token != "" {
		return token
	}

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

	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
}