diff --git a/docs/events/cognito-user-pool.md b/docs/events/cognito-user-pool.md index 9bdda2eaeb..f161839109 100644 --- a/docs/events/cognito-user-pool.md +++ b/docs/events/cognito-user-pool.md @@ -140,6 +140,45 @@ function handler(event, context, callback) { } ``` +### PreTokenGeneration Trigger + +The PreTokenGeneration trigger supports multiple lambda versions for enhanced token customization: + +```yml +functions: + preTokenGenerationV1: + handler: preToken.handler + events: + - cognitoUserPool: + pool: MyUserPool + trigger: PreTokenGeneration + # No lambdaVersion = V1_0 behavior (ID token customization only) + + preTokenGenerationV2: + handler: preToken.handler + events: + - cognitoUserPool: + pool: MyUserPool + trigger: PreTokenGeneration + lambdaVersion: V2_0 # ID + Access token customization + + preTokenGenerationV3: + handler: preToken.handler + events: + - cognitoUserPool: + pool: MyUserPool + trigger: PreTokenGeneration + lambdaVersion: V3_0 # Includes M2M client-credentials grants +``` + +**Lambda Version Support:** + +- `V1_0` (default): ID token customization only +- `V2_0`: ID and access token customization +- `V3_0`: Includes machine-to-machine (M2M) client-credentials grants + +**NOTE:** V2_0 and V3_0 require your user pool to be on the Essentials or Plus feature plan, as documented [by AWS](https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-pre-token-generation.html). + ### Custom Message Trigger Handlers For custom messages, you will need to check `event.triggerSource` type inside your handler function: diff --git a/lib/plugins/aws/custom-resources/resources/cognito-user-pool/lib/user-pool.js b/lib/plugins/aws/custom-resources/resources/cognito-user-pool/lib/user-pool.js index 724911a15d..c9b02cb752 100644 --- a/lib/plugins/aws/custom-resources/resources/cognito-user-pool/lib/user-pool.js +++ b/lib/plugins/aws/custom-resources/resources/cognito-user-pool/lib/user-pool.js @@ -91,6 +91,11 @@ async function updateConfiguration(config) { LambdaVersion: poolConfig.LambdaVersion, }; LambdaConfig.KMSKeyID = poolConfig.KMSKeyID; + } else if (poolConfig.Trigger === 'PreTokenGeneration' && poolConfig.LambdaVersion) { + LambdaConfig.PreTokenGenerationConfig = { + LambdaArn: lambdaArn, + LambdaVersion: poolConfig.LambdaVersion, + }; } else { LambdaConfig[poolConfig.Trigger] = lambdaArn; } @@ -131,6 +136,9 @@ async function removeConfiguration(config) { function removeExistingLambdas(lambdaConfig, lambdaArn) { const res = Object.fromEntries( Object.entries(lambdaConfig).filter(([key, value]) => { + if (key === 'PreTokenGenerationConfig' && value && value.LambdaArn === lambdaArn) { + return false; + } return ( !(customSenderSources.includes(key) && value.LambdaArn === lambdaArn) && value !== lambdaArn ); diff --git a/lib/plugins/aws/package/compile/events/cognito-user-pool.js b/lib/plugins/aws/package/compile/events/cognito-user-pool.js index aaeb4cceb5..9ab132d7a9 100644 --- a/lib/plugins/aws/package/compile/events/cognito-user-pool.js +++ b/lib/plugins/aws/package/compile/events/cognito-user-pool.js @@ -20,7 +20,8 @@ const validTriggerSources = [ 'UserMigration', ].concat(customSenderSources); -const validLambdaVersions = ['V1_0']; +const customSenderValidVersions = ['V1_0']; +const preTokenGenerationValidVersions = ['V1_0', 'V2_0', 'V3_0']; class AwsCompileCognitoUserPoolEvents { constructor(serverless, options) { @@ -154,6 +155,8 @@ class AwsCompileCognitoUserPoolEvents { const { pool, trigger, forceDeploy, kmsKeyId, lambdaVersion } = event.cognitoUserPool; usesExistingCognitoUserPool = funcUsesExistingCognitoUserPool = true; + this.validateLambdaVersion(trigger, lambdaVersion); + if (!currentPoolName) { currentPoolName = pool; } @@ -190,12 +193,14 @@ class AwsCompileCognitoUserPoolEvents { userPoolConfig = { ...userPoolConfig, ...{ - LambdaVersion: lambdaVersion || validLambdaVersions[0], + LambdaVersion: lambdaVersion || customSenderValidVersions[0], }, }; this.checkKmsArn(kmsKeyId, poolKmsIdMap, pool); userPoolConfig.KMSKeyID = kmsKeyId; + } else if (trigger === 'PreTokenGeneration' && lambdaVersion) { + userPoolConfig.LambdaVersion = lambdaVersion; } if (numEventsForFunc === 1) { @@ -310,6 +315,33 @@ class AwsCompileCognitoUserPoolEvents { poolKmsIdMap.set(currentPoolName, kmsKeyId); } + validateLambdaVersion(trigger, lambdaVersion) { + if (customSenderSources.includes(trigger)) { + if (lambdaVersion && !customSenderValidVersions.includes(lambdaVersion)) { + throw new ServerlessError( + `Invalid lambdaVersion "${lambdaVersion}" for trigger "${trigger}". Custom sender triggers only support: ${customSenderValidVersions.join( + ', ' + )}`, + 'COGNITO_INVALID_LAMBDA_VERSION' + ); + } + } else if (trigger === 'PreTokenGeneration') { + if (lambdaVersion && !preTokenGenerationValidVersions.includes(lambdaVersion)) { + throw new ServerlessError( + `Invalid lambdaVersion "${lambdaVersion}" for trigger "${trigger}". PreTokenGeneration supports: ${preTokenGenerationValidVersions.join( + ', ' + )}`, + 'COGNITO_INVALID_LAMBDA_VERSION' + ); + } + } else if (lambdaVersion) { + throw new ServerlessError( + `lambdaVersion is not supported for trigger "${trigger}". It's only supported for: CustomSMSSender, CustomEmailSender, and PreTokenGeneration`, + 'COGNITO_LAMBDA_VERSION_NOT_SUPPORTED' + ); + } + } + findUserPoolsAndFunctions() { const userPools = []; const cognitoUserPoolTriggerFunctions = []; @@ -323,6 +355,11 @@ class AwsCompileCognitoUserPoolEvents { if (event.cognitoUserPool) { if (event.cognitoUserPool.existing) return; + this.validateLambdaVersion( + event.cognitoUserPool.trigger, + event.cognitoUserPool.lambdaVersion + ); + // Save trigger functions so we can use them to generate // IAM permissions later cognitoUserPoolTriggerFunctions.push({ @@ -354,11 +391,18 @@ class AwsCompileCognitoUserPoolEvents { triggerObject = { [value.triggerSource]: { LambdaArn: resolveLambdaTarget(value.functionName, functionObj), - LambdaVersion: value.lambdaVersion || validLambdaVersions[0], + LambdaVersion: value.lambdaVersion || customSenderValidVersions[0], }, }; this.checkKmsArn(value.kmsKeyId, poolKmsIdMap, poolName); triggerObject.KMSKeyID = value.kmsKeyId; + } else if (value.triggerSource === 'PreTokenGeneration' && value.lambdaVersion) { + triggerObject = { + PreTokenGenerationConfig: { + LambdaArn: resolveLambdaTarget(value.functionName, functionObj), + LambdaVersion: value.lambdaVersion, + }, + }; } else { triggerObject = { [value.triggerSource]: resolveLambdaTarget(value.functionName, functionObj), diff --git a/test/unit/lib/plugins/aws/package/compile/events/cognito-user-pool.test.js b/test/unit/lib/plugins/aws/package/compile/events/cognito-user-pool.test.js index b87ed43c89..b4855005b9 100644 --- a/test/unit/lib/plugins/aws/package/compile/events/cognito-user-pool.test.js +++ b/test/unit/lib/plugins/aws/package/compile/events/cognito-user-pool.test.js @@ -162,6 +162,49 @@ const serverlessConfigurationExtension = { }, }; +const preTokenGenerationConfigurationExtension = { + configValidationMode: 'off', + functions: { + preTokenGenerationV2: { + handler: 'index.js', + events: [ + { + cognitoUserPool: { + pool: 'PreTokenGenerationV2Pool', + trigger: 'PreTokenGeneration', + lambdaVersion: 'V2_0', + }, + }, + ], + }, + preTokenGenerationV3: { + handler: 'index.js', + events: [ + { + cognitoUserPool: { + pool: 'PreTokenGenerationV3Pool', + trigger: 'PreTokenGeneration', + lambdaVersion: 'V3_0', + }, + }, + ], + }, + preTokenGenerationV2Existing: { + handler: 'index.js', + events: [ + { + cognitoUserPool: { + pool: 'PreTokenGenerationV2PoolExisting', + trigger: 'PreTokenGeneration', + lambdaVersion: 'V2_0', + existing: true, + }, + }, + ], + }, + }, +}; + describe('AwsCompileCognitoUserPoolEvents', () => { let serverless; let awsCompileCognitoUserPoolEvents; @@ -1385,4 +1428,64 @@ describe('lib/plugins/aws/package/compile/events/cognito-user-pool.test.js', () }); }); }); + + describe('PreTokenGeneration Lambda Versions', () => { + describe('New Pools', () => { + it('should create PreTokenGenerationConfig for V2_0', async () => { + const { cfTemplate } = await runServerless({ + fixture: 'cognito-user-pool', + configExt: preTokenGenerationConfigurationExtension, + command: 'package', + }); + + const userPoolResource = cfTemplate.Resources.CognitoUserPoolPreTokenGenerationV2Pool; + expect(userPoolResource.Properties.LambdaConfig).to.have.property( + 'PreTokenGenerationConfig' + ); + expect( + userPoolResource.Properties.LambdaConfig.PreTokenGenerationConfig.LambdaVersion + ).to.equal('V2_0'); + expect(userPoolResource.Properties.LambdaConfig).to.not.have.property('PreTokenGeneration'); + }); + + it('should create PreTokenGenerationConfig for V3_0', async () => { + const { cfTemplate } = await runServerless({ + fixture: 'cognito-user-pool', + configExt: preTokenGenerationConfigurationExtension, + command: 'package', + }); + + const userPoolResource = cfTemplate.Resources.CognitoUserPoolPreTokenGenerationV3Pool; + expect(userPoolResource.Properties.LambdaConfig).to.have.property( + 'PreTokenGenerationConfig' + ); + expect( + userPoolResource.Properties.LambdaConfig.PreTokenGenerationConfig.LambdaVersion + ).to.equal('V3_0'); + expect(userPoolResource.Properties.LambdaConfig).to.not.have.property('PreTokenGeneration'); + }); + }); + + describe('Existing Pools', () => { + it('should create correct UserPoolConfigs for V2_0', async () => { + const { cfTemplate } = await runServerless({ + fixture: 'cognito-user-pool', + configExt: preTokenGenerationConfigurationExtension, + command: 'package', + }); + + const customResources = Object.values(cfTemplate.Resources).filter( + (resource) => resource.Type === 'Custom::CognitoUserPool' + ); + const v2Resource = customResources.find( + (resource) => resource.Properties.UserPoolName === 'PreTokenGenerationV2PoolExisting' + ); + + expect(v2Resource.Properties.UserPoolConfigs[0]).to.deep.include({ + Trigger: 'PreTokenGeneration', + LambdaVersion: 'V2_0', + }); + }); + }); + }); });