query-orchestration/internal/cognitoauth/middleware.go

package cognitoauth

import (
	"net/http"
	"strings"

	"github.com/labstack/echo/v4"
)

// TokenValidationMiddleware verifies JWT tokens efficiently
// This is the primary middleware that handles both authentication and authorization
// Applied to all routes to enforce security requirements
func TokenValidationMiddleware(config *Config) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			requestPath := c.Request().URL.Path
			config.Logger.Debug("Processing request", "path", requestPath, "method", c.Request().Method)

			// Skip token validation for specified public paths
			publicPaths := []string{config.HomePath, config.LogoutPath}
			for _, path := range publicPaths {
				if requestPath == path {
					config.Logger.Debug("Skipping token validation for public path", "path", requestPath)
					return next(c)
				}
			}

			// Special case for the OAuth callback path
			if requestPath == config.CallbackPath && (c.QueryParam("code") != "" || c.QueryParam("error") != "") {
				config.Logger.Debug("Handling OAuth callback",
					"has_code", c.QueryParam("code") != "",
					"has_error", c.QueryParam("error") != "")
				// Skip token validation for OAuth callback - it will be handled by the callback handler
				return handleOAuthCallback(c, config)
			}

			// Initialize login flow if this is a login request
			if requestPath == config.LoginPath {
				return initiateLoginWithPKCE(c, config)
			}

			// Get authorization header
			authHeader := c.Request().Header.Get("Authorization")
			if authHeader == "" {
				config.Logger.Warn("No authorization header provided")
				return c.JSON(http.StatusUnauthorized, map[string]string{"error": "Authorization required"})
			}

			// Extract token
			tokenStr := authHeader
			if strings.HasPrefix(authHeader, "Bearer ") {
				tokenStr = authHeader[7:]
			}

			// Get JWKS (cached if possible)
			keySet, err := GetJWKS(config.JwksURL, config.Logger)
			if err != nil {
				config.Logger.Error("Failed to get JWKS", "error", err)
				return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Failed to verify token"})
			}

			// Verify token
			claims, err := verifyToken(tokenStr, keySet, *config)
			if err != nil {
				config.Logger.Warn("Token verification failed", "error", err)
				return c.JSON(http.StatusUnauthorized, map[string]string{"error": "Invalid token"})
			}

			// Extract user groups
			userGroups, err := GetUserGroups(claims)
			if err != nil {
				config.Logger.Warn("Failed to extract user groups", "error", err)
				userGroups = []string{} // Empty array if no groups found
			}

			// Store in context
			c.Set("user_claims", claims)
			c.Set("user_groups", userGroups)

			// Create and store user info in context
			userInfo := ExtractUserInfo(claims)
			c.Set("user_info", userInfo)

			// Check authorization
			authorized, requiredGroups := checkPermissions(requestPath, userGroups, config.RoutePermissions, config.Logger)
			if !authorized {
				config.Logger.Warn("Access denied",
					"path", requestPath,
					"user", claims["cognito:username"],
					"groups", userGroups,
					"required", requiredGroups)

				return c.JSON(http.StatusForbidden, map[string]interface{}{
					"error":           "Insufficient permissions",
					"message":         "User doesn't have the required group membership",
					"username":        claims["cognito:username"],
					"groups":          userGroups,
					"required_groups": requiredGroups,
				})
			}

			// User is authenticated and authorized
			return next(c)
		}
	}
}

// JWTAuthMiddleware checks for JWT token in cookies and adds it to the Authorization header
func JWTAuthMiddleware(config *Config) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// Skip for login and callback paths
			if c.Path() == config.LoginPath || c.Path() == config.CallbackPath {
				return next(c)
			}

			// Skip for home and logout as they should be accessible without auth
			if c.Path() == config.HomePath || c.Path() == config.LogoutPath {
				return next(c)
			}

			// Check if Authorization header is already set
			authHeader := c.Request().Header.Get("Authorization")
			if authHeader != "" && strings.HasPrefix(authHeader, "Bearer ") {
				// Authorization header is already set, proceed to the next handler
				return next(c)
			}

			// Try to get token from cookie
			tokenCookie, err := c.Cookie("auth_token")
			if err != nil || tokenCookie.Value == "" {
				// No token cookie found
				return c.Redirect(http.StatusFound, config.LoginPath)
			}

			// Add token to Authorization header
			c.Request().Header.Set("Authorization", "Bearer "+tokenCookie.Value)

			return next(c)
		}
	}
}