Get access to the request headers from Axum IntoResponse

Question

Get access to the request headers from Axum IntoResponse

3.2k views Asked by wallefan At 16 September 2024 at 15:07

I have a custom IntoResponse type that needs to return different data to the client depending on the value of the Accept header in the request. I would prefer, if possible, to have this be seamless (i.e. I don't have to pass the header value into the IntoResponse type's constructor). I thought the most straightforward way to accomplish this would be to write a layer using Axum's middleware system that all functions returning that type are required to be wrapped with, not have the type implement IntoResponse, and instead perform the conversion inside the middleware so that it could access the request headers, but it seems that Axum's middleware system requires the service function to return an IntoResponse. Could I whip something up with a tower::Layer that does the same thing?

Original Q&A

There are 1 answers

**kmdreko** · Accepted Answer · 2023-05-07 02:48:02

Firstly, lets get it out of the way that there is not a built-in mechanism to support the Accept header impacting the response. See this brief discussion.

You've probably noticed the IntoResponse trait that handlers must yield does not provide access to the original request:

pub trait IntoResponse {
    fn into_response(self) -> Response<UnsyncBoxBody<Bytes, Error>>;
}

So you'd have to get the header from either the handler itself or some middleware.

Even worse, it requires the body be a sequence of bytes (which makes sense but loses type information). So a middleware would have to decode data from the encoded response body to modify it for a different mime type. This could probably work if your data is really simple, but for a lot of cases this would probably be an annoying and unnecessary cost. This data flow also impacts how underlying tower services can affect the response, so the distinction doesn't help.

Fortunately, there is another way: you can provide any data to your middleware via an Extension. Through .extensions()/.extensions_mut() on a response, you can store arbitrary data, so your IntoResponse implementation can store itself in an extension that a middeware could then pull out and re-format as necessary.

Here's a mockup of how that could work:

use axum::http::header::ACCEPT;
use axum::http::{Request, StatusCode};
use axum::middleware::{from_fn, Next};
use axum::response::{Html, IntoResponse, Json, Response};
use axum::{routing::get, Router};
use serde::Serialize;

#[derive(Serialize)]
struct MyData {
    data: String,
}

impl IntoResponse for MyData {
    fn into_response(self) -> Response {
        let mut response = StatusCode::NOT_IMPLEMENTED.into_response();
        response.extensions_mut().insert(self);
        response
    }
}

async fn get_accept_aware() -> MyData {
    MyData {
        data: "test".to_owned(),
    }
}

async fn my_data_applicator<B>(request: Request<B>, next: Next<B>) -> Response {
    let accept_header = request
        .headers()
        .get(&ACCEPT)
        .map(|value| value.as_ref().to_owned());

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

    if let Some(my_data) = response.extensions_mut().remove::<MyData>() {
        match accept_header.as_deref() {
            Some(b"application/json") => return Json(my_data).into_response(),
            Some(b"html") => return Html(format!("<body>{}</body>", my_data.data)).into_response(),
            _ => { /* yield original 501 response */ }
        }
    }

    response
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/", get(get_accept_aware))
        .layer(from_fn(my_data_applicator));

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

TechQA.

Get access to the request headers from Axum IntoResponse

There are 1 answers

Related Questions in RUST

Related Questions in RUST-AXUM

Related Questions in RUST-TOWER

Popular Questions

Popular Tags

Trending Questions