core/server/middleware/auth.go

package middleware

import (
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt/v5"

	"armature/config"
	"armature/database"
	"armature/store"
)

// Claims represents the JWT payload. Must match handlers.Claims.
type Claims struct {
	UserID string `json:"user_id"`
	Email  string `json:"email"`
	jwt.RegisteredClaims
}

// ─── User Status Cache ──────────────────────────────────────
// Short-TTL cache so deactivation takes effect within seconds,
// without a DB query on every single request.

const userCacheTTL = 30 * time.Second

type userCacheEntry struct {
	isActive  bool
	fetchedAt time.Time
}

// UserStatusCache is a concurrency-safe TTL cache for user active/role state.
type UserStatusCache struct {
	mu      sync.RWMutex
	entries map[string]userCacheEntry
}

// NewUserStatusCache creates a cache and starts a background cleanup goroutine.
func NewUserStatusCache() *UserStatusCache {
	c := &UserStatusCache{entries: make(map[string]userCacheEntry)}
	go func() {
		for {
			time.Sleep(2 * time.Minute)
			c.evictExpired()
		}
	}()
	return c
}

func (c *UserStatusCache) get(userID string) (userCacheEntry, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	e, ok := c.entries[userID]
	if !ok || time.Since(e.fetchedAt) > userCacheTTL {
		return userCacheEntry{}, false
	}
	return e, true
}

func (c *UserStatusCache) set(userID string, active bool) {
	c.mu.Lock()
	c.entries[userID] = userCacheEntry{isActive: active, fetchedAt: time.Now()}
	c.mu.Unlock()
}

// Evict removes a specific user from the cache (call on role/active change).
func (c *UserStatusCache) Evict(userID string) {
	c.mu.Lock()
	delete(c.entries, userID)
	c.mu.Unlock()
}

func (c *UserStatusCache) evictExpired() {
	c.mu.Lock()
	defer c.mu.Unlock()
	cutoff := time.Now().Add(-2 * userCacheTTL)
	for k, v := range c.entries {
		if v.fetchedAt.Before(cutoff) {
			delete(c.entries, k)
		}
	}
}

// ─── Shared user verification ───────────────────────────────

// verifyUser checks that the user exists and is active.
// Returns true on success or false (with error already sent) on failure.
func verifyUser(c *gin.Context, claims *Claims, users store.UserStore, cache *UserStatusCache) bool {
	userID := claims.UserID
	if userID == "" {
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "invalid token (empty user_id)"})
		return false
	}
	return verifyUserByID(c, userID, users, cache)
}

// verifyUserByID checks is_active for a user ID.
// Used by both JWT auth (via verifyUser) and ticket auth.
func verifyUserByID(c *gin.Context, userID string, users store.UserStore, cache *UserStatusCache) bool {
	// Cache hit
	if entry, ok := cache.get(userID); ok {
		if !entry.isActive {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "account deactivated"})
			return false
		}
		return true
	}

	// Cache miss → DB lookup
	user, err := users.GetByID(c.Request.Context(), userID)
	if err != nil {
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "user not found"})
		return false
	}

	cache.set(userID, user.IsActive)

	if !user.IsActive {
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "account deactivated"})
		return false
	}
	return true
}

// parseAndValidateJWT parses a raw JWT string and returns claims if valid.
func parseAndValidateJWT(tokenString string, jwtSecret string) (*Claims, bool) {
	claims := &Claims{}
	token, err := jwt.ParseWithClaims(tokenString, claims, func(t *jwt.Token) (interface{}, error) {
		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, jwt.ErrSignatureInvalid
		}
		return []byte(jwtSecret), nil
	})
	if err != nil || !token.Valid {
		return nil, false
	}
	return claims, true
}

// UserIDFromCookie extracts the user ID from the sb_token cookie without
// requiring authentication. Returns "" if no valid token is found.
// Used by unauthenticated routes that want optional user context.
func UserIDFromCookie(c *gin.Context, jwtSecret string) string {
	cookie, err := c.Cookie("sb_token")
	if err != nil || cookie == "" {
		return ""
	}
	claims, ok := parseAndValidateJWT(cookie, jwtSecret)
	if !ok {
		return ""
	}
	return claims.UserID
}

// ─── Auth middleware ─────────────────────────────────────────

// Auth returns a Gin middleware that validates JWT bearer tokens and
// verifies the user is active with their current DB role.
func Auth(cfg *config.Config, users store.UserStore, cache *UserStatusCache) gin.HandlerFunc {
	return func(c *gin.Context) {
		// Skip auth when running without a database (unmanaged mode)
		if !database.IsConnected() {
			c.Next()
			return
		}

		header := c.GetHeader("Authorization")
		if header == "" {
			// WebSocket connections can't set headers from browser.
			// Fall back to ?token= query parameter.
			if qToken := c.Query("token"); qToken != "" {
				header = "Bearer " + qToken
			}
		}
		if header == "" {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"error": "missing authorization header",
			})
			return
		}

		tokenString := strings.TrimPrefix(header, "Bearer ")
		if tokenString == header {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"error": "invalid authorization format",
			})
			return
		}

		claims, ok := parseAndValidateJWT(tokenString, cfg.JWTSecret)
		if !ok {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"error": "invalid or expired token",
			})
			return
		}

		// Verify user is active
		if !verifyUser(c, claims, users, cache) {
			return
		}

		c.Set("user_id", claims.UserID)
		c.Set("email", claims.Email)
		c.Next()
	}
}

// ─── CORS middleware ─────────────────────────────────────────

// CORS returns a middleware that sets cross-origin headers.
// In production, restricts to the configured allowed origin(s).
// In development/test, allows all origins.
func CORS(cfg *config.Config) gin.HandlerFunc {
	allowedOrigin := getAllowedOrigin(cfg)

	return func(c *gin.Context) {
		origin := c.GetHeader("Origin")

		if allowedOrigin == "*" {
			c.Header("Access-Control-Allow-Origin", "*")
		} else if origin != "" && originAllowed(origin, allowedOrigin) {
			c.Header("Access-Control-Allow-Origin", origin)
			c.Header("Vary", "Origin")
		}

		c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS")
		c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization")
		c.Header("Access-Control-Allow-Credentials", "true")
		c.Header("Access-Control-Max-Age", "86400")

		if c.Request.Method == "OPTIONS" {
			c.AbortWithStatus(http.StatusNoContent)
			return
		}
		c.Next()
	}
}

// GetAllowedOrigins returns the resolved CORS origin string for use by
// other subsystems (e.g. WebSocket upgrader CheckOrigin). Call after CORS
// middleware is initialized.
func GetAllowedOrigins(cfg *config.Config) string {
	return getAllowedOrigin(cfg)
}

func getAllowedOrigin(cfg *config.Config) string {
	// Explicit env var overrides everything
	if v := os.Getenv("CORS_ALLOWED_ORIGINS"); v != "" {
		if v == "*" {
			log.Println("[CORS] WARNING: CORS_ALLOWED_ORIGINS is explicitly set to '*' — all origins allowed. " +
				"Set to a comma-separated list of allowed origins for production deployments.")
		}
		return v
	}
	// Production: restrict. Dev/test: allow all.
	if cfg.Environment == "production" {
		log.Println("[CORS] production mode with no CORS_ALLOWED_ORIGINS — defaulting to same-origin only")
		return "" // empty = no Access-Control-Allow-Origin header → same-origin only
	}
	return "*"
}

func originAllowed(origin, allowed string) bool {
	if allowed == "*" {
		return true
	}
	// allowed can be comma-separated
	for _, a := range strings.Split(allowed, ",") {
		if strings.TrimSpace(a) == origin {
			return true
		}
	}
	return false
}

// ── WebSocket Auth (ticket exchange) ─────────────────────────

// TicketValidator validates a single-use WebSocket ticket.
// Implemented by events.TicketStore. Interface avoids circular import.
type TicketValidator interface {
	Validate(ticketID string) (userID string, ok bool)
}

// WsAuth returns a Gin middleware for the WebSocket endpoint.
// It authenticates via three methods in priority order:
//
//  1. ?ticket=<opaque>  — single-use ticket (preferred)
//  2. ?token=<jwt>      — legacy JWT in query param (deprecated)
//  3. Authorization header — standard Bearer JWT
//
// When ?token= is used, a deprecation notice is logged. The ticket
// path avoids exposing the JWT in server logs, proxy logs, and
// browser history.
func WsAuth(cfg *config.Config, users store.UserStore, cache *UserStatusCache, tickets TicketValidator) gin.HandlerFunc {
	return func(c *gin.Context) {
		if !database.IsConnected() {
			c.Next()
			return
		}

		// Path 1: Ticket exchange (preferred)
		if ticketID := c.Query("ticket"); ticketID != "" {
			userID, ok := tickets.Validate(ticketID)
			if !ok {
				c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
					"error": "invalid or expired ticket",
				})
				return
			}

			if !verifyUserByID(c, userID, users, cache) {
				return
			}

			c.Set("user_id", userID)
			c.Set("ws_auth", "ticket")
			c.Next()
			return
		}

		// Path 2: Legacy ?token= (deprecated — log warning)
		if qToken := c.Query("token"); qToken != "" {
			log.Printf("[ws] DEPRECATED: client using ?token= query param for WebSocket auth — migrate to ticket exchange")

			claims, ok := parseAndValidateJWT(qToken, cfg.JWTSecret)
			if !ok {
				c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
					"error": "invalid or expired token",
				})
				return
			}

			if !verifyUser(c, claims, users, cache) {
				return
			}

			c.Set("user_id", claims.UserID)
			c.Set("email", claims.Email)
			c.Set("ws_auth", "token_legacy")
			c.Next()
			return
		}

		// Path 3: Authorization header (non-browser clients)
		header := c.GetHeader("Authorization")
		if header == "" {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"error": "missing authentication — use ticket exchange or Authorization header",
			})
			return
		}

		tokenString := strings.TrimPrefix(header, "Bearer ")
		if tokenString == header {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"error": "invalid authorization format",
			})
			return
		}

		claims, ok := parseAndValidateJWT(tokenString, cfg.JWTSecret)
		if !ok {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"error": "invalid or expired token",
			})
			return
		}

		if !verifyUser(c, claims, users, cache) {
			return
		}

		c.Set("user_id", claims.UserID)
		c.Set("email", claims.Email)
		c.Set("ws_auth", "bearer")
		c.Next()
	}
}