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

package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"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"
)

// 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
	ClientSecret string // Client secret for authenticated clients
}

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

// PkceTokenValidationMiddleware verifies JWT tokens efficiently
// This is the primary middleware that handles both authentication and authorization
// Applied to all routes to enforce security requirements
func PkceTokenValidationMiddleware(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 - first thing after we arrive at /login
			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
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
}

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

func getJwksURL(region, userPoolID string) string {
	// original logic
	//       :    fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json", region, userPoolID),
	return fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json", region, userPoolID)
}

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

	fmt.Printf("Using domain %s\n", domain)

	// 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"),
		ClientSecret: os.Getenv("COGNITO_CLIENT_SECRET"),
		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:    getJwksURL(region, userPoolID),
		UserPoolID: userPoolID,
	}

	// Route permissions - matching the provided example
	routePermissions := map[string][]string{
		"/users":                {"exporters", "uploaders", "querybuilders"},
		"/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":                {}, // No permissions required for login
		"/home":                 {}, // No permissions required for home page
		"/logout":               {}, // 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(PkceTokenValidationMiddleware(config, routePermissions, logger))

	// Login route entirely handled by middleware - nothing to do here.
	e.GET("/login", func(c echo.Context) error {
		// The PkceTokenValidationMiddleware 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 - just for demonstration
	e.GET("/users", func(c echo.Context) error {
		return c.String(http.StatusOK, "Users endpoint (OK)")
	})

	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 (OK): %s", id))
	})

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

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

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

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

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