Axum avanzado: middleware, manejo de errores tipado y capas con Tower

En una aplicación web real, hay tareas que no pertenecen a ningún handler concreto: registrar el tiempo de respuesta, añadir cabeceras CORS, verificar que el token JWT es válido o comprimir la respuesta. Si las escribes dentro de cada handler, repites código y mezclas responsabilidades. La solución en Axum, como en el resto del ecosistema Tokio, es Tower.

Tower define los traits Service y Layer. Un Layer envuelve un Service añadiendo comportamiento antes, después o alrededor de él. Axum usa esta abstracción directamente: cualquier middleware de Tower funciona en Axum sin adaptadores.

TraceLayer: logging de peticiones

El crate tower-http incluye TraceLayer, que añade spans de tracing a cada petición. Para usarlo:

[dependencies]
axum = "0.7"
tower-http = { version = "0.5", features = ["trace"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }

use tower_http::trace::TraceLayer;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

#[tokio::main]
async fn main() {
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::new(
            std::env::var("RUST_LOG").unwrap_or_else(|_| "info".into()),
        ))
        .with(tracing_subscriber::fmt::layer())
        .init();

    let app = Router::new()
        .route("/", get(handler))
        .layer(TraceLayer::new_for_http());

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

Con RUST_LOG=debug verás el span de inicio de petición, el span con el código de respuesta y la duración. Puedes personalizar qué datos se registran pasando funciones de callback a TraceLayer.

CORS configurable

Otro middleware incluido en tower-http es CorsLayer:

use tower_http::cors::{Any, CorsLayer};
use http::{HeaderValue, Method};

let cors = CorsLayer::new()
    .allow_origin("https://mi-frontend.com".parse::<HeaderValue>().unwrap())
    .allow_methods([Method::GET, Method::POST, Method::DELETE])
    .allow_headers(Any);

let app = Router::new()
    .route("/api/datos", get(obtener_datos))
    .layer(cors);

Si necesitas permitir múltiples orígenes, usa allow_origin con un AllowOrigin::list(). En desarrollo puedes usar Any para los orígenes, pero en producción siempre especifica los dominios explícitos.

Middleware de autenticación con FromRequestParts

Para autenticación, la forma idiomática en Axum es implementar el trait FromRequestParts en un tipo propio. Esto permite usar el tipo directamente como extractor en los handlers que requieran autenticación.

use axum::{
    async_trait,
    extract::FromRequestParts,
    http::{request::Parts, StatusCode},
};

struct UsuarioAutenticado {
    id: u64,
    email: String,
}

#[async_trait]
impl<S> FromRequestParts<S> for UsuarioAutenticado
where
    S: Send + Sync,
{
    type Rejection = StatusCode;

    async fn from_request_parts(
        parts: &mut Parts,
        _state: &S,
    ) -> Result<Self, Self::Rejection> {
        let token = parts
            .headers
            .get("Authorization")
            .and_then(|v| v.to_str().ok())
            .and_then(|s| s.strip_prefix("Bearer "))
            .ok_or(StatusCode::UNAUTHORIZED)?;

        // En producción validarías el JWT aquí
        if token == "token-valido" {
            Ok(UsuarioAutenticado {
                id: 1,
                email: "[email protected]".to_string(),
            })
        } else {
            Err(StatusCode::UNAUTHORIZED)
        }
    }
}

Ahora cualquier handler que declare UsuarioAutenticado como parámetro exigirá autenticación automáticamente:

async fn perfil(usuario: UsuarioAutenticado) -> Json<serde_json::Value> {
    Json(serde_json::json!({
        "id": usuario.id,
        "email": usuario.email
    }))
}

// Ruta protegida, sin código extra de autenticación en el handler
let app = Router::new().route("/perfil", get(perfil));

Errores tipados con thiserror e IntoResponse

La gestión de errores es donde muchas APIs acaban con código repetido. La combinación de thiserror para definir el tipo de error e IntoResponse para convertirlo en una respuesta HTTP es la solución estándar en Axum:

[dependencies]
thiserror = "1"

use thiserror::Error;
use axum::{
    http::StatusCode,
    response::{IntoResponse, Response, Json},
};
use serde_json::json;

#[derive(Debug, Error)]
enum ApiError {
    #[error("Recurso no encontrado")]
    NoEncontrado,

    #[error("No autorizado")]
    NoAutorizado,

    #[error("Error de validación: {0}")]
    Validacion(String),

    #[error("Error interno: {0}")]
    Interno(#[from] anyhow::Error),
}

impl IntoResponse for ApiError {
    fn into_response(self) -> Response {
        let (status, mensaje) = match &self {
            ApiError::NoEncontrado    => (StatusCode::NOT_FOUND, self.to_string()),
            ApiError::NoAutorizado    => (StatusCode::UNAUTHORIZED, self.to_string()),
            ApiError::Validacion(msg) => (StatusCode::UNPROCESSABLE_ENTITY, msg.clone()),
            ApiError::Interno(_)      => (StatusCode::INTERNAL_SERVER_ERROR, "Error interno".into()),
        };

        let body = Json(json!({ "error": mensaje }));
        (status, body).into_response()
    }
}

Con este tipo, cualquier handler puede devolver Result<T, ApiError> y las respuestas de error serán JSON consistentes:

async fn obtener_usuario(
    Path(id): Path<u64>,
    State(db): State<DB>,
    _usuario: UsuarioAutenticado,  // exige autenticación
) -> Result<Json<Usuario>, ApiError> {
    let db = db.read().unwrap();
    db.get(&id)
        .cloned()
        .map(Json)
        .ok_or(ApiError::NoEncontrado)
}

Routers anidados con estado diferente

En una aplicación grande, suele ser útil dividir las rutas en módulos. Axum permite anidar routers y que cada uno tenga su propio estado:

// rutas/usuarios.rs
pub fn router() -> Router<AppState> {
    Router::new()
        .route("/", get(listar).post(crear))
        .route("/:id", get(obtener).put(actualizar).delete(eliminar))
}

// rutas/productos.rs
pub fn router() -> Router<AppState> {
    Router::new()
        .route("/", get(listar_productos).post(crear_producto))
        .route("/:id", get(obtener_producto))
}

// main.rs
let app = Router::new()
    .nest("/usuarios", rutas::usuarios::router())
    .nest("/productos", rutas::productos::router())
    .with_state(estado_global);

Si un sub-router necesita un estado diferente al global, usa .with_state() dentro del router anidado antes de pasarlo a .nest().

ServiceBuilder: componer múltiples capas

Cuando necesitas apilar varios middlewares, ServiceBuilder de Tower hace el código más legible que encadenar .layer() repetidamente:

use tower::ServiceBuilder;
use tower_http::{
    trace::TraceLayer,
    cors::CorsLayer,
    compression::CompressionLayer,
    timeout::TimeoutLayer,
};
use std::time::Duration;

let middlewares = ServiceBuilder::new()
    .layer(TraceLayer::new_for_http())
    .layer(TimeoutLayer::new(Duration::from_secs(10)))
    .layer(CompressionLayer::new())
    .layer(cors);

let app = Router::new()
    .route("/api/datos", get(handler))
    .layer(middlewares);

El orden importa: en ServiceBuilder el primer .layer() es el más externo (el primero en recibir la petición y el último en procesar la respuesta). Poner TimeoutLayer antes del resto garantiza que el timeout se aplica incluyendo el tiempo de todos los middlewares internos.

Middleware personalizado con axum::middleware::from_fn

Si no quieres implementar el trait Layer directamente, Axum incluye middleware::from_fn para escribir middlewares como funciones async simples:

use axum::{
    middleware::{self, Next},
    extract::Request,
    response::Response,
};

async fn registrar_tiempo(req: Request, next: Next) -> Response {
    let inicio = std::time::Instant::now();
    let path = req.uri().path().to_string();
    let metodo = req.method().clone();

    let respuesta = next.run(req).await;

    let duracion = inicio.elapsed();
    println!("{} {} - {}ms - {}",
        metodo, path,
        duracion.as_millis(),
        respuesta.status()
    );

    respuesta
}

let app = Router::new()
    .route("/", get(handler))
    .layer(middleware::from_fn(registrar_tiempo));

Esta función recibe la petición, puede inspeccionarla o modificarla, llama a next.run(req) para pasar el control al siguiente middleware o handler, y después puede procesar la respuesta antes de devolverla.

La siguiente pieza del rompecabezas es la base de datos. El artículo sobre SQLx en esta misma serie cubre cómo conectar Axum con PostgreSQL usando queries tipadas en tiempo de compilación, pool de conexiones y migraciones automáticas.