gochat/backend/internal/middleware/auth.go

package middleware

import (
	"net/http"
	"strconv"
	"strings"

	"github.com/gin-gonic/gin"

	"github.com/gochat/gochat/internal/auth"
	"github.com/gochat/gochat/internal/config"
	"github.com/gochat/gochat/internal/model"
	"gorm.io/gorm"
)

// AuthMiddleware validates JWT tokens on protected routes.
// Uses JWTService.ParseClaims to extract the typed Claims struct
// instead of raw MapClaims, ensuring Provider and CustomRoleID are
// available for downstream Policy middleware and permission checks.
// Reference: Chatwoot DeviseTokenAuth replacement (P2E §1.2).
//
// Auth flow:
//  1. Extract JWT from Authorization Bearer header → set user_id/account_id from claims
//  2. If no Authorization header, fall back to X-User-ID / X-Account-ID headers (dev/test mode)
//  3. Abort with 401 if neither source provides a valid identity
func AuthMiddleware(cfg *config.JWTConfig) gin.HandlerFunc {
	jwtSvc := auth.NewJWTService(cfg)
	return AuthMiddlewareWithService(jwtSvc)
}

// AuthMiddlewareWithService accepts a pre-built JWTService (for DI in tests/bootstrap).
func AuthMiddlewareWithService(jwtSvc *auth.JWTService) gin.HandlerFunc {
	return AuthMiddlewareWithServiceAndDB(jwtSvc, nil)
}

// AuthMiddlewareWithServiceAndDB additionally enforces Chatwoot client-session
// revocation for tokens carrying a client_id claim. Legacy/API JWTs without a
// client id retain the existing stateless behavior.
func AuthMiddlewareWithServiceAndDB(jwtSvc *auth.JWTService, db *gorm.DB) gin.HandlerFunc {
	return func(c *gin.Context) {
		authHeader := c.GetHeader("Authorization")
		chatwootAccessToken := strings.TrimSpace(c.GetHeader("access-token"))
		if authHeader != "" || chatwootAccessToken != "" {
			tokenString := strings.TrimPrefix(authHeader, "Bearer ")
			if authHeader == "" {
				tokenString = chatwootAccessToken
			} else if tokenString == authHeader {
				c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "bearer token required"})
				return
			}

			claims, _, err := auth.ValidateUserAccessToken(c.Request.Context(), jwtSvc, db, tokenString)
			if err != nil {
				c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "invalid token"})
				return
			}

			// Set typed claims values in Gin context for downstream middleware/handlers.
			// PolicyMiddleware reads these to build PolicyContext.
			c.Set("user_id", claims.UserID)
			c.Set("account_id", claims.AccountID)
			c.Set("role", claims.Role)
			c.Set("user_type", claims.UserType)
			c.Set("provider", claims.Provider)
			c.Set("custom_role_id", claims.CustomRoleID)
			c.Set("claims", claims) // full Claims struct for handlers that need it
			c.Next()
			return
		}

		// Fallback: X-User-ID / X-Account-ID headers for development and testing.
		// This allows integration tests and dev environments to bypass JWT while
		// still exercising the same middleware → handler pipeline.
		headerUserID := c.GetHeader("X-User-ID")
		if headerUserID != "" {
			userID, err := strconv.ParseUint(headerUserID, 10, 32)
			if err != nil {
				c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "invalid X-User-ID header"})
				return
			}
			c.Set("user_id", uint(userID))

			// Also extract account_id from X-Account-ID header if present
			headerAccountID := c.GetHeader("X-Account-ID")
			if headerAccountID != "" {
				accountID, err := strconv.ParseUint(headerAccountID, 10, 32)
				if err == nil && accountID != 0 {
					c.Set("account_id", uint(accountID))
				}
			}

			// Set defaults for role/provider so downstream PolicyMiddleware doesn't break
			c.Set("role", "agent")
			c.Set("provider", "dev_header")
			c.Set("custom_role_id", uint(0))
			c.Next()
			return
		}

		// No Authorization header and no X-User-ID header → unauthorized
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "authorization header required"})
	}
}

// GenerateToken creates a JWT token for a user.
// DEPRECATED: Prefer JWTService.GenerateTokenPair which produces proper typed Claims.
// Kept for backward compatibility with existing test helpers.
func GenerateToken(cfg *config.JWTConfig, userID uint, accountID uint, role string) (string, error) {
	pair, err := auth.NewJWTService(cfg).GenerateTokenPair(&model.User{
		Base:     model.Base{ID: userID},
		Provider: "email",
	}, accountID, role)
	if err != nil {
		return "", err
	}
	return pair.AccessToken, nil
}