個人的備忘録：Go × Amazon Cognito で JWT 認証を実装していたので、勝手に体系的にまとめてみた

## はじめに
本記事では、Amazon Cognito を利用した JWT 認証を行う Go のパッケージについて解説します。

:::note info
**このコードは、私が作成したものではなく、受講している IT スクールのハッカソンでバックエンド担当の方が作成したものです。**
:::

実装の中で理解が難しい部分があったため、自分の備忘録として記事にまとめることにしました。

## コードの流れ

このコードは、主に以下の機能を実装しています。

- Cognito のユーザー情報を取得する機能
- JWT トークンの検証
- 認証ミドルウェアの提供

## 実際に解説してみた

### 構造体の概要

**(1) CognitoUserService**
```go
// Cognito API クライアントを保持する構造体
type CognitoUserService struct {
	client *cognitoidentityprovider.Client
}
```
- AWS SDK for Go を使用して Cognito API クライアントを作成。
- ユーザー情報取得 (`GetUserInfo`) に利用する。

**(2) CognitoAuth**
```go
type CognitoAuth struct {
	verifier CognitoJwtVerifier
}
```
- JWT トークンの検証を担当する `CognitoJwtVerifier` を保持。

**(3) CognitoJwtVerifier**
```go
type CognitoJwtVerifier struct {
	issuer   string
	jwksUri  string
	tokenUse string
	cache    *utils.Cache
}
```
- Cognito の JWT を検証するための情報を管理。
- `issuer`（発行者）、`jwksUri`（公開鍵の取得 URL）、`tokenUse`（トークンの用途）を保持。

**(4) CognitoClaims**
```go
type CognitoClaims struct {
	jwt.RegisteredClaims
	Sub      string `json:"sub"`
	Iss      string `json:"iss"`
	Version  int    `json:"version"`
	ClientID string `json:"client_id"`
	Username string `json:"username"`
	TokenUse string `json:"token_use"`
}
```
- JWT のペイロード部分（クレーム）を Go の構造体として定義。
- `sub`（ユーザーの一意な ID）、`username`（Cognito のユーザー名） などを含む。

### ユーザー情報の取得
```go
func (s *CognitoUserService) GetUserInfo(ctx context.Context, accessToken string) (*cognitoidentityprovider.GetUserOutput, error) {
	input := &cognitoidentityprovider.GetUserInput{
		AccessToken: &accessToken,
	}
	return s.client.GetUser(ctx, input)
}
```
- `accessToken` を使って Cognito の `GetUser` API を呼び出し、ユーザー情報を取得。

### JWT 認証の処理

### **(1) トークンの検証**
```go
func (c CognitoJwtVerifier) Verify(token string) (jwt.Claims, error) {
	decomposeUnverifiedJwt, err := utils.DecomposeUnverifiedJwt(token)
	if err != nil {
		return nil, err
	}

	jwk, err := utils.GetJwk(decomposeUnverifiedJwt, c.jwksUri, c.cache)
	if err != nil {
		return nil, err
	}

	err = utils.VerifyDecomposedJwt(decomposeUnverifiedJwt, c.issuer, c.tokenUse, jwk.Alg)
	if err != nil {
		return nil, err
	}

	validToken, err := utils.ValidateJwt(token, jwk)
	if err != nil {
		return nil, err
	}

	return validToken.Claims, nil
}
```
- `utils.DecomposeUnverifiedJwt(token)` で JWT の構造を分解。
- `utils.GetJwk()` で Cognito の JWK（公開鍵）を取得。
- `utils.VerifyDecomposedJwt()` で、発行者 (`issuer`) やトークン用途 (`tokenUse`) をチェック。
- `utils.ValidateJwt()` で署名の検証を行い、正しいトークンなら `Claims` を返す。

### 認証ミドルウェア

### **(1) `AuthMiddleware` の処理**
```go
func (a *CognitoAuth) AuthMiddleware(cognitoService *CognitoUserService) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// Authorization ヘッダーからトークンを取得
			auth := c.Request().Header.Get("Authorization")
			if auth == "" {
				return echo.ErrUnauthorized
			}

			// "Bearer " を除去
			if len(auth) < 8 || auth[:7] != "Bearer " {
				return echo.ErrUnauthorized
			}
			token := auth[7:]

			// トークンの検証
			claims, err := a.verifier.Verify(token)
			if err != nil {
				return echo.ErrUnauthorized
			}

			// コンテキストにユーザー情報を保存
			c.Set("user", claims)

			return next(c)
		}
	}
}
```

#### **(2) 処理の流れ**
1. `Authorization` ヘッダーからトークンを取得。
2. `Bearer ` を取り除いて JWT のみを抽出。
3. `Verify()` を実行し、トークンの検証を行う。
4. `GetUserInfo()` を呼び出して Cognito からユーザー情報を取得。
5. 認証成功時にはログを出力し、コンテキストに情報を保存。

## 実際のコード

ここまで紹介した部分的なコードをまとめたものが以下になります。必要に応じて参考にしてください。

```go
// Cognito認証関連
package auth

import (
    "context"
    "fmt"
    "log"
    "os"

    // "github.com/99designs/gqlgen/codegen/config"
    "github.com/aws/aws-sdk-go-v2/config"
    "github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider"
    "github.com/golang-jwt/jwt/v5"
    "github.com/jhosan7/cognito-jwt-verify/utils"
    "github.com/labstack/echo/v4"
)

// Cognito API クライアントを保持する構造体
type CognitoUserService struct {
    client *cognitoidentityprovider.Client
}

// ユーザー情報取得メソッド
func (s *CognitoUserService) GetUserInfo(ctx context.Context, accessToken string) (*cognitoidentityprovider.GetUserOutput, error) {
    input := &cognitoidentityprovider.GetUserInput{
        AccessToken: &accessToken,
    }
    return s.client.GetUser(ctx, input)
}

// Cognitoの設定やクレームを扱う構造体
type CognitoAuth struct {
    verifier CognitoJwtVerifier
}

// JWTクレームの構造体
type CognitoClaims struct {
    jwt.RegisteredClaims
    Sub string `json:"sub"`
    Iss string `json:"iss"`
    Version int `json:"version"`
    ClientID string `json:"client_id"`
    OriginJti string `json:"origin_jti"`
    EventID string `json:"event_id"`
    TokenUse string `json:"token_use"` // トークンタイプ（"access" of "id"）
    Scope string `json:"scope"`
    AuthTime int64 `json:"auth_time"`
    Username string `json:"username"`
}

// CognitoJwtVerifierの構造体
type CognitoJwtVerifier struct {
    issuer string
    jwksUri string
    tokenUse string
    cache *utils.Cache
}

type Config struct {
    UserPoolId string
    TokenUse string
    ClientId string
}

func NewCognitoUserService() (*CognitoUserService, error) {
    // SDkの設定を初期化
    cfg, err := config.LoadDefaultConfig(context.TODO(),
        config.WithRegion(os.Getenv("AWS_REGION")), 
    )
    if err != nil {
        return nil, fmt.Errorf("unable to load SDK config: %v", err)
    }
    // Cognitoクライアントを作成
    client := cognitoidentityprovider.NewFromConfig(cfg)

    return &CognitoUserService{
        client: client,
    }, nil
}

// 新しいCognitoAuth インスタンス作成
func NewCognitoAuth() (*CognitoAuth, error) {
    config := Config{
        UserPoolId: os.Getenv("COGNITO_USER_POOL_ID"),
        TokenUse: "access",
        ClientId: os.Getenv("COGNITO_CLIENT_ID"),
        Scope: "aws.cognito.signin.user.admin openid email phone",
    }

    verifier, err := Create(config)
    if err != nil {
        return nil, fmt.Errorf("failed to create cognito verifier: %w", err)
    }

    return &CognitoAuth{
        verifier: verifier,
    }, nil
}

// JWTVerifierの作成
func Create(config Config) (CognitoJwtVerifier, error) {
    issuer, jwksUri, err := utils.ParseUserPoolId(config.UserPoolId)
    if err != nil {
        return CognitoJwtVerifier{}, err
    }
    return CognitoJwtVerifier{
        issuer: issuer,
        jwksUri: jwksUri,
        tokenUse: config.TokenUse,
        cache: utils.NewCache(),
    }, nil
}

// トークン検証
func (c CognitoJwtVerifier) Verify(token string) (jwt.Claims, error) {
    decomposeUnverifiedJwt, err := utils.DecomposeUnverifiedJwt(token)
    if err != nil {
        return nil, err
    }

    jwk, err := utils.GetJwk(decomposeUnverifiedJwt, c.jwksUri, c.cache)
    if err != nil {
        return nil, err
    }

    err = utils.VerifyDecomposedJwt(decomposeUnverifiedJwt, c.issuer, c.tokenUse, jwk.Alg)
    if err != nil {
        return nil, err
    }

    validToken, err := utils.ValidateJwt(token, jwk)
    if err != nil {
        return nil, err
    }

    return validToken.Claims, nil
}

// 認証ミドルウェア
func (a *CognitoAuth) AuthMiddleware(cognitoService *CognitoUserService) echo.MiddlewareFunc {
    return func(next echo.HandlerFunc) echo.HandlerFunc {
        return func(c echo.Context) error {
            // Authorizationヘッダーからトークンを取得
            auth := c.Request().Header.Get("Authorization")
            if auth == "" {
                return echo.ErrUnauthorized
            }

            // "Bearer"を除去
            if len(auth) < 8 || auth[:7] != "Bearer " {
                return echo.ErrUnauthorized
            }
            token := auth[7:]

            // トークンを検証
            claims, err := a.verifier.Verify(token)
            if err != nil {
                return echo.ErrUnauthorized
            }

            mapClaims, ok := claims.(jwt.MapClaims)
            if !ok {
                return echo.ErrInternalServerError
            }

            // クレームをCognitoClaimsにマッピング
            username, _ := mapClaims["username"].(string)
            cognitoClaims := &CognitoClaims{
                Sub: mapClaims["sub"].(string),
                TokenUse: mapClaims["token_use"].(string),
                Username: username,
            }

            // Cognitoからユーザー情報を取得してログ出力
            if cognitoClaims.TokenUse != "access" {
                log.Printf("Invalid token use: expected access token, got %s", cognitoClaims.TokenUse)
                return echo.ErrUnauthorized
            }
            
            userInfo, err := cognitoService.GetUserInfo(c.Request().Context(), token)
            if err != nil {
                log.Printf("Failed to get user info from Cognito: %v", err)
                return echo.ErrUnauthorized
            }
            if err != nil {
                log.Printf("Failed to get user info from Cognito: %v", err)
            } else {
                log.Printf("Cognito User Attributes:")
                for _, attr :=range userInfo.UserAttributes {
                    if attr.Name != nil && attr.Value != nil {
                        log.Printf(" %s: %s", *attr.Name, *attr.Value)
                    }
                }
            }

            log.Printf("Authentication successful - User: %s", cognitoClaims.Username)

            // コンテキストにユーザー情報を保存
            c.Set("user", cognitoClaims)

            return next(c)
        }
    }
}
```

## まとめ

**なにをしているのか**
- Cognito の JWT トークンを検証し、ユーザー認証を行うミドルウェアを提供。
- Cognito からユーザー情報を取得し、ログ出力する処理を実装。

**どんな場面で使うのか**
- API のリクエストごとに、Cognito の JWT を検証する必要がある場合。
- ユーザー情報を取得し、認証後にアプリケーションで利用したい場合。

**ポイント**
- JWT の検証を `CognitoJwtVerifier` で実施。
- `GetUserInfo()` で Cognito API を使ってユーザー情報を取得。
- `AuthMiddleware` で API ルートごとに認証処理を適用。

このコードを使えば、Cognito を活用した API 認証の仕組みを簡単に組み込むことが可能です！