query-orchestration/internal/cognitoauth/middleware.go

package cognitoauth

import (
	"net/http"
	"time"

	"queryorchestration/internal/serviceconfig/auth"

	"strings"

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

// TokenValidationMiddleware verifies JWT tokens efficiently
// This is the primary middleware that handles both authentication and authorization
// Applied to all routes to enforce security requirements
func TokenValidationMiddleware(config auth.ConfigProvider) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			requestPath := c.Request().URL.Path
			config.GetAuthLogger().Debug("Processing request", "path", requestPath, "method", c.Request().Method)

			// Skip token validation for specified public paths
			publicPaths := []string{config.GetAuthHomePath(), config.GetAuthLogoutPath()}
			for _, path := range publicPaths {
				if requestPath == path {
					config.GetAuthLogger().Debug("Skipping token validation for public path", "path", requestPath)
					return next(c)
				}
			}

			// Special case for the OAuth callback path
			if requestPath == config.GetAuthCallbackPath() && (c.QueryParam("code") != "" || c.QueryParam("error") != "") {
				config.GetAuthLogger().Debug("Handling OAuth callback",
					"has_code", c.QueryParam("code") != "",
					"has_error", c.QueryParam("error") != "")
				// Skip token validation for OAuth callback - it will be handled by the callback handler
				return handleOAuthCallback(c, config)
			}

			// Initialize login flow if this is a login request
			if requestPath == config.GetAuthLoginPath() {
				return initiateLoginWithPKCE(c, config)
			}

			// Get authorization header
			authHeader := c.Request().Header.Get("Authorization")
			if authHeader == "" {
				config.GetAuthLogger().Warn("No authorization header provided")
				return c.JSON(http.StatusUnauthorized, map[string]string{"error": "Authorization required"})
			}

			// Extract token
			tokenStr := authHeader
			if strings.HasPrefix(authHeader, "Bearer ") {
				tokenStr = authHeader[7:]
			}

			// Get JWKS (cached if possible)
			keySet, err := GetJWKS(config.GetAuthJwksURL(), config.GetAuthLogger())
			if err != nil {
				config.GetAuthLogger().Error("Failed to get JWKS", "error", err)
				return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Failed to verify token"})
			}

			// Verify token
			claims, err := verifyToken(tokenStr, keySet, config)
			if err != nil {
				config.GetAuthLogger().Warn("Token verification failed", "error", err)
				return c.JSON(http.StatusUnauthorized, map[string]string{"error": "Invalid token"})
			}

			// Extract user groups
			userGroups, err := GetUserGroups(claims)
			if err != nil {
				config.GetAuthLogger().Warn("Failed to extract user groups", "error", err)
				userGroups = []string{} // Empty array if no groups found
			}

			// Store in context
			c.Set("user_claims", claims)
			c.Set("user_groups", userGroups)

			// Create and store user info in context
			userInfo := ExtractUserInfo(claims)
			c.Set("user_info", userInfo)

			// Check authorization
			authorized, requiredGroups := checkPermissions(requestPath, userGroups, config.GetAuthRoutePermissions(), config.GetAuthLogger())
			if !authorized {
				config.GetAuthLogger().Warn("Access denied",
					"path", requestPath,
					"user", claims["cognito:username"],
					"groups", userGroups,
					"required", requiredGroups)

				return c.JSON(http.StatusForbidden, map[string]interface{}{
					"error":           "Insufficient permissions",
					"message":         "User doesn't have the required group membership",
					"username":        claims["cognito:username"],
					"groups":          userGroups,
					"required_groups": requiredGroups,
				})
			}

			// User is authenticated and authorized
			return next(c)
		}
	}
}

// JWTAuthMiddleware checks for JWT token in cookies and adds it to the Authorization header
func JWTAuthMiddleware(config auth.ConfigProvider) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// Skip for login and callback paths
			if c.Path() == config.GetAuthLoginPath() || c.Path() == config.GetAuthCallbackPath() {
				return next(c)
			}

			// if logout then remove the auth cookie
			if c.Path() == config.GetAuthLogoutPath() {
				cookie := new(http.Cookie)
				cookie.Name = "auth_token"
				cookie.Value = ""
				cookie.Path = "/"
				cookie.Expires = time.Now().Add(-1 * time.Hour) // Set expiry in the past
				cookie.HttpOnly = true
				c.SetCookie(cookie)
			}

			// Skip for home and logout as they should be accessible without auth
			if c.Path() == config.GetAuthHomePath() || c.Path() == config.GetAuthLogoutPath() {
				return next(c)
			}

			// Check if Authorization header is already set
			authHeader := c.Request().Header.Get("Authorization")
			if authHeader != "" && strings.HasPrefix(authHeader, "Bearer ") {
				// Authorization header is already set, proceed to the next handler
				return next(c)
			}

			// Try to get token from cookie
			tokenCookie, err := c.Cookie("auth_token")
			if err != nil || tokenCookie.Value == "" {
				// No token cookie found
				return c.Redirect(http.StatusFound, config.GetAuthLoginPath())
			}

			// Add token to Authorization header
			c.Request().Header.Set("Authorization", "Bearer "+tokenCookie.Value)

			return next(c)
		}
	}
}