query-orchestration/internal/cognitoauth/auth.go

package cognitoauth

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"time"

	"queryorchestration/internal/serviceconfig/auth"

	"strings"

	"github.com/labstack/echo/v4"
	"github.com/lestrrat-go/jwx/v2/jwt"
)

// RegisterRoutes registers all authentication-related routes to the Echo engine.
// It configures login and callback endpoints, and applies JWT authentication middleware.
// If the DISABLE_AUTH environment variable is set to "true", authentication is bypassed.
//
// Parameters:
//   - e: The Echo instance to register routes on
//   - config: Configuration provider for authentication settings
func RegisterRoutes(e *echo.Echo, config auth.ConfigProvider) {
	// check the auth feature flag
	fmt.Println("Registering routes for Auth.")
	disableAuth := os.Getenv("DISABLE_AUTH")
	if strings.ToLower(disableAuth) == "true" {
		fmt.Println("Auth is disabled. Skipping auth routes.")
		return // Do nothing if auth is not enabled
	}

	// Register the login route - the middleware will handle this
	e.GET(config.GetAuthLoginPath(), func(c echo.Context) error {
		// This is handled by the middleware
		return nil
	})

	// Register callback route - the middleware will handle this
	e.GET(config.GetAuthCallbackPath(), 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,
		})
	})

	// Apply the JWT Auth middleware first
	e.Use(JWTAuthMiddleware(config))

	// Then apply the token validation middleware
	e.Use(TokenValidationMiddleware(config))
}

// GetTokenFromRequest extracts the JWT token from the HTTP request.
// It first attempts to retrieve the token from the Authorization header
// in Bearer token format. If not found, it then tries to retrieve the token
// from the "auth_token" cookie. Returns an empty string if no token is found.
//
// Parameters:
//   - c: The Echo context containing the HTTP request
//
// Returns:
//   - string: The extracted JWT token, or an empty string if not found
func GetTokenFromRequest(c echo.Context) string {
	// First try from Authorization header
	authHeader := c.Request().Header.Get("Authorization")
	if authHeader != "" && strings.HasPrefix(authHeader, "Bearer ") {
		return authHeader[7:]
	}

	// Then try from cookie
	tokenCookie, err := c.Cookie("auth_token")
	if err == nil && tokenCookie.Value != "" {
		return tokenCookie.Value
	}

	return ""
}

// GetUserInfo retrieves user information from the Echo context.
// It first attempts to retrieve user info directly from the context.
// If not found there, it tries to extract and parse the JWT token,
// then extract user information from the token claims.
//
// Parameters:
//   - c: The Echo context containing the HTTP request and context values
//
// Returns:
//   - UserInfo: Structure containing username, email, and groups
//   - bool: True if user information was successfully retrieved, false otherwise
func GetUserInfo(c echo.Context) (UserInfo, bool) {
	// Try to get user info from context
	userInfo, ok := c.Get("user_info").(UserInfo)
	if ok {
		return userInfo, true
	}

	// Try to get token and extract user info
	token := GetTokenFromRequest(c)
	if token == "" {
		return UserInfo{}, false
	}

	// Parse token to extract claims
	parsedToken, err := jwt.Parse([]byte(token), jwt.WithVerify(false))
	if err != nil {
		return UserInfo{}, false
	}

	claims, err := parsedToken.AsMap(context.Background())
	if err != nil {
		return UserInfo{}, false
	}

	// Extract user info
	userInfo = ExtractUserInfo(claims)
	return userInfo, true
}

// GetUserSubject extracts the subject ID from JWT claims for Permit.io
// The subject ('sub') claim uniquely identifies the user in Cognito
//
// Parameters:
//   - c: The Echo context containing the HTTP request and context values
//
// Returns:
//   - string: The user's subject ID from the JWT token
//   - bool: True if subject was successfully extracted, false otherwise
func GetUserSubject(c echo.Context) (string, bool) {
	userClaims, ok := c.Get("user_claims").(map[string]interface{})
	if !ok {
		return "", false
	}

	subject, ok := userClaims["sub"].(string)
	if !ok {
		return "", false
	}

	return subject, true
}

// setTokenCookies sets both access and refresh token cookies.
//
// This function creates HTTP-only cookies for both the access token and refresh token.
// The access token cookie expires with the token, while the refresh token cookie
// has a longer expiration (30 days) to allow for token refresh operations.
//
// Parameters:
//   - c: The Echo context for setting cookies
//   - tokens: TokenResponse containing both access and refresh tokens
//   - accessTokenExpiresAt: Expiration time for the access token
//   - config: Configuration provider for auth settings
func setTokenCookies(c echo.Context, tokens *TokenResponse, accessTokenExpiresAt time.Time, config auth.ConfigProvider) {
	// Log key cookie setting information for debugging
	config.GetAuthLogger().Debug("Setting token cookies",
		"access_token_len", len(tokens.AccessToken),
		"refresh_token_len", len(tokens.RefreshToken))

	// Set access token cookie (shorter expiration)
	// Ensure proper UTC time handling
	accessTokenExpiresUTC := accessTokenExpiresAt.UTC()
	accessTokenCookie := &http.Cookie{
		Name:     "auth_token",
		Value:    tokens.AccessToken,
		Path:     "/",
		Expires:  accessTokenExpiresUTC,
		HttpOnly: true,
		Secure:   isSecureContext(config),
		SameSite: http.SameSiteLaxMode, // Modern browser compatibility
	}
	c.SetCookie(accessTokenCookie)

	// Set refresh token cookie (longer expiration - 30 days)
	// Use UTC time and ensure proper expiration
	refreshTokenExpires := time.Now().UTC().Add(30 * 24 * time.Hour)
	refreshTokenCookie := &http.Cookie{
		Name:     "refresh_token",
		Value:    tokens.RefreshToken,
		Path:     "/",
		Expires:  refreshTokenExpires,
		HttpOnly: true,
		Secure:   isSecureContext(config),
		SameSite: http.SameSiteLaxMode, // Modern browser compatibility
	}
	c.SetCookie(refreshTokenCookie)

	config.GetAuthLogger().Debug("Token cookies set successfully")

	config.GetAuthLogger().Debug("Token cookies set",
		"access_token_expires", accessTokenExpiresAt,
		"refresh_token_expires", refreshTokenCookie.Expires)
}

// clearTokenCookies removes both access and refresh token cookies.
//
// This function clears authentication cookies by setting them to empty values
// with past expiration dates, effectively removing them from the client.
//
// Parameters:
//   - c: The Echo context for setting cookies
//   - config: Configuration provider for auth settings
func clearTokenCookies(c echo.Context, config auth.ConfigProvider) {
	// Log when clearTokenCookies is called for debugging
	config.GetAuthLogger().Debug("Clearing auth cookies", "path", c.Request().URL.Path)

	// Use UTC time for consistent cookie clearing
	pastTime := time.Now().UTC().Add(-1 * time.Hour)

	// Clear access token cookie - must match original attributes exactly
	accessTokenCookie := &http.Cookie{
		Name:     "auth_token",
		Value:    "",
		Path:     "/",
		Expires:  pastTime,
		HttpOnly: true,
		Secure:   isSecureContext(config),
		SameSite: http.SameSiteLaxMode, // Must match original
	}
	c.SetCookie(accessTokenCookie)

	// Clear refresh token cookie - must match original attributes exactly
	refreshTokenCookie := &http.Cookie{
		Name:     "refresh_token",
		Value:    "",
		Path:     "/",
		Expires:  pastTime,
		HttpOnly: true,
		Secure:   isSecureContext(config),
		SameSite: http.SameSiteLaxMode, // Must match original
	}
	c.SetCookie(refreshTokenCookie)

	config.GetAuthLogger().Debug("Auth cookies cleared successfully")
	config.GetAuthLogger().Debug("Token cookies cleared")
}

// getRefreshTokenFromCookie extracts the refresh token from the request cookies.
//
// Parameters:
//   - c: The Echo context containing the HTTP request
//   - config: Configuration provider for auth settings
//
// Returns:
//   - string: The refresh token value, or empty string if not found
func getRefreshTokenFromCookie(c echo.Context, config auth.ConfigProvider) string {
	// Check if we have cookies
	cookies := c.Cookies()
	config.GetAuthLogger().Debug("Checking for refresh token cookie",
		"cookie_count", len(cookies),
		"path", c.Request().URL.Path)

	refreshCookie, err := c.Cookie("refresh_token")
	if err != nil {
		config.GetAuthLogger().Debug("Failed to get refresh_token cookie", "error", err)
		return ""
	}
	if refreshCookie.Value == "" {
		config.GetAuthLogger().Debug("Refresh token cookie is empty")
		return ""
	}

	config.GetAuthLogger().Debug("Found refresh token cookie", "token_length", len(refreshCookie.Value))
	return refreshCookie.Value
}

// isSecureContext determines if cookies should be marked as Secure.
//
// In production environments, cookies should be marked as Secure to ensure
// they are only sent over HTTPS connections.
//
// Parameters:
//   - config: Configuration provider for auth settings
//
// Returns:
//   - bool: True if cookies should be marked as Secure
func isSecureContext(config auth.ConfigProvider) bool {
	// For now, return false to match existing comment in original code
	// In production, this should return true based on environment detection
	return false
}