feat: add support for PreTokenGeneration trigger with multiple lambda versions in Cognito User Pool

Author: admir-slateteams

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,29 @@ 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 +351,8 @@ 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 +384,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),