query-orchestration/internal/cognitoauth/token.go

package cognitoauth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"time"

	"queryorchestration/internal/serviceconfig/auth"

	"github.com/lestrrat-go/jwx/v2/jwk"
	"github.com/lestrrat-go/jwx/v2/jwt"
)

const PKCE_SESSION_TTL = 5 * time.Minute

// verifyToken verifies the JWT token and returns its claims.
//
// This function validates the provided JWT token against a JWK key set,
// verifying signature, expiration, issuer, and other standard JWT claims.
// It's used during both OAuth callback and subsequent API requests with bearer tokens.
// When NO_JWT_VALIDATION is enabled, it bypasses signature validation for testing.
//
// Parameters:
//   - tokenString: The raw JWT token string to verify
//   - keySet: The JWK set containing keys to verify the token signature
//   - config: Configuration provider containing auth settings like region and user pool ID
//
// Returns:
//   - map[string]interface{}: The verified token claims as a map
//   - error: Error if token verification fails for any reason
func verifyToken(tokenString string, keySet jwk.Set, config auth.ConfigProvider) (map[string]interface{}, error) {
	// Check if JWT validation is disabled for testing
	if config.GetNoJWTValidation() {
		return verifyTokenTestMode(tokenString, config)
	}

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

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

// verifyTokenTestMode parses JWT without signature validation for testing.
//
// This function is used when NO_JWT_VALIDATION=true is set. It parses the JWT
// without verifying the signature, validates basic structure, and extracts claims.
// This allows testing with unsigned tokens while still extracting user information.
//
// Parameters:
//   - tokenString: The raw JWT token string to parse
//   - config: Configuration provider containing auth settings
//
// Returns:
//   - map[string]interface{}: The token claims as a map
//   - error: Error if token parsing fails
func verifyTokenTestMode(tokenString string, config auth.ConfigProvider) (map[string]interface{}, error) {
	logger := config.GetAuthLogger()
	logger.Warn("JWT validation disabled - using test mode. This should only be used for testing!")

	// Handle RFC-compliant unsecured JWTs that may have only 2 parts (header.payload.)
	// The JWT library expects 3 parts, so we need to ensure proper format
	parts := strings.Split(tokenString, ".")
	if len(parts) == 2 {
		// RFC 7519 Section 6 compliant unsecured JWT (header.payload.)
		// Add empty signature part to make it parseable by the JWT library
		tokenString = tokenString + "."
	} else if len(parts) != 3 {
		return nil, fmt.Errorf("invalid JWT format: expected 2 or 3 parts, got %d", len(parts))
	}

	// Parse token without signature verification
	parsedToken, err := jwt.Parse(
		[]byte(tokenString),
		jwt.WithVerify(false),   // Skip signature verification
		jwt.WithValidate(false), // Skip expiration and other validations
	)
	if err != nil {
		return nil, fmt.Errorf("failed to parse token in test mode: %w", err)
	}

	// Extract claims to a map
	claims, err := parsedToken.AsMap(context.Background())
	if err != nil {
		return nil, fmt.Errorf("failed to extract claims in test mode: %w", err)
	}

	return claims, nil
}

// GetUserGroups extracts user groups from the token claims.
//
// This function attempts to find groups in different claim names depending
// on the Cognito setup. It checks several common claim names where group
// information might be stored, handling different data formats.
//
// Parameters:
//   - claims: A map of token claims extracted from a verified JWT token
//
// Returns:
//   - []string: An array of group names that the user belongs to
//   - error: Error if no groups are found in any of the expected claim fields
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")
}

// createCodeChallenge creates a code challenge from a code verifier.
//
// This function implements the S256 code challenge method as specified in
// PKCE (Proof Key for Code Exchange) RFC 7636. It creates a SHA256 hash
// of the verifier and encodes it as base64url without padding.
//
// Parameters:
//   - verifier: The PKCE code verifier string to transform
//
// Returns:
//   - string: The generated code challenge for use in the authorization request
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
}

// generateRandomString creates a cryptographically secure random string.
//
// This function generates a cryptographically secure random byte sequence
// of the specified length and encodes it as a URL-safe base64 string.
// It's used for generating the state parameter and code verifier in the
// OAuth/PKCE flow.
//
// Parameters:
//   - length: The desired length of the random string
//
// Returns:
//   - string: The generated random string
//   - error: Error if random bytes generation fails
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
}

// storePKCESession stores the PKCE session for later verification.
//
// This function creates and stores a new PKCE session in the in-memory session map,
// associating it with the provided state parameter. It also performs cleanup of
// expired sessions. In a production environment, this should be replaced with
// a proper persistent session store.
//
// Parameters:
//   - state: The OAuth state parameter that serves as the session identifier
//   - codeVerifier: The PKCE code verifier to store for later validation
//   - logger: A structured logger for recording debug information
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(PKCE_SESSION_TTL),
	}

	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.
//
// This function looks up a PKCE session by its state parameter and returns
// the associated code verifier if the session exists and hasn't expired.
// It's used during the token exchange step of the OAuth/PKCE flow.
//
// Parameters:
//   - state: The OAuth state parameter that serves as the session identifier
//   - logger: A structured logger for recording debug information
//
// Returns:
//   - string: The code verifier associated with the state
//   - error: Error if no session is found or if the session has expired
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
}

// ExtractUserInfo extracts user information from JWT claims.
//
// This function parses a verified JWT claims map and extracts relevant user
// information such as username, email, and group memberships. It creates and
// returns a UserInfo struct with the extracted data.
//
// Parameters:
//   - claims: A map of token claims extracted from a verified JWT token
//
// Returns:
//   - UserInfo: A struct containing the extracted user information
func ExtractUserInfo(claims map[string]interface{}) UserInfo {
	userInfo := UserInfo{}

	// Extract username
	if username, ok := claims["cognito:username"].(string); ok {
		userInfo.Username = username
	}

	// Extract email
	if email, ok := claims["email"].(string); ok {
		userInfo.Email = email
	}

	// Extract groups
	groups, _ := GetUserGroups(claims)
	userInfo.Groups = groups

	return userInfo
}

// isTokenExpired checks if a JWT token is expired or will expire soon.
//
// This function parses the JWT token without verification to extract the
// expiration claim and checks if the token is expired or will expire within
// a buffer period (5 minutes). This allows for proactive token refresh.
// For malformed tokens, it returns false to let the normal validation handle them.
//
// Parameters:
//   - tokenString: The JWT token string to check
//
// Returns:
//   - bool: True if the token is expired or will expire soon, false otherwise
func isTokenExpired(tokenString string) bool {
	if tokenString == "" {
		return false // Let normal validation handle empty tokens
	}

	// Parse token without signature verification to get claims
	parsedToken, err := jwt.Parse(
		[]byte(tokenString),
		jwt.WithVerify(false),   // Skip signature verification
		jwt.WithValidate(false), // Skip expiration validation (we'll check manually)
	)
	if err != nil {
		return false // Let normal validation handle malformed tokens
	}

	// Extract expiration time
	exp := parsedToken.Expiration()
	if exp.IsZero() {
		return false // Let normal validation handle tokens without expiration
	}

	// Check if token is expired or will expire within 5 minutes (buffer)
	bufferTime := 5 * time.Minute
	return time.Now().Add(bufferTime).After(exp)
}