Skip to content

Commit

Permalink
Added support for Cognito pre token generation with access token cust…
Browse files Browse the repository at this point in the history 
…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 <bmoffatt@users.noreply.github.com>
Co-authored-by: Bryan Moffatt <moffattb@amazon.com>
  • Loading branch information
@AMZ-brandon @bmoffatt
3 people committed Dec 20, 2023
1 parent 30d7c7e commit 326401a
Show file tree
Hide file tree
Showing 3 changed files with 150 additions and 15 deletions.
72 changes: 57 additions & 15 deletions events/cognito.go
View file
Expand Up @@ -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"`
Expand All @@ -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
Expand All @@ -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"`
Expand Down Expand Up @@ -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"`
Expand Down Expand Up @@ -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"`
Expand Down Expand Up @@ -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"`
Expand All @@ -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"`
Expand All @@ -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
Expand Down
22 changes: 22 additions & 0 deletions events/cognito_test.go
View file
Expand Up @@ -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)
Expand Down
71 changes: 71 additions & 0 deletions events/testdata/cognito-event-userpools-pretokengen-v2.json
View file
@@ -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": "testuser@example.com",
"cognito:user_status": "CONFIRMED",
"email_verified": "true",
"email": "testuser@example.com"
},
"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"
}
}
}
}

0 comments on commit 326401a

Please sign in to comment.