query-orchestration/internal/cognitoauth/auth.go

package cognitoauth

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

	"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
func RegisterRoutes(e *echo.Echo, config auth.ConfigProvider) {
	// check the auth feature flag
	enableAuth := os.Getenv("ENABLE_AUTH")
	if enableAuth == "" || strings.ToLower(enableAuth) == "false" {
		return // Do nothing if auth is not enabled
	}

	// Only proceed if ENABLE_AUTH is "true" or "auth"
	if strings.ToLower(enableAuth) != "true" {
		return
	}

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

	// Register logout route
	e.GET(config.GetAuthLogoutPath(), func(c echo.Context) error {
		return LogoutHandler(c, config)
	})

	// Default home page handler if requested
	e.GET(config.GetAuthHomePath(), func(c echo.Context) error {
		return HomeHandler(c, config)
	})

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

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

// HomeHandler implements a simple home page that shows auth status and available endpoints
func HomeHandler(c echo.Context, config auth.ConfigProvider) error {
	// Create a list of all registered routes
	var endpoints []string

	// Add the auth routes
	endpoints = append(endpoints, []string{
		config.GetAuthLoginPath(),
		config.GetAuthCallbackPath(),
		config.GetAuthHomePath(),
		config.GetAuthLogoutPath(),
	}...)

	// Add routes from the permission map
	for route := range config.GetAuthRoutePermissions() {
		// Skip routes that are already in the list
		alreadyAdded := false
		for _, endpoint := range endpoints {
			if route == endpoint {
				alreadyAdded = true
				break
			}
		}
		if !alreadyAdded {
			endpoints = append(endpoints, route)
		}
	}

	// Check if user is authenticated by looking for token in cookie
	tokenCookie, err := c.Cookie("auth_token")
	isAuthenticated := (err == nil && tokenCookie.Value != "")

	// Variables for user info
	username := ""
	email := ""
	var userGroups []string

	// If authenticated, try to decode the JWT token to get user info
	if isAuthenticated {
		// Parse the JWT token without verification (just to extract info for display)
		token, _ := jwt.Parse([]byte(tokenCookie.Value), jwt.WithVerify(false))
		if token != nil {
			claims, _ := token.AsMap(context.Background())
			username, _ = claims["cognito:username"].(string)
			email, _ = claims["email"].(string)

			// Try to extract groups
			userGroups, _ = GetUserGroups(claims)
		}
	}

	// Create HTML for the endpoints list
	var linksHTML string
	baseURL := c.Scheme() + "://" + c.Request().Host

	// Create list items for each endpoint
	for _, endpoint := range endpoints {
		// Skip endpoints with path parameters for direct linking
		if strings.Contains(endpoint, ":") {
			displayPath := strings.Replace(endpoint, ":id", "{id}", -1)
			linksHTML += fmt.Sprintf("<li>%s (requires parameter)</li>\n", displayPath)
		} else {
			linksHTML += fmt.Sprintf("<li><a href=\"%s%s\">%s</a></li>\n", baseURL, endpoint, endpoint)
		}
	}

	// Create authentication status section
	var authStatusHTML string
	if isAuthenticated {
		authStatusHTML = fmt.Sprintf(`
		<div class="auth-status authenticated">
			<h2>Authentication Status: Authenticated</h2>
			<p><strong>Username:</strong> %s</p>
			<p><strong>Email:</strong> %s</p>
			<p><strong>Groups:</strong> %s</p>
			<p><a href="%s/logout" class="logout-btn">Logout</a></p>
		</div>
		`, username, email, strings.Join(userGroups, ", "), baseURL)
	} else {
		authStatusHTML = fmt.Sprintf(`
		<div class="auth-status unauthenticated">
			<h2>Authentication Status: Not Authenticated</h2>
			<p>You are not currently logged in.</p>
			<p><a href="%s/login" class="login-btn">Login</a></p>
		</div>
		`, baseURL)
	}

	// Create the complete HTML page
	html := fmt.Sprintf(`
<!DOCTYPE html>
<html>
<head>
    <title>Authentication Home</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            line-height: 1.6;
            margin: 20px;
            max-width: 800px;
            margin: 0 auto;
            padding: 20px;
        }
        h1, h2 {
            color: #333;
            border-bottom: 1px solid #ddd;
            padding-bottom: 10px;
        }
        ul {
            margin-top: 20px;
        }
        li {
            margin-bottom: 8px;
        }
        a {
            color: #0066cc;
            text-decoration: none;
        }
        a:hover {
            text-decoration: underline;
        }
        .note {
            background-color: #f8f9fa;
            border-left: 4px solid #5bc0de;
            padding: 10px 15px;
            margin-top: 20px;
            font-size: 0.9em;
        }
        .auth-status {
            margin: 20px 0;
            padding: 15px;
            border-radius: 5px;
        }
        .authenticated {
            background-color: #dff0d8;
            border: 1px solid #d6e9c6;
        }
        .unauthenticated {
            background-color: #f2dede;
            border: 1px solid #ebccd1;
        }
        .login-btn, .logout-btn {
            display: inline-block;
            padding: 8px 16px;
            background-color: #0066cc;
            color: white;
            border-radius: 4px;
            text-decoration: none;
            margin-top: 10px;
        }
        .logout-btn {
            background-color: #d9534f;
        }
    </style>
</head>
<body>
    <h1>Authentication Home</h1>
    
    %s
    
    <h2>Available Endpoints</h2>
    <ul>
        %s
    </ul>
    
    <div class="note">
        <p><strong>Note:</strong> This is a debugging page. Some endpoints require authentication or specific permissions.</p>
    </div>
</body>
</html>
	`, authStatusHTML, linksHTML)

	return c.HTML(http.StatusOK, html)
}

// GetTokenFromRequest extracts the token from the request
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 gets user information from the context
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
}

// RequireGroups creates a middleware that checks if the user is in any of the specified groups
func RequireGroups(groups ...string) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			userInfo, ok := GetUserInfo(c)
			if !ok {
				return c.JSON(http.StatusUnauthorized, map[string]string{
					"error": "Authentication required",
				})
			}

			// Check if user has any of the required groups
			for _, requiredGroup := range groups {
				for _, userGroup := range userInfo.Groups {
					if userGroup == requiredGroup {
						return next(c)
					}
				}
			}

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