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/") || strings.HasPrefix(requestPath, "/metrics")
}

// 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
}

// isAuthOnlyPath returns true for paths that require authentication but NOT authorization.
// These endpoints are accessible to all authenticated users regardless of their Permit.io roles.
// This enables users to discover their own roles/permissions before making other API calls.
func isAuthOnlyPath(requestPath string) bool {
	authOnlyPaths := []string{
		"/identity", // Users need to discover their roles without having any specific role
	}
	for _, path := range authOnlyPaths {
		if requestPath == path {
			return true
		}
	}
	return false
}

// 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
	}

	// Skip authorization for auth-only paths (authentication required, but no specific permissions needed)
	if isAuthOnlyPath(requestPath) {
		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
}

// HTTPError represents an HTTP error response with status code and response data
type HTTPError struct {
	StatusCode int
	Response   interface{}
}

func (e *HTTPError) Error() string {
	return fmt.Sprintf("HTTP %d error", e.StatusCode)
}

// attemptTokenRefresh handles the complete token refresh flow including non-rotating token preservation.
//
// This function consolidates the duplicate refresh logic that was scattered throughout the middleware.
// It attempts to refresh tokens using the provided refresh token and handles the case where Cognito
// returns empty refresh tokens (non-rotating tokens).
//
// Parameters:
//   - c: The Echo context for setting cookies and headers
//   - refreshToken: The refresh token to use for getting new tokens
//   - config: Configuration provider for auth settings
//
// Returns:
//   - *TokenResponse: New tokens if refresh succeeds
//   - error: Error if refresh fails
func attemptTokenRefresh(c echo.Context, refreshToken string, config auth.ConfigProvider) (*TokenResponse, error) {
	config.GetAuthLogger().Debug("Attempting token refresh with refresh token")

	newTokens, err := refreshTokens(refreshToken, config)
	if err != nil {
		config.GetAuthLogger().Warn("Token refresh failed", "error", err)
		clearTokenCookies(c, config)
		return nil, err
	}

	// Handle non-rotating refresh tokens - preserve original if new one is empty
	if newTokens.RefreshToken == "" {
		config.GetAuthLogger().Debug("Refresh response missing refresh token, preserving original")
		newTokens.RefreshToken = refreshToken
	}

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

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

	config.GetAuthLogger().Debug("Token refresh successful, proceeding with new token")
	return newTokens, nil
}

// acquireValidToken gets a valid access token through extraction or refresh.
//
// This function handles the complete token acquisition flow:
//   - Attempts to extract token from header/cookie
//   - Falls back to refresh if no token found
//   - Checks for expiration and refreshes if needed
//   - Returns appropriate HTTP error responses
//
// Parameters:
//   - c: The Echo context for cookies, headers, and responses
//   - config: Configuration provider for auth settings
//
// Returns:
//   - string: Valid access token
//   - error: HTTP error response if token acquisition fails
func acquireValidToken(c echo.Context, config auth.ConfigProvider) (string, error) {
	// Extract JWT token from header or cookie
	tokenStr, err := extractToken(c, config)

	// If token extraction fails, try refresh before returning 401
	if err != nil {
		config.GetAuthLogger().Debug("No access token found, checking for refresh token")
		config.GetAuthLogger().Debug("extractToken error", "error", err)

		// Check if we have a refresh token
		config.GetAuthLogger().Debug("About to call getRefreshTokenFromCookie")
		refreshToken := getRefreshTokenFromCookie(c, config)
		config.GetAuthLogger().Debug("getRefreshTokenFromCookie returned", "token_length", len(refreshToken))

		if refreshToken == "" {
			config.GetAuthLogger().Debug("No refresh token found, redirecting to login")
			return "", &HTTPError{
				StatusCode: http.StatusUnauthorized,
				Response: map[string]interface{}{
					"error":     "Authorization required",
					"message":   "This endpoint requires authentication. Please obtain a valid JWT token.",
					"login_url": config.GetAuthLoginPath(),
				},
			}
		}

		// Attempt to refresh tokens
		newTokens, err := attemptTokenRefresh(c, refreshToken, config)
		if err != nil {
			return "", &HTTPError{
				StatusCode: http.StatusUnauthorized,
				Response: map[string]interface{}{
					"error":     "Session expired",
					"message":   "Unable to refresh your session. Please log in again.",
					"login_url": config.GetAuthLoginPath(),
				},
			}
		}

		// Use the new token for this request
		tokenStr = newTokens.AccessToken
	}

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

		refreshToken := getRefreshTokenFromCookie(c, config)
		if refreshToken == "" {
			config.GetAuthLogger().Debug("No refresh token found, redirecting to login")
			clearTokenCookies(c, config)
			return "", &HTTPError{
				StatusCode: http.StatusUnauthorized,
				Response: 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 := attemptTokenRefresh(c, refreshToken, config)
		if err != nil {
			return "", &HTTPError{
				StatusCode: http.StatusUnauthorized,
				Response: map[string]interface{}{
					"error":     "Session expired",
					"message":   "Unable to refresh your session. Please log in again.",
					"login_url": config.GetAuthLoginPath(),
				},
			}
		}

		// Update the token string for subsequent verification
		tokenStr = newTokens.AccessToken
		config.GetAuthLogger().Info("Token refreshed successfully")
	}

	return tokenStr, nil
}

// verifyTokenAndSetContext verifies token signature and populates request context.
//
// This function handles token verification and context setup:
//   - Retrieves JWKS for token signature verification (unless in test mode)
//   - Verifies the token against the key set
//   - Extracts claims and user info
//   - Populates the Echo context with user information
//
// Parameters:
//   - c: The Echo context for setting context values and responses
//   - tokenStr: The JWT token string to verify
//   - config: Configuration provider for auth settings
//
// Returns:
//   - map[string]interface{}: Token claims if verification succeeds
//   - error: HTTP error response if verification fails
func verifyTokenAndSetContext(c echo.Context, tokenStr string, config auth.ConfigProvider) (map[string]interface{}, error) {
	// Get JWKS (skip in test mode)
	var keySet jwk.Set
	var err error
	if !config.GetNoJWTValidation() {
		keySet, err = GetJWKS(config.GetAuthJwksURL(), config.GetAuthLogger())
		if err != nil {
			config.GetAuthLogger().Error("Failed to get JWKS", "error", err)
			return nil, &HTTPError{
				StatusCode: http.StatusInternalServerError,
				Response:   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 nil, &HTTPError{
			StatusCode: http.StatusUnauthorized,
			Response:   map[string]string{"error": "Invalid token"},
		}
	}

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

	return claims, 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
			}

			// Acquire valid token (handles extraction, refresh, expiration)
			tokenStr, err := acquireValidToken(c, config)
			if err != nil {
				if httpErr, ok := err.(*HTTPError); ok {
					return c.JSON(httpErr.StatusCode, httpErr.Response)
				}
				return err
			}

			// Verify token and populate context
			_, err = verifyTokenAndSetContext(c, tokenStr, config)
			if err != nil {
				if httpErr, ok := err.(*HTTPError); ok {
					return c.JSON(httpErr.StatusCode, httpErr.Response)
				}
				return err
			}

			// 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)
		}
	}
}