feat: implement signIn with TOTP

frouriojs · Dec 24, 2024 · 81dda4c · 81dda4c
1 parent c3a92ba
commit 81dda4c
Show file tree

Hide file tree

Showing 14 changed files with 253 additions and 85 deletions.
diff --git a/server/common/types/signIn.ts b/server/common/types/signIn.ts
@@ -1,3 +1,4 @@
+import type { MFA_SETTING_LIST } from 'common/constants';
 import type { TargetBody } from './auth';
 import type { MaybeId } from './brandedId';
 
@@ -39,24 +40,42 @@ export type RefreshTokenAuthTarget = TargetBody<
 export type InitiateAuthTarget = UserSrpAuthTarget | RefreshTokenAuthTarget;
 
 export type RespondToAuthChallengeTarget = TargetBody<
-  {
-    ChallengeName: 'PASSWORD_VERIFIER';
-    ChallengeResponses: {
-      PASSWORD_CLAIM_SECRET_BLOCK: string;
-      PASSWORD_CLAIM_SIGNATURE: string;
-      TIMESTAMP: string;
-      USERNAME: string;
-    };
-    ClientId: MaybeId['userPoolClient'];
-  },
-  {
-    AuthenticationResult: {
-      AccessToken: string;
-      ExpiresIn: number;
-      IdToken: string;
-      RefreshToken: string;
-      TokenType: 'Bearer';
-    };
-    ChallengeParameters: Record<string, never>;
-  }
+  | {
+      ChallengeName: 'PASSWORD_VERIFIER';
+      ChallengeResponses: {
+        PASSWORD_CLAIM_SECRET_BLOCK: string;
+        PASSWORD_CLAIM_SIGNATURE: string;
+        TIMESTAMP: string;
+        USERNAME: string;
+      };
+      ClientId: MaybeId['userPoolClient'];
+      Session?: undefined;
+    }
+  | {
+      ChallengeName: 'SOFTWARE_TOKEN_MFA';
+      ChallengeResponses: {
+        SOFTWARE_TOKEN_MFA_CODE: string;
+        USERNAME: string;
+      };
+      ClientId: MaybeId['userPoolClient'];
+      Session: string;
+    },
+  | {
+      AuthenticationResult: {
+        AccessToken: string;
+        ExpiresIn: number;
+        IdToken: string;
+        RefreshToken: string;
+        TokenType: 'Bearer';
+      };
+      ChallengeName?: undefined;
+      Session?: undefined;
+      ChallengeParameters: Record<string, never>;
+    }
+  | {
+      AuthenticationResult?: undefined;
+      ChallengeName: (typeof MFA_SETTING_LIST)[number];
+      Session: string;
+      ChallengeParameters: Record<string, never>;
+    }
 >;
diff --git a/server/common/types/user.ts b/server/common/types/user.ts
@@ -35,6 +35,7 @@ export type SocialUserEntity = {
   createdTime: number;
   updatedTime: number;
   challenge?: undefined;
+  srpAuth?: undefined;
   preferredMfaSetting?: undefined;
   mfaSettingList?: undefined;
   totpSecretCode?: undefined;
@@ -60,6 +61,7 @@ export type CognitoUserEntity = {
   createdTime: number;
   updatedTime: number;
   challenge?: ChallengeVal;
+  srpAuth?: { timestamp: string; clientSignature: string };
   preferredMfaSetting?: (typeof MFA_SETTING_LIST)[number];
   mfaSettingList?: (typeof MFA_SETTING_LIST)[number][];
   totpSecretCode?: string;

diff --git a/server/domain/user/model/signInMethod.ts b/server/domain/user/model/signInMethod.ts
@@ -66,4 +66,8 @@ export const signInMethod = {
       user: params.user,
     });
   },
+  challengeMfa: (
+    user: CognitoUserEntity,
+    srpAuth: { timestamp: string; clientSignature: string },
+  ): CognitoUserEntity => ({ ...user, srpAuth }),
 };
diff --git a/server/domain/user/repository/toUserEntity.ts b/server/domain/user/repository/toUserEntity.ts
@@ -51,6 +51,17 @@ export const toCognitoUserEntity = (
         value: attr.value,
       }),
     ),
+    srpAuth: z
+      .object({ timestamp: z.string(), clientSignature: z.string() })
+      .optional()
+      .parse(
+        prismaUser.srpAuthTimestamp
+          ? {
+              timestamp: prismaUser.srpAuthTimestamp,
+              clientSignature: prismaUser.srpAuthClientSignature,
+            }
+          : undefined,
+      ),
     preferredMfaSetting: z
       .enum(MFA_SETTING_LIST)
       .optional()

diff --git a/server/domain/user/repository/userCommand.ts b/server/domain/user/repository/userCommand.ts
@@ -27,6 +27,8 @@ export const userCommand = {
         pubA: user.challenge?.pubA,
         pubB: user.challenge?.pubB,
         secB: user.challenge?.secB,
+        srpAuthTimestamp: user.srpAuth?.timestamp,
+        srpAuthClientSignature: user.srpAuth?.clientSignature,
         preferredMfaSetting: user.preferredMfaSetting,
         enabledTotp: user.mfaSettingList?.some((setting) => setting === 'SOFTWARE_TOKEN_MFA'),
         totpSecretCode: user.totpSecretCode,
@@ -48,9 +50,6 @@ export const userCommand = {
         confirmationCode: user.confirmationCode,
         authorizationCode: user.authorizationCode,
         codeChallenge: user.codeChallenge,
-        preferredMfaSetting: user.preferredMfaSetting,
-        enabledTotp: user.mfaSettingList?.some((setting) => setting === 'SOFTWARE_TOKEN_MFA'),
-        totpSecretCode: user.totpSecretCode,
         userPoolId: user.userPoolId,
         attributes: { createMany: { data: user.attributes } },
         createdAt: new Date(user.createdTime),

diff --git a/server/domain/user/useCase/signInUseCase.ts b/server/domain/user/useCase/signInUseCase.ts
@@ -11,6 +11,7 @@ import { userPoolQuery } from 'domain/userPool/repository/userPoolQuery';
 import { cognitoAssert } from 'service/cognitoAssert';
 import { EXPIRES_SEC } from 'service/constants';
 import { transaction } from 'service/prismaClient';
+import { mfaMethod } from '../model/mfaMethod';
 import { signInMethod } from '../model/signInMethod';
 import { isEmailVerified } from '../service/isEmailVerified';
 
@@ -64,27 +65,69 @@ export const signInUseCase = {
       const jwks = await userPoolQuery.findJwks(tx, user.userPoolId);
 
       assert(pool.id === poolClient.userPoolId);
-      assert(user.challenge?.secretBlock === req.ChallengeResponses.PASSWORD_CLAIM_SECRET_BLOCK);
 
-      cognitoAssert(isEmailVerified(user), 'User is not confirmed.');
+      if (req.ChallengeName === 'PASSWORD_VERIFIER') {
+        assert(user.challenge);
+        assert(user.challenge.secretBlock === req.ChallengeResponses.PASSWORD_CLAIM_SECRET_BLOCK);
 
-      const tokens = signInMethod.srpAuth({
-        user,
-        timestamp: req.ChallengeResponses.TIMESTAMP,
-        clientSignature: req.ChallengeResponses.PASSWORD_CLAIM_SIGNATURE,
-        jwks,
-        pool,
-        poolClient,
-      });
+        cognitoAssert(isEmailVerified(user), 'User is not confirmed.');
 
-      return {
-        AuthenticationResult: {
-          ...tokens,
-          ExpiresIn: 3600,
-          RefreshToken: user.refreshToken,
-          TokenType: 'Bearer',
-        },
-        ChallengeParameters: {},
-      };
+        if (user.mfaSettingList?.some((s) => s === 'SOFTWARE_TOKEN_MFA')) {
+          const updated = signInMethod.challengeMfa(user, {
+            timestamp: req.ChallengeResponses.TIMESTAMP,
+            clientSignature: req.ChallengeResponses.PASSWORD_CLAIM_SIGNATURE,
+          });
+
+          await userCommand.save(tx, updated);
+
+          return {
+            ChallengeName: 'SOFTWARE_TOKEN_MFA',
+            Session: 'magnito_dummy_session',
+            ChallengeParameters: {},
+          };
+        }
+
+        const tokens = signInMethod.srpAuth({
+          user,
+          timestamp: req.ChallengeResponses.TIMESTAMP,
+          clientSignature: req.ChallengeResponses.PASSWORD_CLAIM_SIGNATURE,
+          jwks,
+          pool,
+          poolClient,
+        });
+
+        return {
+          AuthenticationResult: {
+            ...tokens,
+            ExpiresIn: 3600,
+            RefreshToken: user.refreshToken,
+            TokenType: 'Bearer',
+          },
+          ChallengeParameters: {},
+        };
+      } else {
+        const updated = mfaMethod.verify(user, req.ChallengeResponses.SOFTWARE_TOKEN_MFA_CODE);
+
+        assert(updated.srpAuth);
+
+        const tokens = signInMethod.srpAuth({
+          user,
+          timestamp: updated.srpAuth.timestamp,
+          clientSignature: updated.srpAuth.clientSignature,
+          jwks,
+          pool,
+          poolClient,
+        });
+
+        return {
+          AuthenticationResult: {
+            ...tokens,
+            ExpiresIn: 3600,
+            RefreshToken: user.refreshToken,
+            TokenType: 'Bearer',
+          },
+          ChallengeParameters: {},
+        };
+      }
     }),
 };
diff --git a/server/prisma/migrations/20241224100030_/migration.sql b/server/prisma/migrations/20241224100030_/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "User" ADD COLUMN "srpAuthClientSignature" TEXT;
+ALTER TABLE "User" ADD COLUMN "srpAuthTimestamp" TEXT;
diff --git a/server/prisma/schema.prisma b/server/prisma/schema.prisma
@@ -8,32 +8,34 @@ generator client {
 }
 
 model User {
-  id                  String          @id
-  kind                String?
-  name                String
-  email               String
-  enabled             Boolean
-  status              String
-  password            String?
-  confirmationCode    String?
-  salt                String?
-  verifier            String?
-  refreshToken        String
-  createdAt           DateTime
-  updatedAt           DateTime
-  provider            String?
-  authorizationCode   String?
-  codeChallenge       String?
-  secretBlock         String?
-  pubA                String?
-  pubB                String?
-  secB                String?
-  preferredMfaSetting String?
-  enabledTotp         Boolean?
-  totpSecretCode      String?
-  attributes          UserAttribute[]
-  UserPool            UserPool        @relation(fields: [userPoolId], references: [id])
-  userPoolId          String
+  id                     String          @id
+  kind                   String?
+  name                   String
+  email                  String
+  enabled                Boolean
+  status                 String
+  password               String?
+  confirmationCode       String?
+  salt                   String?
+  verifier               String?
+  refreshToken           String
+  createdAt              DateTime
+  updatedAt              DateTime
+  provider               String?
+  authorizationCode      String?
+  codeChallenge          String?
+  secretBlock            String?
+  pubA                   String?
+  pubB                   String?
+  secB                   String?
+  srpAuthTimestamp       String?
+  srpAuthClientSignature String?
+  preferredMfaSetting    String?
+  enabledTotp            Boolean?
+  totpSecretCode         String?
+  attributes             UserAttribute[]
+  UserPool               UserPool        @relation(fields: [userPoolId], references: [id])
+  userPoolId             String
 }
 
 model UserAttribute {

diff --git a/server/tests/api/apiClient.ts b/server/tests/api/apiClient.ts
@@ -3,7 +3,6 @@ import api from 'api/$api';
 import axios from 'axios';
 import { COOKIE_NAME } from 'service/constants';
 import { PORT } from 'service/envValues';
-import { createCognitoUserAndToken } from './utils';
 
 const baseURL = `http://127.0.0.1:${PORT}`;
 
@@ -15,9 +14,9 @@ export const testUserName = 'test-user';
 
 export const testPassword = 'Test-user-password1';
 
-export const createUserClient = async (): Promise<typeof noCookieClient> => {
-  const token = await createCognitoUserAndToken();
-
+export const createUserClient = async (token: {
+  AccessToken: string;
+}): Promise<typeof noCookieClient> => {
   const agent = axios.create({
     baseURL,
     headers: { cookie: `${COOKIE_NAME}=${token.AccessToken}`, 'Content-Type': 'text/plain' },

diff --git a/server/tests/api/changePassword.test.ts b/server/tests/api/changePassword.test.ts
@@ -4,10 +4,12 @@ import { calculateSrpA } from 'domain/user/service/srp/calcSrpA';
 import { fromBuffer } from 'domain/user/service/srp/util';
 import { DEFAULT_USER_POOL_CLIENT_ID } from 'service/envValues';
 import { test } from 'vitest';
-import { createUserClient, noCookieClient, testPassword, testUserName } from './apiClient';
+import { noCookieClient, testPassword, testUserName } from './apiClient';
+import { createCognitoUserAndToken } from './utils';
 
 test('changePassword', async () => {
-  await createUserClient();
+  await createCognitoUserAndToken();
+
   const { a, A } = calculateSrpA();
   const res1 = await noCookieClient.$post({
     headers: { 'x-amz-target': 'AWSCognitoIdentityProviderService.InitiateAuth' },

diff --git a/server/tests/api/private.test.ts b/server/tests/api/private.test.ts
@@ -1,9 +1,9 @@
 import { expect, test } from 'vitest';
 import { createUserClient, noCookieClient } from './apiClient';
-import { GET } from './utils';
+import { createCognitoUserAndToken, GET } from './utils';
 
 test(GET(noCookieClient.private), async () => {
-  const userClient = await createUserClient();
+  const userClient = await createCognitoUserAndToken().then(createUserClient);
   const res = await userClient.private.$get();
 
   expect(res).toEqual('');
@@ -12,7 +12,7 @@ test(GET(noCookieClient.private), async () => {
 });
 
 test(GET(noCookieClient.private.me), async () => {
-  const userClient = await createUserClient();
+  const userClient = await createCognitoUserAndToken().then(createUserClient);
   const res = await userClient.private.me.get();
 
   expect(res.status).toBe(200);