query-orchestration/cmd/cognito_test/cognito.pkce/main.go

package main

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

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

// CognitoTokenResponse represents the response from the token endpoint
// Returned by Cognito when exchanging an authorization code for tokens
type CognitoTokenResponse struct {
	IDToken      string `json:"id_token"`      // OpenID Connect ID token containing user claims
	AccessToken  string `json:"access_token"`  // OAuth2 access token for accessing resources
	RefreshToken string `json:"refresh_token"` // Token used to get new access tokens without re-authentication
	ExpiresIn    int    `json:"expires_in"`    // Token validity period in seconds
	TokenType    string `json:"token_type"`    // Type of token, typically "Bearer"
}

// CognitoConfig holds the configuration for AWS Cognito
// Contains all necessary parameters to interact with Cognito endpoints
type CognitoConfig struct {
	ClientID    string // OAuth2 client ID registered with Cognito
	RedirectURI string // URL where Cognito redirects after authentication
	TokenURL    string // Cognito endpoint for token operations
	JwksURL     string // URL for JSON Web Key Set (for token verification)
	UserPoolID  string // Cognito User Pool ID
	AuthURL     string // Cognito authorization endpoint for initiating login
}

// PKCESession stores PKCE code verifier for verification
// We need to store the code verifier temporarily to use during token exchange
type PKCESession struct {
	CodeVerifier string
	CreatedAt    time.Time
	ExpiresAt    time.Time
}

// Global map to store PKCE sessions by state parameter
// In a production environment, this should be replaced with a proper session store
var pkceSessionMap = struct {
	sync.RWMutex
	sessions map[string]*PKCESession
}{
	sessions: make(map[string]*PKCESession),
}

// JWKSCache caches the JWKS to avoid frequent fetches
// Implements a thread-safe caching mechanism with expiration
type JWKSCache struct {
	KeySet    jwk.Set      // The cached JSON Web Key Set
	ExpiresAt time.Time    // Expiration time for the cache
	mutex     sync.RWMutex // Mutex for thread-safe access
}

// Global JWKS cache instance
var jwksCache = &JWKSCache{
	ExpiresAt: time.Now(), // Initial state is expired, forcing a fetch on first use
}

const LOGIN_CALLBACK_PATH = "/login-callback"

// GetJWKS returns the cached JWKS or fetches a new one if needed
// This function implements caching logic to minimize external requests
// Used during token validation to get the key set for signature verification
func GetJWKS(jwksURL string, logger *slog.Logger) (jwk.Set, error) {
	// Try to use cached JWKS first (read lock)
	jwksCache.mutex.RLock()
	if jwksCache.KeySet != nil && time.Now().Before(jwksCache.ExpiresAt) {
		defer jwksCache.mutex.RUnlock()
		logger.Debug("Using cached JWKS")
		return jwksCache.KeySet, nil
	}
	jwksCache.mutex.RUnlock()

	// Need to fetch new JWKS (write lock)
	jwksCache.mutex.Lock()
	defer jwksCache.mutex.Unlock()

	// Double-check expiration after acquiring the write lock
	if jwksCache.KeySet != nil && time.Now().Before(jwksCache.ExpiresAt) {
		return jwksCache.KeySet, nil
	}

	logger.Info("Fetching new JWKS", "url", jwksURL)
	jwksJSON, err := fetchJWKS(jwksURL)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch JWKS: %w", err)
	}

	keySet, err := jwk.ParseString(jwksJSON)
	if err != nil {
		return nil, fmt.Errorf("failed to parse JWKS: %w", err)
	}

	jwksCache.KeySet = keySet
	// Cache for 24 hours - you can adjust this based on your needs
	jwksCache.ExpiresAt = time.Now().Add(24 * time.Hour)

	return keySet, nil
}

// 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 *CognitoConfig, routePermissions map[string][]string, logger *slog.Logger) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			requestPath := c.Request().URL.Path
			logger.Debug("Processing request", "path", requestPath, "method", c.Request().Method)

			// Skip token validation for specific public paths
			if requestPath == "/home" || requestPath == "/logout" {
				logger.Debug("Skipping token validation for public path", "path", requestPath)
				return next(c)
			}

			// Special case for the OAuth callback path
			if requestPath == LOGIN_CALLBACK_PATH && (c.QueryParam("code") != "" || c.QueryParam("error") != "") {
				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, routePermissions, logger)
			}

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

			// Get authorization header
			authHeader := c.Request().Header.Get("Authorization")
			if authHeader == "" {
				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, logger)
			if err != nil {
				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 {
				logger.Warn("Token verification failed", "error", err)
				return c.JSON(http.StatusUnauthorized, map[string]string{"error": "Invalid token"})
			}

			// Debug logging if enabled
			if os.Getenv("DEBUG") == "true" {
				logger.Info("Token verified successfully", "claims", claims)
			}

			// Extract user groups
			userGroups, err := GetUserGroups(claims)
			if err != nil {
				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)

			// Check authorization
			authorized, requiredGroups := checkPermissions(requestPath, userGroups, routePermissions, logger)
			if !authorized {
				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
// This middleware is simpler than full session management and aligns with JWT's stateless nature
func JWTAuthMiddleware(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		// Skip for login and callback paths
		if c.Path() == "/login" || c.Path() == LOGIN_CALLBACK_PATH {
			return next(c)
		}

		// Skip for home and logout as they should be accessible without auth
		if c.Path() == "/home" || c.Path() == "/logout" {
			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, "/login")
		}

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

		return next(c)
	}
}

// logoutHandler for JWT-based authentication simply clears the token cookie
func logoutHandler(c echo.Context) error {
	// Clear the token cookie
	cookie := new(http.Cookie)
	cookie.Name = "auth_token"
	cookie.Value = ""
	cookie.Path = "/"
	cookie.Expires = time.Now().Add(-1 * time.Hour) // Set expiry in the past
	cookie.HttpOnly = true
	c.SetCookie(cookie)

	// Redirect to home page
	return c.Redirect(http.StatusFound, "/home")
}

// checkPermissions checks if the user has the required permissions
// Used by both the main token validation middleware and the OAuth callback handler
// Returns true if authorized, false if not, along with the required groups
func checkPermissions(path string, userGroups []string, routePermissions map[string][]string, logger *slog.Logger) (bool, []string) {
	// Find matching route pattern
	var requiredGroups []string
	var matched bool

	for pattern, groups := range routePermissions {
		if matchRoute(path, pattern) {
			requiredGroups = groups
			matched = true
			logger.Debug("Found matching route pattern", "pattern", pattern, "requiredGroups", requiredGroups)
			break
		}
	}

	if !matched {
		// No matching pattern found - could either allow or deny by default
		logger.Info("No matching route pattern found for path", "path", path)
		return true, nil // Allowing by default
	}

	// Check if user has any of the required groups
	for _, requiredGroup := range requiredGroups {
		for _, userGroup := range userGroups {
			if userGroup == requiredGroup {
				logger.Debug("User has required group", "group", requiredGroup)
				return true, requiredGroups
			}
		}
	}

	return false, requiredGroups
}

// initiateLoginWithPKCE generates PKCE code challenge and redirects to Cognito
// This function starts the OAuth authorization flow with PKCE
func initiateLoginWithPKCE(c echo.Context, config *CognitoConfig, logger *slog.Logger) error {
	// Generate random state parameter to prevent CSRF
	state, err := generateRandomString(32)
	if err != nil {
		logger.Error("Failed to generate state parameter", "error", err)
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Authentication initialization failed"})
	}

	// Generate code verifier (random string between 43-128 chars)
	codeVerifier, err := generateRandomString(64)
	if err != nil {
		logger.Error("Failed to generate code verifier", "error", err)
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Authentication initialization failed"})
	}

	// Create code challenge from verifier (SHA256 + Base64URL without padding)
	codeChallenge := createCodeChallenge(codeVerifier)

	// Store in session for later verification (with expiration time)
	storePKCESession(state, codeVerifier, logger)

	// Build authorization URL with PKCE parameters
	params := url.Values{}
	params.Set("client_id", config.ClientID)
	params.Set("response_type", "code")
	params.Set("redirect_uri", config.RedirectURI)
	params.Set("scope", "openid email profile")
	params.Set("state", state)
	params.Set("code_challenge", codeChallenge)
	params.Set("code_challenge_method", "S256")

	authURL := fmt.Sprintf("%s?%s", config.AuthURL, params.Encode())

	if os.Getenv("DEBUG") == "true" {
		logger.Debug("Initiating login with PKCE",
			"code_verifier", codeVerifier,
			"code_challenge", codeChallenge,
			"state", state,
			"redirect_uri", config.RedirectURI)
	}

	// Redirect user to Cognito login page
	return c.Redirect(http.StatusFound, authURL)
}

// handleOAuthCallback handles the OAuth 2.0 authorization code flow callback with PKCE
// Called when Cognito redirects back to our application with an authorization code
// Exchanges the code for tokens, verifies them, and redirects to home page with token in cookie
func handleOAuthCallback(c echo.Context, config *CognitoConfig, routePermissions map[string][]string, logger *slog.Logger) error {
	// Extract the authorization code and state
	code := c.QueryParam("code")
	if code == "" {
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Missing code parameter"})
	}

	state := c.QueryParam("state")
	if state == "" {
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Missing state parameter"})
	}

	// Check for OAuth errors
	errorMsg := c.QueryParam("error")
	if errorMsg != "" {
		errorDesc := c.QueryParam("error_description")
		return c.JSON(http.StatusBadRequest, map[string]string{
			"error":             errorMsg,
			"error_description": errorDesc,
		})
	}

	// Retrieve stored code verifier
	codeVerifier, err := getCodeVerifier(state, logger)
	if err != nil {
		logger.Error("Failed to retrieve code verifier", "error", err, "state", state)
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Invalid or expired session"})
	}

	// Exchange the code for tokens using PKCE
	tokens, err := exchangeCodeForTokensWithPKCE(code, codeVerifier, *config, logger)
	if err != nil {
		logger.Error("Failed to exchange code for tokens", "error", err)
		return c.JSON(http.StatusUnauthorized, map[string]string{"error": err.Error()})
	}

	// Get the JWKS for token verification
	keySet, err := GetJWKS(config.JwksURL, logger)
	if err != nil {
		logger.Error("Failed to fetch JWKS", "error", err)
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Failed to verify token"})
	}

	// Verify ID token
	idTokenClaims, err := verifyToken(tokens.IDToken, keySet, *config)
	if err != nil {
		logger.Error("Failed to verify ID token", "error", err)
		return c.JSON(http.StatusUnauthorized, map[string]string{"error": "Invalid token"})
	}

	// Extract user groups
	userGroups, err := GetUserGroups(idTokenClaims)
	if err != nil {
		logger.Warn("Failed to extract user groups", "error", err)
		userGroups = []string{} // Initialize as empty array
	}

	// Debug logging for tokens
	if os.Getenv("DEBUG") == "true" {
		// Log decoded token for debugging
		logger.Info("ID Token Claims", "claims", idTokenClaims)
		logger.Info("Raw ID Token", "token", tokens.IDToken)
		logger.Info("Raw Access Token", "token", tokens.AccessToken)
	}

	// Check authorization
	authorized, requiredGroups := checkPermissions(LOGIN_CALLBACK_PATH, userGroups, routePermissions, logger)
	if !authorized {
		logger.Warn("OAuth callback - access denied",
			"user", idTokenClaims["cognito:username"],
			"groups", userGroups,
			"required", requiredGroups)

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

	// User is authenticated and authorized

	// Calculate token expiration time based on token's expires_in value
	expiresAt := time.Now().Add(time.Duration(tokens.ExpiresIn) * time.Second)

	// Set a cookie with the access token
	tokenCookie := new(http.Cookie)
	tokenCookie.Name = "auth_token"
	tokenCookie.Value = tokens.AccessToken
	tokenCookie.Path = "/"
	tokenCookie.Expires = expiresAt
	tokenCookie.HttpOnly = true // Not accessible via JavaScript
	// In production, set Secure to true
	// tokenCookie.Secure = true
	c.SetCookie(tokenCookie)

	// Extract username for display purposes (optional)
	username, _ := idTokenClaims["cognito:username"].(string)

	logger.Info("User authenticated successfully", "username", username, "groups", userGroups)

	// Redirect to home page
	return c.Redirect(http.StatusFound, "/home")
}

// generateRandomString creates a cryptographically secure random string
// Used for generating state parameter and code verifier
func generateRandomString(length int) (string, error) {
	b := make([]byte, length)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(b)[:length], nil
}

// createCodeChallenge creates a code challenge from a code verifier
// Implements S256 code challenge method as specified in PKCE RFC 7636
func createCodeChallenge(verifier string) string {
	// Create SHA256 hash of the verifier
	hash := sha256.Sum256([]byte(verifier))

	// Base64URL encode the hash without padding
	challenge := base64.RawURLEncoding.EncodeToString(hash[:])

	return challenge
}

// storePKCESession stores the PKCE session for later verification
// In a production environment, this should be replaced with a proper session store
func storePKCESession(state, codeVerifier string, logger *slog.Logger) {
	pkceSessionMap.Lock()
	defer pkceSessionMap.Unlock()

	// Create new session with expiration (5 minutes is common)
	session := &PKCESession{
		CodeVerifier: codeVerifier,
		CreatedAt:    time.Now(),
		ExpiresAt:    time.Now().Add(5 * time.Minute),
	}

	pkceSessionMap.sessions[state] = session

	// Clean up expired sessions
	for k, v := range pkceSessionMap.sessions {
		if time.Now().After(v.ExpiresAt) {
			logger.Debug("Cleaning up expired PKCE session", "state", k)
			delete(pkceSessionMap.sessions, k)
		}
	}
}

// getCodeVerifier retrieves and validates the code verifier for a state
func getCodeVerifier(state string, logger *slog.Logger) (string, error) {
	pkceSessionMap.RLock()
	defer pkceSessionMap.RUnlock()

	session, ok := pkceSessionMap.sessions[state]
	if !ok {
		return "", errors.New("no session found for state")
	}

	if time.Now().After(session.ExpiresAt) {
		return "", errors.New("session expired")
	}

	return session.CodeVerifier, nil
}

// exchangeCodeForTokensWithPKCE exchanges the authorization code for tokens using PKCE
// Called during the OAuth callback flow to get tokens from the authorization code
// Makes an HTTP request to Cognito's token endpoint to perform this exchange
func exchangeCodeForTokensWithPKCE(authCode, codeVerifier string, config CognitoConfig, logger *slog.Logger) (*CognitoTokenResponse, error) {
	data := url.Values{}
	data.Set("grant_type", "authorization_code")
	data.Set("client_id", config.ClientID)
	data.Set("code", authCode)
	data.Set("redirect_uri", config.RedirectURI)
	data.Set("code_verifier", codeVerifier) // Include code verifier for PKCE

	logger.Debug("Token request details",
		"url", config.TokenURL,
		"client_id", config.ClientID,
		"redirect_uri", config.RedirectURI)

	req, err := http.NewRequest("POST", config.TokenURL, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, fmt.Errorf("failed to create token request: %w", err)
	}

	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")

	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("token request failed: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read token response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("token endpoint returned status %d: %s", resp.StatusCode, string(body))
	}

	var tokenResponse CognitoTokenResponse
	if err := json.Unmarshal(body, &tokenResponse); err != nil {
		return nil, fmt.Errorf("failed to parse token response: %w", err)
	}

	return &tokenResponse, nil
}

// verifyToken verifies the JWT token and returns its claims
// Used during both OAuth callback and subsequent API requests with bearer token
// Verifies signature, expiration, issuer, and other JWT claims
func verifyToken(tokenString string, keySet jwk.Set, config CognitoConfig) (map[string]interface{}, error) {
	region := "us-east-2"
	if strings.Contains(config.JwksURL, "amazonaws.com/") {
		parts := strings.Split(config.JwksURL, "amazonaws.com/")
		if len(parts) > 1 {
			regionParts := strings.Split(parts[0], ".")
			if len(regionParts) > 0 {
				lastPart := regionParts[len(regionParts)-1]
				if lastPart != "" {
					region = lastPart
				}
			}
		}
	}

	issuer := fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s", region, config.UserPoolID)

	// Verify the token with the keySet
	verifiedToken, err := jwt.Parse(
		[]byte(tokenString),
		jwt.WithKeySet(keySet),
		jwt.WithValidate(true),
		jwt.WithIssuer(issuer),
		// Add other validation options as needed
	)
	if err != nil {
		return nil, err
	}

	// Extract claims to a map
	claims, err := verifiedToken.AsMap(context.Background())
	if err != nil {
		return nil, err
	}

	return claims, nil
}

// GetUserGroups extracts user groups from the token claims
// Attempts to find groups in different claim names depending on Cognito setup
// Returns an array of group names or an error if no groups found
func GetUserGroups(claims map[string]interface{}) ([]string, error) {
	// The claim containing the groups might have different names depending on your Cognito setup
	// Common names are "cognito:groups", "groups", or a custom attribute
	possibleGroupFields := []string{"cognito:groups", "groups", "custom:groups"}

	for _, field := range possibleGroupFields {
		if rawGroups, ok := claims[field]; ok {
			switch groups := rawGroups.(type) {
			case []interface{}:
				result := make([]string, len(groups))
				for i, g := range groups {
					result[i] = fmt.Sprintf("%v", g)
				}
				return result, nil
			case []string:
				return groups, nil
			case string:
				return []string{groups}, nil
			}
		}
	}

	return nil, errors.New("no groups found in token claims")
}

// fetchJWKS retrieves the JSON Web Key Set from Cognito
// Used to get the public keys needed to verify token signatures
// Makes an HTTP request to the JWKS URL and returns the raw JSON
func fetchJWKS(jwksURL string) (string, error) {
	client := &http.Client{Timeout: 10 * time.Second}
	req, err := http.NewRequest("GET", jwksURL, nil)
	if err != nil {
		return "", err
	}

	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	jwksData, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	return string(jwksData), nil
}

// matchRoute checks if a request path matches a route pattern
// Converts Echo route patterns to regex for matching
// Supports parameter patterns like '/users/:id' and wildcard patterns
func matchRoute(requestPath, routePattern string) bool {
	// Convert Echo route pattern to regex pattern
	// e.g., "/users/:id" to "/users/([^/]+)"
	regexPattern := "^"
	patternParts := strings.Split(routePattern, "/")

	for i, part := range patternParts {
		if i > 0 {
			regexPattern += "/"
		}

		if strings.HasPrefix(part, ":") {
			// Parameter part (e.g., :id)
			regexPattern += "([^/]+)"
		} else if part == "*" {
			// Wildcard - match anything
			regexPattern += ".*"
		} else {
			// Literal part - escape any regex metacharacters
			regexPattern += regexp.QuoteMeta(part)
		}
	}

	regexPattern += "$"

	// Compile and match
	regex, err := regexp.Compile(regexPattern)
	if err != nil {
		return false
	}

	return regex.MatchString(requestPath)
}

// Helper function to mask sensitive values
// Used for logging sensitive information like client IDs
func maskString(s string) string {
	if s == "" {
		return "<not set>"
	}
	if len(s) <= 8 {
		return "****"
	}
	return s[:4] + "..." + s[len(s)-4:]
}

// main is the entry point of the application
// Sets up the Echo server, middleware, routes, and starts listening
func main() {
	// Create a new Echo instance
	e := echo.New()

	// Basic middleware
	e.Use(middleware.Logger())
	e.Use(middleware.Recover())

	// Initialize logger with appropriate level based on DEBUG env var
	var logLevel slog.Level
	if os.Getenv("DEBUG") == "true" {
		logLevel = slog.LevelDebug
	} else {
		logLevel = slog.LevelInfo
	}

	logHandler := slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
		Level: logLevel,
	})
	logger := slog.New(logHandler)

	// Get Cognito configuration
	region := os.Getenv("AWS_REGION")
	if region == "" {
		region = "us-east-2" // Fallback to hardcoded value
	}

	userPoolID := os.Getenv("COGNITO_USER_POOL_ID")
	if userPoolID == "" {
		userPoolID = "us-east-2_1y6po8rR8" // Fallback to hardcoded value
	}

	domain := os.Getenv("COGNITO_DOMAIN")
	// In main function
	if domain == "" {
		logger.Warn("COGNITO_DOMAIN environment variable is not set, using default")
		// Use the user pool ID directly as part of the domain
		domain = fmt.Sprintf("%s.auth.%s.amazoncognito.com", userPoolID, region)
	}

	// Ensure domain format is correct (no protocol prefix)
	domain = strings.TrimPrefix(domain, "https://")
	domain = strings.TrimPrefix(domain, "http://")

	config := &CognitoConfig{
		ClientID:    os.Getenv("COGNITO_CLIENT_ID"),
		RedirectURI: "http://localhost:8080" + LOGIN_CALLBACK_PATH,
		TokenURL:    fmt.Sprintf("https://%s/oauth2/token", domain),
		AuthURL:     fmt.Sprintf("https://%s/oauth2/authorize", domain),
		// actual https://cognito-idp.us-east-2.amazonaws.com/us-east-2_1y6po8rR8/.well-known/jwks.json
		JwksURL:    fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json", region, userPoolID),
		UserPoolID: userPoolID,
	}

	// Route permissions - matching the provided example
	routePermissions := map[string][]string{
		"/users":                {"exporters"},
		"/users/:id":            {"exporters", "querybuilders"},
		"/orders":               {"exporters", "exporters"},
		"/reports/sales":        {"exporters", "uploaders", "querybuilders"},
		"/settings":             {"exporters"},
		LOGIN_CALLBACK_PATH:     {"exporters", "uploaders", "querybuilders"}, // Only exporters can access
		"/api/inventory/update": {"exporters", "uploaders"},
		"/login":                {"exporters", "uploaders", "querybuilders"}, // No permissions required for login
		"/home":                 {"exporters", "uploaders", "querybuilders"}, // No permissions required for home page
		"/logout":               {"exporters", "uploaders", "querybuilders"}, // No permissions required for logout
	}

	// Apply JWT auth middleware first
	// This will extract the JWT token from cookies and add it to the Authorization header
	e.Use(JWTAuthMiddleware)

	// Then apply token validation middleware
	// This will validate the token from the Authorization header
	e.Use(TokenValidationMiddleware(config, routePermissions, logger))

	// Login route
	e.GET("/login", func(c echo.Context) error {
		// The TokenValidationMiddleware will handle this
		return nil
	})

	// Logout route
	e.GET("/logout", logoutHandler)

	// Callback endpoint - handled by middleware
	e.GET(LOGIN_CALLBACK_PATH, func(c echo.Context) error {
		// This handler will only be called for requests that don't have a code parameter
		// (those are handled directly in the middleware)

		// Get user info from context
		userClaims, _ := c.Get("user_claims").(map[string]interface{})
		userGroups, _ := c.Get("user_groups").([]string)

		return c.JSON(http.StatusOK, map[string]interface{}{
			"authenticated": true,
			"username":      userClaims["cognito:username"],
			"email":         userClaims["email"],
			"groups":        userGroups,
		})
	})

	// Other endpoint handlers
	e.GET("/users", func(c echo.Context) error {
		return c.String(http.StatusOK, "Users endpoint")
	})

	e.GET("/users/:id", func(c echo.Context) error {
		id := c.Param("id")
		return c.String(http.StatusOK, fmt.Sprintf("User details endpoint for ID: %s", id))
	})

	e.GET("/orders", func(c echo.Context) error {
		return c.String(http.StatusOK, "Orders endpoint")
	})

	e.GET("/reports/sales", func(c echo.Context) error {
		return c.String(http.StatusOK, "Sales reports endpoint")
	})

	e.GET("/settings", func(c echo.Context) error {
		return c.String(http.StatusOK, "Settings endpoint")
	})

	e.PUT("/api/inventory/update", func(c echo.Context) error {
		return c.String(http.StatusOK, "Inventory update endpoint")
	})

	// Home endpoint for debugging - shows all available endpoints
	e.GET("/home", HomeHandler)

	// Print startup information
	logger.Info("Server configuration",
		"client_id", maskString(config.ClientID),
		"domain", domain,
		"user_pool_id", userPoolID,
		"region", region,
		"debug_mode", os.Getenv("DEBUG") == "true",
		"auth_method", "PKCE with JWT cookie")

	// Configure server
	server := &http.Server{
		Addr:              ":8080",
		ReadHeaderTimeout: 3 * time.Second,
		ReadTimeout:       20 * time.Second,
		WriteTimeout:      20 * time.Second,
		IdleTimeout:       120 * time.Second,
	}

	// Start server with custom server config
	logger.Info("Server starting on :8080")
	e.Logger.Fatal(e.StartServer(server))
}