Week 1 : AWS Cognito Pre token generation Lambda trigger in .NET


Sometimes when working with Cognito you want to add, remove, or modify claims on both ID and Access tokens before returning them to an application. This week I had to work through creating such a trigger using .Net so here is a small tutorial on getting a simple one up and running to hopefully make another developer out theres day easier.

Things you will need (assuming you already have dotnet installed)

First you can create a empty lambda with the following command.

mkdir tokenLambdaDemo
cd tokenLambdaDemo
dotnet new lambda.EmptyFunction

Open this project up in your editor of choice and open to src/tokenLambdaDemo/Function.cs. This is the function that gets executed when the lambda is run, and the one we are going to make some changes to for our own end.

The first change we are going to make is adding the package Amazon.Lambda.CognitoEvents to our project dependencies. This package contains classes that we can make use of as input types for our function handler.
from src/tokenLambdaDemo/ run the following :

dotnet add package Amazon.Lambda.CognitoEvents

Now you can edit the function to the following :

using Amazon.Lambda.Core;
using Amazon.Lambda.CognitoEvents;

// Assembly attribute to enable the Lambda function's JSON input to be converted into a .NET class.
[assembly: LambdaSerializer(typeof(Amazon.Lambda.Serialization.SystemTextJson.DefaultLambdaJsonSerializer))]

namespace tokenLambdaDemo;

public class Function
{

    /// <summary>
    /// A simple function that takes a CognitoPreTokenGenerationV2Event and add a claim to the access token
    /// </summary>
    /// <param name="input">The event for the Lambda function handler to process.</param>
    /// <param name="context">The ILambdaContext that provides methods for logging and describing the Lambda environment.</param>
    /// <returns></returns>
    public CognitoPreTokenGenerationV2Event FunctionHandler(CognitoPreTokenGenerationV2Event input, ILambdaContext context)
    {
        return input;
    }
}

Our function now both accepts as input and returns a CognitoPreTokenGenerationV2Event. We use a V2 event as V1 only allows making modifications to the ID token.

Next we need to add the claims we want in our access token to the Response property of this input event. First by creating a new CognitoPreTokenGenerationV2Response and then adding the claim. This event is documented at https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-pre-token-generation.html

   public CognitoPreTokenGenerationV2Event FunctionHandler(CognitoPreTokenGenerationV2Event input, ILambdaContext context)
    {
        var resp = new CognitoPreTokenGenerationV2Response();
        resp.ClaimsAndScopeOverrideDetails.AccessTokenGeneration.ClaimsToAddOrOverride.Add("demoClaim", "demoValue");
        input.Response = resp;
        return input;
    }

You can then deploy this trigger (make sure you are authenticated with the aws cli!)

dotnet lambda deploy-function tokenLambdaDemo

And then you can run this sample event against it.

{
    "request": {
        "userAttributes": {
            "string": "string"
        },
        "scopes": ["string", "string"],
        "groupConfiguration": {
            "groupsToOverride": ["string", "string"],
            "iamRolesToOverride": ["string", "string"],
            "preferredRole": "string"
        },
        "clientMetadata": {
            "string": "string"
        }
    }
}

You will get a response that looks simillar to this:

{
  "callerContext": {},
  "request": {
    "groupConfiguration": {
      "groupsToOverride": [
        "string",
        "string"
      ],
      "iamRolesToOverride": [
        "string",
        "string"
      ],
      "preferredRole": "string"
    },
    "clientMetadata": {
      "string": "string"
    },
    "scopes": [
      "string",
      "string"
    ],
    "userAttributes": {
      "string": "string"
    }
  },
  "response": {
    "claimsAndScopeOverrideDetails": {
      "idTokenGeneration": {
        "claimsToAddOrOverride": {},
        "claimsToSuppress": []
      },
      "accessTokenGeneration": {
        "claimsToAddOrOverride": {
          "demoClaim": "demoValue"
        },
        "claimsToSuppress": [],
        "scopesToAdd": [],
        "scopesToSuppress": []
      },
      "groupOverrideDetails": {
        "groupsToOverride": [],
        "iamRolesToOverride": []
      }
    }
  }
}

And that’s it!

This basic example just adds a new claim demoClaim with the values demoValue. You can add whatever logic you see fit for modifying the token in whatever way you need. Want to use the subclaim to call a service and enrich the token with permissions from a database somewhere? Go for it! Want to remove some claims because you don’t want the client to know some things for whatever reason? You can do it with this!


Leave a Reply

Your email address will not be published. Required fields are marked *