Merge pull request #98 from admirkadriu/feat/cognito-pretoken-generation-v2-v3-support

Add Support for V2_0 and V3_0 Lambda Versions in PreTokenGeneration Cognito Triggers

Author: mnapoli

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:

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
       );

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),

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',
+        });
+      });
+    });
+  });
 });