axum-layout.rs

use std::convert::Infallible;

use axum::{
    async_trait,
    extract::FromRequestParts,
    http::{request::Parts, Request, header::CONTENT_LENGTH},
    middleware::{from_fn, Next},
    response::{Html, IntoResponse, IntoResponseParts, Response},
    routing::get,
    Router, body::{Full, self},
};
use maud::{html, Markup};

#[tokio::main]
async fn main() {
    let one = Router::new().route("/one", get(one));

    let two = Router::new()
        .route("/two", get(two))
        .layer(from_fn(wrap_in_layout));

    let app = one.merge(two);

    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

// Option 1: A layout extractor

struct Layout {
    current_user: String,
}

#[async_trait]
impl<S> FromRequestParts<S> for Layout
where
    S: Send + Sync,
{
    type Rejection = std::convert::Infallible;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        // extract whatever your layout needs
        Ok(Self {
            current_user: "Bob".to_owned(),
        })
    }
}

impl Layout {
    fn render(self, template: Markup) -> Response {
        let with_layout = html! {
            div {
                h1 {
                    (format!("Welcome! {}", self.current_user))
                }

                div {
                    (template)
                }
            }
        }.into_string();

        Html(with_layout).into_response()
    }
}

async fn one(layout: Layout) -> Response {
    layout.render(html! { "Hello, World!" })
}

// Option 2: Middleware

async fn two() -> Template {
    Template(html! {
        "Hello, World!"
    })
}

struct Template(Markup);

impl IntoResponseParts for Template {
    type Error = Infallible;

    fn into_response_parts(
        self,
        mut res: axum::response::ResponseParts,
    ) -> Result<axum::response::ResponseParts, Self::Error> {
        res.extensions_mut().insert(self);
        Ok(res)
    }
}

impl IntoResponse for Template {
    fn into_response(self) -> Response {
        (self, ()).into_response()
    }
}

async fn wrap_in_layout<B>(
    // whatever extractors you need can go here
    request: Request<B>,
    next: Next<B>,
) -> Response {
    // maybe this is extracted from a JWT or something
    let current_user = "Bob".to_owned();

    let mut response = next.run(request).await;

    if let Some(Template(template)) = response.extensions_mut().remove::<Template>() {
        let with_layout = html! {
            div {
                h1 {
                    (format!("Welcome! {current_user}"))
                }

                div {
                    (template)
                }
            }
        }.into_string();

        let new_body = Html(with_layout);

        let (mut parts, _) = response.into_parts();
        parts.headers.remove(CONTENT_LENGTH);

        (parts, new_body).into_response()
    } else {
        response
    }
}