Spring SecurityとAmazon Cognitoで実現する二段階認証(≠二要素認証)

public SecurityFilterChain securityFilterChain(HttpSecurity http) throws Exception {
        return http
                // 認可の設定
                .authorizeHttpRequests(c -> c
                        // ログイン画面は誰でも見られる
                        .requestMatchers("/login").permitAll()
                        // 二段階認証画面はログイン済のユーザーのみ見られる
                        .requestMatchers("/login/step").access(new TwoStepVerificationManager())
                        // エラー画面は誰でも見られる
                        .requestMatchers("/error").permitAll()
                        // CSSとJSは誰でも参照できる
                        .requestMatchers("/css/**", "/js/**").permitAll()
                        // 上記以外の画面は二段階認証が必要
                        .anyRequest().authenticated())

                // ログインの設定
                .formLogin(c -> c
                        // ログイン認証を行うパスを設定する
                        .loginPage("/login")
                        // ユーザー名のパラメーター名を設定する
                        .usernameParameter("loginId")
                        // パスワードのパラメーター名を設定する
                        .passwordParameter("userPassword")
                        // パスワード認証後に２段階認証の画面を表示するように設定
                        .successHandler(new TwoStepVerificationSuccessHandler("/login/step")))

                .authenticationManager(cognitoAuthenticationManager())

public class TwoStepVerificationManager implements AuthorizationManager {

	// ここではauthentication.get()がTwoFactorAuthentication型であること→一段階目の認証を突破していることのみcheckする。
	@Override
	public AuthorizationDecision check(Supplier authentication, RequestAuthorizationContext object) {
		return new AuthorizationDecision(authentication.get() instanceof TwoStepVerification);
	}
}

@PostMapping(path = "step")
public String executeTwoStepVerification(
        @Validated CognitoAuthenticationForm form,
        TwoStepVerification twoStepVerification,
        HttpServletRequest request,
        HttpServletResponse response,
        SecurityContextLogoutHandler logoutHandler,
        RedirectAttributes redirectAttributes) {

    Authentication authentication = twoStepVerification.getAuthentication();
    UserDetailsImpl accountUserDetails = (UserDetailsImpl) authentication.getPrincipal();

    // DBからセッション情報を取得する
    SystemAccount account = systemAccountRepository.findByLoginId(accountUserDetails.getUsername());

    // カスタム認証チャレンジに対する応答を行う
    RespondToAuthChallengeResult authChallengeResult;
    try {
        authChallengeResult = securityLoginService.respondToAuthCustomChallenge(
                account.getUserId(),
                form.getAuthCode(),
                account.getCognitoSession());
    } catch (NotAuthorizedException e) {
        // 二段階認証コードの有効期限が切れた場合、一度ログアウトしてログイン画面に返す。
        logoutHandler.logout(request, response, authentication);

        // (中略。エラーメッセージをログイン画面に設定する)
        return "redirect:/login";
    }

    // DBにセッション情報を保存
    securityLoginService.saveCognitoSession(account.getLoginId(), authChallengeResult.getSession());

    if (authChallengeResult.getAuthenticationResult() != null &&
            authChallengeResult.getAuthenticationResult().getAccessToken() != null) {
        // SecurityContextの内容を書き換えることで、isAuthenticated()をtrueにする。
        SecurityContext securityContext = SecurityContextHolder.getContext();
        securityContext.setAuthentication(authentication);
        SecurityContextHolder.setContext(securityContext);

        // ログイン成功時には、最終ログイン日時を更新する
        securityLoginService.updateSystemAccount(account.getLoginId());
        return "redirect:/portal";
    } else {
        redirectAttributes.addFlashAttribute("authError",
                messageSource.getMessage("E00012", null, LocaleContextHolder.getLocale()));
        return "redirect:/login/step";
    }
}

public RespondToAuthChallengeResult respondToAuthCustomChallenge(String userSub, String verificationCode, String session) {
    return cognitoService.respondToAuthCustomChallenge(new RespondCustomAuthDTO(userSub, verificationCode, session));
}

def lambda_handler(event, context):
    session = event['request']['session']
    latest_session = session[-1] if session else None
    def set_response(issue_tokens, fail_authentication, challenge_name=None):
        event['response']['issueTokens'] = issue_tokens
        event['response']['failAuthentication'] = fail_authentication
        if challenge_name:
            event['response']['challengeName'] = challenge_name
    if latest_session:
        challenge_name = latest_session['challengeName']
        challenge_result = latest_session['challengeResult']

        if challenge_name == "SRP_A" and challenge_result:
            set_response(False, False, "PASSWORD_VERIFIER")
        elif challenge_name == "PASSWORD_VERIFIER" and challenge_result:
            set_response(False, False, "CUSTOM_CHALLENGE")
        elif challenge_name == "CUSTOM_CHALLENGE" and challenge_result:
            set_response(True, False)
        elif challenge_name == "CUSTOM_CHALLENGE" and not challenge_result:
            set_response(False, False, "CUSTOM_CHALLENGE")
        else:
            set_response(False, True)
    else:
        set_response(False, True)
    return event

import random
import string
import boto3
import os
def lambda_handler(event, context):
    if event['request']['challengeName'] == "CUSTOM_CHALLENGE":
        # セッションが存在し、CUSTOM_CHALLENGEの場合
        if event['request']['session'] and event['request']['session'][-1]['challengeName'] == "CUSTOM_CHALLENGE":
            verification_code = event['request']['session'][-1]['challengeMetadata']
        else:
            ses = boto3.client('ses')
            # ランダムな数字6桁を生成
            verification_code = ''.join(random.choice(string.digits) for _ in range(6))
            to_email_address = event['request']['userAttributes']['email']
            send_email(ses, to_email_address, verification_code)
        event['response']['privateChallengeParameters'] = {'answer': verification_code}
        event['response']['challengeMetadata'] = verification_code
    return event
def send_email(ses, email, code):
    params = {
        'Destination': {
            'ToAddresses': [email],
        },
        'Message': {
            'Body': {
                'Text': {
                    'Data': f'Your verification code is {code}.',
                },
            },
            'Subject': {
                'Data': 'Your verification code',
            },
        },
        'Source': os.getenv('FROM_MAIL_ADDRESS'),
    }
    ses.send_email(**params)

def lambda_handler(event, context):
    # ユーザーが入力した値(challengeAnswer)が正しいか
    if event['request']['privateChallengeParameters']['answer'] == event['request']['challengeAnswer']:
        event['response']['answerCorrect'] = True
    else:
        event['response']['answerCorrect'] = False
    return event