query-orchestration/internal/cognitoauth/token.go

package cognitoauth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"log/slog"
	"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
// 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 auth.ConfigProvider) (map[string]interface{}, error) {
	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
}

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

// createCodeChallenge creates a code challenge from a code verifier
// Implements S256 code challenge method as specified in PKCE RFC 7636
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
// Used for generating state parameter and code verifier
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
// In a production environment, this should be replaced with a proper session store
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
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
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
}