Axum Layers によるオブザーバビリティとセキュリティのためのモジュラーWebサービス構築

はじめに

急速に進化するWebサービス開発の状況において、堅牢でスケーラブル、かつ保守性の高いアプリケーションを構築することは最優先事項です。オブザーバビリティ（ロギングとトレース）とセキュリティ（認証）は、単なる後付けではなく、あらゆる本番グレードのシステムの基本的な柱です。サービスが複雑化するにつれて、これらの横断的懸念を効果的かつ非侵襲的に統合することは、大きな課題となります。そこでミドルウェアが活躍し、これらの懸念をコアビジネスロジックから分離してカプセル化する強力なパターンを提供します。Rustエコシステムで人気のあるWebフレームワークであるAxumは、Towerプロジェクトに大きく影響を受けた「Layer」システムを通じて、ミドルウェアを実装するためのエレガントで柔軟なメカニズムを提供します。Axum Layersを理解し活用することで、開発者は非常にモジュラーでオブザーバブルなサービスを構築し、ロギング、認証、トレースなどの不可欠な機能を驚くほど簡単かつ再利用可能に注入できます。この記事では、Axum Layersの内部を探り、これらの重要な側面のカスタムミドルウェアを構築するためのガイドを提供し、それらがRust Webアプリケーションの構造と保守性をどのように劇的に改善できるかを実証します。

Axum Layersの理解

実装に入る前に、Axum Layersに関連するいくつかのコア用語を明確にしましょう。

Service: TowerとAxumのコンテキストでは、Serviceはリクエストを受け取り、レスポンスまたはエラーのFutureを返す非同期関数を定義するトレイトです。これはリクエストを処理するための基本的なビルディングブロックです。Axumハンドラは効果的にサービスです。

Layer: Layerは、Serviceを受け取り、それをラップして新しいServiceを返すトレイトです。この新しいサービスは、基盤となるサービスを呼び出す前、後、または場合によっては代わりに操作を実行できます。Layersはコンポーザブルであり、複数のレイヤーをスタックして処理ステップのパイプラインを作成できます。

Middleware: Layerは特定のTowerトレイトですが、「ミドルウェア」はそれが実装するより一般的な概念です。ミドルウェアはリクエストとレスポンスをインターセプトし、ロギング、認証、キャッシングなどの補助的なタスクを実行します。

Tower: Towerは、堅牢なクライアントおよびサーバーアプリケーションを構築するためのモジュラーで再利用可能なコンポーネントのライブラリです。AxumはTowerのServiceおよびLayerトレイトを広範に活用し、その機能を拡張するための強力でイディオマティックな方法を提供します。

Axum Layersの原則

Axum Layersの強力さは composability にあります。各レイヤーはデコレータとして機能し、ラップするサービスに機能を追加します。リクエストが到着すると、適用された順序で各レイヤーを通過し、最も内側のサービス（ハンドラ）に到達し、最終的にレスポンスは逆順でレイヤーから戻ります。この「オニオンスキン」モデルにより、懸念事項を明確に分離できます。

カスタムLayerは通常、2つの主要なコンポーネントを定義します。

1. Layer struct: このstructはtower::Layerトレイトを実装します。そのlayerメソッドは、内部サービスを受け取り、それをラップする新しいServiceを返します。
2. Service struct: このstructはtower::Serviceトレイトを実装します。内部サービスへの参照を保持し、内部サービスを呼び出す前または後に実行されるロジックをカプセル化します。

これを、ロギング、認証、トレースの実用的な例で説明しましょう。

カスタムロギングレイヤー

ロギングレイヤーは、本番環境でのアプリケーションの動作を理解し、エラーを追跡し、リクエストフローを監視するために不可欠です。

use axum::{
    body::{Body, BoxBody},
    http::Request,
    middleware::Next,
    response::IntoResponse,
    routing::get,
    Router,
};
use std::{
    task::{Context, Poll},
    time::Instant,
};
tower::Layer;
tracing::{info, instrument};

// 1. ロギングレイヤーの定義
#[derive(Debug, Clone)]
struct LogLayer;

impl<S> Layer<S> for LogLayer {
    type Service = LogService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        LogService { inner }
    }
}

// 2. ロギングサービスの定義
#[derive(Debug, Clone)]
struct LogService<S> {
    inner: S,
}

impl<S, B> Service<Request<B>> for LogService<S>
where
    S: Service<Request<B>, Response = Response<BoxBody>> + Send + 'static,
    S::Future: Send + 'static,
    B: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = futures::future::BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, request: Request<B>) -> Self::Future {
        let path = request.uri().path().to_owned();
        let method = request.method().to_string();
        
        info!(%method, %path, "Incoming request");

        let start = Instant::now();
        let future = self.inner.call(request);

        Box::pin(async move {
            let response = future.await?;
            let latency = start.elapsed();
            let status = response.status();
            info!(%method, %path, %status, "Request finished in {}ms", latency.as_millis());
            Ok(response)
        })
    }
}

// レイヤーを作成するためのヘルパー関数（オプションだが便利）
fn log_layer() -> LogLayer {
    LogLayer
}

この例では：

• LogLayerは状態を持たず、Layerトレイトを実装するだけです。そのlayerメソッドは、内部サービスをラップするLogServiceインスタンスを作成します。
• LogService<S>はinnerサービスを保持します。そのcallメソッドは、実際のロギングロジックが存在する場所です。内部サービスを呼び出す前にログを記録し、レスポンスを受信した後、リクエストメソッド、パス、ステータス、レイテンシを含めて再度ログを記録します。
• 構造化ロギングにはtracingを使用しています。これはRustでは強く推奨されます。

カスタム認証レイヤー

認証レイヤーは、許可されたリクエストのみがコアビジネスロジックに到達することを保証し、不正アクセスを防ぎます。

use axum::http::HeaderValue;
use std::collections::HashMap;

// デモンストレーションのためのシンプルなインメモリユーザーーストア
lazy_static::lazy_static! {
    static ref USERS: HashMap<&'static str, &'static str> = {
        let mut m = HashMap::new();
        m.insert("admin", "password123");
        m.insert("user", "mysecret");
        m
    };
}

// 認証失敗のためのエラー型を定義
#[derive(Debug)]
enum AuthError {
    InvalidCredentials,
    MissingCredentials,
}

impl IntoResponse for AuthError {
    fn into_response(self) -> Response<BoxBody> {
        let (status, msg) = match self {
            AuthError::InvalidCredentials => (StatusCode::UNAUTHORIZED, "Invalid credentials"),
            AuthError::MissingCredentials => (StatusCode::UNAUTHORIZED, "Missing Authorization header"),
        };
        (status, msg).into_response()
    }
}

// 1. 認証レイヤーの定義
#[derive(Debug, Clone)]
struct AuthLayer;

impl<S> Layer<S> for AuthLayer {
    type Service = AuthService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        AuthService { inner }
    }
}

// 2. 認証サービスの定義
#[derive(Debug, Clone)]
struct AuthService<S> {
    inner: S,
}

impl<S, B> Service<Request<B>> for AuthService<S>
where
    S: Service<Request<B>, Response = Response<BoxBody>, Error = AuthError> + Send + 'static, // Inner service can return AuthError
    S::Future: Send + 'static,
    B: Send + 'static,
{
    type Response = S::Response;
    type Error = AuthError; // This service can also return AuthError
    type Future = futures::future::BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, mut request: Request<B>) -> Self::Future {
        let auth_header = request.headers().get(axum::http::header::AUTHORIZATION);

        let authenticated = match auth_header {
            Some(header_value) => {
                let header_str = header_value.to_str().unwrap_or_default();
                if header_str.starts_with("Basic ") {
                    let encoded_credentials = header_str["Basic ".len()..].trim();
                    let decoded_bytes = base64::decode(encoded_credentials).unwrap_or_default();
                    let decoded_str = String::from_utf8(decoded_bytes).unwrap_or_default();
                    let parts: Vec<&str> = decoded_str.split(':').collect();

                    if parts.len() == 2 {
                        let (username, password) = (parts[0], parts[1]);
                        USERS.get(username) == Some(&password)
                    } else {
                        false
                    }
                } else {
                    false
                }
            }
            None => false,
        };

        if !authenticated {
            return Box::pin(async { Err(AuthError::InvalidCredentials) });
        }

        // If authenticated, we could potentially extract user ID and insert it into request extensions
        // request.extensions_mut().insert(UserId("some-user-id".to_string()));

        let future = self.inner.call(request);

        Box::pin(async move {
            future.await // Pass through to the inner service
        })
    }
}

// レイヤーを作成するためのヘルパー関数
fn auth_layer() -> AuthLayer {
    AuthLayer
}

認証の主なポイント：

• カスタムAuthErrorを定義し、IntoResponseを実装するため、Axumは認証失敗を適切に処理できます。
• AuthServiceはAuthorizationヘッダーを抽出し、Basic認証資格情報を解析しようとし、USERSマップに対してチェックします。
• 認証に失敗した場合、すぐにErr(AuthError::InvalidCredentials)を返します。それ以外の場合は、リクエストを内部サービスに渡します。
• 重要: AuthServiceとinnerサービス制約の両方でError = AuthErrorを定義していることに注意してください。これは、内部サービスがAuthErrorで失敗した場合でも、正しく伝播することを意味します。

カスタムトレースレイヤー

トレースは、単一サービス内および異なるサービス間のリクエストフローを視覚化するのに役立ち、分散システムのデバッグやパフォーマンス分析に不可欠です。Axumはすでにtracingとよく統合されていますが、カスタムレイヤーはトレースをリッチ化したり、特定のコンテキストを追加したりできます。実際には、tower_http::trace::TraceLayerのような専用のトレースレイヤーを使用することが多いですが、教育目的で、単純なものをいくつか作成してデモンストレーションしましょう。

use axum_extra::extract::Extension;
use uuid::Uuid;

// トレース用にリクエストに注入されるデータ
#[derive(Clone)]
struct RequestId(String);

// 1. トレースレイヤーの定義
#[derive(Debug, Clone)]
struct TraceLayer;

impl<S> Layer<S> for TraceLayer {
    type Service = TraceService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        TraceService { inner }
    }
}

// 2. トレースサービスの定義
#[derive(Debug, Clone)]
struct TraceService<S> {
    inner: S,
}

impl<S, B> Service<Request<B>> for TraceService<S>
where
    S: Service<Request<B>, Response = Response<BoxBody>> + Send + 'static,
    S::Future: Send + 'static,
    B: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = futures::future::BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, mut request: Request<B>) -> Self::Future {
        // 一意のリクエストIDを生成
        let request_id = Uuid::new_v4().to_string();
        info!(request_id = %request_id, "Assigning request ID");

        // ハンドラがアクセスできるように、リクエスト拡張機能にリクエストIDを挿入
        request.extensions_mut().insert(RequestId(request_id.clone()));

        // ダウンストリームサービスのためにヘッダーとしても追加
        request.headers_mut().insert("X-Request-ID", HeaderValue::from_str(&request_id).unwrap());

        // リクエストのトレーススパンを作成
        let span = tracing::info_span!("request", request_id = %request_id, method = %request.method(), path = %request.uri().path());
        let _guard = span.enter(); // リクエストの期間スパンに入る

        let future = self.inner.call(request);

        Box::pin(async move {
            let response = future.await?;
            info!(request_id = %request_id, status = %response.status(), "Request processed");
            Ok(response)
        })
    }
}

// レイヤーを作成するためのヘルパー関数
fn trace_layer() -> TraceLayer {
    TraceLayer
}

トレースの場合：

• 一意のrequest_idを生成するためにuuidを使用します。
• これはrequest.extensions_mut()に挿入され、Extension<RequestId>を使用して後続のレイヤーとAxumハンドラからアクセス可能になります。
• tracing::info_span!が作成され、このリクエストコンテキスト内のすべてのログが自動的にrequest_idを含むようになります。
• X-Request-IDヘッダーが追加され、サービス境界を越えてトレースIDを伝播するのに役立ちます。

レイヤーをAxumルーターに構成する

これらのレイヤーを組み立て、Axumアプリケーションに適用してみましょう。

use axum::{
    extract::State,
    middleware,
    response::Html,
    routing::{get, post},
    Router,
};
use std::sync::Arc;
tokio::net::TcpListener;
tower_http::trace::TraceLayer as TowerTraceLayer; // より良い機能のためのtower-httpのトレースを使用
tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

// 認証が必要で、リクエストIDにアクセスできるハンドラ
#[instrument(skip(app_state))]
async fn protected_handler(
    State(app_state): State<Arc<String>>,
    Extension(request_id): Extension<RequestId>,
) -> Html<String> {
    info!("Accessing protected handler with request ID: {}", request_id.0);
    Html(format!(
        "<h1>Hello from protected handler!</h1><p>App State: {}</p><p>Request ID: {}</p>",
        app_state,
        request_id.0
    ))
}

// 認証を必要としないハンドラ
async fn public_handler() -> Html<String> {
    info!("Accessing public handler");
    Html("<h1>Hello from public handler!</h1>".to_string())
} 

#[tokio::main]
async fn main() {
    // トレーシングの初期化
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "info,tower_http=debug".into()),
        )
        .with(tracing_subscriber::fmt::layer())
        .init();

    info!("Starting server...");

    let app_state = Arc::new("My Awesome App".to_string());

    // カスタムレイヤーを構築
    let custom_layers = ServiceBuilder::new()
        // カスタムロギングレイヤー（すべてをログに記録するため、外側にあるべき）
        .layer(log_layer())
        // カスタムトレースレイヤー（リクエストをインストゥルメントするために隣にあるべき）
        .layer(trace_layer());

    // レイヤーを特定のルートまたはルーター全体に適用
    let app = Router::new()
        .route("/public", get(public_handler))
        .route("/protected", get(protected_handler))
        .route_layer(middleware::from_fn(|req, next| async {
            // 特定のルートにのみ認証を適用
            // 注：カスタム認証レイヤーはTower Layerとして動作します
            // 特定のルート認証のためのよりイディオマティックなAxumの方法は、
            // カスタムエクストラクタまたはネストされたルーターにレイヤーを適用することかもしれません。
            
            let fut = next.call(req);
            // この例では、ミドルウェア列挙関数に直接AuthServiceを適用しています。
            // しかし、より一般的なアプローチはTower Layerを使用することです。
            // ここでは、Axumのミドルウェア関数としてAuthServiceをシミュレートします。
            // 実際には、AuthLayer.layer(next) を使用します。
            // fut.await
            // 本来ならAuthLayerのインスタンスを作成し、nextを渡す。
            // しかし、ここでは直接実装します。
            let mut auth_service = AuthService { inner: next };
            auth_service.call(req).await
        }))
        .layer(TowerTraceLayer::new()) // 包括的なトレースのためのtower-httpのトレースレイヤーを使用
        // .layer(custom_layers) // ここでカスタムレイヤーを適用
        // レイヤーは定義の逆順に適用されます。
        // *最後*に追加されたレイヤーが*ルーター全体/サービス全体*をラップし、最も外側のロジックになります。
        // 例：.layer(MyOuterLayer).layer(MyInnerLayer) -> リクエストはOuter、次にInner、次にハンドラを通過します。
        .layer(auth_layer()) // カスタム認証レイヤー
        .layer(custom_layers) // カスタムロギングとトレースを適用（最も外側）
        .with_state(app_state);

    let listener = TcpListener::bind("127.0.0.1:3000").await.unwrap();
    info!("Listening on {}", listener.local_addr().unwrap());
    axum::serve(listener, app).await.unwrap();
}

この例を実行するには、Cargo.tomlに次の依存関係が必要です：

[dependencies]
axum = { version = "0.7", features = ["macros"] }
tokio = { version = "1.36", features = ["full"] }
tower = { version = "0.4", features = ["full"] }
tower-http = { version = "0.5", features = ["full"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] }
futures = "0.3"
uuid = { version = "1.7", features = ["v4", "fast-rng"] }
base64 = "0.21"
lazy_static = "1.4"
axum-extra = { version = "0.9", features = ["extract-intervals"] } # For Extension

main関数では：

• ログ出力のためにtracing_subscriberを初期化します。
• ServiceBuilder::new().layer(log_layer()).layer(trace_layer())は、カスタムロギングおよびトレースサービスから結合されたレイヤーを作成します。ServiceBuilderチェーンのトップダウンでレイヤーが適用されます。つまり、log_layerはtrace_layerの外側になります。
• app.layer(auth_layer())は、認証レイヤーをルーター全体に適用します。これは、すべてのルートが認証を通過することを意味します。選択的な認証が必要な場合は、ネストされたRouter::with_no_routes()インスタンスにレイヤーを適用できます。
• protected_handlerは、リクエスト拡張機能からRequestIdを抽出する方法を示しています。
• サーバーは127.0.0.1:3000でリッスンします。

テストする：

1. パブリックエンドポイント： curl http://127.0.0.1:3000/public（パブリックHTMLを返すはずです）。401 Unauthorized with "Invalid credentials".
2. 保護されたエンドポイント（認証なし）： curl http://127.0.0.1:3000/protected（401 Unauthorized with "Invalid credentials" を返すはずです）。
3. 保護されたエンドポイント（認証あり）： curl -H "Authorization: Basic YWRtaW46cGFzc3dvcmQxMjM=" http://127.0.0.1:3000/protected（保護されたHTMLを返すはずです）。 YWRtaW46cGFzc3dvcmQxMjM= は base64 エンコードされた "admin" です。

コンソールに詳細なログが表示され、リクエストID、メソッド、パス、ステータス、レイテンシが表示され、カスタムレイヤーの効果が実証されます。

アプリケーションシナリオ

• ロギング： すべてのAPIリクエストは、監査、デバッグ、監視のためにログに記録する必要があります。
• 認証/認可： ユーザーロールまたは権限に基づいて、特定のとは言え、APIのセグメント全体を保護します。
• トレース： 分散トレースのためにリクエストIDを伝播させ、マイクロサービス全体のリクエストのエンドツーエンドの可視性を可能にします。
• レート制限： クライアントが一定期間内に送信できるリクエスト数を制限して、乱用を防ぎます。
• リクエスト/レスポンス変換： ヘッダーを変更したり、ボディを圧縮したり、リクエスト/レスポンスに共通のデータを注入したりします。
• メトリクス： 監視ダッシュボードのために、リクエスト数、エラー率、レイテンシなどのメトリクスを収集して公開します。
• CORS： クロスオリジンリソース共有ヘッダーを処理します。

結論

Towerエコシステムを基盤とするAxumのLayerシステムは、Rust Webアプリケーションにカスタムミドルウェアを実装するための、非常に柔軟でイディオマティックな方法を提供します。LayerとServiceトレイトを理解することで、開発者はロギング、認証、トレースのような横断的懸念をモジュール化し、よりクリーンで保守性が高く、堅牢なコードベースにつながります。このアプローチは、サービスのオブザーバビリティとセキュリティを強化するだけでなく、再利用性と懸念事項の分離を促進し、スケーラブルで満足のいくWebアプリケーションの構築に不可欠です。Axum Layersを活用することで、信頼性の高い、本番対応のRust Webサービスを構築できます。