feat: Implement session store fallback for token validation in the authentication middleware.

api/api.go

@@ -252,6 +252,10 @@ func RegisterRoutes(r *gin.Engine, opts ...Option) {
 		opt(h)
 	}
 
+	if h.tokenService != nil && h.store != nil {
+		h.tokenService.SetStore(h.store)
+	}
+
 	r.GET("/healthz", func(c *gin.Context) {
 		c.JSON(http.StatusOK, gin.H{"status": "ok"})
 	})

internal/auth/middleware.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/gin-gonic/gin"
 
@@ -52,6 +53,7 @@ const (
 )
 
 // AuthMiddleware is a middleware that validates JWT access tokens
+// with a fallback to database-backed session tokens.
 func (s *TokenService) AuthMiddleware() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		var token string
@@ -74,24 +76,45 @@ func (s *TokenService) AuthMiddleware() gin.HandlerFunc {
 			return
 		}
 
+		// 1. Try JWT validation first.
 		claims, err := s.ValidateAccessToken(token)
-		if err != nil {
-			c.JSON(http.StatusUnauthorized, gin.H{
-				"error":  "invalid or expired token",
-				"detail": err.Error(),
-			})
-			c.Abort()
+		if err == nil {
+			// JWT is valid, populate context from claims.
+			ctx := context.WithValue(c.Request.Context(), userIDKey, claims.UserID)
+			ctx = context.WithValue(ctx, emailKey, claims.Email)
+			ctx = context.WithValue(ctx, rolesKey, claims.Roles)
+			ctx = context.WithValue(ctx, mfaKey, claims.MFA)
+			c.Request = c.Request.WithContext(ctx)
+			c.Next()
 			return
 		}
 
-		// Store claims in context
-		ctx := context.WithValue(c.Request.Context(), userIDKey, claims.UserID)
-		ctx = context.WithValue(ctx, emailKey, claims.Email)
-		ctx = context.WithValue(ctx, rolesKey, claims.Roles)
-		ctx = context.WithValue(ctx, mfaKey, claims.MFA)
-		c.Request = c.Request.WithContext(ctx)
+		// 2. Fallback to database session store if JWT fails and store is available.
+		if s.store != nil {
+			userID, expiresAt, err := s.store.GetSession(c.Request.Context(), token)
+			if err == nil && time.Now().Before(expiresAt) {
+				// Valid session found in store.
+				user, err := s.store.GetUserByID(c.Request.Context(), userID)
+				if err == nil {
+					ctx := context.WithValue(c.Request.Context(), userIDKey, user.ID)
+					ctx = context.WithValue(ctx, emailKey, user.Email)
+					ctx = context.WithValue(ctx, rolesKey, []string{user.Role})
+					// Assume MFA verified for active sessions found in DB for now,
+					// or we can refine this if session store tracks MFA state.
+					ctx = context.WithValue(ctx, mfaKey, true)
+					c.Request = c.Request.WithContext(ctx)
+					c.Next()
+					return
+				}
+			}
+		}
 
-		c.Next()
+		// Both JWT and Session lookups failed.
+		c.JSON(http.StatusUnauthorized, gin.H{
+			"error":  "invalid or expired token",
+			"detail": "token could not be validated as JWT or session",
+		})
+		c.Abort()
 	}
 }
 

internal/auth/token_service.go

@@ -5,6 +5,8 @@ import (
 	"time"
 
 	"github.com/golang-jwt/jwt/v5"
+
+	"account/internal/store"
 )
 
 // TokenPair represents a pair of Public and Access tokens
@@ -32,6 +34,7 @@ type TokenService struct {
 	accessSecret  string
 	accessExpiry  time.Duration
 	refreshExpiry time.Duration
+	store         store.Store
 }
 
 // TokenConfig holds configuration for token service
@@ -41,6 +44,7 @@ type TokenConfig struct {
 	AccessSecret  string
 	AccessExpiry  time.Duration
 	RefreshExpiry time.Duration
+	Store         store.Store
 }
 
 // NewTokenService creates a new TokenService instance
@@ -51,9 +55,15 @@ func NewTokenService(config TokenConfig) *TokenService {
 		accessSecret:  config.AccessSecret,
 		accessExpiry:  config.AccessExpiry,
 		refreshExpiry: config.RefreshExpiry,
+		store:         config.Store,
 	}
 }
 
+// SetStore sets the store for the token service.
+func (s *TokenService) SetStore(st store.Store) {
+	s.store = st
+}
+
 // ValidatePublicToken validates the public token
 func (s *TokenService) ValidatePublicToken(publicToken string) bool {
 	return publicToken == s.publicToken