feat: implement SetUserMFAPreferenceCommand

IMPLEMENTATION_COVERAGE.md

@@ -3,7 +3,7 @@
 ## cognito-idp
 
 <details>
-<summary> 21% implemented  </summary>
+<summary> 24% implemented  </summary>
 
 - [ ] AddCustomAttributes
 - [ ] AdminAddUserToGroup
@@ -32,7 +32,7 @@
 - [ ] AdminUpdateDeviceStatus
 - [x] AdminUpdateUserAttributes
 - [ ] AdminUserGlobalSignOut
-- [ ] AssociateSoftwareToken
+- [x] AssociateSoftwareToken
 - [x] ChangePassword
 - [ ] ConfirmDevice
 - [x] ConfirmForgotPassword
@@ -89,7 +89,7 @@
 - [ ] SetLogDeliveryConfiguration
 - [ ] SetRiskConfiguration
 - [ ] SetUiCustomization
-- [ ] SetUserMfaPreference
+- [x] SetUserMfaPreference
 - [ ] SetUserPoolMfaConfig
 - [ ] SetUserSettings
 - [x] SignUp
@@ -106,7 +106,7 @@
 - [ ] UpdateUserPool
 - [ ] UpdateUserPoolClient
 - [ ] UpdateUserPoolDomain
-- [ ] VerifySoftwareToken
+- [x] VerifySoftwareToken
 - [x] VerifyUserAttribute
 
 </details>

server/api/controller.ts

@@ -41,6 +41,7 @@ const useCases: {
   'AWSCognitoIdentityProviderService.DeleteUserAttributes': authUseCase.deleteUserAttributes,
   'AWSCognitoIdentityProviderService.AssociateSoftwareToken': mfaUseCase.associateSoftwareToken,
   'AWSCognitoIdentityProviderService.VerifySoftwareToken': mfaUseCase.verifySoftwareToken,
+  'AWSCognitoIdentityProviderService.SetUserMFAPreference': mfaUseCase.setUserMFAPreference,
 };
 
 const main = <T extends keyof AmzTargets>(target: T, body: AmzTargets[T]['reqBody']) => {

server/common/constants.ts

@@ -13,6 +13,8 @@ export const USER_KIND_LIST = ['social', 'cognito'] as const;
 
 export const PROVIDER_LIST = ['Google', 'Apple', 'Amazon', 'Facebook'] as const;
 
+export const MFA_SETTING_LIST = ['SOFTWARE_TOKEN_MFA'] as const;
+
 const listToDict = <T extends readonly [string, ...string[]]>(list: T): { [U in T[number]]: U } =>
   list.reduce((dict, type) => ({ ...dict, [type]: type }), {} as { [U in T[number]]: U });
 

server/common/types/auth.ts

@@ -22,6 +22,8 @@ import type {
   ListUserPoolsResponse,
   ListUsersRequest,
   ListUsersResponse,
+  SetUserMFAPreferenceRequest,
+  SetUserMFAPreferenceResponse,
   SignUpRequest,
   SignUpResponse,
   UpdateUserAttributesRequest,
@@ -130,6 +132,11 @@ export type VerifySoftwareTokenTarget = TargetBody<
   VerifySoftwareTokenResponse
 >;
 
+export type SetUserMFAPreferenceTarget = TargetBody<
+  SetUserMFAPreferenceRequest,
+  SetUserMFAPreferenceResponse
+>;
+
 export type AmzTargets = {
   'AWSCognitoIdentityProviderService.SignUp': SignUpTarget;
   'AWSCognitoIdentityProviderService.ConfirmSignUp': ConfirmSignUpTarget;
@@ -155,4 +162,5 @@ export type AmzTargets = {
   'AWSCognitoIdentityProviderService.DeleteUserAttributes': DeleteUserAttributesTarget;
   'AWSCognitoIdentityProviderService.AssociateSoftwareToken': AssociateSoftwareTokenTarget;
   'AWSCognitoIdentityProviderService.VerifySoftwareToken': VerifySoftwareTokenTarget;
+  'AWSCognitoIdentityProviderService.SetUserMFAPreference': SetUserMFAPreferenceTarget;
 };

server/common/types/user.ts

@@ -1,5 +1,5 @@
 import type { UserStatusType } from '@aws-sdk/client-cognito-identity-provider';
-import type { PROVIDER_LIST, USER_KINDS } from 'common/constants';
+import type { MFA_SETTING_LIST, PROVIDER_LIST, USER_KINDS } from 'common/constants';
 import type { EntityId, MaybeId } from './brandedId';
 
 export type ChallengeVal = {
@@ -35,6 +35,8 @@ export type SocialUserEntity = {
   createdTime: number;
   updatedTime: number;
   challenge?: undefined;
+  preferredMfaSetting?: undefined;
+  mfaSettingList?: undefined;
   totpSecretCode?: undefined;
 };
 
@@ -58,6 +60,8 @@ export type CognitoUserEntity = {
   createdTime: number;
   updatedTime: number;
   challenge?: ChallengeVal;
+  preferredMfaSetting?: (typeof MFA_SETTING_LIST)[number];
+  mfaSettingList?: (typeof MFA_SETTING_LIST)[number][];
   totpSecretCode?: string;
 };
 

server/domain/user/model/mfaMethod.ts

@@ -1,3 +1,4 @@
+import type { SetUserMFAPreferenceTarget } from 'common/types/auth';
 import type { CognitoUserEntity } from 'common/types/user';
 import { authenticator } from 'otplib';
 import { cognitoAssert } from 'service/cognitoAssert';
@@ -12,6 +13,20 @@ export const mfaMethod = {
       'Invalid verification code provided, please try again.',
     );
 
-    return user;
+    return { ...user, mfaSettingList: ['SOFTWARE_TOKEN_MFA'] };
+  },
+  setPreference: (
+    user: CognitoUserEntity,
+    req: SetUserMFAPreferenceTarget['reqBody'],
+  ): CognitoUserEntity => {
+    return {
+      ...user,
+      preferredMfaSetting: req.SoftwareTokenMfaSettings?.PreferredMfa
+        ? 'SOFTWARE_TOKEN_MFA'
+        : user.preferredMfaSetting,
+      mfaSettingList: req.SoftwareTokenMfaSettings?.Enabled
+        ? ['SOFTWARE_TOKEN_MFA']
+        : user.mfaSettingList,
+    };
   },
 };

server/domain/user/repository/toUserEntity.ts

@@ -1,6 +1,6 @@
 import { UserStatusType } from '@aws-sdk/client-cognito-identity-provider';
 import type { User, UserAttribute } from '@prisma/client';
-import { PROVIDER_LIST, USER_KIND_LIST, USER_KINDS } from 'common/constants';
+import { MFA_SETTING_LIST, PROVIDER_LIST, USER_KIND_LIST, USER_KINDS } from 'common/constants';
 import type {
   CognitoUserEntity,
   SocialUserEntity,
@@ -51,6 +51,11 @@ export const toCognitoUserEntity = (
         value: attr.value,
       }),
     ),
+    preferredMfaSetting: z
+      .enum(MFA_SETTING_LIST)
+      .optional()
+      .parse(prismaUser.preferredMfaSetting ?? undefined),
+    mfaSettingList: prismaUser.enabledTotp ? ['SOFTWARE_TOKEN_MFA'] : undefined,
     totpSecretCode: prismaUser.totpSecretCode ?? undefined,
     createdTime: prismaUser.createdAt.getTime(),
     updatedTime: prismaUser.updatedAt.getTime(),

server/domain/user/repository/userCommand.ts

@@ -3,6 +3,7 @@ import type { EntityId } from 'common/types/brandedId';
 import type { UserAttributeEntity, UserEntity } from 'common/types/user';
 
 export const userCommand = {
+  // eslint-disable-next-line complexity
   save: async (tx: Prisma.TransactionClient, user: UserEntity): Promise<void> => {
     await tx.userAttribute.deleteMany({ where: { userId: user.id } });
 
@@ -26,6 +27,8 @@ export const userCommand = {
         pubA: user.challenge?.pubA,
         pubB: user.challenge?.pubB,
         secB: user.challenge?.secB,
+        preferredMfaSetting: user.preferredMfaSetting,
+        enabledTotp: user.mfaSettingList?.some((setting) => setting === 'SOFTWARE_TOKEN_MFA'),
         totpSecretCode: user.totpSecretCode,
         attributes: { createMany: { data: user.attributes } },
         updatedAt: new Date(user.updatedTime),
@@ -45,6 +48,8 @@ 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 } },

server/domain/user/useCase/authUseCase.ts

@@ -29,7 +29,12 @@ export const authUseCase = {
       const decoded = jwtDecode<AccessTokenJwt>(req.AccessToken);
       const user = await userQuery.findById(tx, decoded.sub);
 
-      return { UserAttributes: toAttributeTypes(user), Username: user.name };
+      return {
+        UserAttributes: toAttributeTypes(user),
+        Username: user.name,
+        PreferredMfaSetting: user.preferredMfaSetting,
+        UserMFASettingList: user.mfaSettingList,
+      };
     }),
   listUsers: (req: ListUsersTarget['reqBody']): Promise<ListUsersTarget['resBody']> =>
     transaction(async (tx) => {

server/domain/user/useCase/mfaUseCase.ts

@@ -1,6 +1,10 @@
 import { VerifySoftwareTokenResponseType } from '@aws-sdk/client-cognito-identity-provider';
 import assert from 'assert';
-import type { AssociateSoftwareTokenTarget, VerifySoftwareTokenTarget } from 'common/types/auth';
+import type {
+  AssociateSoftwareTokenTarget,
+  SetUserMFAPreferenceTarget,
+  VerifySoftwareTokenTarget,
+} from 'common/types/auth';
 import { jwtDecode } from 'jwt-decode';
 import { transaction } from 'service/prismaClient';
 import type { AccessTokenJwt } from 'service/types';
@@ -43,4 +47,21 @@ export const mfaUseCase = {
 
       return { Status: VerifySoftwareTokenResponseType.SUCCESS, Session: req.Session };
     }),
+  setUserMFAPreference: (
+    req: SetUserMFAPreferenceTarget['reqBody'],
+  ): Promise<SetUserMFAPreferenceTarget['resBody']> =>
+    transaction(async (tx) => {
+      assert(req.AccessToken);
+
+      const decoded = jwtDecode<AccessTokenJwt>(req.AccessToken);
+      const user = await userQuery.findById(tx, decoded.sub);
+
+      assert(user.kind === 'cognito');
+
+      const updated = mfaMethod.setPreference(user, req);
+
+      await userCommand.save(tx, updated);
+
+      return {};
+    }),
 };

server/prisma/migrations/20241223120649_/migration.sql

@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "User" ADD COLUMN "enabledTotp" BOOLEAN;
+ALTER TABLE "User" ADD COLUMN "preferredMfaSetting" TEXT;

server/prisma/schema.prisma

@@ -8,30 +8,32 @@ 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?
-  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?
+  preferredMfaSetting String?
+  enabledTotp         Boolean?
+  totpSecretCode      String?
+  attributes          UserAttribute[]
+  UserPool            UserPool        @relation(fields: [userPoolId], references: [id])
+  userPoolId          String
 }
 
 model UserAttribute {

server/tests/sdk/mfa.test.ts

@@ -1,9 +1,12 @@
 import {
   AssociateSoftwareTokenCommand,
+  GetUserCommand,
+  SetUserMFAPreferenceCommand,
   VerifySoftwareTokenCommand,
   VerifySoftwareTokenResponseType,
 } from '@aws-sdk/client-cognito-identity-provider';
 import assert from 'assert';
+import { MFA_SETTING_LIST } from 'common/constants';
 import { authenticator } from 'otplib';
 import { cognitoClient } from 'service/cognito';
 import { createCognitoUserAndToken } from 'tests/api/utils';
@@ -41,3 +44,25 @@ test(VerifySoftwareTokenCommand.name, async () => {
 
   expect(res.Status).toBe(VerifySoftwareTokenResponseType.SUCCESS);
 });
+
+test(SetUserMFAPreferenceCommand.name, async () => {
+  const token = await createCognitoUserAndToken();
+  const session = 'dummySession';
+
+  const { SecretCode } = await cognitoClient.send(
+    new AssociateSoftwareTokenCommand({ AccessToken: token.AccessToken, Session: session }),
+  );
+
+  assert(SecretCode);
+
+  await cognitoClient.send(
+    new SetUserMFAPreferenceCommand({
+      AccessToken: token.AccessToken,
+      SoftwareTokenMfaSettings: { PreferredMfa: true, Enabled: true },
+    }),
+  );
+
+  const user = await cognitoClient.send(new GetUserCommand(token));
+
+  expect(user.PreferredMfaSetting).toBe(MFA_SETTING_LIST['0']);
+});