Tuesday, December 3, 2024

Configuring Cognito with React and WebSockets

 


Cognito Authorization with React and AWS WS Gateway

M 917 536 3378
maksim_kozyarchuk@yahoo.com





Introduction

Configuring AWS Cognito for the Kupala Nich application proved to be more challenging than anticipated. This document outlines the essential steps and key considerations for setting up Cognito with a focus on authentication and its integration with React and AWS Lambda for seamless user management.

AWS Cognito provides a robust and feature-rich system for user authentication and management, but its extensive capabilities can be overwhelming. Focus of this article is achieving a minimally sufficient configuration to meet the requirements without unnecessary complexity. Along the way, it introduces foundational concepts of OpenID Connect (OIDC) and provides practical examples of:

  • CloudFormation templates for automating the setup of Cognito resources.

  • React code for integrating Cognito authentication with the frontend.

  • AWS Lambda functions for handling authentication workflows on the backend.



Authentication vs. Authorization

AWS Cognito offers comprehensive capabilities for both authentication and authorization, each with a rich feature set and associated complexity. However, the focus of this article is on authentication, as it forms the foundation of secure user management in the Kupala Nich application.

In the Kupala Nich application, authorization is implemented at the data level—specifically, by controlling access to user portfolios. This approach differs from Cognito's native authorization framework, which is designed for API-level or resource-level access control. By shifting authorization to the data layer, the application avoids the need to use Identity Pools or API Gateway Authorizers.



Configuring Authentication with Cognito

To set up authentication with AWS Cognito, two primary components need to be configured: the User Pool and the User Pool Client. These can be configured either through the Cognito UI or via CloudFormation. For the Kupala Nich application, all configurations are managed through CloudFormation to streamline deployment across different environments.

User Pool

The User Pool defines and manages user accounts and their attributes for your application. Key aspects of the User Pool configuration include:

  • Fields defined for the account, determining which standard fields (e.g., email, phone_number) are required or optional, and allowing for custom fields.

  • Password policies to enforce rules such as complexity and expiration.

  • Multi-Factor Authentication (MFA) settings, specifying whether MFA is required, optional, or disabled.

  • Email and phone verification procedures for confirming user accounts.

  • Account recovery policies that define how users regain access to their accounts (e.g., using verified email).

Important Note: Once created, the User Pool configuration cannot be modified. To make changes, you must export users to a CSV file, delete and recreate the User Pool with the updated attributes, and re-import the users.

For the Kupala Nich application, the following configuration choices were made:

  • All User Pool configurations are defined and managed via CloudFormation (template included in the Appendix).

  • A minimal setup is used and MFA setup has been deferred until the product launch strategy is finalized. This choice acknowledges the potential need for at least one user account migration in the future.

User Pool Client

The User Pool Client defines how applications authenticate against the User Pool. Its configuration depends heavily on the frontend design and the libraries chosen for integration.

To integrate Cognito into the Kupala Nich application, three React libraries were evaluated:

  • amazon-cognito-identity-js: Provides direct API-level integration with Cognito. However, this library requires developers to build custom screens for workflows such as login, logout, password reset, and user registration. It is suitable for applications needing a highly customized user experience.

  • react-oidc-context: Handles user session and token management within the application while delegating workflows like login and password reset to Cognito's Hosted UI. This library is featured in examples on the Cognito User Pool Client welcome page. It was selected for Kupala Nich due to its simplicity and ease of integration.

  • aws-amplify: Offers a hybrid approach, combining API-level integration with built-in user workflows and native integration with other AWS services like S3 and DynamoDB. While not explored for Kupala Nich, this library is worth considering for applications requiring broader AWS service integration.

The following configuration choices were made for the Kupala Nich User Pool Client:

  • Hosted UIs are used for authentication and user management, enabling workflows such as sign-in, sign-up, password reset, and email verification.

  • Users can maintain an open session for 30 days without re-authentication, supported by Refresh Tokens for automatic renewal.

  • ID and Access Tokens expire every 119 minutes, matching the WebSocket session duration, which is re-established every two hours.

  • Only Code-based Authorization Flows are enabled to ensure security. These flows are also automatically handled by react-oidc-context.

  • Initially, only Cognito-based authentication is supported, with plans to add Google and other social identity providers in the future.

Callback URLs were defined to support authentication workflows:

  • https://kupala-nich.com/callback is used for login and token exchange.

  • https://kupala-nich.com/silent-callback handles silent token refresh.

  • https://kupala-nich.com is the redirect after logout.

  • Localhost versions of these callbacks are included for debugging and testing during development.



Tokens in AWS Cognito and Their Role in Kupala Nich

Before jumping into the implementation, it’s worth discussing the tokens provided by AWS Cognito as part of the OpenID Connect (OIDC) protocol. There are three main types of tokens:

  • ID Token: Confirms the user’s identity. It is used by the Kupala Nich application to establish a trusted identity for WebSocket connections.

  • Access Token: Used as part of the authorization flow to grant access to specific resources or APIs. It does not apply to Kupala Nich at this point.

  • Refresh Token: Since ID and Access tokens are short-lived, the Refresh Token provides a long-lived mechanism to renew these tokens upon expiration without requiring the user to re-authenticate.

Tokens in the Kupala Nich Application

The Kupala Nich application primarily uses the ID Token to verify the user’s identity. This token is passed from the client to the API Gateway and backend Lambda function. For WebSocket-based interactions, the Lambda function that handles the $connect method validates the ID Token and links the user’s identity to the WebSocket connection.

Within AWS WebSocket Gateway, a WebSocket session has a maximum lifespan of 2 hours. This means that the $connect function—and, by extension, the ID Token—is validated every 2 hours.

Managing Token Expiry

ID Tokens are short-lived and, by default, expire after 60 minutes. To maintain an active connection between the React app and the backend, these tokens need to be renewed periodically.

The Kupala Nich application uses Refresh Tokens to handle this renewal process. The frontend React application leverages the react-oidc-context library to automatically refresh the ID Token when it is close to expiration. This ensures that:

  • The ID Token remains valid for as long as the user’s session is active.

  • The WebSocket connection between the React app and the backend is seamlessly maintained, without requiring the user to re-authenticate.



React Integration with react-oidc-context

Implementing authentication, token management, and session handling within a React app can be complex. This is one area where I do not recommend asking guidance from ChatGPT and looking at source documentation instead.  A good starting point is the AWS Cognito console documentation for your User Pool Client. The console provides a basic code snippet that establishes the connection and serves as a helpful starting point.  After the initial setup, review the examples available in the github repository( https://github.com/authts/react-oidc-context) for react-oidc-context. These examples are particularly useful for setting up event listeners for token lifecycle events and maintaining user authentication across browser sessions.

The available documentation, however, falls short in addressing three important areas: handling redirects, managing token refresh, and implementing logout. Below, I describe these areas in detail with examples from Kupala Nich.

Handling Redirects

The react-oidc-context library requires handling three types of redirects to implement the authentication flow effectively:

  • Login Redirect: This occurs after the user is authenticated with the Cognito Hosted UI. It establishes the session in the app. During this redirect, the user can be informed about the redirection process or any authentication errors that occurred.

  • Silent Redirect: This is used to refresh the id_token in the background. The library launches a hidden iframe to process the token refresh. The silent redirect page extracts the new token information while ensuring that the process does not interfere with other parts of the application.

  • Logout Redirect: After logout, the user is redirected to a specific page. In the Kupala Nich application, users are sent back to the main page, though this can be customized as needed.

In the Kupala Nich application, the react-router-dom library is used to handle these callbacks within the React app. The application is deployed as a Single Page Application (SPA) to enable smooth handling of redirects. Refer to the appendix for the relevant code extract from index.js.

Handling Token Refresh

Token refresh is straightforward but not effectively documented. To enable token refresh:

  • Add the automaticSilentRenew and silent_redirect_uri flags to the cognitoAuthConfig configuration.

  • Ensure that your React app handles redirects to the silent_redirect_uri internally.

  • Include the silent_redirect_uri in the list of allowed redirect URIs in the User Pool Client configuration.

These steps allow tokens to be refreshed in the background without user intervention, ensuring uninterrupted sessions.

Handling Logout

Implementing logout requires combining recommendations from AWS Cognito User Pool Client example and from react-oidc-context library documentation:

  • Use the logout route provided by the Cognito domain to invalidate the server-side session.

  • Call auth.removeUser() from the react-oidc-context library to clear the local session and remove tokens stored in the browser.

In the Kupala Nich application, the user is redirected to the main page after logout. Refer to the appendix for an example demonstrating logout integration.

Refer to the appendix for detailed code examples from the index.js file, demonstrating the setup of react-oidc-context and react-router-dom in the Kupala Nich application.



Server-Side Token Validation

As discussed in the previous section, user identity in the Kupala Nich application is established during the $connect Lambda handler. This identity is then referenced by subsequent Lambda calls for the duration of the WebSocket connection’s two-hour lifespan.

This process relies on the pyjwt library to validate and decode the id_token passed by the React app during the WebSocket connection request. The validation is performed against the OpenID Connect (OIDC) configuration associated with the Cognito User Pool Client.

The id_token validation ensures:

  • The token’s integrity by verifying its signature.

  • The token’s authenticity by comparing claims (e.g., issuer, audience).

  • That the token has not expired.

Important part of token validation is logging, which can then be analyzed to determine attempts to break security. The implementation details for this mechanism, including code examples, can be found in the Appendix.



Takeaways and Next Steps

I encourage you to critically review the content above and provide comments or suggestions to further enhance the security and scalability of the Kupala Nich application. Given that the application will host personal investment strategy details in a multi-tenant environment, it is crucial to implement security measures correctly and comprehensively.

The next step in the security journey is the implementation of robust authorization logic.

Your feedback is invaluable in refining these approaches and ensuring that the Kupala Nich application adheres to best practices for security and scalability.



Appendix 1: Cloudformation template 

  KupalaNichUserPool:

    Type: AWS::Cognito::UserPool

    Properties:

      UserPoolName: KupalaNichUserPool

      Policies:

        PasswordPolicy:

          MinimumLength: 8

          RequireUppercase: true

          RequireLowercase: true

          RequireNumbers: true

      MfaConfiguration: 'OFF'

      Schema:

        - Name: email

          Required: true  

          Mutable: true

      AutoVerifiedAttributes:

        - email

      UsernameAttributes:  

        - email

      AccountRecoverySetting:

        RecoveryMechanisms:

          - Name: verified_email

            Priority: 1


  KupalaNichUserPoolDomain:

    Type: AWS::Cognito::UserPoolDomain

    Properties:

      Domain: !Ref CognitoDomain

      UserPoolId: !Ref KupalaNichUserPool


  KupalaNichUserPoolClient:

    Type: AWS::Cognito::UserPoolClient

    Properties:

      UserPoolId: !Ref KupalaNichUserPool

      ClientName: KupalaNich

      GenerateSecret: false

      AllowedOAuthFlows:

        - code

      AllowedOAuthScopes:

        - openid

        - profile

        - email

      CallbackURLs:

        - https://kupala-nich.com/callback

        - https://kupala-nich.com/silent-callback

        - http://localhost:3000/callback

        - http://localhost:3000/silent-callback

      LogoutURLs:  

        - https://kupala-nich.com

        - http://localhost:3000

      SupportedIdentityProviders:

        - COGNITO

      TokenValidityUnits:

        AccessToken: minutes

        IdToken: minutes

        RefreshToken: days

      AccessTokenValidity: 119

      IdTokenValidity: 119

      RefreshTokenValidity: 30

      EnableTokenRevocation: true

      ExplicitAuthFlows:  

        - ALLOW_USER_SRP_AUTH

        - ALLOW_REFRESH_TOKEN_AUTH

        - ALLOW_USER_PASSWORD_AUTH



Appendix 2: Configuring react-oidc-context inside index.js

const cognitoAuthConfig = {

  authority: process.env.REACT_APP_AUTHORITY,

  client_id: process.env.REACT_APP_CLIENT_ID,

  redirect_uri: process.env.REACT_APP_REDIRECT_URI,

  response_type: "code",

  scope: "openid profile",

  automaticSilentRenew: true,

  silent_redirect_uri: process.env.REACT_APP_SILENT_CALLBACK,

  userStore: new WebStorageStateStore({ store: window.localStorage }),

};


const Callback = () => {

  const auth = useAuth();


  if (auth.isLoading) {

      return <div>Processing login...</div>;

  }


  if (auth.error) {

    return (

      <div>

        <div>Error: {auth.error.message}</div>

        <a href="/">Go back to the main page</a>

      </div>

    );

  }

  return null; // Redirect to the dashboard

  };


const SilentCallback = () => null;


const root = ReactDOM.createRoot(document.getElementById('root'));

root.render(

  <React.StrictMode>

    <AuthProvider {...cognitoAuthConfig}

      onSigninCallback={() => {

        window.location.href = "/";

      }}

      onSigninError={(error) => {

        console.error("Signin error:", error);

      }}

    >

    <Router>

          <Routes>

              <Route path="/silent-callback" element={<SilentCallback />} />

              <Route path="/callback" element={<Callback />} />

              <Route path="/" element={<Dashboard />} />

          </Routes>

      </Router>

    </AuthProvider>

  </React.StrictMode>

);


 

Appendix 3: Configuring $connect lambda to validate JWT token

from common import logger, connection_table

import os

import jwt

import urllib.request

import json

import time

from typing import Dict, Any


CACHE: Dict[str, Dict[str, Any]] = {

    "oidc_config": {},

    "jwks": {},

    "signing_keys": {},

    "expiration": {}

}


ACCOUNT_CACHE: Dict[str, Dict[str, Any]] = {

    "expiration": {}

}


def fetch_with_cache(url, cache_key, cache_duration=3600):

    current_time = time.time()

    if cache_key in CACHE and current_time < CACHE["expiration"].get(cache_key, 0):

        return CACHE[cache_key]


    try:

        with urllib.request.urlopen(url) as response:

            data = json.loads(response.read().decode())

    except Exception as e:

        raise ValueError(f"Failed to fetch data from {url}: {e}")

   

    CACHE[cache_key] = data

    CACHE["expiration"][cache_key] = current_time + cache_duration

    return data


def get_signing_key_with_cache(id_token, jwks_url):

    unverified_header = jwt.get_unverified_header(id_token)

    kid = unverified_header["kid"]

    if kid in CACHE["signing_keys"]:

        return CACHE["signing_keys"][kid]


    jwks_client = jwt.PyJWKClient(jwks_url)

    signing_key = jwks_client.get_signing_key_from_jwt(id_token)

    CACHE["signing_keys"][kid] = signing_key

    return signing_key


def validate_and_parse_cognito_id_token( id_token ):

    user_pool_id = os.getenv("COGNITO_USER_POOL_ID")

    region = os.getenv("AWS_REGION")

    client_id = os.getenv("COGNITO_APP_CLIENT_ID")

    logger.info(f"COGNITO_USER_POOL_ID: {user_pool_id}, AWS_REGION: {region}, COGNITO_APP_CLIENT_ID: {client_id}")


    if not all([id_token, user_pool_id, region, client_id]):

        logger.error(f'Not all parameters are provided. id_token: {id_token}, user_pool_id: {user_pool_id}, region: {region}, client_id: {client_id}')

        raise ValueError(

            "All parameters (id_token, access_token, user_pool_id, region, client_id) are required."

        )


    try:

        oidc_config_url = f"https://cognito-idp.{region}.amazonaws.com/{user_pool_id}/.well-known/openid-configuration"

        oidc_config = fetch_with_cache(oidc_config_url, "oidc_config")

        signing_algos = oidc_config.get("id_token_signing_alg_values_supported", ["RS256"])

        jwks_url = oidc_config["jwks_uri"]

        signing_key = get_signing_key_with_cache(id_token, jwks_url)


        data = jwt.decode_complete(

            id_token,

            key=signing_key.key,

            audience=client_id,

            algorithms=signing_algos,

            issuer=f"https://cognito-idp.{region}.amazonaws.com/{user_pool_id}",

        )

        logger.info(f"Decoded id_token: {data}")

        return data["payload"]


    except ValueError as e:

        raise ValueError(f"Validation error: {e}")    

    except jwt.ExpiredSignatureError:

        raise ValueError("The token has expired.")

    except jwt.InvalidTokenError as e:

        raise ValueError(f"Invalid token: {str(e)}")

    except Exception as e:

        raise ValueError(f"An error occurred during token validation: {str(e)}")


def get_user_account_from_connection_id(connection_id):

    current_time = time.time()

    if connection_id in ACCOUNT_CACHE and current_time < CACHE["expiration"].get(connection_id, 0):

        account = ACCOUNT_CACHE[connection_id]

    else:

        account = connection_table.get_item(Key={'ConnectionId': connection_id}).get('Item', {}).get('account', None)

        set_user_account_to_connection_id(connection_id, account)


    return account


def set_user_account_to_connection_id(connection_id, account, cache_duration=3600):

    ACCOUNT_CACHE[connection_id] = account

    ACCOUNT_CACHE["expiration"][connection_id] = time.time() + cache_duration


def handle_connect(connection_id,event):

    requestContext = event['requestContext']


    sub = None

    id_token  = event.get('queryStringParameters',{}).get('id_token', None)

    if id_token:

        parsed  = validate_and_parse_cognito_id_token(id_token)

        sub = parsed.get("sub")


    logger.debug(f"Handling connect event for connection ID: {connection_id}, sub: {sub}")

    connection_table.put_item(Item={'ConnectionId': connection_id, 'requestContext':requestContext, 'account':sub})

    set_user_account_to_connection_id(connection_id, sub)

    logger.info(f"Successfully connected: {connection_id}")


No comments: