Skip to content

Update: [AEA-5840] - Removed state and session state dynamo tables from auth flow #1412

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MatthewPopat-NHS merged 4 commits into from
Oct 16, 2025
Merged

Update: [AEA-5840] - Removed state and session state dynamo tables from auth flow #1412

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9ab1e6a
Removed state and session state dynamo tables from auth flow
MatthewPopat-NHS Oct 7, 2025
1b00e7c
Fixed duplication
MatthewPopat-NHS Oct 8, 2025
66de7b7
Fixed object syntax
MatthewPopat-NHS Oct 16, 2025
2d9c5fa
Merge branch 'main' into AEA-5840
MatthewPopat-NHS Oct 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 4 additions & 35 deletions packages/cognito/src/authorize.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,11 @@ import {Logger} from "@aws-lambda-powertools/logger"
import {APIGatewayProxyEvent, APIGatewayProxyResult} from "aws-lambda"
import {injectLambdaContext} from "@aws-lambda-powertools/logger/middleware"

import {DynamoDBClient} from "@aws-sdk/client-dynamodb"
import {DynamoDBDocumentClient} from "@aws-sdk/lib-dynamodb"

import {MiddyErrorHandler} from "@cpt-ui-common/middyErrorHandler"

import middy from "@middy/core"
import inputOutputLogger from "@middy/input-output-logger"

import {createHash} from "crypto"

import {insertStateMapping} from "@cpt-ui-common/dynamoFunctions"

/*
* Expects the following environment variables to be set:
*
Expand All @@ -24,24 +17,18 @@ import {insertStateMapping} from "@cpt-ui-common/dynamoFunctions"
*
* FULL_CLOUDFRONT_DOMAIN
*
* StateMappingTableName
*
*/

// Environment variables
const authorizeEndpoint = process.env["IDP_AUTHORIZE_PATH"] as string
const cis2ClientId = process.env["OIDC_CLIENT_ID"] as string
const userPoolClientId = process.env["COGNITO_CLIENT_ID"] as string
const cloudfrontDomain = process.env["FULL_CLOUDFRONT_DOMAIN"] as string
const stateMappingTableName = process.env["StateMappingTableName"] as string

const logger = new Logger({serviceName: "authorize"})
const errorResponseBody = {message: "A system error has occurred"}
const middyErrorHandler = new MiddyErrorHandler(errorResponseBody)

const dynamoClient = new DynamoDBClient()
const documentClient = DynamoDBDocumentClient.from(dynamoClient)

const lambdaHandler = async (
event: APIGatewayProxyEvent
): Promise<APIGatewayProxyResult> => {
Expand All @@ -50,14 +37,12 @@ const lambdaHandler = async (
authorizeEndpoint,
cis2ClientId,
userPoolClientId,
cloudfrontDomain,
stateMappingTableName
cloudfrontDomain
}})

// Validate required environment variables
if (!authorizeEndpoint) throw new Error("Authorize endpoint environment variable not set")
if (!cloudfrontDomain) throw new Error("Cloudfront domain environment variable not set")
if (!stateMappingTableName) throw new Error("State mapping table name environment variable not set")
if (!userPoolClientId) throw new Error("Cognito user pool client ID environment variable not set")
if (!cis2ClientId) throw new Error("OIDC client ID environment variable not set")

Expand All @@ -74,34 +59,18 @@ const lambdaHandler = async (
queryParams.scope = "openid profile nhsperson nationalrbacaccess associatedorgs"

// Ensure the state parameter is provided
const originalState = queryParams.state
if (!originalState) throw new Error("Missing state parameter")

// Generate the hashed state value
const cis2State = createHash("sha256").update(originalState).digest("hex")

// Set TTL for 5 minutes from now
const stateTtl = Math.floor(Date.now() / 1000) + 300
const state = queryParams.state
if (!state) throw new Error("Missing state parameter")

// Build the callback URI for redirection
const callbackUri = `https://${cloudfrontDomain}/oauth2/callback`

// Store original state mapping in DynamoDB
const item = {
State: cis2State,
CognitoState: originalState,
ExpiryTime: stateTtl
}

await insertStateMapping(documentClient, stateMappingTableName, item, logger)
logger.debug("State mapping inserted", {item})

// Build the redirect parameters for CIS2
const responseParameters = {
response_type: queryParams.response_type as string,
scope: queryParams.scope as string,
client_id: cis2ClientId,
state: cis2State,
state,
redirect_uri: callbackUri,
prompt: "login"
}
Expand Down
32 changes: 1 addition & 31 deletions packages/cognito/src/authorizeMock.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,11 @@ import {Logger} from "@aws-lambda-powertools/logger"
import {APIGatewayProxyEvent, APIGatewayProxyResult} from "aws-lambda"
import {injectLambdaContext} from "@aws-lambda-powertools/logger/middleware"

import {DynamoDBClient} from "@aws-sdk/client-dynamodb"
import {DynamoDBDocumentClient} from "@aws-sdk/lib-dynamodb"

import {MiddyErrorHandler} from "@cpt-ui-common/middyErrorHandler"

import middy from "@middy/core"
import inputOutputLogger from "@middy/input-output-logger"

import {createHash} from "crypto"

import {insertStateMapping} from "@cpt-ui-common/dynamoFunctions"

/*
* Expects the following environment variables to be set:
*
Expand All @@ -23,25 +16,19 @@ import {insertStateMapping} from "@cpt-ui-common/dynamoFunctions"
*
* FULL_CLOUDFRONT_DOMAIN
*
* StateMappingTableName
*
*/

// Environment variables
const authorizeEndpoint = process.env["IDP_AUTHORIZE_PATH"] as string
const cis2ClientId = process.env["OIDC_CLIENT_ID"] as string
const userPoolClientId = process.env["COGNITO_CLIENT_ID"] as string
const cloudfrontDomain = process.env["FULL_CLOUDFRONT_DOMAIN"] as string
const stateMappingTableName = process.env["StateMappingTableName"] as string
const apigeeApiKey = process.env["APIGEE_API_KEY"] as string

const logger = new Logger({serviceName: "authorize"})
const errorResponseBody = {message: "A system error has occurred"}
const middyErrorHandler = new MiddyErrorHandler(errorResponseBody)

const dynamoClient = new DynamoDBClient()
const documentClient = DynamoDBDocumentClient.from(dynamoClient)

const lambdaHandler = async (
event: APIGatewayProxyEvent
): Promise<APIGatewayProxyResult> => {
Expand All @@ -56,14 +43,12 @@ const lambdaHandler = async (
cis2ClientId,
userPoolClientId,
cloudfrontDomain,
stateMappingTableName,
apigeeApiKey
}})

// Validate required environment variables
if (!authorizeEndpoint) throw new Error("Authorize endpoint environment variable not set")
if (!cloudfrontDomain) throw new Error("Cloudfront domain environment variable not set")
if (!stateMappingTableName) throw new Error("State mapping table name environment variable not set")
if (!userPoolClientId) throw new Error("Cognito user pool client ID environment variable not set")
if (!cis2ClientId) throw new Error("OIDC client ID environment variable not set")
if (!apigeeApiKey) throw new Error("apigee api key environment variable not set")
Expand All @@ -81,12 +66,6 @@ const lambdaHandler = async (
const originalState = queryParams.state
if (!originalState) throw new Error("Missing state parameter")

// Generate the hashed state value
const cis2State = createHash("sha256").update(originalState).digest("hex")

// Set TTL for 5 minutes from now
const stateTtl = Math.floor(Date.now() / 1000) + 300

// Build the callback URI for redirection
// for pull requests we pack the real callback url for this pull request into the state
// the callback lambda then decodes this and redirects to the callback url for this pull request
Expand All @@ -96,19 +75,10 @@ const lambdaHandler = async (
const newStateJson = {
isPullRequest: true,
redirectUri: realCallbackUri,
originalState: cis2State
originalState: originalState
}
const newState = Buffer.from(JSON.stringify(newStateJson)).toString("base64")

// Store original state mapping in DynamoDB
const item = {
State: cis2State,
CognitoState: originalState,
ExpiryTime: stateTtl
}

await insertStateMapping(documentClient, stateMappingTableName, item, logger)

// Build the redirect parameters for CIS2
const responseParameters = {
redirect_uri: callbackUri,
Expand Down
32 changes: 2 additions & 30 deletions packages/cognito/src/callback.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,17 +1,14 @@
import {Logger} from "@aws-lambda-powertools/logger"
import {APIGatewayProxyEvent, APIGatewayProxyResult} from "aws-lambda"
import {injectLambdaContext} from "@aws-lambda-powertools/logger/middleware"
import {DynamoDBClient} from "@aws-sdk/client-dynamodb"
import {DynamoDBDocumentClient} from "@aws-sdk/lib-dynamodb"
import {MiddyErrorHandler} from "@cpt-ui-common/middyErrorHandler"
import middy from "@middy/core"
import inputOutputLogger from "@middy/input-output-logger"
import {deleteStateMapping, getStateMapping} from "@cpt-ui-common/dynamoFunctions"
import {buildCallbackRedirect} from "./helpers"

/*
* Expects the following environment variables to be set:
*
* StateMappingTableName
* COGNITO_CLIENT_ID
* COGNITO_DOMAIN
* PRIMARY_OIDC_ISSUER
Expand All @@ -23,17 +20,12 @@ const errorResponseBody = {message: "A system error has occurred"}
const middyErrorHandler = new MiddyErrorHandler(errorResponseBody)

// Environment variables
const stateMappingTableName = process.env["StateMappingTableName"] as string
const fullCognitoDomain = process.env["COGNITO_DOMAIN"] as string

const dynamoClient = new DynamoDBClient()
const documentClient = DynamoDBDocumentClient.from(dynamoClient)

const lambdaHandler = async (event: APIGatewayProxyEvent): Promise<APIGatewayProxyResult> => {
logger.appendKeys({"apigw-request-id": event.requestContext?.requestId})

// Destructure and validate required query parameters
// TODO: investigate if session_state is needed at all for this function
const {state, code, session_state} = event.queryStringParameters || {}
if (!state || !code) {
logger.error(
Expand All @@ -44,27 +36,7 @@ const lambdaHandler = async (event: APIGatewayProxyEvent): Promise<APIGatewayPro
}
logger.info("Incoming query parameters", {state, code})

const cognitoStateItem = await getStateMapping(documentClient, stateMappingTableName, state, logger)
await deleteStateMapping(documentClient, stateMappingTableName, state, logger)

// Build response parameters for redirection
const responseParams = {
state: cognitoStateItem.CognitoState,
session_state: session_state || "",
code
}

const redirectUri = `https://${fullCognitoDomain}/oauth2/idpresponse` +
`?${new URLSearchParams(responseParams).toString()}`

logger.info("Redirecting to Cognito", {redirectUri})

return {
statusCode: 302,
headers: {Location: redirectUri},
isBase64Encoded: false,
body: JSON.stringify({})
}
return buildCallbackRedirect(logger, state, code, session_state, fullCognitoDomain)
}

export const handler = middy(lambdaHandler)
Expand Down
62 changes: 2 additions & 60 deletions packages/cognito/src/callbackMock.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,18 +1,14 @@
import {Logger} from "@aws-lambda-powertools/logger"
import {APIGatewayProxyEvent, APIGatewayProxyResult} from "aws-lambda"
import {injectLambdaContext} from "@aws-lambda-powertools/logger/middleware"
import {DynamoDBClient} from "@aws-sdk/client-dynamodb"
import {DynamoDBDocumentClient} from "@aws-sdk/lib-dynamodb"
import {MiddyErrorHandler} from "@cpt-ui-common/middyErrorHandler"
import middy from "@middy/core"
import inputOutputLogger from "@middy/input-output-logger"
import {createHash, randomBytes} from "crypto"
import {deleteStateMapping, getStateMapping, insertSessionState} from "@cpt-ui-common/dynamoFunctions"
import {buildCallbackRedirect} from "./helpers"

/*
* Expects the following environment variables to be set:
*
* StateMappingTableName
* COGNITO_CLIENT_ID
* COGNITO_DOMAIN
* MOCK_OIDC_ISSUER
Expand All @@ -25,13 +21,8 @@ const errorResponseBody = {message: "A system error has occurred"}
const middyErrorHandler = new MiddyErrorHandler(errorResponseBody)

// Environment variables
const stateMappingTableName = process.env["StateMappingTableName"] as string
const SessionStateMappingTableName = process.env["SessionStateMappingTableName"] as string
const fullCognitoDomain = process.env["COGNITO_DOMAIN"] as string

const dynamoClient = new DynamoDBClient()
const documentClient = DynamoDBDocumentClient.from(dynamoClient)

const lambdaHandler = async (event: APIGatewayProxyEvent): Promise<APIGatewayProxyResult> => {
logger.appendKeys({"apigw-request-id": event.requestContext?.requestId})

Expand Down Expand Up @@ -72,56 +63,7 @@ const lambdaHandler = async (event: APIGatewayProxyEvent): Promise<APIGatewayPro
// Continue with regular flow
}

// TODO: remove session state mapping we can just use the state mapping
//TODO: make sure to update the logic that currently points to the session state mapping table
// If not a PR redirect, continue with the standard Cognito flow
// Get the original Cognito state from DynamoDB
logger.debug("trying to get data from session state table", {
stateMappingTableName,
state
})
const cognitoStateItem = await getStateMapping(documentClient, stateMappingTableName, state, logger)
await deleteStateMapping(documentClient, stateMappingTableName, state, logger)

// we need to generate a session state param and store it along with code returned
// as that will be used in the token lambda
// Generate the hashed state value
const sessionState = createHash("sha256").update(state).digest("hex")
const localCode = randomBytes(20).toString("hex")

const sessionStateExpiryTime = Math.floor(Date.now() / 1000) + 300

const item = {
LocalCode: localCode,
SessionState: sessionState,
ApigeeCode: code,
ExpiryTime: sessionStateExpiryTime
}

logger.debug("going to insert into session state mapping table", {
SessionStateMappingTableName,
item
})
await insertSessionState(documentClient, SessionStateMappingTableName, item, logger)

// Build response parameters for redirection
const responseParams = {
state: cognitoStateItem.CognitoState,
session_state: sessionState,
code: localCode
}

const redirectUri = `https://${fullCognitoDomain}/oauth2/idpresponse` +
`?${new URLSearchParams(responseParams).toString()}`

logger.info("Redirecting to Cognito", {redirectUri})

return {
statusCode: 302,
headers: {Location: redirectUri},
isBase64Encoded: false,
body: JSON.stringify({})
}
return buildCallbackRedirect(logger, state, code, session_state, fullCognitoDomain)
}

export const handler = middy(lambdaHandler)
Expand Down
26 changes: 26 additions & 0 deletions packages/cognito/src/helpers.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,3 +39,29 @@ export function rewriteRequestBody(
delete objectBodyParameters.client_secret
return objectBodyParameters
}

export function buildCallbackRedirect(
logger: Logger,
state: string,
code: string,
session_state: string | undefined,
fullCognitoDomain: string
) {
const responseParams = {
state,
session_state: session_state || "",
code
}

const redirectUri = `https://${fullCognitoDomain}/oauth2/idpresponse` +
`?${new URLSearchParams(responseParams).toString()}`

logger.info("Redirecting to Cognito", {redirectUri})

return {
statusCode: 302,
headers: {Location: redirectUri},
isBase64Encoded: false,
body: JSON.stringify({})
}
}
Loading