jay brown
500 lines
15 KiB
Go
500 lines
15 KiB
Go
package cognitotest
|
|
|
|
import (
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"log/slog"
|
|
"net/http"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/labstack/echo/v4"
|
|
"github.com/labstack/echo/v4/middleware"
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"queryorchestration/internal/cognitoauth"
|
|
"queryorchestration/internal/serviceconfig/auth"
|
|
|
|
"github.com/chromedp/cdproto/network"
|
|
"github.com/chromedp/chromedp"
|
|
)
|
|
|
|
// setupEnvironmentForRefresh ensures all required environment variables are set for refresh testing
|
|
func setupEnvironmentForRefresh(t *testing.T) {
|
|
// Define all required environment variables with test values
|
|
envVars := map[string]string{
|
|
"USER": os.Getenv("USER"),
|
|
"PASSWORD": os.Getenv("PASSWORD"),
|
|
"COGNITO_DOMAIN": os.Getenv("COGNITO_DOMAIN"),
|
|
"AWS_REGION": os.Getenv("AWS_REGION"),
|
|
"COGNITO_CLIENT_SECRET": os.Getenv("COGNITO_CLIENT_SECRET"),
|
|
"COGNITO_USER_POOL_ID": os.Getenv("COGNITO_USER_POOL_ID"),
|
|
"COGNITO_CLIENT_ID": os.Getenv("COGNITO_CLIENT_ID"),
|
|
"DEBUG": "false",
|
|
"HEADLESS": os.Getenv("HEADLESS"),
|
|
}
|
|
|
|
// Verify all required variables are set
|
|
for key, value := range envVars {
|
|
if value == "" {
|
|
// HEADLESS is optional with a default of "true"
|
|
if key == "HEADLESS" {
|
|
os.Setenv(key, "true")
|
|
continue
|
|
}
|
|
t.Fatalf("Required environment variable %s is not set", key)
|
|
}
|
|
// Set the environment variable for the test
|
|
os.Setenv(key, value)
|
|
}
|
|
}
|
|
|
|
// setupTestServerForRefresh creates and configures an Echo server for refresh testing
|
|
func setupTestServerForRefresh(t *testing.T) (*http.Server, *auth.CognitoConfig) {
|
|
// Create a new Echo instance
|
|
e := echo.New()
|
|
|
|
// Add basic middleware
|
|
e.Use(middleware.Logger())
|
|
e.Use(middleware.Recover())
|
|
|
|
// Initialize logger
|
|
logHandler := slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
|
|
Level: slog.LevelDebug,
|
|
})
|
|
logger := slog.New(logHandler)
|
|
|
|
// Create auth config with correct base URL that matches Cognito settings
|
|
baseURL := "http://localhost:8080"
|
|
config := auth.InitializeConfig(baseURL, logger)
|
|
|
|
// Set route permissions based on test requirements
|
|
routePermissions := map[string][]string{
|
|
"/query": {"exporters", "uploaders", "querybuilders"},
|
|
}
|
|
config.SetAuthRoutePermissions(routePermissions)
|
|
|
|
// Register Cognito auth routes and middleware
|
|
cognitoauth.RegisterRoutes(e, config)
|
|
|
|
// Add test endpoints
|
|
e.GET("/home", func(c echo.Context) error {
|
|
// Check if user is authenticated
|
|
userInfo, authenticated := cognitoauth.GetUserInfo(c)
|
|
|
|
response := map[string]interface{}{
|
|
"authenticated": authenticated,
|
|
}
|
|
|
|
if authenticated {
|
|
response["username"] = userInfo.Username
|
|
response["email"] = userInfo.Email
|
|
response["groups"] = userInfo.Groups
|
|
}
|
|
|
|
return c.JSON(http.StatusOK, response)
|
|
})
|
|
|
|
e.GET("/query", func(c echo.Context) error {
|
|
return c.JSON(http.StatusOK, map[string]string{
|
|
"status": "success",
|
|
"message": "You have access to the query endpoint",
|
|
})
|
|
})
|
|
|
|
// Create server with fixed address to match Cognito config
|
|
server := &http.Server{
|
|
Addr: ":8080",
|
|
Handler: e,
|
|
ReadHeaderTimeout: 30 * time.Second,
|
|
}
|
|
|
|
// Log the redirect URL for debugging
|
|
logger.Info("Using redirect URI", "uri", config.GetAuthRedirectURI())
|
|
|
|
return server, config
|
|
}
|
|
|
|
// parseJWTTokenExpiration extracts the expiration time from a JWT token without verification
|
|
func parseJWTTokenExpiration(tokenString string) (time.Time, error) {
|
|
// Split JWT into parts
|
|
parts := strings.Split(tokenString, ".")
|
|
if len(parts) != 3 {
|
|
return time.Time{}, fmt.Errorf("invalid JWT format")
|
|
}
|
|
|
|
// Decode payload (second part)
|
|
payload := parts[1]
|
|
// Add padding if needed
|
|
for len(payload)%4 != 0 {
|
|
payload += "="
|
|
}
|
|
|
|
decoded, err := base64.URLEncoding.DecodeString(payload)
|
|
if err != nil {
|
|
return time.Time{}, fmt.Errorf("failed to decode JWT payload: %v", err)
|
|
}
|
|
|
|
// Parse JSON payload
|
|
var claims map[string]interface{}
|
|
if err := json.Unmarshal(decoded, &claims); err != nil {
|
|
return time.Time{}, fmt.Errorf("failed to parse JWT claims: %v", err)
|
|
}
|
|
|
|
// Extract expiration time
|
|
exp, ok := claims["exp"]
|
|
if !ok {
|
|
return time.Time{}, fmt.Errorf("no exp claim found in JWT")
|
|
}
|
|
|
|
// Convert to time.Time
|
|
switch expValue := exp.(type) {
|
|
case float64:
|
|
return time.Unix(int64(expValue), 0), nil
|
|
default:
|
|
return time.Time{}, fmt.Errorf("invalid exp claim type")
|
|
}
|
|
}
|
|
|
|
// TestAutoRefreshFlow runs the complete end-to-end test for auto refresh functionality
|
|
func TestAutoRefreshFlow(t *testing.T) {
|
|
if os.Getenv("ENABLE_REFRESH_TESTS") != "true" {
|
|
t.Skip("Skipping auto refresh tests - set ENABLE_REFRESH_TESTS=true to run")
|
|
}
|
|
|
|
// Skip if running in CI environment or short test mode
|
|
if testing.Short() {
|
|
t.Skip("Skipping integration test in short mode")
|
|
}
|
|
|
|
// Setup test environment
|
|
setupEnvironmentForRefresh(t)
|
|
|
|
// Get credentials from environment
|
|
username := os.Getenv("USER")
|
|
password := os.Getenv("PASSWORD")
|
|
|
|
// Setup test server
|
|
server, config := setupTestServerForRefresh(t)
|
|
|
|
// Start the server in a goroutine
|
|
go func() {
|
|
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
|
t.Logf("Server error: %v", err)
|
|
}
|
|
}()
|
|
|
|
// Ensure server shuts down at the end of the test
|
|
defer func() {
|
|
ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
|
|
defer cancel()
|
|
if err := server.Shutdown(ctx); err != nil {
|
|
t.Logf("Server shutdown error: %v", err)
|
|
}
|
|
}()
|
|
|
|
// Give the server time to start
|
|
time.Sleep(1 * time.Second)
|
|
|
|
// Log important URLs for debugging
|
|
baseURL := "http://localhost:8080"
|
|
t.Logf("Server URL: %s", baseURL)
|
|
t.Logf("Login URL: %s%s", baseURL, config.GetAuthLoginPath())
|
|
t.Logf("Callback URL: %s", config.GetAuthRedirectURI())
|
|
|
|
// Get headless mode from environment variable
|
|
headless := true
|
|
if os.Getenv("HEADLESS") == "false" {
|
|
headless = false
|
|
}
|
|
|
|
// Setup Chrome options
|
|
opts := append(chromedp.DefaultExecAllocatorOptions[:],
|
|
chromedp.Flag("headless", headless),
|
|
chromedp.Flag("disable-gpu", true),
|
|
chromedp.Flag("no-sandbox", true),
|
|
chromedp.Flag("disable-dev-shm-usage", true),
|
|
chromedp.Flag("host-blocking-patterns", "*prometheus*"),
|
|
chromedp.Flag("disable-cache", true),
|
|
)
|
|
|
|
allocCtx, cancel := chromedp.NewExecAllocator(t.Context(), opts...)
|
|
defer cancel()
|
|
|
|
// Create Chrome context with extended timeout for token expiration testing
|
|
ctx, cancel := chromedp.NewContext(allocCtx, chromedp.WithLogf(log.Printf))
|
|
defer cancel()
|
|
ctx, cancel = context.WithTimeout(ctx, 15*time.Minute) // Extended timeout for token expiration
|
|
defer cancel()
|
|
|
|
// Enable request interception for debugging
|
|
if err := chromedp.Run(ctx, network.Enable()); err != nil {
|
|
t.Fatalf("Failed to enable network: %v", err)
|
|
}
|
|
|
|
// Store cookies to extract JWT tokens
|
|
var cookies []*network.Cookie
|
|
|
|
// Phase 1: Initial Authentication
|
|
t.Run("Initial Authentication", func(t *testing.T) {
|
|
// Navigate to login page
|
|
loginURL := fmt.Sprintf("%s%s", baseURL, config.GetAuthLoginPath())
|
|
err := chromedp.Run(ctx,
|
|
chromedp.Navigate(loginURL),
|
|
chromedp.WaitVisible(`input[name="username"]`, chromedp.ByQuery),
|
|
)
|
|
assert.NoError(t, err, "Failed to navigate to login page")
|
|
|
|
// Submit login credentials
|
|
err = chromedp.Run(ctx,
|
|
chromedp.WaitVisible(`input[name="username"]`, chromedp.ByQuery),
|
|
chromedp.Clear(`input[name="username"]`, chromedp.ByQuery),
|
|
chromedp.SendKeys(`input[name="username"]`, username, chromedp.ByQuery),
|
|
chromedp.Click(`button[type="submit"], button.amplify-button--primary, input[type="Submit"]`, chromedp.ByQuery),
|
|
)
|
|
assert.NoError(t, err, "Failed to submit username")
|
|
|
|
// Enter password
|
|
err = chromedp.Run(ctx,
|
|
chromedp.WaitVisible(`input[name="password"]`, chromedp.ByQuery),
|
|
chromedp.Clear(`input[name="password"]`, chromedp.ByQuery),
|
|
chromedp.SendKeys(`input[name="password"]`, password, chromedp.ByQuery),
|
|
chromedp.Sleep(1*time.Second),
|
|
chromedp.Click(`button[type="submit"], button.amplify-button--primary, input[type="Submit"]`, chromedp.ByQuery),
|
|
)
|
|
assert.NoError(t, err, "Failed to submit password")
|
|
|
|
// Wait for redirection to home page
|
|
err = chromedp.Run(ctx, chromedp.ActionFunc(func(ctx context.Context) error {
|
|
startTime := time.Now()
|
|
timeout := 45 * time.Second
|
|
|
|
for {
|
|
if time.Since(startTime) > timeout {
|
|
return fmt.Errorf("timed out waiting for redirect")
|
|
}
|
|
|
|
var currentURL string
|
|
if err := chromedp.Location(¤tURL).Do(ctx); err != nil {
|
|
return err
|
|
}
|
|
|
|
if strings.Contains(currentURL, "/home") {
|
|
return nil
|
|
}
|
|
|
|
if strings.Contains(currentURL, "/login-callback") {
|
|
time.Sleep(2 * time.Second)
|
|
continue
|
|
}
|
|
|
|
time.Sleep(1 * time.Second)
|
|
}
|
|
}))
|
|
assert.NoError(t, err, "Failed to wait for redirect after login")
|
|
|
|
// Get cookies
|
|
err = chromedp.Run(ctx,
|
|
chromedp.Sleep(2*time.Second),
|
|
chromedp.ActionFunc(func(ctx context.Context) error {
|
|
cookies, err = network.GetCookies().Do(ctx)
|
|
return err
|
|
}),
|
|
)
|
|
assert.NoError(t, err, "Failed to get cookies")
|
|
|
|
t.Log("✅ Initial authentication completed successfully")
|
|
})
|
|
|
|
// Phase 2: Extract and Log Token Expiration
|
|
var authToken string
|
|
var tokenExpiration time.Time
|
|
t.Run("Extract Token and Log Expiration", func(t *testing.T) {
|
|
for _, cookie := range cookies {
|
|
if cookie.Name == "auth_token" {
|
|
authToken = cookie.Value
|
|
break
|
|
}
|
|
}
|
|
assert.NotEmpty(t, authToken, "JWT token cookie not found")
|
|
|
|
var err error
|
|
tokenExpiration, err = parseJWTTokenExpiration(authToken)
|
|
assert.NoError(t, err, "Failed to parse token expiration")
|
|
|
|
t.Logf("✅ Token extracted successfully")
|
|
t.Logf("📅 Token expires at: %s", tokenExpiration.Format(time.RFC3339))
|
|
t.Logf("⏰ Current time: %s", time.Now().Format(time.RFC3339))
|
|
t.Logf("⌛ Time until expiration: %s", time.Until(tokenExpiration))
|
|
t.Logf("🔄 Auto refresh should trigger when token has < 5 minutes remaining")
|
|
})
|
|
|
|
// Phase 3: Wait for Token Expiration
|
|
t.Run("Wait for Token Expiration", func(t *testing.T) {
|
|
// Calculate when the token will expire (plus 1 second buffer)
|
|
timeToWait := time.Until(tokenExpiration) + 1*time.Second
|
|
|
|
// If token expires in more than 10 minutes, something might be wrong
|
|
if timeToWait > 10*time.Minute {
|
|
t.Logf("⚠️ Token expires in %s, which seems longer than expected for a 5-minute token", timeToWait)
|
|
t.Logf("🤔 Proceeding anyway, but you may want to check Cognito token settings")
|
|
}
|
|
|
|
if timeToWait > 0 {
|
|
t.Logf("⏳ Waiting %s for token to expire...", timeToWait)
|
|
|
|
// Wait in chunks and log progress
|
|
start := time.Now()
|
|
for timeToWait > 0 {
|
|
if timeToWait > 30*time.Second {
|
|
time.Sleep(30 * time.Second)
|
|
timeToWait = time.Until(tokenExpiration) + 1*time.Second
|
|
elapsed := time.Since(start)
|
|
t.Logf("⏳ Still waiting... elapsed: %s, remaining: %s", elapsed.Round(time.Second), timeToWait.Round(time.Second))
|
|
} else {
|
|
time.Sleep(timeToWait)
|
|
timeToWait = 0
|
|
}
|
|
}
|
|
}
|
|
|
|
t.Logf("✅ Wait period completed - token should now be expired")
|
|
t.Logf("⏰ Current time: %s", time.Now().Format(time.RFC3339))
|
|
})
|
|
|
|
// Phase 4: Test Auto Refresh
|
|
t.Run("Test Auto Refresh Functionality", func(t *testing.T) {
|
|
// Get cookies before the request
|
|
var cookiesBefore []*network.Cookie
|
|
err := chromedp.Run(ctx,
|
|
chromedp.ActionFunc(func(ctx context.Context) error {
|
|
var err error
|
|
cookiesBefore, err = network.GetCookies().Do(ctx)
|
|
return err
|
|
}),
|
|
)
|
|
assert.NoError(t, err, "Failed to get cookies before refresh test")
|
|
|
|
// Log all cookies for debugging
|
|
cookieNames := make([]string, 0, len(cookiesBefore))
|
|
for _, cookie := range cookiesBefore {
|
|
cookieNames = append(cookieNames, cookie.Name)
|
|
}
|
|
t.Logf("🍪 Cookies before refresh test: %v", cookieNames)
|
|
|
|
var originalAuthToken string
|
|
for _, cookie := range cookiesBefore {
|
|
if cookie.Name == "auth_token" {
|
|
originalAuthToken = cookie.Value
|
|
break
|
|
}
|
|
}
|
|
|
|
if originalAuthToken == "" {
|
|
t.Logf("⚠️ No auth token found in cookies before refresh test")
|
|
} else {
|
|
t.Logf("🔍 Original token (first 20 chars): %s...", originalAuthToken[:20])
|
|
}
|
|
|
|
// Make a request to a protected endpoint to trigger auto refresh
|
|
queryURL := fmt.Sprintf("%s/query", baseURL)
|
|
var queryResponseBody string
|
|
|
|
t.Logf("🚀 Making request to %s to trigger auto refresh...", queryURL)
|
|
|
|
err = chromedp.Run(ctx,
|
|
chromedp.Navigate(queryURL),
|
|
chromedp.Sleep(5*time.Second), // Give time for auto refresh to complete
|
|
chromedp.OuterHTML("body", &queryResponseBody),
|
|
)
|
|
assert.NoError(t, err, "Failed to access protected endpoint")
|
|
|
|
// Get cookies after the request
|
|
var cookiesAfter []*network.Cookie
|
|
err = chromedp.Run(ctx,
|
|
chromedp.ActionFunc(func(ctx context.Context) error {
|
|
var err error
|
|
cookiesAfter, err = network.GetCookies().Do(ctx)
|
|
return err
|
|
}),
|
|
)
|
|
assert.NoError(t, err, "Failed to get cookies after refresh test")
|
|
|
|
// Log all cookies after refresh for debugging
|
|
cookieNamesAfter := make([]string, 0, len(cookiesAfter))
|
|
for _, cookie := range cookiesAfter {
|
|
cookieNamesAfter = append(cookieNamesAfter, cookie.Name)
|
|
}
|
|
t.Logf("🍪 Cookies after refresh test: %v", cookieNamesAfter)
|
|
|
|
// Log the response body for debugging
|
|
t.Logf("📄 Response body: %s", queryResponseBody)
|
|
|
|
// Check if we got a successful response (not a redirect to login)
|
|
assert.Contains(t, queryResponseBody, "success", "Auto refresh should have succeeded and allowed access to protected endpoint")
|
|
|
|
// Find the new auth token
|
|
var newAuthToken string
|
|
for _, cookie := range cookiesAfter {
|
|
if cookie.Name == "auth_token" {
|
|
newAuthToken = cookie.Value
|
|
break
|
|
}
|
|
}
|
|
|
|
assert.NotEmpty(t, newAuthToken, "Auth token should be present after successful auto refresh")
|
|
|
|
if newAuthToken != "" {
|
|
// Verify the token was refreshed (tokens should be different)
|
|
if originalAuthToken == newAuthToken {
|
|
t.Logf("⚠️ Token appears to be the same - auto refresh may not have occurred")
|
|
t.Logf("🤔 This could happen if the token wasn't actually expired yet")
|
|
} else {
|
|
t.Logf("✅ Token was refreshed successfully!")
|
|
if len(newAuthToken) >= 20 {
|
|
t.Logf("🔍 New token (first 20 chars): %s...", newAuthToken[:20])
|
|
} else {
|
|
t.Logf("🔍 New token: %s", newAuthToken)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Parse new token expiration
|
|
if newAuthToken != "" {
|
|
newExpiration, err := parseJWTTokenExpiration(newAuthToken)
|
|
if err == nil {
|
|
t.Logf("📅 New token expires at: %s", newExpiration.Format(time.RFC3339))
|
|
t.Logf("⌛ New token valid for: %s", time.Until(newExpiration))
|
|
|
|
// Verify new token has a later expiration than original
|
|
if newExpiration.After(tokenExpiration) {
|
|
t.Logf("✅ New token expiration is later than original - auto refresh confirmed!")
|
|
}
|
|
}
|
|
}
|
|
|
|
t.Logf("🎉 Auto refresh test completed successfully!")
|
|
})
|
|
|
|
// Phase 5: Verify Continued Access to Protected Endpoint
|
|
t.Run("Verify Continued Access After Refresh", func(t *testing.T) {
|
|
// Make another request to a protected endpoint to ensure the refreshed token works
|
|
queryURL := fmt.Sprintf("%s/query", baseURL)
|
|
var secondResponseBody string
|
|
|
|
err := chromedp.Run(ctx,
|
|
chromedp.Navigate(queryURL),
|
|
chromedp.Sleep(3*time.Second),
|
|
chromedp.OuterHTML("body", &secondResponseBody),
|
|
)
|
|
assert.NoError(t, err, "Failed to access protected endpoint after refresh")
|
|
|
|
assert.Contains(t, secondResponseBody, "success", "Should continue to have access to protected endpoints after auto refresh")
|
|
t.Logf("✅ Continued access to protected endpoint verified")
|
|
})
|
|
} |