【Cognito】カスタム認証チャレンジでトークン認証を実装して、SAMでデプロイしてみる。

## 柔軟性のある認証フローを作成できる

Cognitoを触ってみて、最初は色々不便だなあと思っていました。
しかし、ちゃんとドキュメントを読んでいくうちに、拡張性が高く、結構柔軟な実装が可能だと気づきます。
その柔軟性を担保してくれいている仕組みの一つが「カスタム認証チャレンジ」です。

![カスタム認証チャレンジフロー.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/632997/68970d48-a8e7-e239-0e05-0373418f10d5.png)

https://docs.aws.amazon.com/ja_jp/cognito/latest/developerguide/user-pool-lambda-challenge.html より引用

## 何が便利なのか

ログイン方法を複数指定することができます。
一般的な認証方法の１つにユーザ名とパスワードを入力するパスワード認証があります。
例えば、カスタム認証フローを使えば、あるロールのユーザにはパスワード認証を、あるロールのユーザにはトークン認証を使うなど、**相手によって認証方法を変える**ことができます。
また、パスワード認証の前後にMFAを設定したり、合言葉認証を複数設定することも実装できます。

## どうやって実装するのか

Lambdaを使って認証を定義していきます。
上記の図のように３つのトリガーがあり、そのトリガーごとにLambdaを設置します。

## 各Lambdaの役割

### Define Auth Challenge

ここでは、次にユーザに要求するチャレンジを選択し、認証するかどうかを決めます。

レスポンスを `return` する際、レスポンスの中に `challengeName` の値を以下の

`CUSTOM_CHALLENGE`、`SRP_A`、`PASSWORD_VERIFIER`、`SMS_MFA`, `DEVICE_SRP_AUTH`、`DEVICE_PASSWORD_VERIFIER`、`ADMIN_NO_SRP_AUTH` 

のうち、いずれかを指定します。
`return` された` challengeName` に指定されたチャレンジをユーザに要求します。
ユーザは `respond-to-auth-challenge` APIなどを使って、そのチャレンジに答えます。
`CUSTOM_CHALLENGE`以外のチャレンジは、ユーザからの答えをCognitoが内部的に処理をし、
**Define Auth Challenge** を再度呼び出します。

`return`された`challengeName`に`CUSTOM_CHALLENGE`が返されると次のラムダ関数である **Create Auth Challenge** を呼び出します。
ユーザがこの **Create Auth Challenge** からのチャレンジに応え **Verify Auth Challenge** が呼び出された後、再度この **Define Auth Challenge** が呼び出されます。


### Create Auth Challenge

ここでは、ユーザに独自のチャレンジを返します。

例えば、ユーザに特定の数字「100」を返し、ユーザからその2倍の数が送られてきたらチャレンジをクリアしたと判定させることもできます。
それを実装する場合、この関数ですることは、
ユーザに「100」を返し、答えとする値「200」を保存することだけです。
ユーザから返ってきた値が望む値かどうかを判定するのは次のラムダ関数の役割です。

### Verify Auth Challenge

ここではユーザから返されたチャレンジに対する答えを処理し、チャレンジが通過したかどうかを判定します。

レスポンス中の `answerCorrect` を `true` にして `return` することで、この`CUSTOM_CHALLENGE`が通過したかどうかを **Define Auth Challenge** に返します。

## よくわかんない！

と思う方は多いと思います。
私も複雑怪奇な公式の日本語ドキュメントを読んで、最初はそういう感想がでました。

とりあえずコードを見た方がイメージがつきやすいと思うので、ひとまずコードをば。

## コード

**「アプリでJWTを発行し、それをCognitoに投げたユーザを認証する」**という場合を考えてみます。

注意：

- ここではJWTシークレットをパラメーターストアに保存してそれを取得するようにしています。
    - JWTについてよくわからない方は別途調べてください。
- jsonwebtokenが必要なので、レイヤーを作る必要があります。
    - https://zenn.dev/nekoniki/articles/6a30b75da0fac5 がわかりやすいです。
- 簡易的な実装なので、実際に実装をする場合にはJWTの検証をより適切に行う必要があります。
- AWS CLI, SAM CLI が必要です。


``` define_auth_challenge.js
exports.handler = async (event, context, callback) => {
    console.log(context);
    console.log(event);  

    if (event.request.session.length > 0 && event.request.session[event.request.session.length - 1].challengeResult){
        event.response.issueTokens = true;
        event.response.failAuthentication = false;
    } else {
        event.response.issueTokens = false;
        event.response.failAuthentication = false;
        event.response.challengeName = 'CUSTOM_CHALLENGE';
    }

    callback(null, event);
};
```

``` create_auth_challenge.js
exports.handler = (event, context, callback) => {
    console.log(event);

    event.response.publicChallengeParameters = {};
    event.response.publicChallengeParameters.NEXT_ACTION = 'respond-to-auth-challenge';

    callback(null, event);
}
```

``` verify_auth_challenge.js
const jwt = require('jsonwebtoken');
const ssm = new (require('aws-sdk/clients/ssm'))();
const env = process.env['ENV'];

exports.handler = async (event, context, callback) => {
    console.log(event);
    const token = event.request?.challengeAnswer

    if (verifyJwt(token)) {
        event.response.answerCorrect = true;
    } else {
        event.response.answerCorrect = false;
    }

    callback(null, event);
};

const verifyJwt = async function(token){
    const ssmData = await ssm.getParameter({
        Name: `/${env}/secure/JWT_SECRET`,
        WithDecryption: true
    }).promise();

    const jwtSecret = ssmData.Parameter.Value;

    return jwt.verify(token, jwtSecret, (err, decoded) => {
        if (err) {
            throw new Error("This JWT token was invalid.");
        } else {
            console.log("verifyJwt() was successful.");
            return true;
        }
    });
};

```

``` template.yml
AWSTemplateFormatVersion: '2010-09-09'
Transform: 'AWS::Serverless-2016-10-31'
Description: Custom auth sample
Parameters:
  Env:
    Type: String
    AllowedValues:
      - prod
      - stg
      - dev
  Layer:
    Description: ARN of Lambda Layer which has jwt package
    Type: String
Resources:
  DefineAuthCallengeLambda:
    Type: 'AWS::Serverless::Function'
    Properties:
      Handler: define_auth_challenge.handler
      Runtime: nodejs14.x
      FunctionName: define-auth-challenge
      CodeUri: ./define_auth_challenge.js
      Description: ''
      MemorySize: 128
      Timeout: 10
      Role: !GetAtt DefineAuthChallengeRole.Arn
      Environment:
        Variables:
          ENV: !Ref Env
  CreateAuthChallengeLambda:
    Type: 'AWS::Serverless::Function'
    Properties:
      Handler: create_auth_challenge.handler
      Runtime: nodejs14.x
      FunctionName: create-auth-challenge
      CodeUri: ./create_auth_challenge.js
      Description: ''
      MemorySize: 128
      Timeout: 10
      Role: !GetAtt CreateAuthChallengeRole.Arn
      Environment:
        Variables:
          ENV: !Ref Env
  VerifyAuthChallengeLambda:
    Type: 'AWS::Serverless::Function'
    Properties:
      Handler: verify_auth_challenge.handler
      Runtime: nodejs14.x
      FunctionName: verify-auth-challenge
      CodeUri: ./verify_auth_challenge.js
      Description: ''
      MemorySize: 128
      Timeout: 30
      Role: !GetAtt VerifyAuthChallengeRole.Arn
      Layers:
        - !Sub ${Layer}
      Environment:
        Variables:
          ENV: !Ref Env
  DefineAuthChallengeRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: define-auth-challenge-role
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Action: "sts:AssumeRole"
            Principal:
              Service: lambda.amazonaws.com
      Policies:
        - PolicyName: define-auth-challenge-policy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "logs:CreateLogGroup"
                Resource: !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:*"
              - Effect: "Allow"
                Action:
                  - "logs:CreateLogStream"
                  - "logs:PutLogEvents"
                Resource: !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/define-auth-challenge:*"
  CreateAuthChallengeRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: create-auth-challenge-role
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Action: "sts:AssumeRole"
            Principal:
              Service: lambda.amazonaws.com
      Policies:
        - PolicyName: create-auth-challenge-policy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "logs:CreateLogGroup"
                Resource: !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:*"
              - Effect: "Allow"
                Action:
                  - "logs:CreateLogStream"
                  - "logs:PutLogEvents"
                Resource: !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/create-auth-challenge:*"
  VerifyAuthChallengeRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: verify-auth-challenge-role
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Action: "sts:AssumeRole"
            Principal:
              Service: lambda.amazonaws.com
      Policies:
        - PolicyName: verify-auth-challenge-policy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "logs:CreateLogGroup"
                Resource: !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:*"
              - Effect: "Allow"
                Action:
                  - "logs:CreateLogStream"
                  - "logs:PutLogEvents"
                Resource: !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/verify-auth-challenge:*"
              - Effect: "Allow"
                Action:
                  - "ssm:GetParameter"
                Resource: !Sub "arn:aws:ssm:${AWS::Region}:${AWS::AccountId}:parameter/${Env}/secure/JWT_SECRET"
```

### デプロイしてみよう

上記のファイルを全て同じディレクトリに置いた状態でデプロイします。

SAMのデプロイコマンド例(SAM CLIのインストールが必要)

```
sam deploy --region ap-northeast-1 --capabilities CAPABILITY_NAMED_IAM \ 
--stack-name lambda-custom-auth --s3-bucket templateを保存するS3バケット名 \
--parameter-overrides Env=dev Layer=レイヤーのARN --s3-prefix custom-auth
```

### 動作確認をしよう

`initiate-auth` がログインのためのAPIとなります。
ここの `--auth-flow` を `CUSTOM_AUTH` とすることで カスタム認証チャレンジ を開始できます。
(`--auth-flow` を `USER_PASSWORD_AUTH` とすると、一般的なパスワード認証で認証させることができます)

このコマンドを打つと

```
aws cognito-idp initiate-auth --auth-flow CUSTOM_AUTH \ 
--auth-parameters "USERNAME=+81000000000" --client-id クライアントID
```

以下のように

```
{
    "ChallengeName": "CUSTOM_CHALLENGE",
    "Session": "session情報",
    "ChallengeParameters": {
        "NEXT_ACTION": "respond-to-auth-challenge",
        "USERNAME": "uuid"
    }
}
```
こんな値が返ってくるので、`Session`の値を用いて、

```
aws cognito-idp respond-to-auth-challenge --client-id クライアントID --challenge-name CUSTOM_CHALLENGE \ 
--challenge-responses USERNAME=+81000000000,ANSWER=ここにJWTを入力 --session "session情報"
```

上記のようにJWTを送ってあげます。

```
{
    "ChallengeParameters": {},
    "AuthenticationResult": {
        "AccessToken": "アクセストークン",
        "ExpiresIn": 86400,
        "TokenType": "Bearer",,
```

こんな感じでアクセストークンが返ってきたら成功です！！

当然ですが、JWT生成時と、Lambdaで同じJWT_SECRETを利用していなかったり、JWTの期限が切れていると認証は失敗します。

### 参考

https://awscli.amazonaws.com/v2/documentation/api/latest/reference/cognito-idp/index.html