query-orchestration/internal/cognitoauth/handler.go

package cognitoauth

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"

	"queryorchestration/internal/serviceconfig/auth"

	"strings"
	"time"

	"github.com/labstack/echo/v4"
)

// initiateLoginWithPKCE initiates the OAuth 2.0 authorization flow with PKCE (Proof Key for Code Exchange)
// for Amazon Cognito authentication.
//
// This function:
// 1. Generates a random state parameter to prevent CSRF attacks
// 2. Creates a code verifier and its corresponding code challenge for PKCE
// 3. Stores the PKCE parameters in a session for later verification
// 4. Builds the authorization URL with all required OAuth and PKCE parameters
// 5. Redirects the user to the Cognito login page
//
// Parameters:
//   - c: The Echo context for handling the HTTP request and response
//   - config: Configuration provider that supplies auth-related settings
//
// Returns:
//   - An error if any step of the process fails, or nil on successful redirect
func initiateLoginWithPKCE(c echo.Context, config auth.ConfigProvider) error {
	// Generate random state parameter to prevent CSRF
	state, err := generateRandomString(32)
	if err != nil {
		config.GetAuthLogger().Error("Failed to generate state parameter", "error", err)
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Authentication initialization failed"})
	}

	// Generate code verifier (random string between 43-128 chars)
	codeVerifier, err := generateRandomString(64)
	if err != nil {
		config.GetAuthLogger().Error("Failed to generate code verifier", "error", err)
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Authentication initialization failed"})
	}

	// Create code challenge from verifier (SHA256 + Base64URL without padding)
	codeChallenge := createCodeChallenge(codeVerifier)

	// Store in session for later verification (with expiration time)
	storePKCESession(state, codeVerifier, config.GetAuthLogger())

	// Build authorization URL with PKCE parameters
	params := url.Values{}
	params.Set("client_id", config.GetAuthClientID())
	params.Set("response_type", "code")
	params.Set("redirect_uri", config.GetAuthRedirectURI())
	params.Set("scope", "openid email profile")
	params.Set("state", state)
	params.Set("code_challenge", codeChallenge)
	params.Set("code_challenge_method", "S256")

	authURL := fmt.Sprintf("%s?%s", config.GetAuthURL(), params.Encode())

	if os.Getenv("DEBUG") == "true" {
		config.GetAuthLogger().Debug("Initiating login with PKCE",
			"code_verifier", codeVerifier,
			"code_challenge", codeChallenge,
			"state", state,
			"redirect_uri", config.GetAuthRedirectURI())
	}

	// Redirect user to Cognito login page
	return c.Redirect(http.StatusFound, authURL)
}

// exchangeCodeForTokensWithPKCE exchanges an authorization code for OAuth tokens using PKCE.
//
// This function performs the second part of the OAuth 2.0 authorization code flow with PKCE:
// 1. Builds a token request with the authorization code, code verifier, and other required parameters
// 2. Sends the request to the Cognito token endpoint
// 3. Processes the response to extract the tokens
//
// It handles both confidential clients (with a client secret) and public clients (without a secret)
// by adapting the authentication method used with the token endpoint.
//
// Parameters:
//   - authCode: The authorization code received from the OAuth provider (Cognito)
//   - codeVerifier: The original code verifier that was used to generate the code challenge
//   - config: Configuration provider that supplies auth-related settings
//
// Returns:
//   - A pointer to a TokenResponse containing the ID token, access token, and refresh token
//   - An error if the token exchange fails for any reason
func exchangeCodeForTokensWithPKCE(authCode, codeVerifier string, config auth.ConfigProvider) (*TokenResponse, error) {
	data := url.Values{}
	data.Set("grant_type", "authorization_code")
	data.Set("client_id", config.GetAuthClientID())
	data.Set("code", authCode)
	data.Set("redirect_uri", config.GetAuthRedirectURI())
	data.Set("code_verifier", codeVerifier) // Include code verifier for PKCE

	config.GetAuthLogger().Debug("Token request details",
		"url", config.GetAuthTokenURL(),
		"client_id", config.GetAuthClientID(),
		"redirect_uri", config.GetAuthRedirectURI())

	req, err := http.NewRequest("POST", config.GetAuthTokenURL(), 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.GetAuthClientSecret() != "" {
		req.SetBasicAuth(config.GetAuthClientID(), config.GetAuthClientSecret())
		config.GetAuthLogger().Debug("Using Basic Auth authentication")
	} else {
		config.GetAuthLogger().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 TokenResponse
	if err := json.Unmarshal(body, &tokenResponse); err != nil {
		return nil, fmt.Errorf("failed to parse token response: %w", err)
	}

	return &tokenResponse, nil
}

// handleOAuthCallback processes the callback from Cognito after a user has authenticated.
//
// This function completes the OAuth 2.0 authorization code flow with PKCE by:
// 1. Extracting and validating the authorization code and state parameters
// 2. Retrieving the stored code verifier associated with the state parameter
// 3. Exchanging the authorization code for tokens using the code verifier
// 4. Verifying the received tokens using the JSON Web Key Set (JWKS)
// 5. Extracting user information and group memberships from the tokens
// 6. Checking if the user has the required permissions based on their group memberships
// 7. Setting the access token as a cookie and redirecting to the home page if authorized
//
// Parameters:
//   - c: The Echo context for handling the HTTP request and response
//   - config: Configuration provider that supplies auth-related settings and permissions
//
// Returns:
//   - Various HTTP status codes and error messages based on the outcome of each step
//   - Redirects to the home page on successful authentication and authorization
func handleOAuthCallback(c echo.Context, config auth.ConfigProvider) error {
	// Extract the authorization code and state
	code := c.QueryParam("code")
	if code == "" {
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Missing code parameter"})
	}

	state := c.QueryParam("state")
	if state == "" {
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Missing state 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,
		})
	}

	// Retrieve stored code verifier
	codeVerifier, err := getCodeVerifier(state, config.GetAuthLogger())
	if err != nil {
		config.GetAuthLogger().Error("Failed to retrieve code verifier", "error", err, "state", state)
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Invalid or expired session"})
	}

	// Exchange the code for tokens using PKCE
	tokens, err := exchangeCodeForTokensWithPKCE(code, codeVerifier, config)
	if err != nil {
		config.GetAuthLogger().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
	keySet, err := GetJWKS(config.GetAuthJwksURL(), config.GetAuthLogger())
	if err != nil {
		config.GetAuthLogger().Error("Failed to fetch 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 {
		config.GetAuthLogger().Error("Failed to verify ID token", "error", err)
		return c.JSON(http.StatusUnauthorized, map[string]string{"error": "Invalid token"})
	}

	// Log decoded token for debugging
	config.GetAuthLogger().Debug("ID Token Claims", "claims", idTokenClaims)
	config.GetAuthLogger().Debug("Raw ID Token", "token", tokens.IDToken)
	config.GetAuthLogger().Debug("Raw Access Token", "token", tokens.AccessToken)

	// OAuth callback path is typically allowed for all authenticated users
	// No additional authorization check needed here as the callback handles authentication
	config.GetAuthLogger().Debug("OAuth callback - user authenticated successfully",
		"user", idTokenClaims["cognito:username"],
		"email", idTokenClaims["email"])

	// User is authenticated and authorized

	// Calculate token expiration time based on token's expires_in value
	expiresAt := time.Now().Add(time.Duration(tokens.ExpiresIn) * time.Second)

	// Set a cookie with the access token
	tokenCookie := new(http.Cookie)
	tokenCookie.Name = "auth_token"
	tokenCookie.Value = tokens.AccessToken
	tokenCookie.Path = "/"
	tokenCookie.Expires = expiresAt
	tokenCookie.HttpOnly = true // Not accessible via JavaScript
	// In production, set Secure to true
	// tokenCookie.Secure = true
	c.SetCookie(tokenCookie)

	// Extract username for display purposes (optional)
	username, _ := idTokenClaims["cognito:username"].(string)

	config.GetAuthLogger().Info("User authenticated successfully", "username", username)

	// Redirect to home page
	return c.Redirect(http.StatusFound, config.GetAuthHomePath())
}