feat(token): add refresh token, revoke token and initiate auth

Signed-off-by: James Gregory <[email protected]>
michaelruocco · Dec 6, 2021 · 0d46ed7 · 0d46ed7
1 parent 65be7e9
commit 0d46ed7
Show file tree

Hide file tree

Showing 17 changed files with 429 additions and 59 deletions.
diff --git a/README.md b/README.md
@@ -41,7 +41,7 @@ A _Good Enough_ offline emulator for [Amazon Cognito](https://aws.amazon.com/cog
 | AdminForgetDevice                | ❌                   |
 | AdminGetDevice                   | ❌                   |
 | AdminGetUser                     | ✅                   |
-| AdminInitiateAuth                | ❌                   |
+| AdminInitiateAuth                | 🕒 (partial support) |
 | AdminLinkProviderForUser         | ❌                   |
 | AdminListDevices                 | ❌                   |
 | AdminListGroupsForUser           | ❌                   |
@@ -108,7 +108,7 @@ A _Good Enough_ offline emulator for [Amazon Cognito](https://aws.amazon.com/cog
 | ListUsersInGroup                 | ❌                   |
 | ResendConfirmationCode           | ❌                   |
 | RespondToAuthChallenge           | 🕒 (partial support) |
-| RevokeToken                      | ❌                   |
+| RevokeToken                      | 🕒 (partial support) |
 | SetRiskConfiguration             | ❌                   |
 | SetUICustomization               | ❌                   |
 | SetUserMFAPreference             | ❌                   |

diff --git a/integration-tests/userPoolService.test.ts b/integration-tests/userPoolService.test.ts
@@ -61,6 +61,7 @@ describe("User Pool Service", () => {
           UserLastModifiedDate: now,
           UserCreateDate: now,
           Enabled: true,
+          RefreshTokens: [],
         });
 
         const file = JSON.parse(
@@ -79,6 +80,7 @@ describe("User Pool Service", () => {
             UserLastModifiedDate: now.toISOString(),
             UserCreateDate: now.toISOString(),
             Enabled: true,
+            RefreshTokens: [],
           },
         });
       });
@@ -99,6 +101,7 @@ describe("User Pool Service", () => {
           UserLastModifiedDate: now,
           UserCreateDate: now,
           Enabled: true,
+          RefreshTokens: [],
         });
 
         let file = JSON.parse(
@@ -118,6 +121,7 @@ describe("User Pool Service", () => {
             UserLastModifiedDate: now.toISOString(),
             UserCreateDate: now.toISOString(),
             Enabled: true,
+            RefreshTokens: [],
           },
         });
 
@@ -132,6 +136,7 @@ describe("User Pool Service", () => {
           UserLastModifiedDate: now,
           UserCreateDate: now,
           Enabled: true,
+          RefreshTokens: [],
         });
 
         file = JSON.parse(
@@ -150,6 +155,7 @@ describe("User Pool Service", () => {
             UserLastModifiedDate: now.toISOString(),
             UserCreateDate: now.toISOString(),
             Enabled: true,
+            RefreshTokens: [],
           },
         });
       });
@@ -174,6 +180,7 @@ describe("User Pool Service", () => {
           UserCreateDate: new Date(),
           UserLastModifiedDate: new Date(),
           Enabled: true,
+          RefreshTokens: [],
         });
       });
 
@@ -212,6 +219,7 @@ describe("User Pool Service", () => {
         UserCreateDate: now,
         UserLastModifiedDate: now,
         Enabled: true,
+        RefreshTokens: [],
       });
 
       await userPool.saveUser({
@@ -222,6 +230,7 @@ describe("User Pool Service", () => {
         UserCreateDate: now,
         UserLastModifiedDate: now,
         Enabled: true,
+        RefreshTokens: [],
       });
     });
 
@@ -241,6 +250,7 @@ describe("User Pool Service", () => {
           UserCreateDate: now,
           UserLastModifiedDate: now,
           Enabled: true,
+          RefreshTokens: [],
         },
         {
           Username: "2",
@@ -250,6 +260,7 @@ describe("User Pool Service", () => {
           UserCreateDate: now,
           UserLastModifiedDate: now,
           Enabled: true,
+          RefreshTokens: [],
         },
       ]);
     });

diff --git a/src/__tests__/testDataBuilder.ts b/src/__tests__/testDataBuilder.ts
@@ -27,6 +27,7 @@ export const user = (partial?: Partial<User>): User => ({
   UserLastModifiedDate: partial?.UserLastModifiedDate ?? new Date(),
   Username: partial?.Username ?? id("User"),
   UserStatus: partial?.UserStatus ?? "CONFIRMED",
+  RefreshTokens: [],
 });
 
 export const userPool = (partial?: Partial<UserPool>): UserPool => {

diff --git a/src/services/dataStore.ts b/src/services/dataStore.ts
@@ -60,6 +60,7 @@ export const createDataStore: CreateDataStore = async (
       `Deserialize: Expected ${key} to contain a String or Number, received a ${typeof value}`
     );
   };
+
   const engine = new StormDB.localFileEngine(`${directory}/${id}.json`, {
     async: true,
     serialize: (obj: unknown) =>

diff --git a/src/services/messageDelivery/messageDelivery.test.ts b/src/services/messageDelivery/messageDelivery.test.ts
@@ -11,6 +11,7 @@ describe("Message Delivery", () => {
     Enabled: true,
     UserCreateDate: new Date(),
     UserLastModifiedDate: new Date(),
+    RefreshTokens: [],
   };
 
   describe("when delivery method is EMAIL", () => {

diff --git a/src/services/tokens.ts b/src/services/tokens.ts
@@ -31,6 +31,16 @@ export function generateTokens(
   const authTime = Math.floor(clock.get().getTime() / 1000);
   const sub = attributeValue("sub", user.Attributes);
 
+  const customAttributes = user.Attributes.reduce<{
+    [key: string]: string | undefined;
+  }>((acc, attr) => {
+    if (attr.Name.startsWith("custom:")) {
+      acc[attr.Name] = attr.Value;
+    }
+
+    return acc;
+  }, {});
+
   const issuer = `${tokenConfig.IssuerDomain}/${userPoolId}`;
   return {
     AccessToken: jwt.sign(
@@ -61,6 +71,7 @@ export function generateTokens(
         auth_time: authTime,
         "cognito:username": user.Username,
         email: attributeValue("email", user.Attributes),
+        ...customAttributes,
       },
       PrivateKey.pem,
       {
@@ -71,6 +82,21 @@ export function generateTokens(
         keyid: "CognitoLocal",
       }
     ),
-    RefreshToken: "<< TODO >>",
+    // this content is for debugging purposes only
+    // in reality token payload is encrypted and uses different algorithm
+    RefreshToken: jwt.sign(
+      {
+        "cognito:username": user.Username,
+        email: attributeValue("email", user.Attributes),
+        // something unique for each token
+        unique: uuid.v4(),
+      },
+      PrivateKey.pem,
+      {
+        algorithm: "RS256",
+        issuer,
+        expiresIn: "7d",
+      }
+    ),
   };
 }
diff --git a/src/services/triggers/userMigration.ts b/src/services/triggers/userMigration.ts
@@ -41,61 +41,64 @@ interface UserMigrationServices {
   lambda: Lambda;
 }
 
-export const UserMigration = ({
-  lambda,
-  cognitoClient,
-  clock,
-}: UserMigrationServices): UserMigrationTrigger => async ({
-  clientId,
-  clientMetadata,
-  password,
-  userAttributes,
-  username,
-  userPoolId,
-  validationData,
-}): Promise<User> => {
-  const userPool = await cognitoClient.getUserPoolForClientId(clientId);
-  if (!userPool) {
-    throw new ResourceNotFoundError();
-  }
+export const UserMigration =
+  ({
+    lambda,
+    cognitoClient,
+    clock,
+  }: UserMigrationServices): UserMigrationTrigger =>
+  async ({
+    clientId,
+    clientMetadata,
+    password,
+    userAttributes,
+    username,
+    userPoolId,
+    validationData,
+  }): Promise<User> => {
+    const userPool = await cognitoClient.getUserPoolForClientId(clientId);
+    if (!userPool) {
+      throw new ResourceNotFoundError();
+    }
 
-  let result: CognitoUserPoolResponse;
+    let result: CognitoUserPoolResponse;
 
-  try {
-    result = await lambda.invoke("UserMigration", {
-      clientId,
-      clientMetadata,
-      password,
-      triggerSource: "UserMigration_Authentication",
-      userAttributes: attributesToRecord(userAttributes),
-      username,
-      userPoolId,
-      validationData,
-    });
-  } catch (ex) {
-    throw new NotAuthorizedError();
-  }
+    try {
+      result = await lambda.invoke("UserMigration", {
+        clientId,
+        clientMetadata,
+        password,
+        triggerSource: "UserMigration_Authentication",
+        userAttributes: attributesToRecord(userAttributes),
+        username,
+        userPoolId,
+        validationData,
+      });
+    } catch (ex) {
+      throw new NotAuthorizedError();
+    }
 
-  const now = clock.get();
-  const user: User = {
-    Attributes: attributesFromRecord(result.userAttributes ?? {}),
-    Enabled: true,
-    Password: password,
-    UserCreateDate: now,
-    UserLastModifiedDate: now,
-    Username: uuid.v4(),
-    UserStatus: result.finalUserStatus ?? "CONFIRMED",
-  };
+    const now = clock.get();
+    const user: User = {
+      Attributes: attributesFromRecord(result.userAttributes ?? {}),
+      Enabled: true,
+      Password: password,
+      UserCreateDate: now,
+      UserLastModifiedDate: now,
+      Username: uuid.v4(),
+      UserStatus: result.finalUserStatus ?? "CONFIRMED",
+      RefreshTokens: [],
+    };
 
-  if (result.forceAliasCreation) {
-    // TODO: do something with aliases?
-  }
+    if (result.forceAliasCreation) {
+      // TODO: do something with aliases?
+    }
 
-  await userPool.saveUser(user);
+    await userPool.saveUser(user);
 
-  if (result.messageAction !== "SUPPRESS") {
-    // TODO: send notification when not suppressed?
-  }
+    if (result.messageAction !== "SUPPRESS") {
+      // TODO: send notification when not suppressed?
+    }
 
-  return user;
-};
+    return user;
+  };
diff --git a/src/services/userPoolService.ts b/src/services/userPoolService.ts
@@ -55,6 +55,7 @@ export interface User {
   Password: string;
   ConfirmationCode?: string;
   MFACode?: string;
+  RefreshTokens: string[];
 }
 
 export interface Group {