Added support for Cognito pre token generation with access token cust…

…omization (#538) * Fixed Cognito service name and added event structure for Cognito pre token generation event V2 * gofmt -s -w . * Add sample data and serde test * Tweak the test data to avoid having to break response marshaling compatability of the GroupConfiguration struct --------- Co-authored-by: Bryan Moffatt <[email protected]> Co-authored-by: Bryan Moffatt <[email protected]>
aws · Dec 20, 2023 · 326401a · 326401a
1 parent 30d7c7e
commit 326401a
Show file tree

Hide file tree

Showing 3 changed files with 150 additions and 15 deletions.
diff --git a/events/cognito.go b/events/cognito.go
@@ -2,7 +2,7 @@
 
 package events
 
-// CognitoEvent contains data from an event sent from AWS Cognito Sync
+// CognitoEvent contains data from an event sent from Amazon Cognito Sync
 type CognitoEvent struct {
 	DatasetName    string                          `json:"datasetName"`
 	DatasetRecords map[string]CognitoDatasetRecord `json:"datasetRecords"`
@@ -13,54 +13,62 @@ type CognitoEvent struct {
 	Version        int                             `json:"version"`
 }
 
-// CognitoDatasetRecord represents a record from an AWS Cognito Sync event
+// CognitoDatasetRecord represents a record from an Amazon Cognito Sync event
 type CognitoDatasetRecord struct {
 	NewValue string `json:"newValue"`
 	OldValue string `json:"oldValue"`
 	Op       string `json:"op"`
 }
 
-// CognitoEventUserPoolsPreSignup is sent by AWS Cognito User Pools when a user attempts to register
+// CognitoEventUserPoolsPreSignup is sent by Amazon Cognito User Pools when a user attempts to register
 // (sign up), allowing a Lambda to perform custom validation to accept or deny the registration request
 type CognitoEventUserPoolsPreSignup struct {
 	CognitoEventUserPoolsHeader
 	Request  CognitoEventUserPoolsPreSignupRequest  `json:"request"`
 	Response CognitoEventUserPoolsPreSignupResponse `json:"response"`
 }
 
-// CognitoEventUserPoolsPreAuthentication is sent by AWS Cognito User Pools when a user submits their information
+// CognitoEventUserPoolsPreAuthentication is sent by Amazon Cognito User Pools when a user submits their information
 // to be authenticated, allowing you to perform custom validations to accept or deny the sign in request.
 type CognitoEventUserPoolsPreAuthentication struct {
 	CognitoEventUserPoolsHeader
 	Request  CognitoEventUserPoolsPreAuthenticationRequest  `json:"request"`
 	Response CognitoEventUserPoolsPreAuthenticationResponse `json:"response"`
 }
 
-// CognitoEventUserPoolsPostConfirmation is sent by AWS Cognito User Pools after a user is confirmed,
+// CognitoEventUserPoolsPostConfirmation is sent by Amazon Cognito User Pools after a user is confirmed,
 // allowing the Lambda to send custom messages or add custom logic.
 type CognitoEventUserPoolsPostConfirmation struct {
 	CognitoEventUserPoolsHeader
 	Request  CognitoEventUserPoolsPostConfirmationRequest  `json:"request"`
 	Response CognitoEventUserPoolsPostConfirmationResponse `json:"response"`
 }
 
-// CognitoEventUserPoolsPreTokenGen is sent by AWS Cognito User Pools when a user attempts to retrieve
+// CognitoEventUserPoolsPreTokenGen is sent by Amazon Cognito User Pools when a user attempts to retrieve
 // credentials, allowing a Lambda to perform insert, suppress or override claims
 type CognitoEventUserPoolsPreTokenGen struct {
 	CognitoEventUserPoolsHeader
 	Request  CognitoEventUserPoolsPreTokenGenRequest  `json:"request"`
 	Response CognitoEventUserPoolsPreTokenGenResponse `json:"response"`
 }
 
-// CognitoEventUserPoolsPostAuthentication is sent by AWS Cognito User Pools after a user is authenticated,
+// CognitoEventUserPoolsPreTokenGenV2 is sent by Amazon Cognito User Pools when a user attempts to retrieve
+// credentials, allowing a Lambda to perform insert, suppress or override claims and scopes
+type CognitoEventUserPoolsPreTokenGenV2 struct {
+	CognitoEventUserPoolsHeader
+	Request  CognitoEventUserPoolsPreTokenGenV2Request  `json:"request"`
+	Response CognitoEventUserPoolsPreTokenGenV2Response `json:"response"`
+}
+
+// CognitoEventUserPoolsPostAuthentication is sent by Amazon Cognito User Pools after a user is authenticated,
 // allowing the Lambda to add custom logic.
 type CognitoEventUserPoolsPostAuthentication struct {
 	CognitoEventUserPoolsHeader
 	Request  CognitoEventUserPoolsPostAuthenticationRequest  `json:"request"`
 	Response CognitoEventUserPoolsPostAuthenticationResponse `json:"response"`
 }
 
-// CognitoEventUserPoolsMigrateUser is sent by AWS Cognito User Pools when a user does not exist in the
+// CognitoEventUserPoolsMigrateUser is sent by Amazon Cognito User Pools when a user does not exist in the
 // user pool at the time of sign-in with a password, or in the forgot-password flow.
 type CognitoEventUserPoolsMigrateUser struct {
 	CognitoEventUserPoolsHeader
@@ -74,7 +82,7 @@ type CognitoEventUserPoolsCallerContext struct {
 	ClientID      string `json:"clientId"`
 }
 
-// CognitoEventUserPoolsHeader contains common data from events sent by AWS Cognito User Pools
+// CognitoEventUserPoolsHeader contains common data from events sent by Amazon Cognito User Pools
 type CognitoEventUserPoolsHeader struct {
 	Version       string                             `json:"version"`
 	TriggerSource string                             `json:"triggerSource"`
@@ -125,11 +133,24 @@ type CognitoEventUserPoolsPreTokenGenRequest struct {
 	ClientMetadata     map[string]string  `json:"clientMetadata"`
 }
 
-// CognitoEventUserPoolsPreTokenGenResponse containst the response portion of  a PreTokenGen event
+// CognitoEventUserPoolsPreTokenGenV2Request contains request portion of V2 PreTokenGen event
+type CognitoEventUserPoolsPreTokenGenV2Request struct {
+	UserAttributes     map[string]string  `json:"userAttributes"`
+	GroupConfiguration GroupConfiguration `json:"groupConfiguration"`
+	ClientMetadata     map[string]string  `json:"clientMetadata,omitempty"`
+	Scopes             []string           `json:"scopes"`
+}
+
+// CognitoEventUserPoolsPreTokenGenResponse contains the response portion of a PreTokenGen event
 type CognitoEventUserPoolsPreTokenGenResponse struct {
 	ClaimsOverrideDetails ClaimsOverrideDetails `json:"claimsOverrideDetails"`
 }
 
+// CognitoEventUserPoolsPreTokenGenV2Response contains the response portion of a V2 PreTokenGen event
+type CognitoEventUserPoolsPreTokenGenV2Response struct {
+	ClaimsAndScopeOverrideDetails ClaimsAndScopeOverrideDetails `json:"claimsAndScopeOverrideDetails"`
+}
+
 // CognitoEventUserPoolsPostAuthenticationRequest contains the request portion of a PostAuthentication event
 type CognitoEventUserPoolsPostAuthenticationRequest struct {
 	NewDeviceUsed  bool              `json:"newDeviceUsed"`
@@ -157,14 +178,35 @@ type CognitoEventUserPoolsMigrateUserResponse struct {
 	ForceAliasCreation     bool              `json:"forceAliasCreation"`
 }
 
+// ClaimsAndScopeOverrideDetails allows lambda to add, suppress or override V2 claims and scopes in the token
+type ClaimsAndScopeOverrideDetails struct {
+	IDTokenGeneration     IDTokenGeneration     `json:"idTokenGeneration"`
+	AccessTokenGeneration AccessTokenGeneration `json:"accessTokenGeneration"`
+	GroupOverrideDetails  GroupConfiguration    `json:"groupOverrideDetails"`
+}
+
+// IDTokenGeneration allows lambda to modify the ID token
+type IDTokenGeneration struct {
+	ClaimsToAddOrOverride map[string]string `json:"claimsToAddOrOverride"`
+	ClaimsToSuppress      []string          `json:"claimsToSuppress"`
+}
+
+// AccessTokenGeneration allows lambda to modify the access token
+type AccessTokenGeneration struct {
+	ClaimsToAddOrOverride map[string]string `json:"claimsToAddOrOverride"`
+	ClaimsToSuppress      []string          `json:"claimsToSuppress"`
+	ScopesToAdd           []string          `json:"scopesToAdd"`
+	ScopesToSuppress      []string          `json:"scopesToSuppress"`
+}
+
 // ClaimsOverrideDetails allows lambda to add, suppress or override claims in the token
 type ClaimsOverrideDetails struct {
 	GroupOverrideDetails  GroupConfiguration `json:"groupOverrideDetails"`
 	ClaimsToAddOrOverride map[string]string  `json:"claimsToAddOrOverride"`
 	ClaimsToSuppress      []string           `json:"claimsToSuppress"`
 }
 
-// GroupConfiguration allows lambda to override groups, roles and set a perferred role
+// GroupConfiguration allows lambda to override groups, roles and set a preferred role
 type GroupConfiguration struct {
 	GroupsToOverride   []string `json:"groupsToOverride"`
 	IAMRolesToOverride []string `json:"iamRolesToOverride"`
@@ -194,7 +236,7 @@ type CognitoEventUserPoolsDefineAuthChallengeResponse struct {
 	FailAuthentication bool   `json:"failAuthentication"`
 }
 
-// CognitoEventUserPoolsDefineAuthChallenge sent by AWS Cognito User Pools to initiate custom authentication flow
+// CognitoEventUserPoolsDefineAuthChallenge sent by Amazon Cognito User Pools to initiate custom authentication flow
 type CognitoEventUserPoolsDefineAuthChallenge struct {
 	CognitoEventUserPoolsHeader
 	Request  CognitoEventUserPoolsDefineAuthChallengeRequest  `json:"request"`
@@ -216,7 +258,7 @@ type CognitoEventUserPoolsCreateAuthChallengeResponse struct {
 	ChallengeMetadata          string            `json:"challengeMetadata"`
 }
 
-// CognitoEventUserPoolsCreateAuthChallenge sent by AWS Cognito User Pools to create a challenge to present to the user
+// CognitoEventUserPoolsCreateAuthChallenge sent by Amazon Cognito User Pools to create a challenge to present to the user
 type CognitoEventUserPoolsCreateAuthChallenge struct {
 	CognitoEventUserPoolsHeader
 	Request  CognitoEventUserPoolsCreateAuthChallengeRequest  `json:"request"`
@@ -236,15 +278,15 @@ type CognitoEventUserPoolsVerifyAuthChallengeResponse struct {
 	AnswerCorrect bool `json:"answerCorrect"`
 }
 
-// CognitoEventUserPoolsVerifyAuthChallenge sent by AWS Cognito User Pools to verify if the response from the end user
+// CognitoEventUserPoolsVerifyAuthChallenge sent by Amazon Cognito User Pools to verify if the response from the end user
 // for a custom Auth Challenge is valid or not
 type CognitoEventUserPoolsVerifyAuthChallenge struct {
 	CognitoEventUserPoolsHeader
 	Request  CognitoEventUserPoolsVerifyAuthChallengeRequest  `json:"request"`
 	Response CognitoEventUserPoolsVerifyAuthChallengeResponse `json:"response"`
 }
 
-// CognitoEventUserPoolsCustomMessage is sent by AWS Cognito User Pools before a verification or MFA message is sent,
+// CognitoEventUserPoolsCustomMessage is sent by Amazon Cognito User Pools before a verification or MFA message is sent,
 // allowing a user to customize the message dynamically.
 type CognitoEventUserPoolsCustomMessage struct {
 	CognitoEventUserPoolsHeader

diff --git a/events/cognito_test.go b/events/cognito_test.go
@@ -140,6 +140,28 @@ func TestCognitoEventUserPoolsPreTokenGenMarshaling(t *testing.T) {
 	test.AssertJsonsEqual(t, inputJSON, outputJSON)
 }
 
+func TestCognitoEventUserPoolsPreTokenGenV2Marshaling(t *testing.T) {
+	// read json from file
+	inputJSON, err := ioutil.ReadFile("./testdata/cognito-event-userpools-pretokengen-v2.json")
+	if err != nil {
+		t.Errorf("could not open test file. details: %v", err)
+	}
+
+	// de-serialize into CognitoEvent
+	var inputEvent CognitoEventUserPoolsPreTokenGenV2
+	if err := json.Unmarshal(inputJSON, &inputEvent); err != nil {
+		t.Errorf("could not unmarshal event. details: %v", err)
+	}
+
+	// serialize to json
+	outputJSON, err := json.Marshal(inputEvent)
+	if err != nil {
+		t.Errorf("could not marshal event. details: %v", err)
+	}
+
+	test.AssertJsonsEqual(t, inputJSON, outputJSON)
+}
+
 func TestCognitoEventUserPoolsDefineAuthChallengeMarshaling(t *testing.T) {
 	var inputEvent CognitoEventUserPoolsDefineAuthChallenge
 	test.AssertJsonFile(t, "./testdata/cognito-event-userpools-define-auth-challenge.json", &inputEvent)

diff --git a/events/testdata/cognito-event-userpools-pretokengen-v2.json b/events/testdata/cognito-event-userpools-pretokengen-v2.json
@@ -0,0 +1,71 @@
+{
+  "version": "2",
+  "triggerSource": "TokenGeneration_Authentication",
+  "region": "us-west-2",
+  "userPoolId": "us-east-1_EXAMPLE",
+  "userName": "brcotter",
+  "callerContext": {
+    "awsSdkVersion": "aws-sdk-unknown-unknown",
+    "clientId": "1example23456789"
+  },
+  "request": {
+    "userAttributes": {
+      "sub": "a36036a8-9061-424d-a737-56d57dae7bc6",
+      "cognito:email_alias": "[email protected]",
+      "cognito:user_status": "CONFIRMED",
+      "email_verified": "true",
+      "email": "[email protected]"
+    },
+    "groupConfiguration": {
+      "groupsToOverride": [],
+      "iamRolesToOverride": [],
+      "preferredRole": null
+    },
+    "scopes": [
+      "aws.cognito.signin.user.admin"
+    ]
+  },
+  "response": {
+    "claimsAndScopeOverrideDetails": {
+      "idTokenGeneration": {
+        "claimsToAddOrOverride": {
+          "family_name": "xyz"
+        },
+        "claimsToSuppress": [
+          "email",
+          "birthdate"
+        ]
+      },
+      "accessTokenGeneration": {
+        "claimsToAddOrOverride": {
+          "family_name": "xyz"
+        },
+        "claimsToSuppress": [
+          "email",
+          "birthdate"
+        ],
+        "scopesToAdd": [
+          "scope1",
+          "scope2",
+          "scopeLomond"
+        ],
+        "scopesToSuppress": [
+          "phone_number"
+        ]
+      },
+      "groupOverrideDetails": {
+        "groupsToOverride": [
+          "group-A",
+          "group-B",
+          "group-C"
+        ],
+        "iamRolesToOverride": [
+          "arn:aws:iam::123456789012:role/sns_callerA",
+          "arn:aws:iam::123456789012:role/sns_callerB",
+          "arn:aws:iam::123456789012:role/sns_callerC"
+        ],
+        "preferredRole": "arn:aws:iam::123456789012:role/sns_caller"
+      }
+    }
+  }
+}