query-orchestration/cmd/cognito_test/cognitotest/main.go

package main

import (
	"context"
	"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
	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
}

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

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

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

	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")

	// Add Authorization header if client secret is provided
	if config.ClientSecret != "" {
		req.SetBasicAuth(config.ClientID, config.ClientSecret)
		logger.Debug("Using Basic Auth authentication")
	} else {
		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, 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)
}

// 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")
	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
		"/query":                {"exporters"}, // 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
	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{
		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))
}

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