LWA（Lambda Web Adapter）でAurora DSQLに接続する際の認証トークン期限切れに注意

use anyhow::Result;
use aws_config::BehaviorVersion;
use aws_sdk_dsql::auth_token::{AuthTokenGenerator, Config as AuthConfig};
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode};
use axum::routing::get;
use axum::Router;
use sqlx::postgres::{PgConnectOptions, PgPoolOptions, PgSslMode};
use sqlx::PgPool;
use std::net::{Ipv4Addr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpListener;
use tokio::sync::RwLock;

/// トークンの有効期限（秒）
/// ローカル再現時は短くする（60秒）、Lambda 上では本番同等（900秒=15分）
const TOKEN_EXPIRES_IN_SECS: u64 = 60;

/// トークンリフレッシュ間隔（秒）— 有効期限の80%で切り下げ
const TOKEN_REFRESH_INTERVAL_SECS: u64 = TOKEN_EXPIRES_IN_SECS * 8 / 10;

// -------------------------------------------------------------------------
// DSQL 接続ヘルパー
// -------------------------------------------------------------------------

async fn generate_token(endpoint: &str, expires_in: u64) -> Result<String> {
    let sdk_config = aws_config::load_defaults(BehaviorVersion::latest()).await;
    let signer = AuthTokenGenerator::new(
        AuthConfig::builder()
            .hostname(endpoint)
            .expires_in(expires_in)
            .build()
            .map_err(|e| anyhow::anyhow!("Failed to build auth config: {e:?}"))?,
    );
    let token = signer
        .db_connect_admin_auth_token(&sdk_config)
        .await
        .map_err(|e| anyhow::anyhow!("Failed to generate auth token: {e:?}"))?;
    Ok(token.to_string())
}

async fn create_pool(endpoint: &str, token: &str) -> Result<PgPool> {
    let connect_options = PgConnectOptions::new()
        .host(endpoint)
        .port(5432)
        .database("postgres")
        .username("admin")
        .password(token)
        .ssl_mode(PgSslMode::Require);

    let pool = PgPoolOptions::new()
        .max_connections(2)
        // 接続の最大生存時間。トークン有効期限と同じにすることで、
        // 期限切れ前後に接続がリサイクルされるようにする。
        .max_lifetime(Duration::from_secs(TOKEN_EXPIRES_IN_SECS))
        // アイドル接続のタイムアウト。有効期限の半分で設定。
        .idle_timeout(Duration::from_secs(TOKEN_EXPIRES_IN_SECS / 2))
        .connect_with(connect_options)
        .await?;
    Ok(pool)
}

// -------------------------------------------------------------------------
// RefreshablePool（修正版で使用）
// -------------------------------------------------------------------------

#[derive(Clone)]
struct RefreshablePool {
    pool: Arc<RwLock<PgPool>>,
}

impl RefreshablePool {
    async fn new(endpoint: &str, token: &str) -> Result<Self> {
        let pool = create_pool(endpoint, token).await?;
        Ok(Self {
            pool: Arc::new(RwLock::new(pool)),
        })
    }

    async fn get(&self) -> PgPool {
        self.pool.read().await.clone()
    }

    async fn refresh(&self, endpoint: &str) -> Result<()> {
        // 1. 新しいトークンで新しいプールを作成（max_lifetime/idle_timeout は無効化）
        let token = generate_token(endpoint, TOKEN_EXPIRES_IN_SECS).await?;
        let connect_options = PgConnectOptions::new()
            .host(endpoint)
            .port(5432)
            .database("postgres")
            .username("admin")
            .password(&token)
            .ssl_mode(PgSslMode::Require);
        let new_pool = PgPoolOptions::new()
            .max_connections(2)
            .max_lifetime(None)
            .idle_timeout(None)
            .connect_with(connect_options)
            .await?;

        // 2. 書き込みロックを取得し、古いプールを新しいプールに差し替える
        //    ブロック {} でスコープを限定し、ロック保持時間を最小にしている
        //    ブロックを抜けると guard がドロップされてロックが解放される
        let old_pool = {
            let mut guard = self.pool.write().await;
            std::mem::replace(&mut *guard, new_pool)
        };

        // 3. ロック解放後に古いプールを閉じる（既存の実行中クエリは完了を待つ）
        old_pool.close().await;
        tracing::info!("pool refreshed");
        Ok(())
    }
}

// -------------------------------------------------------------------------
// HTTP サーバー
// -------------------------------------------------------------------------

#[derive(Clone)]
struct AppState {
    pool: RefreshablePool,
}

async fn health_handler(headers: HeaderMap, State(state): State<AppState>) -> (StatusCode, String) {
    let request_id = headers
        .get("x-amzn-lambda-context")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| serde_json::from_str::<serde_json::Value>(v).ok())
        .and_then(|v| v.get("request_id")?.as_str().map(String::from))
        .unwrap_or_else(|| "-".to_string());
    let pool = state.pool.get().await;
    match sqlx::query("SELECT 1").fetch_one(&pool).await {
        Ok(_) => {
            tracing::info!(request_id, "SELECT 1 => OK");
            (StatusCode::OK, "OK".to_string())
        }
        Err(e) => {
            tracing::error!(request_id, error = %e, "SELECT 1 => ERROR");
            (StatusCode::INTERNAL_SERVER_ERROR, format!("ERROR: {e}"))
        }
    }
}

async fn serve(pool: RefreshablePool) -> Result<()> {
    let state = AppState { pool };
    let app = Router::new()
        .route("/health", get(health_handler))
        .with_state(state);

    let addr = SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 8080);
    let listener = TcpListener::bind(&addr).await?;
    tracing::info!("Listening on {addr}");
    axum::serve(listener, app).await?;
    Ok(())
}

/// トークンリフレッシュなし — 期限切れ後にエラーになる
async fn run_bug(endpoint: &str) -> Result<()> {
    tracing::info!("=== bug mode: no token refresh ===");

    let token = generate_token(endpoint, TOKEN_EXPIRES_IN_SECS).await?;
    let pool = RefreshablePool::new(endpoint, &token).await?;

    serve(pool).await
}

/// トークンリフレッシュあり — 期限切れ前にプールを差し替える
/// プール丸ごと差し替えで対応するため、max_lifetime/idle_timeout は無効化する。
/// これらを設定すると、リフレッシュ前にプール内部で接続リサイクルが走り、
/// 期限切れトークンで再接続してエラーになる可能性がある。
async fn run_fix(endpoint: &str) -> Result<()> {
    tracing::info!("=== fix mode: with token refresh ===");

    let token = generate_token(endpoint, TOKEN_EXPIRES_IN_SECS).await?;
    let connect_options = PgConnectOptions::new()
        .host(endpoint)
        .port(5432)
        .database("postgres")
        .username("admin")
        .password(&token)
        .ssl_mode(PgSslMode::Require);

    let pg_pool = PgPoolOptions::new()
        .max_connections(2)
        .max_lifetime(None)
        .idle_timeout(None)
        .connect_with(connect_options)
        .await?;

    let pool = RefreshablePool {
        pool: Arc::new(RwLock::new(pg_pool)),
    };

    let refresh_pool = pool.clone();
    let refresh_endpoint = endpoint.to_string();
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(Duration::from_secs(TOKEN_REFRESH_INTERVAL_SECS)).await;
            if let Err(e) = refresh_pool.refresh(&refresh_endpoint).await {
                tracing::error!("refresh failed: {e}");
            }
        }
    });

    serve(pool).await
}

// -------------------------------------------------------------------------
// main
// -------------------------------------------------------------------------

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .json()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()),
        )
        .init();

    let endpoint = std::env::var("DSQL_CLUSTER_ENDPOINT")
        .expect("DSQL_CLUSTER_ENDPOINT 環境変数を設定してください");

    let mode = std::env::var("MODE").unwrap_or_else(|_| "bug".to_string());

    match mode.as_str() {
        "bug" => run_bug(&endpoint).await?,
        "fix" => run_fix(&endpoint).await?,
        _ => {
            eprintln!("MODE=bug|fix を指定してください");
            std::process::exit(1);
        }
    }

    Ok(())
}

import * as path from 'node:path';
import { fileURLToPath } from 'node:url';
import * as cdk from 'aws-cdk-lib';
import * as dsql from 'aws-cdk-lib/aws-dsql';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as logs from 'aws-cdk-lib/aws-logs';
import { Construct } from 'constructs';

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

export class DsqlSampleStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    // DSQL Cluster
    const dsqlCluster = new dsql.CfnCluster(this, 'DsqlCluster', {
      deletionProtectionEnabled: false,
    });

    const dsqlClusterEndpoint = cdk.Fn.join('', [
      dsqlCluster.attrIdentifier,
      '.dsql.',
      cdk.Aws.REGION,
      '.on.aws',
    ]);

    const dsqlClusterArn = cdk.Fn.join('', [
      'arn:aws:dsql:',
      cdk.Aws.REGION,
      ':',
      cdk.Aws.ACCOUNT_ID,
      ':cluster/',
      dsqlCluster.attrIdentifier,
    ]);

    // Docker Lambda with LWA
    const fn = new lambda.DockerImageFunction(this, 'SampleLambda', {
      functionName: 'dsql-sample',
      code: lambda.DockerImageCode.fromImageAsset(
        path.resolve(__dirname, '../..'),
      ),
      memorySize: 1024,
      timeout: cdk.Duration.seconds(30),
      architecture: lambda.Architecture.ARM_64,
      loggingFormat: lambda.LoggingFormat.JSON,
      logGroup: new logs.LogGroup(this, 'SampleLogGroup', {
        logGroupName: '/aws/lambda/dsql-sample',
        retention: logs.RetentionDays.ONE_WEEK,
        removalPolicy: cdk.RemovalPolicy.DESTROY,
      }),
      environment: {
        AWS_LWA_INVOKE_MODE: 'response_stream',
        DSQL_CLUSTER_ENDPOINT: dsqlClusterEndpoint,
        MODE: 'fix',
      },
    });

    // DSQL 接続用 IAM ポリシー
    fn.addToRolePolicy(
      new iam.PolicyStatement({
        effect: iam.Effect.ALLOW,
        actions: ['dsql:DbConnectAdmin'],
        resources: [dsqlClusterArn],
      }),
    );

    // Function URL（API Gateway なしで簡易的にアクセス）
    const fnUrl = fn.addFunctionUrl({
      authType: lambda.FunctionUrlAuthType.NONE,
    });

    new cdk.CfnOutput(this, 'FunctionUrl', {
      value: fnUrl.url,
    });

    new cdk.CfnOutput(this, 'DsqlClusterEndpoint', {
      value: dsqlClusterEndpoint,
    });
  }
}

# syntax=docker/dockerfile:1

ARG RUST_VERSION=1.94

FROM lukemathwalker/cargo-chef:latest-rust-${RUST_VERSION} AS chef
WORKDIR /app

FROM chef AS planner
COPY . .
RUN cargo chef prepare --recipe-path recipe.json

FROM chef AS builder
COPY --from=planner /app/recipe.json recipe.json
RUN --mount=type=cache,target=/usr/local/cargo/registry,sharing=locked \
    --mount=type=cache,target=/usr/local/cargo/git,sharing=locked \
    cargo chef cook --release --recipe-path recipe.json

COPY . .
RUN --mount=type=cache,target=/usr/local/cargo/registry,sharing=locked \
    --mount=type=cache,target=/usr/local/cargo/git,sharing=locked \
    --mount=type=cache,target=/app/target,sharing=locked \
    cargo build --release && \
    cp ./target/release/rust-dsql-connect-error-sample /bin/server

FROM public.ecr.aws/awsguru/aws-lambda-adapter:0.9.0 AS aws-lambda-adapter

FROM gcr.io/distroless/cc-debian13:nonroot
COPY --from=aws-lambda-adapter /lambda-adapter /opt/extensions/lambda-adapter

WORKDIR /app
COPY --from=builder --chown=nonroot:nonroot /bin/server /app
USER nonroot

EXPOSE 8080
CMD ["/app/server"]