Amplify Gen 2のバックエンドでCognito JWTを検証する（第3回）— Lambda Function URLでの認証と監査ログ

// amplify/backend.ts
import { defineBackend } from '@aws-amplify/backend'
import { FunctionUrlAuthType, HttpMethod, InvokeMode } from 'aws-cdk-lib/aws-lambda'
import type { Function as LambdaFunction } from 'aws-cdk-lib/aws-lambda'
import { auth } from './auth/resource'
import { streamingChat } from './functions/streaming-chat/resource'

const backend = defineBackend({ auth, streamingChat })

// Cognito リソースへの参照を取得
const userPool = backend.auth.resources.userPool
const userPoolClient = backend.auth.resources.cfnResources.cfnUserPoolClient

// Lambda 関数への参照
const streamingLambda = backend.streamingChat.resources.lambda as LambdaFunction

// Function URL を設定（レスポンスストリーミング有効）
const fnUrl = streamingLambda.addFunctionUrl({
  authType: FunctionUrlAuthType.NONE,  // JWT検証を自前で行うため
  invokeMode: InvokeMode.RESPONSE_STREAM,
  cors: {
    allowedOrigins: ['https://your-app.amplifyapp.com', 'http://localhost:5173'],
    allowedHeaders: ['content-type', 'authorization'],
    allowedMethods: [HttpMethod.POST],
  },
})

// Cognito 環境変数を動的に注入
streamingLambda.addEnvironment('USER_POOL_ID', userPool.userPoolId)
streamingLambda.addEnvironment('USER_POOL_CLIENT_ID', userPoolClient.ref)

Authorization: Bearer <token>
  ↓
1. JWTをデコード（ヘッダー、ペイロード、署名）
  ↓
2. クレーム検証
  ├─ issuer: Cognito User Pool のURL
  ├─ token_use: "id"（IDトークン）
  ├─ client_id: User Pool Client ID
  └─ exp: 現在時刻より未来
  ↓
3. JWKS取得（Cognito公開鍵）
  ↓
4. JWKから公開鍵をインポート（Web Crypto API）
  ↓
5. RSA-SHA256で署名検証
  ↓
6. 検証成功 → payload.sub（ユーザーID）を返す

interface JwkKey {
  kid: string  // Key ID
  kty: string  // Key Type (RSA)
  n: string    // Modulus
  e: string    // Exponent
  alg: string  // Algorithm (RS256)
  use: string  // Usage (sig)
}

interface JwksResponse {
  keys: JwkKey[]
}

interface JwtHeader {
  kid: string
  alg: string
}

interface JwtPayload {
  sub: string       // Subject（ユーザーID）
  iss: string       // Issuer（Cognito URL）
  client_id?: string
  token_use: string // "id" or "access"
  exp: number       // Expiration（Unix timestamp）
}

let cachedJwks: JwksResponse | null = null

async function getJwks(): Promise<JwksResponse> {
  if (cachedJwks) return cachedJwks

  const userPoolId = getEnv('USER_POOL_ID')
  const region = userPoolId.split('_')[0]  // "ap-northeast-1_xxx" → "ap-northeast-1"
  const url = `https://cognito-idp.${region}.amazonaws.com/${userPoolId}/.well-known/jwks.json`

  const res = await fetch(url)
  if (!res.ok) throw new Error(`Failed to fetch JWKS: ${res.status}`)

  cachedJwks = (await res.json()) as JwksResponse
  return cachedJwks
}

function base64UrlDecode(str: string): Uint8Array {
  const base64 = str.replace(/-/g, '+').replace(/_/g, '/')
  const padded = base64 + '='.repeat((4 - (base64.length % 4)) % 4)
  const binary = Buffer.from(padded, 'base64')
  return new Uint8Array(binary)
}

function decodeJwtParts(token: string): {
  header: JwtHeader
  payload: JwtPayload
  signatureInput: string
  signature: Uint8Array
} {
  const parts = token.split('.')
  if (parts.length !== 3) throw new Error('Invalid JWT format')

  const header = JSON.parse(
    Buffer.from(parts[0], 'base64url').toString()
  ) as JwtHeader

  const payload = JSON.parse(
    Buffer.from(parts[1], 'base64url').toString()
  ) as JwtPayload

  const signatureInput = `${parts[0]}.${parts[1]}`
  const signature = base64UrlDecode(parts[2])

  return { header, payload, signatureInput, signature }
}

async function importJwk(jwk: JwkKey) {
  return crypto.subtle.importKey(
    'jwk',
    { kty: jwk.kty, n: jwk.n, e: jwk.e, alg: jwk.alg, ext: true },
    { name: 'RSASSA-PKCS1-v1_5', hash: 'SHA-256' },
    false,
    ['verify'],
  )
}

async function verifyJwt(token: string): Promise<string> {
  const { header, payload, signatureInput, signature } = decodeJwtParts(token)

  // 1. クレーム検証
  const userPoolId = getEnv('USER_POOL_ID')
  const clientId = getEnv('USER_POOL_CLIENT_ID')
  const region = userPoolId.split('_')[0]
  const expectedIssuer = `https://cognito-idp.${region}.amazonaws.com/${userPoolId}`

  if (payload.iss !== expectedIssuer) {
    throw new Error('Invalid issuer')
  }
  if (payload.token_use !== 'id') {
    throw new Error('Invalid token_use: expected id token')
  }
  if (payload.client_id && payload.client_id !== clientId) {
    throw new Error('Invalid client_id')
  }
  if (payload.exp * 1000 < Date.now()) {
    throw new Error('Token expired')
  }

  // 2. 署名検証
  const jwks = await getJwks()
  const jwk = jwks.keys.find((k) => k.kid === header.kid)
  if (!jwk) throw new Error('No matching JWK found')

  const cryptoKey = await importJwk(jwk)
  const encoder = new TextEncoder()
  const valid = await crypto.subtle.verify(
    'RSASSA-PKCS1-v1_5',
    cryptoKey,
    signature,
    encoder.encode(signatureInput),
  )
  if (!valid) throw new Error('Invalid JWT signature')

  // 検証成功 → ユーザーID を返す
  return payload.sub
}

// Lambda Function URL ハンドラー（抜粋）
export const handler = async (event: FunctionUrlEvent) => {
  // Authorization ヘッダーの取得
  const authHeader =
    event.headers['authorization'] ?? event.headers['Authorization']

  if (!authHeader?.startsWith('Bearer ')) {
    return {
      statusCode: 401,
      body: JSON.stringify({ error: 'Missing or invalid Authorization header' }),
    }
  }

  // JWT 検証
  let userId: string
  try {
    userId = await verifyJwt(authHeader.slice(7))
  } catch (err) {
    return {
      statusCode: 401,
      body: JSON.stringify({
        error: `JWT verification failed: ${(err as Error).message}`,
      }),
    }
  }

  // userId を使ってビジネスロジックを実行
  // セッションの所有権チェックなど
  const session = await getSession(sessionId)
  if (session.userId !== userId) {
    return { statusCode: 403, body: JSON.stringify({ error: 'Unauthorized' }) }
  }

  // ... 以降のビジネスロジック
}

// src/features/chat/hooks/useStreamingChat.ts（抜粋）
import { fetchAuthSession } from 'aws-amplify/auth'

async function callStreamingEndpoint(sessionId: string, content: string) {
  // Cognito IDトークンを取得
  const session = await fetchAuthSession()
  const idToken = session.tokens?.idToken?.toString()

  if (!idToken) {
    throw new Error('Not authenticated')
  }

  const response = await fetch(STREAMING_CHAT_URL, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'Authorization': `Bearer ${idToken}`,
    },
    body: JSON.stringify({ sessionId, content }),
  })

  // ... レスポンス処理
}

// src/shared/utils/auditLogger.ts
import { getAmplifyClient } from '@/lib/amplifyClient'
import type { AuditAction, AuditResourceType } from '@/types'

interface LogAuditEventParams {
  userId: string
  userEmail: string
  action: AuditAction
  resourceType: AuditResourceType
  resourceId?: string
  metadata?: Record<string, string>
}

/**
 * Fire-and-forget 監査ログ。
 * 認証フローをブロックしないよう、エラーは握りつぶしてconsoleに出力するだけ。
 */
export async function logAuditEvent(params: LogAuditEventParams): Promise<void> {
  try {
    const client = getAmplifyClient()
    await client.models.AuditLog.create({
      userId: params.userId,
      userEmail: params.userEmail,
      action: params.action,
      resourceType: params.resourceType,
      resourceId: params.resourceId,
      metadata: params.metadata ? JSON.stringify(params.metadata) : undefined,
      timestamp: new Date().toISOString(),
    })
  } catch (err) {
    console.error('[AuditLogger] Failed to write audit log:', err)
  }
}

// useAuth.ts の signIn 内
const updatedUser = await checkCurrentUser()
if (updatedUser) {
  void logAuditEvent({
    userId: updatedUser.id,
    userEmail: updatedUser.email,
    action: 'LOGIN',
    resourceType: 'Auth',
  })
}

async function verifyJwtWithFallback(token: string): Promise<string> {
  try {
    return await verifyJwt(token)
  } catch (err) {
    // キャッシュをクリアしてリトライ
    cachedJwks = null
    return await verifyJwt(token)
  }
}