query-orchestration/internal/cognitoauth/middleware.go

package cognitoauth

import (
	"fmt"
	"net/http"
	"time"

	"queryorchestration/internal/serviceconfig/auth"

	"strings"

	"github.com/labstack/echo/v4"
	"github.com/lestrrat-go/jwx/v2/jwk"
)

// isPublicPath checks if the request path should skip authentication
func isPublicPath(requestPath string, config auth.ConfigProvider) bool {
	publicPaths := []string{config.GetAuthHomePath(), config.GetAuthLogoutPath(), config.GetHealthPath()}
	for _, path := range publicPaths {
		if requestPath == path {
			return true
		}
	}
	return strings.HasPrefix(requestPath, "/swagger/")
}

// isSpecialAuthPath checks for OAuth callback and login paths that need special handling
func isSpecialAuthPath(requestPath string, c echo.Context, config auth.ConfigProvider) (bool, error) {
	// Special case for the OAuth callback path
	if requestPath == config.GetAuthCallbackPath() && (c.QueryParam("code") != "" || c.QueryParam("error") != "") {
		config.GetAuthLogger().Debug("Handling OAuth callback",
			"has_code", c.QueryParam("code") != "",
			"has_error", c.QueryParam("error") != "")
		return true, handleOAuthCallback(c, config)
	}

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

	return false, nil
}

// extractToken gets the JWT token from Authorization header or cookie
func extractToken(c echo.Context, config auth.ConfigProvider) (string, error) {
	authHeader := c.Request().Header.Get("Authorization")
	if authHeader != "" && strings.HasPrefix(authHeader, "Bearer ") {
		return authHeader[7:], nil
	}

	// Try cookie
	tokenCookie, err := c.Cookie("auth_token")
	if err != nil || tokenCookie.Value == "" {
		config.GetAuthLogger().Warn("No authorization header or cookie provided")
		return "", fmt.Errorf("no authorization header or cookie provided")
	}
	return tokenCookie.Value, nil
}

// performPermitIOAuthorization handles the Permit.io authorization check
func performPermitIOAuthorization(c echo.Context, requestPath string, config auth.ConfigProvider) error {
	// Skip authorization for swagger and metrics
	if strings.HasPrefix(requestPath, "/swagger") || strings.HasPrefix(requestPath, "/metrics") {
		return nil
	}

	// Get user subject from JWT for Permit.io
	userSubject, ok := GetUserSubject(c)
	if !ok {
		config.GetAuthLogger().Warn("No user subject found in request")
		return c.JSON(http.StatusUnauthorized, map[string]string{
			"error": "Authentication required",
		})
	}

	// Initialize Permit.io client
	permitClient := NewPermitIOClient(
		config.GetPermitIOAPIKey(),
		config.GetPermitIOPDPURL(),
		config.GetAuthLogger(),
	)

	// Authorize with Permit.io
	resource := GetResourceFromRoute(requestPath)
	action := GetActionFromMethod(c.Request().Method)

	permitted, err := permitClient.CheckPermission(userSubject, action, resource)
	if err != nil {
		config.GetAuthLogger().Error("Permit.io authorization failed",
			"error", err,
			"user", userSubject,
			"resource", resource,
			"action", action)
		return c.JSON(http.StatusInternalServerError, map[string]string{
			"error": "Authorization service unavailable",
		})
	}

	if !permitted {
		config.GetAuthLogger().Warn("Access denied by Permit.io",
			"user", userSubject,
			"resource", resource,
			"action", action,
			"path", requestPath)

		if err := c.JSON(http.StatusForbidden, map[string]interface{}{
			"error":    "Insufficient permissions",
			"message":  "Access denied by authorization service",
			"user":     userSubject,
			"resource": resource,
			"action":   action,
		}); err != nil {
			config.GetAuthLogger().Error("Failed to send JSON response", "error", err)
		}
		return echo.NewHTTPError(http.StatusForbidden, "access denied")
	}

	config.GetAuthLogger().Debug("Access granted by Permit.io",
		"user", userSubject,
		"resource", resource,
		"action", action)

	return nil
}

// TokenValidationMiddleware verifies JWT tokens and enforces authentication and authorization.
//
// This middleware performs several critical security functions:
//   - Validates JWT tokens and their signatures against the JWKS from the identity provider
//   - Extracts and validates claims from the token
//   - Checks authorization using Permit.io for dynamic permissions
//   - Enforces route-specific permissions
//
// The middleware skips authentication for public paths and handles special cases
// like OAuth callbacks and login initiation. For protected routes, it enforces
// both authentication (valid token) and authorization (Permit.io permissions).
//
// Parameters:
//   - config: A ConfigProvider implementation that supplies all necessary auth configuration
//
// Returns:
//   - echo.MiddlewareFunc: An Echo middleware function that can be added to the middleware chain
func TokenValidationMiddleware(config auth.ConfigProvider) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			requestPath := c.Request().URL.Path
			config.GetAuthLogger().Debug("Processing request", "path", requestPath, "method", c.Request().Method)

			// Check for public paths that skip authentication
			if isPublicPath(requestPath, config) {
				config.GetAuthLogger().Debug("Skipping token validation for public path", "path", requestPath)
				return next(c)
			}

			// Check for special auth paths (OAuth callback, login)
			if isSpecial, err := isSpecialAuthPath(requestPath, c, config); isSpecial {
				return err
			}

			// Extract JWT token from header or cookie
			tokenStr, err := extractToken(c, config)
			if err != nil {
				return c.JSON(http.StatusUnauthorized, map[string]interface{}{
					"error":     "Authorization required",
					"message":   "This endpoint requires authentication. Please obtain a valid JWT token.",
					"login_url": config.GetAuthLoginPath(),
				})
			}

			// Check if token is expired and attempt refresh if needed
			if isTokenExpired(tokenStr) {
				config.GetAuthLogger().Debug("Access token expired, attempting refresh")

				refreshToken := getRefreshTokenFromCookie(c)
				if refreshToken == "" {
					config.GetAuthLogger().Debug("No refresh token found, redirecting to login")
					clearTokenCookies(c, config)
					return c.JSON(http.StatusUnauthorized, map[string]interface{}{
						"error":     "Session expired",
						"message":   "Your session has expired. Please log in again.",
						"login_url": config.GetAuthLoginPath(),
					})
				}

				// Attempt to refresh tokens
				newTokens, err := refreshTokens(refreshToken, config)
				if err != nil {
					config.GetAuthLogger().Warn("Token refresh failed", "error", err)
					clearTokenCookies(c, config)
					return c.JSON(http.StatusUnauthorized, map[string]interface{}{
						"error":     "Session expired",
						"message":   "Unable to refresh your session. Please log in again.",
						"login_url": config.GetAuthLoginPath(),
					})
				}

				// Update cookies with new tokens
				newExpiresAt := time.Now().Add(time.Duration(newTokens.ExpiresIn) * time.Second)
				setTokenCookies(c, newTokens, newExpiresAt, config)

				// Update the token string for subsequent verification
				tokenStr = newTokens.AccessToken

				// Update the Authorization header for this request
				c.Request().Header.Set("Authorization", "Bearer "+tokenStr)

				config.GetAuthLogger().Info("Token refreshed successfully")
			}

			// Get JWKS (skip in test mode)
			var keySet jwk.Set
			if !config.GetNoJWTValidation() {
				keySet, err = GetJWKS(config.GetAuthJwksURL(), config.GetAuthLogger())
				if err != nil {
					config.GetAuthLogger().Error("Failed to get JWKS", "error", err)
					return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Failed to verify token"})
				}
			}

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

			// Store user info in context
			c.Set("user_claims", claims)
			userInfo := ExtractUserInfo(claims)
			c.Set("user_info", userInfo)

			// Perform Permit.io authorization
			if err := performPermitIOAuthorization(c, requestPath, config); err != nil {
				return err
			}

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

// JWTAuthMiddleware extracts JWT tokens from cookies and adds them to the Authorization header.
//
// This middleware serves as a bridge between cookie-based and header-based authentication:
//   - Checks for auth_token cookie and transfers it to the Authorization header
//   - Skips processing for login and callback paths that handle authentication
//   - Manages cookie deletion during logout operations
//   - Redirects unauthenticated requests to the login path
//
// This middleware should be registered before TokenValidationMiddleware in the middleware
// chain to ensure the JWT token is properly placed in the Authorization header before
// validation occurs.
//
// Parameters:
//   - config: A ConfigProvider implementation that supplies auth configuration values
//
// Returns:
//   - echo.MiddlewareFunc: An Echo middleware function that can be added to the middleware chain
func JWTAuthMiddleware(config auth.ConfigProvider) 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.GetAuthLoginPath() || c.Path() == config.GetAuthCallbackPath() {
				return next(c)
			}

			// if logout then remove both auth cookies
			if c.Path() == config.GetAuthLogoutPath() {
				clearTokenCookies(c, config)
			}

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

			// Skip for swagger documentation paths
			if strings.HasPrefix(c.Path(), "/swagger/") {
				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 - let TokenValidationMiddleware handle the error
				return next(c)
			}

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

			return next(c)
		}
	}
}