Bug Fix: Update Cognito Authentication Flow

[yosefkri]

Fix: Update Cognito Authentication Flow for Federated Users

Summary

This PR updates the Cognito authentication flow for federated users by replacing the POST_AUTHENTICATION trigger with PRE_SIGN_UP and POST_CONFIRMATION triggers. This change resolves the issue where users were experiencing 404 errors after logout when using Cognito federation.

Problem

When using Cognito federation, the previous implementation relied on the POST_AUTHENTICATION trigger to assign users to groups. However, this approach caused issues where users would get 404 errors after logout, and the Lambda function would only trigger after the user had already logged into the console.

Solution

The solution implements a more robust authentication flow by:

1. Using PRE_SIGN_UP trigger to validate and prepare for group assignment before the user is created
2. Using POST_CONFIRMATION trigger to assign the user to the appropriate group after the user is confirmed
3. Enhancing the Lambda function to handle different trigger types with appropriate event structure handling

This approach ensures that users are properly assigned to groups earlier in the authentication flow, preventing the 404 errors that occurred with the previous implementation.

Changes

• Modified lib/authentication/index.ts to:

  • Replace POST_AUTHENTICATION trigger with PRE_SIGN_UP and POST_CONFIRMATION triggers
  • Add environment variable for default user group
  • Update permission names to reflect the new trigger types

• Enhanced lib/authentication/lambda/addFederatedUserToUserGroup/index.py to:

  • Support multiple trigger types (PRE_AUTHENTICATION, POST_AUTHENTICATION, PRE_SIGN_UP, POST_CONFIRMATION)
  • Handle different event structures for each trigger type
  • Add default group assignment from environment variables
  • Improve error handling and logging
  • Add special handling for PRE_SIGN_UP where the user doesn't exist yet

Testing

The changes have been tested with Cognito federation enabled, confirming that:

• Users are properly assigned to groups during the authentication flow
• No 404 errors occur after logout
• The Lambda function correctly handles different trigger types
• Group assignments are consistent across login sessions