Fix: User's group memberships passed to PreTokenGeneration Lambda

Dockerfile.dev

@@ -0,0 +1,17 @@
+FROM node:14-alpine
+WORKDIR /app
+
+# dependencies
+ADD package.json yarn.lock ./
+RUN yarn --frozen-lockfile
+
+# library code
+ADD src src
+
+# bindings
+EXPOSE 9229
+ENV HOST 0.0.0.0
+ENV PORT 9229
+VOLUME /app/.cognito
+
+ENTRYPOINT ["yarn", "start:debug"]

src/services/lambda.test.ts

@@ -443,6 +443,11 @@ describe("Lambda function invoker", () => {
           clientMetadata: {
             client: "metadata",
           },
+          groupConfiguration: {
+            groupsToOverride: ["group1", "group2"],
+            iamRolesToOverride: ["role1", "role2"],
+            preferredRole: "preferredRole",
+          },
         });
 
         expect(mockLambdaClient.invoke).toHaveBeenCalledWith({
@@ -461,9 +466,22 @@ describe("Lambda function invoker", () => {
               clientMetadata: {
                 client: "metadata",
               },
-              groupConfiguration: {},
+              groupConfiguration: {
+                groupsToOverride: ["group1", "group2"],
+                iamRolesToOverride: ["role1", "role2"],
+                preferredRole: "preferredRole",
+              },
+            },
+            response: {
+              claimsOverrideDetails: {
+                claimsToAddOrOverride: {},
+                claimsToSuppress: [],
+                groupOverrideDetails: {
+                  groupsToOverride: [],
+                  iamRolesToOverride: [],
+                },
+              },
             },
-            response: { claimsOverrideDetails: {} },
             userName: "username",
           }),
         });

src/services/lambda.ts

@@ -109,12 +109,12 @@ interface PreTokenGenerationEvent extends EventCommonParameters {
     /**
      * A list of the group names that are associated with the user that the identity token is issued for.
      */
-    groupsToOverride: readonly string[] | undefined;
+    groupsToOverride: string[] | undefined;
 
     /**
      * A list of the current IAM roles associated with these groups.
      */
-    iamRolesToOverride: readonly string[] | undefined;
+    iamRolesToOverride: string[] | undefined;
 
     /**
      * A string indicating the preferred IAM role.
@@ -362,11 +362,19 @@ export class LambdaService implements Lambda {
           userName: event.username,
           request: {
             userAttributes: event.userAttributes,
-            groupConfiguration: {},
+            groupConfiguration: event.groupConfiguration,
             clientMetadata: event.clientMetadata,
           },
           response: {
-            claimsOverrideDetails: {},
+            claimsOverrideDetails: {
+              claimsToAddOrOverride: {},
+              claimsToSuppress: [],
+              groupOverrideDetails: {
+                groupsToOverride: [],
+                iamRolesToOverride: [],
+                preferredRole: undefined,
+              },
+            },
           },
         };
       }

src/services/tokenGenerator.test.ts

@@ -68,6 +68,36 @@ describe("JwtTokenGenerator", () => {
       });
     });
 
+    it("can pass user's group memberships to the PreTokenGeneration lambda", async () => {
+      mockTriggers.enabled.mockImplementation((name) => {
+        return name === "PreTokenGeneration";
+      });
+
+      mockTriggers.preTokenGeneration.mockResolvedValue({
+        claimsOverrideDetails: {
+          claimsToAddOrOverride: {
+            newclaim: "value",
+            email: "something else",
+          },
+        },
+      });
+
+      await tokenGenerator.generate(
+        TestContext,
+        user,
+        ["group1", "group2"],
+        TDB.appClient(),
+        { client: "metadata" },
+        "RefreshTokens"
+      );
+
+      expect(mockTriggers.preTokenGeneration.mock.calls[0][1]).toMatchObject({
+        groupConfiguration: {
+          groupsToOverride: ["group1", "group2"],
+        },
+      });
+    });
+
     it("can suppress claims in the id token", async () => {
       mockTriggers.enabled.mockImplementation((name) => {
         return name === "PreTokenGeneration";

src/services/tokenGenerator.ts

@@ -130,7 +130,7 @@ export class JwtTokenGenerator implements TokenGenerator {
   public async generate(
     ctx: Context,
     user: User,
-    userGroups: readonly string[],
+    userGroups: string[],
     userPoolClient: AppClient,
     clientMetadata: Record<string, string> | undefined,
     source:
@@ -183,14 +183,17 @@ export class JwtTokenGenerator implements TokenGenerator {
         userAttributes: user.Attributes,
         username: user.Username,
         groupConfiguration: {
-          // TODO: this should be populated from the user's groups
-          groupsToOverride: undefined,
-          iamRolesToOverride: undefined,
-          preferredRole: undefined,
+          groupsToOverride: userGroups,
+          iamRolesToOverride: [],
+          preferredRole: "",
         },
         userPoolId: userPoolClient.UserPoolId,
       });
 
+      ctx.logger.info(
+        `claimsOverrideDetails: ${JSON.stringify(result.claimsOverrideDetails)}`
+      );
+
       idToken = applyTokenOverrides(idToken, result.claimsOverrideDetails);
     }
 

src/services/triggers/preTokenGeneration.ts

@@ -36,12 +36,12 @@ export type PreTokenGenerationTrigger = Trigger<
       /**
        * A list of the group names that are associated with the user that the identity token is issued for.
        */
-      groupsToOverride: readonly string[] | undefined;
+      groupsToOverride: string[] | undefined;
 
       /**
        * A list of the current IAM roles associated with these groups.
        */
-      iamRolesToOverride: readonly string[] | undefined;
+      iamRolesToOverride: string[] | undefined;
 
       /**
        * A string indicating the preferred IAM role.