working test harness

cmd/cognito_test/cognitotest/main.go

@@ -6,10 +6,13 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"log/slog"
 	"net/http"
 	"net/url"
 	"os"
+	"regexp"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/labstack/echo/v4"
@@ -19,112 +22,310 @@ import (
 )
 
 // 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"`
-	AccessToken  string `json:"access_token"`
-	RefreshToken string `json:"refresh_token"`
-	ExpiresIn    int    `json:"expires_in"`
-	TokenType    string `json:"token_type"`
+	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 Cognito
+// CognitoConfig holds the configuration for AWS Cognito
+// Contains all necessary parameters to interact with Cognito endpoints
 type CognitoConfig struct {
-	ClientID     string
-	ClientSecret string
-	RedirectURI  string
-	TokenURL     string
-	JwksURL      string
-	UserPoolID   string
+	ClientID     string // OAuth2 client ID registered with Cognito
+	ClientSecret string // Client secret for authenticated clients
+	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
 }
 
-// VerifyCognitoCode exchanges the authorization code for tokens and verifies them
-func VerifyCognitoCode(authCode string, jwksJSON string) (map[string]interface{}, error) {
-	// Get client credentials from environment variables
-	clientID := os.Getenv("COGNITO_CLIENT_ID")
-	clientSecret := os.Getenv("COGNITO_CLIENT_SECRET")
-	userPoolID := os.Getenv("COGNITO_USER_POOL_ID")
-	region := os.Getenv("AWS_REGION")
+// 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
+}
 
-	if clientID == "" {
-		return nil, errors.New("COGNITO_CLIENT_ID environment variable is not set")
+// Global JWKS cache instance
+var jwksCache = &JWKSCache{
+	ExpiresAt: time.Now(), // Initial state is expired, forcing a fetch on first use
+}
+
+// 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
 	}
 
-	if userPoolID == "" {
-		userPoolID = "us-east-2_1y6po8rR8" // Fallback to the hardcoded value
-	}
-
-	if region == "" {
-		region = "us-east-2" // Fallback to the hardcoded value
-	}
-
-	// Get the domain name from environment variable or use a default
-	domain := os.Getenv("COGNITO_DOMAIN")
-	if domain == "" {
-		// Log warning about missing domain
-		fmt.Println("Warning: COGNITO_DOMAIN environment variable is not set, using default")
-		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:     clientID,
-		ClientSecret: clientSecret,
-		RedirectURI:  "http://localhost:8080/query",
-		TokenURL:     fmt.Sprintf("https://%s/oauth2/token", domain),
-		JwksURL:      fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json", region, userPoolID),
-		UserPoolID:   userPoolID,
-	}
-
-	// Debug output
-	fmt.Printf("Configuration:\n")
-	fmt.Printf("  Client ID: %s\n", config.ClientID)
-	fmt.Printf("  Using client secret: %v\n", config.ClientSecret != "")
-	fmt.Printf("  Domain: %s\n", domain)
-	fmt.Printf("  Token URL: %s\n", config.TokenURL)
-	fmt.Printf("  JWKS URL: %s\n", config.JwksURL)
-
-	// Exchange auth code for tokens
-	tokens, err := exchangeAuthCodeForTokens(authCode, config)
+	logger.Info("Fetching new JWKS", "url", jwksURL)
+	jwksJSON, err := fetchJWKS(jwksURL)
 	if err != nil {
-		return nil, fmt.Errorf("failed to exchange auth code for tokens: %w", err)
+		return nil, fmt.Errorf("failed to fetch JWKS: %w", err)
 	}
 
-	// Parse JWKS JSON
 	keySet, err := jwk.ParseString(jwksJSON)
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse JWKS: %w", err)
 	}
 
-	// Verify the ID token
-	idTokenClaims, err := verifyToken(tokens.IDToken, keySet, config)
-	if err != nil {
-		return nil, fmt.Errorf("failed to verify ID token: %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 idTokenClaims, nil
+	return keySet, nil
 }
 
-// exchangeAuthCodeForTokens exchanges the authorization code for tokens
-func exchangeAuthCodeForTokens(authCode string, config CognitoConfig) (*CognitoTokenResponse, error) {
+// 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)
+
+			// Special case for the OAuth callback path
+			if requestPath == "/query" && c.QueryParam("code") != "" {
+				// Skip token validation for OAuth callback - it will be handled by the callback handler
+				return handleOAuthCallback(c, config, routePermissions, 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)
+		}
+	}
+}
+
+// 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
+}
+
+// handleOAuthCallback handles the OAuth 2.0 authorization code flow callback
+// Called when Cognito redirects back to our application with an authorization code
+// Exchanges the code for tokens, verifies them, and checks authorization
+func handleOAuthCallback(c echo.Context, config *CognitoConfig, routePermissions map[string][]string, logger *slog.Logger) error {
+	// Extract the authorization code
+	code := c.QueryParam("code")
+	if code == "" {
+		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Missing code 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,
+		})
+	}
+
+	// Exchange the code for tokens
+	tokens, err := exchangeCodeForTokens(code, *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
+	jwksJSON, err := fetchJWKS(config.JwksURL)
+	if err != nil {
+		logger.Error("Failed to fetch JWKS", "error", err)
+		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Failed to verify token"})
+	}
+
+	// Parse the JWKS
+	keySet, err := jwk.ParseString(jwksJSON)
+	if err != nil {
+		logger.Error("Failed to parse 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 - this was missing and caused the bug
+	authorized, requiredGroups := checkPermissions("/query", 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, return tokens and info
+	return c.JSON(http.StatusOK, map[string]interface{}{
+		"authenticated": true,
+		"username":      idTokenClaims["cognito:username"],
+		"email":         idTokenClaims["email"],
+		"groups":        userGroups,
+		"id_token":      tokens.IDToken,
+		"access_token":  tokens.AccessToken,
+		"token_type":    "Bearer",
+		"expires_in":    tokens.ExpiresIn,
+	})
+}
+
+// exchangeCodeForTokens exchanges the authorization code for tokens
+// 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 exchangeCodeForTokens(authCode 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)
 
-	// Debug output
-	fmt.Printf("Token request details:\n")
-	fmt.Printf("  URL: %s\n", config.TokenURL)
-	fmt.Printf("  client_id: %s\n", config.ClientID)
-	fmt.Printf("  redirect_uri: %s\n", config.RedirectURI)
-	fmt.Printf("  grant_type: authorization_code\n")
+	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, err
+		return nil, fmt.Errorf("failed to create token request: %w", err)
 	}
 
 	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
@@ -132,21 +333,21 @@ func exchangeAuthCodeForTokens(authCode string, config CognitoConfig) (*CognitoT
 	// Add Authorization header if client secret is provided
 	if config.ClientSecret != "" {
 		req.SetBasicAuth(config.ClientID, config.ClientSecret)
-		fmt.Println("  Using Basic Auth authentication")
+		logger.Debug("Using Basic Auth authentication")
 	} else {
-		fmt.Println("  Using public client authentication (no secret)")
+		logger.Debug("Using public client authentication (no secret)")
 	}
 
 	client := &http.Client{Timeout: 10 * time.Second}
 	resp, err := client.Do(req)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("token request failed: %w", err)
 	}
 	defer resp.Body.Close()
 
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to read token response: %w", err)
 	}
 
 	if resp.StatusCode != http.StatusOK {
@@ -155,13 +356,15 @@ func exchangeAuthCodeForTokens(authCode string, config CognitoConfig) (*CognitoT
 
 	var tokenResponse CognitoTokenResponse
 	if err := json.Unmarshal(body, &tokenResponse); err != nil {
-		return nil, err
+		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/") {
@@ -178,7 +381,6 @@ func verifyToken(tokenString string, keySet jwk.Set, config CognitoConfig) (map[
 	}
 
 	issuer := fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s", region, config.UserPoolID)
-	fmt.Printf("Using issuer: %s\n", issuer)
 
 	// Verify the token with the keySet
 	verifiedToken, err := jwt.Parse(
@@ -202,6 +404,8 @@ func verifyToken(tokenString string, keySet jwk.Set, config CognitoConfig) (map[
 }
 
 // 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
@@ -227,105 +431,184 @@ func GetUserGroups(claims map[string]interface{}) ([]string, error) {
 	return nil, errors.New("no groups found in token claims")
 }
 
-// HandleCallback is the Echo handler for OAuth callback
-func HandleCallback(c echo.Context) error {
-	// Extract the authorization code from the query parameters
-	code := c.QueryParam("code")
-	if code == "" {
-		return c.String(http.StatusBadRequest, "Missing code parameter")
-	}
-
-	// Extract error if present
-	errorMsg := c.QueryParam("error")
-	errorDescription := c.QueryParam("error_description")
-	if errorMsg != "" {
-		return c.String(http.StatusBadRequest, fmt.Sprintf("Authorization error: %s - %s", errorMsg, errorDescription))
-	}
-
-	// Get user pool ID and region
-	userPoolID := os.Getenv("COGNITO_USER_POOL_ID")
-	region := os.Getenv("AWS_REGION")
-
-	if userPoolID == "" {
-		userPoolID = "us-east-2_1y6po8rR8" // Fallback to hardcoded value
-	}
-
-	if region == "" {
-		region = "us-east-2" // Fallback to hardcoded value
-	}
-
-	// Fetch the JWKS from Cognito
-	jwksURL := fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json", region, userPoolID)
-	fmt.Printf("Fetching JWKS from: %s\n", jwksURL)
-
-	// Fix G107: Use a safe, validated URL for the HTTP request
-	if !strings.HasPrefix(jwksURL, "https://cognito-idp.") || !strings.Contains(jwksURL, ".amazonaws.com/") {
-		return c.String(http.StatusInternalServerError, "Invalid JWKS URL")
-	}
-
-	// Create a custom client with a timeout to avoid gosec G107 warning
-	client := &http.Client{
-		Timeout: 10 * time.Second,
-	}
+// 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 c.String(http.StatusInternalServerError, fmt.Sprintf("Failed to create request: %v", err))
+		return "", err
 	}
+
 	resp, err := client.Do(req)
 	if err != nil {
-		return c.String(http.StatusInternalServerError, fmt.Sprintf("Failed to fetch JWKS: %v", err))
+		return "", err
 	}
 	defer resp.Body.Close()
 
-	jwksJSON, err := io.ReadAll(resp.Body)
+	jwksData, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return c.String(http.StatusInternalServerError, fmt.Sprintf("Failed to read JWKS response: %v", err))
+		return "", err
 	}
 
-	// Verify the code and get the token claims
-	claims, err := VerifyCognitoCode(code, string(jwksJSON))
-	if err != nil {
-		return c.String(http.StatusUnauthorized, fmt.Sprintf("Failed to verify code: %v", err))
-	}
-
-	// Extract user groups
-	groups, err := GetUserGroups(claims)
-	if err != nil {
-		// Handle the case where groups can't be found, but the token is valid
-		// You might still want to proceed with authenticated but ungrouped user
-		fmt.Printf("Warning: %v\n", err)
-	}
-
-	// Prepare response
-	response := map[string]interface{}{
-		"authenticated": true,
-		"username":      claims["cognito:username"],
-		"email":         claims["email"],
-		"groups":        groups,
-	}
-
-	return c.JSON(http.StatusOK, response)
+	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)
+}
+
+// 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()
 
-	// Middleware
+	// Basic middleware
 	e.Use(middleware.Logger())
 	e.Use(middleware.Recover())
 
-	// Routes
-	e.GET("/query", HandleCallback)
+	// 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")
+	if domain == "" {
+		logger.Warn("COGNITO_DOMAIN environment variable is not set, using default")
+		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"),
+		ClientSecret: os.Getenv("COGNITO_CLIENT_SECRET"),
+		RedirectURI:  "http://localhost:8080/query",
+		TokenURL:     fmt.Sprintf("https://%s/oauth2/token", domain),
+		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"},
+		"/query":                {"exporters", "querybuilders"}, // Only exporters can access /query
+		"/api/inventory/update": {"exporters", "uploaders"},
+	}
+
+	// Apply token validation middleware
+	e.Use(TokenValidationMiddleware(config, routePermissions, logger))
+
+	// Endpoint handlers - much simpler now that auth is fully handled in middleware
+	e.GET("/query", 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")
+	})
 
 	// Print startup information
-	fmt.Println("Server starting on :8080")
-	fmt.Println("Using environment variables:")
-	fmt.Printf("  COGNITO_CLIENT_ID: %s\n", maskString(os.Getenv("COGNITO_CLIENT_ID")))
-	fmt.Printf("  COGNITO_CLIENT_SECRET: %s\n", maskString(os.Getenv("COGNITO_CLIENT_SECRET")))
-	fmt.Printf("  COGNITO_DOMAIN: %s\n", os.Getenv("COGNITO_DOMAIN"))
-	fmt.Printf("  COGNITO_USER_POOL_ID: %s\n", os.Getenv("COGNITO_USER_POOL_ID"))
-	fmt.Printf("  AWS_REGION: %s\n", os.Getenv("AWS_REGION"))
+	logger.Info("Server configuration",
+		"client_id", maskString(config.ClientID),
+		"domain", domain,
+		"user_pool_id", userPoolID,
+		"region", region,
+		"debug_mode", os.Getenv("DEBUG") == "true")
 
 	// Configure server
 	server := &http.Server{
@@ -337,10 +620,12 @@ func main() {
 	}
 
 	// Start server with custom server config
+	logger.Info("Server starting on :8080")
 	e.Logger.Fatal(e.StartServer(server))
 }
 
 // Helper function to mask sensitive values
+// Used for logging sensitive information like client IDs
 func maskString(s string) string {
 	if s == "" {
 		return "<not set>"

cmd/cognito_test/cognitotest/readme.md

@@ -1,198 +1,155 @@
-# Cognito harness for auth tests
-The code in main.go is a harness for running the auth tests. 
+# Cognito Auth Test Harness
 
-To test it set the following environment variables: 
-```json
-export COGNITO_CLIENT_ID="552cqkf3640t39ncehkmgpce31";
-export COGNITO_CLIENT_SECRET="omitted";
-export COGNITO_DOMAIN="us-east-21y6po8rr8.auth.us-east-2.amazoncognito.com";
-export COGNITO_USER_POOL_ID="us-east-2_1y6po8rR8";export AWS_REGION="us-east-2"
+This repository contains a test harness for validating AWS Cognito authentication and authorization with Go and the Echo web framework. It demonstrates a complete OAuth 2.0 flow with role-based access control (RBAC).
+
+## Features
+
+- Complete OAuth 2.0 authentication flow with AWS Cognito
+- Role-based access control (RBAC) for API endpoints
+- JWT token validation with signature verification
+- Efficient JWKS caching to minimize external requests
+- Debug mode for token inspection
+
+## Setup
+
+### Prerequisites
+
+- Go 1.18 or higher
+- AWS Cognito User Pool with configured app client
+- Test user(s) in the Cognito User Pool with assigned groups
+
+### Environment Variables
+
+Set the following environment variables before running the application:
+
+```bash
+export COGNITO_CLIENT_ID="552cqkf3640t39ncehkmgpce31"
+export COGNITO_CLIENT_SECRET="your-client-secret"
+export COGNITO_DOMAIN="us-east-21y6po8rr8.auth.us-east-2.amazoncognito.com"
+export COGNITO_USER_POOL_ID="us-east-2_1y6po8rR8"
+export AWS_REGION="us-east-2"
+
+# Optional: Enable debug mode to see token details
+export DEBUG="true"
 ```
 
-Run the code with `go run main.go` with these ^ environment variables set.
-Try to login to the user group login test page (this will redirect to localhost:8080/query) 
- https://us-east-21y6po8rr8.auth.us-east-2.amazoncognito.com/login?client_id=552cqkf3640t39ncehkmgpce31&response_type=code&scope=email+openid+phone&redirect_uri=http%3A%2F%2Flocalhost%3A8080%2Fquery
+## Running the Application
 
-If it works you should see a response like this:
-```json
-{"authenticated":true,"email":"betot75403@isorax.com","groups":["uploaders","querybuilders"],"username":"testuser"}
+1. Build and run the server:
+
+```bash
+go run main.go
 ```
 
+2. The server will start on `http://localhost:8080`
 
-Now that this code is verified it will be extracted and moved into an echo middleware that will handle the auth for the API.
+## Testing the Authentication Flow
 
-# AWS Cognito Authentication Flow with Localhost - explainer for the code
+### Step 1: Initial Login
 
-This rest of this document explains the step-by-step flow that occurs when a user authenticates through AWS Cognito and gets redirected to your localhost server 
-in this test harness. 
-
-## Prerequisites
-
-- You have an AWS Cognito User Pool set up
-- You have a test user created in this pool
-- Your Cognito app client is configured with `http://localhost:8080/query` as an allowed callback URL
-- The provided Go server is running on your local machine
-
-## Authentication Flow
-
-### 1. Initial Authentication Request
-
-When a user wants to authenticate, they're first directed to the Cognito hosted UI. This happens outside of the code you provided, typically through a link like:
+1. Open the following URL in your browser (update domain/client_id if different):
 
 ```
-https://{domain}.auth.{region}.amazoncognito.com/login?client_id={clientId}&response_type=code&scope=email+openid+profile&redirect_uri=http://localhost:8080/query
+https://us-east-21y6po8rr8.auth.us-east-2.amazoncognito.com/login?client_id=552cqkf3640t39ncehkmgpce31&response_type=code&scope=email+openid+phone&redirect_uri=http%3A%2F%2Flocalhost%3A8080%2Fquery
 ```
 
-Where:
-- `{domain}` is your Cognito domain (from the `COGNITO_DOMAIN` environment variable)
-- `{region}` is your AWS region (from the `AWS_REGION` environment variable)
-- `{clientId}` is your app client ID (from the `COGNITO_CLIENT_ID` environment variable)
+2. Log in with your Cognito user credentials
 
-### 2. User Login
+3. After successful authentication, Cognito will redirect to your local server (`/query` endpoint)
 
-The user logs in with their username and password on the Cognito hosted UI.
-
-### 3. Redirect to Localhost
-
-After successful authentication, Cognito redirects the user to the specified callback URL:
-```
-http://localhost:8080/query?code={authorization_code}
-```
-
-The authorization code is a temporary token that your application can exchange for actual access tokens.
-
-### 4. Server Handles the Callback
-
-When the user's browser loads the redirect URL, your Go server processes the request through the `HandleCallback` function:
-
-1. **Extract Authorization Code**:
-   ```go
-   code := r.URL.Query().Get("code")
-   ```
-
-2. **Check for Errors**:
-   The function checks if Cognito returned any error messages.
-
-3. **Get Configuration**:
-   It retrieves your Cognito User Pool ID and AWS region from environment variables or falls back to hardcoded values.
-
-4. **Fetch JWKS (JSON Web Key Set)**:
-   ```go
-   jwksURL := fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json", region, userPoolID)
-   ```
-
-   The function fetches the JWKS, which contains the public keys needed to verify the JWT tokens that Cognito issues.
-
-### 5. Exchange Authorization Code for Tokens
-
-The `VerifyCognitoCode` function handles the token exchange:
-
-1. **Prepare Request to Token Endpoint**:
-   It creates a request to the Cognito token endpoint with:
-    - `grant_type=authorization_code`
-    - Your client ID
-    - The authorization code
-    - Your redirect URI
-
-2. **Add Authentication**:
-   If you have a client secret, it adds Basic Authentication.
-
-3. **Send the Request**:
-   It sends the request to Cognito's token endpoint:
-   ```
-   https://{domain}.auth.{region}.amazoncognito.com/oauth2/token
-   ```
-
-4. **Parse the Response**:
-   Cognito returns a response containing:
-    - ID token (contains user information)
-    - Access token (for API access)
-    - Refresh token (for getting new tokens)
-    - Token expiration time
-
-### 6. Verify the ID Token
-
-The `verifyToken` function verifies the authenticity of the ID token:
-
-1. **Determine the Issuer**:
-   ```go
-   issuer := fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s", region, config.UserPoolID)
-   ```
-
-2. **Verify Token Signature**:
-   It uses the JWKS (fetched earlier) to verify the token's cryptographic signature.
-
-3. **Extract Claims**:
-   After verification, it extracts all the claims (user information) from the token.
-
-### 7. Extract User Groups
-
-The `GetUserGroups` function attempts to extract the user's group memberships from the token claims:
-
-1. **Check Multiple Possible Fields**:
-   ```go
-   possibleGroupFields := []string{"cognito:groups", "groups", "custom:groups"}
-   ```
-
-   It looks for group information in various possible claim names.
-
-2. **Handle Different Data Types**:
-   It handles cases where groups might be represented as arrays, strings, or other formats.
-
-### 8. Return Response to User
-
-Finally, the server returns a JSON response to the user's browser:
+4. You'll see a JSON response with:
+   - Authentication status
+   - User information (username, email)
+   - User groups
+   - JWT tokens (id_token and access_token)
 
+Example response:
 ```json
 {
   "authenticated": true,
-  "username": "user123",
+  "username": "testuser",
   "email": "user@example.com",
-  "groups": ["admin", "users"]
+  "groups": ["uploaders", "querybuilders"],
+  "id_token": "eyJhbGciOiJSUzI1NiIs...",
+  "access_token": "eyJhbGciOiJSUzI1NiIs...",
+  "token_type": "Bearer",
+  "expires_in": 3600
 }
 ```
 
-This confirms successful authentication and provides basic user information.
+### Step 2: Testing Protected Endpoints
 
-## Debugging Information
+Use the obtained token to access protected endpoints:
 
-Throughout this process, the server outputs various debugging information:
+1. Copy the `id_token` value from the response
 
-- Configuration details (client ID, domain, token URL, etc.)
-- Token request details
-- Authentication method (Basic Auth or public client)
-- JWKS URL
-- Issuer URL
+2. Use it as a Bearer token in subsequent requests:
 
-This information helps in troubleshooting any authentication issues.
+```bash
+# Using curl to access a protected endpoint
+curl -H "Authorization: Bearer eyJhbGciOiJSUzI1NiIs..." http://localhost:8080/users
+```
 
-## Security Features
+3. Test authorization with different endpoints:
 
-The implementation includes several security features:
+| Endpoint | Required Groups |
+|----------|----------------|
+| /users | exporters |
+| /users/:id | exporters, querybuilders |
+| /orders | exporters |
+| /reports/sales | exporters, uploaders, querybuilders |
+| /settings | exporters |
+| /query | exporters |
+| /api/inventory/update | exporters, uploaders |
 
-1. **Token Verification**:
-   It cryptographically verifies the tokens using the JWKS.
+4. If your user doesn't have the required group for an endpoint, you'll receive a 403 Forbidden response:
 
-2. **URL Validation**:
-   It validates the JWKS URL to prevent potential security issues.
+```json
+{
+  "error": "Insufficient permissions",
+  "message": "User doesn't have the required group membership",
+  "username": "testuser",
+  "groups": ["uploaders", "querybuilders"],
+  "required_groups": ["exporters"]
+}
+```
 
-3. **Timeouts**:
-   HTTP requests use timeouts to prevent hanging connections.
+## Authentication Flow Explained
 
-4. **Server Timeouts**:
-   The HTTP server uses timeouts to protect against slow client attacks.
+### Initial Login (OAuth 2.0 Flow)
 
-5. **Masking Sensitive Values**:
-   The `maskString` function masks sensitive values in logs.
+1. **User Authentication**: User authenticates with Cognito's hosted UI
+2. **Authorization Code**: Cognito redirects to `/query?code=...` with an authorization code
+3. **Token Exchange**: The server exchanges this code for OAuth tokens by calling Cognito's token endpoint
+4. **Token Verification**: The server verifies the JWT token signature using Cognito's JWKS
+5. **Group Verification**: The server checks if the user belongs to the required groups for the `/query` endpoint
+6. **Response**: If authorized, the server returns the tokens and user information; if not, it returns a 403 error
 
-## Error Handling
+### Subsequent API Requests
 
-The server handles various error scenarios:
+1. **Token Submission**: Client includes the ID token as a Bearer token in the Authorization header
+2. **Token Validation**: Server validates the token's signature, expiration, and other claims
+3. **Group Extraction**: Server extracts the user's groups from the token
+4. **Permission Check**: Server checks if the user has any of the required groups for the requested endpoint
+5. **Access Control**: If authorized, the request proceeds to the handler; if not, a 403 error is returned
 
-- Missing authorization code
-- Authorization errors from Cognito
-- Failed token requests
-- Invalid or expired tokens
-- Missing user groups
+## Advanced Features
 
-Each error is properly logged and an appropriate HTTP status code is returned.
+### JWKS Caching
+
+The server caches the JSON Web Key Set (JWKS) used to verify token signatures, reducing the number of requests to AWS.
+
+### Debug Mode
+
+Enable debug mode by setting the `DEBUG` environment variable to `true`. This will:
+- Log decoded token claims
+- Log raw tokens
+- Provide more verbose logging throughout the authentication process
+
+## Integration into Your Own Application
+
+This test harness is for proving the design of the echo middleware before attempting integration:
+
+1. The `TokenValidationMiddleware` handles both authentication and authorization
+2. The route permissions map can be configured to match our application's security requirements
+3. The JWKS caching mechanism ensures efficient token validation