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

// IssuerKeysets maps an `iss` claim value (Cognito issuer URL) to the
// JWKS that should validate tokens carrying it. The verifier looks the
// token's issuer up in this map; an `iss` not present is rejected as
// "untrusted issuer". Building the map at startup gives us multi-pool
// support (two-pool plan §4.2) without weakening signature checks.
type IssuerKeysets map[string]jwk.Set

// AuthSourceUser tags JWTs issued by the human Cognito pool (pool A).
const AuthSourceUser = "user"

// AuthSourceService tags JWTs issued by the service-accounts pool
// (pool B). Downstream gates (CheckUserEnvironmentAccess, EULA skip)
// dispatch on this value.
const AuthSourceService = "service"

// MainPoolIssuer returns the iss claim value that tokens minted by the
// human user pool (pool A) must carry. We rely on Cognito's documented
// `https://cognito-idp.<region>.amazonaws.com/<pool-id>` format.
func MainPoolIssuer(config auth.ConfigProvider) string {
	return fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s",
		config.GetAuthRegion(), config.GetAuthUserPoolID())
}

// ServicePoolIssuer returns the iss value for tokens minted by the
// service-accounts pool (pool B). Returns empty string when the
// feature is disabled (COGNITO_SVC_USER_POOL_ID unset), so callers can
// detect the disabled case with a simple emptiness check.
func ServicePoolIssuer(config auth.ConfigProvider) string {
	pool := config.GetAuthSvcUserPoolID()
	if pool == "" {
		return ""
	}
	return fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s",
		config.GetAuthRegion(), pool)
}

// AuthSourceForIssuer maps a verified iss claim to the auth_source tag.
// Returns "" when the issuer matches neither configured pool.
func AuthSourceForIssuer(iss string, config auth.ConfigProvider) string {
	switch iss {
	case MainPoolIssuer(config):
		return AuthSourceUser
	case ServicePoolIssuer(config):
		return AuthSourceService
	default:
		return ""
	}
}

// verifyToken verifies the JWT token against one of the configured
// issuers' JWKS and returns its claims.
//
// We support two issuers (the human pool and the service-accounts
// pool) by:
//  1. Parsing the token without verification to read the iss claim.
//     The peek is intentionally untrusted — its only job is to select
//     a keyset.
//  2. Re-parsing with jwt.WithKeySet(keySet) and jwt.WithIssuer(iss).
//     A forged iss therefore fails because the chosen pool's JWKS
//     does not contain the signer's key.
//
// When NO_JWT_VALIDATION is enabled, this delegates to the test-mode
// parser as before.
//
// Parameters:
//   - tokenString: The raw JWT token string to verify
//   - keysets:     iss -> JWKS lookup built at startup
//   - config:      auth configuration (used only for test-mode hook
//     and for logger access)
//
// Returns the verified claims map, or an error explaining the failure.
func verifyToken(tokenString string, keysets IssuerKeysets, config auth.ConfigProvider) (map[string]interface{}, error) {
	if config.GetNoJWTValidation() {
		return verifyTokenTestMode(tokenString, config)
	}

	// Peek at the issuer without trusting it. We require the token
	// to parse into a real JWT shape but skip signature + claim
	// validation here.
	unverified, err := jwt.Parse(
		[]byte(tokenString),
		jwt.WithVerify(false),
		jwt.WithValidate(false),
	)
	if err != nil {
		return nil, fmt.Errorf("malformed token: %w", err)
	}
	iss := unverified.Issuer()
	if iss == "" {
		return nil, errors.New("token missing iss claim")
	}

	keySet, ok := keysets[iss]
	if !ok {
		return nil, fmt.Errorf("untrusted issuer: %s", iss)
	}

	verified, err := jwt.Parse(
		[]byte(tokenString),
		jwt.WithKeySet(keySet),
		jwt.WithValidate(true),
		jwt.WithIssuer(iss),
	)
	if err != nil {
		return nil, err
	}

	claims, err := verified.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. ID tokens use "cognito:username", access tokens use "username".
	if username, ok := claims["cognito:username"].(string); ok {
		userInfo.Username = username
	} else if username, ok := claims["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)
}