Implement IntoResponse for Extension<T>

[Foleychris]

The pattern of passing state into handlers via Request.extensions works very well across Axum, including passing state from middleware into a handler. Similarly, I would like to pass state back from a handler to a middleware. My thought was have Extension populate Response.extension, which could then be extracted in middleware, and combinations similar to Headers today.

For a more concrete example, in the case of session handling today, individual handlers are required to handle session cookies (https://github.com/tokio-rs/axum/blob/main/examples/sessions/src/main.rs#L54-L58). With this approach, they could return (StatusCode, Extension<User>) from the handler, and the middleware would responsible for the cookie lifecycle end to end (extracting if present, setting if needed, etc).

I would be willing to attempt the implementation of this, but may require some mentoring through the process.

EDIT: There's a long personal journey in this thread, but the tl;dr is in #770 (comment)

[Foleychris]

Just capturing my own thoughts here -- it looks like the approach used for the tuple type with headers won't work here. They implement IntoResponse like:

impl<H, T> IntoResponse for (StatusCode, H, T)
where
    H: IntoResponseHeaders,
    T: IntoResponse,
{
    fn into_response(self) -> Response {
       [...]
    }
}

and

impl<H, T> IntoResponse for (StatusCode, H, T)
where
    H: IntoResponseHeaders,
    T: IntoResponse,
{
    fn into_response(self) -> Response {
       [...]
    }
}

Ignoring the specifics of the implementation, I first tried doing something similar, defining an IntoExtension trait in axum-core, that I could then implement in axum:

pub trait IntoExtensions {
    fn into_extensions(self) -> http::Extensions;
}

and then attempt to implement for a tuple:

impl<E, T> IntoResponse for (E, T)
where
    E: IntoExtensions,
    T: IntoResponse,
{
    fn into_response(self) -> Response {
        todo!()
    }
}

but this quite reasonably fails:

conflicting implementation for (_, _)

which makes sense, a type could implement both IntoResponseHeaders and IntoExtensions (it would probably be nonsense but...) and the complier rejects this.

I'm definitely not familiar enough with the code base to say if this is a blocker. The only thing that immediately comes to mind is pulling Extension<T> into axum-core and using it directly in the impl with something like:

impl<I,T> IntoResponse for (Extension<I>, T)
where
  I: Clone + Send + Sync + 'static
  T: IntoResponse
{
  [...]
}

This does pass the complier, but that feels like a pretty major change for this feature. Its not immediately clear if it would be possible to move the definition of Extension<T> into axum-core even if we wanted to, given its use super::{rejection::*, FromRequest, RequestParts};. I still do like the idea of this feature, so I'll leave this open for someone more familiar with the codebase to hopefully chime in with an alternative approach.

EDIT:

Ok, one approach that might be a middle ground between "ideal ergonomics" and "major refactoring" is:

// in axum-core/response/mod.rs
impl IntoResponse for http::Extensions {
    fn into_response(self) -> Response {
        let mut res = Response::new(boxed(Empty::new()));
        *res.extensions_mut() = self;
        res
    }
}

impl<T> IntoResponse for (http::Extensions, T)
where
    T: IntoResponse,
{
    fn into_response(self) -> Response {
        let mut res = self.1.into_response();
        *res.extensions_mut() = self.0;
        res
    }
}

impl<T> IntoResponse for (StatusCode, http::Extensions, T)
where
    T: IntoResponse,
{
    fn into_response(self) -> Response {
        let mut res = self.2.into_response();
        *res.status_mut() = self.0;
        *res.extensions_mut() = self.1;
        res
    }
}

and

// in axum/src/extract/extension.rs
impl<T> From<Extension<T>> for http::Extensions
where
    T: Send + Sync + 'static,
{
    fn from(Extension(value): Extension<T>) -> Self {
        let mut extensions = http::Extensions::new();
        extensions.insert(value);
        extensions
    }
}

This allows for handlers like:

async fn handler() -> (http::Extensions, Html<&'static str>) {
    (
        Extension(String::from("Extension string!")).into(),
        Html("Hello world!"),
    )
}

which isn't the worst by a long shot. But the possibly nicer -> impl IntoResponse doesn't really work:

async fn handler() -> impl IntoResponse {
    (
        Extension(String::from("Extension string!")).into(),
        Html("Hello world!"),
    )
}

gives

error[E0283]: type annotations needed

Maybe something like Extension<T>.into_extensions() instead?

impl<T> Extension<T>
where
    T: Send + Sync + 'static,
{
    pub fn into_extensions(self) -> http::Extensions {
        let Self(value) = self;
        let mut extensions = http::Extensions::new();
        extensions.insert(value);
        extensions
    }
}

giving a handler of:

async fn handler() -> impl IntoResponse {
    (
        Extension(String::from("Extension string!")).into_extensions(),
        Html("Hello world!"),
    )
}

[Foleychris]

I might have been too attached to the idea of using the original extractor Extension type, I think this accomplishes what I want without any modification to axum:

struct RespExtension<E, T>(E, T);

impl<E, T> axum::response::IntoResponse for RespExtension<E, T>
where
    E: Send + Sync + 'static,
    T: axum::response::IntoResponse,
{
    fn into_response(self) -> axum::response::Response {
        let mut resp = self.1.into_response();
        let extensions = resp.extensions_mut();
        extensions.insert(self.0);
        resp
    }
}
async fn handler() -> impl IntoResponse {
    RespExtension("Extension String!", "Hello, world!")
}

Sorry for all the rambling, hopefully there's some value here somewhere!

[jplatte]

I think it might be possible to change IntoResponseHeaders (which wasn't released yet) to IntoResponseParts or sth. like that, and then implement it for Extension.

[davidpdrsn]

Another thing to consider is that you cannot merge two http::Extensions.

[jplatte]

With the anymap crate approaching 1.0 maybe http could switch to that in its next release. It looks like that one allows merging two maps via map1.extend(map2.into_raw());.

[davidpdrsn]

I've done some experiments now.

To work around not being able to merge http::Extensions I figured we could make a trait like:

pub trait IntoResponseExtensions: Send + Sync + 'static {
    fn into_extensions(self, extensions: &mut Extensions);
}

Which we can the implement for tuples of different sizes. However you cannot implement a trait for T and (T1, T2). So that requires the one tuple impl being (T,). Due to the coherence stuff that @Foleychris mentions we'd have to use some wrapper type like Extension. Assuming we did that and moved it into axum-core the usage would be something like:

// set status, one extension, and a body
(StatusCode::OK, Extension((extension,)), "Hello, World!")

I quite dislike Extension((extension,)). That is very foreign and would probably trip up a lot of people.

I'm not sure how to get around that with a consistent API. I mean one could imagine uses Extension(foo) for adding one extension and something like (foo, bar, baz).into_extensions() for multiple. However I don't like that inconsistency.

The fact that the extensions are different types is making this quite hard. With headers is easy stick things in an array because they're all the same type.

---

I suppose a different solution would be changing IntoResponse to

pub trait IntoResponse {
    fn into_response(self) -> Response;

    fn with_extension<T>(self, extension: T) -> Response
    where
        T: Send + Sync + 'static,
        Self: Sized,
    {
        let mut res = self.into_response();
        res.extensions_mut().insert(extension);
        res
    }
}

Usage would be

async fn handler() -> Response {
    Html("<h1>Hello, World!</h1>")
        .with_extension(1)
        .with_extension(2)
}

It can be chained because Response implements IntoResponse. In practice I'd probably make an IntoResponseExt trait and put that in axum (or maybe event axum-extra), not axum-core but not sure about that.

@jplatte @Foleychris What do you think about this?

[jplatte]

As I was saying above, I think a good solution here could be to change IntoResponseHeaders to a more general IntoResponseParts or similar. Then, in the same way (StatusCode::OK, TypedHeader(…), TypedHeader(…), "Hello, World!") is allowed by that, you could add multiple extensions through (StatusCode::Ok, Extension(1), Extension(2), "Hello, World!").

[davidpdrsn]

Ah I missed that. Thats a great idea! I think #797 should do it!

[Foleychris]

Sorry for going quiet on the issue, work ate up all my free time. With the caveat that I'm super junior to both Rust and Axum (so weigh my opinion accordingly!) I really like this new API from #797 on first read -- I'll definitely try to make some time to try it out this week.

I'll be honest -- if I had thought implementing this would require reworking fundamental traits I probably wouldn't have opened this issue. I'm currently just using Axum for a toy app in a portfolio project. Using extensions in a response to interact with middleware made sense to me, but I definitely don't want to introduce any additional complexity to the project for my sake. If as maintainers this API makes sense to you, then I'm all in favor (and makes me think I'm starting to understand the ecosystem better 😄 )

[davidpdrsn]

Don't worry about it 😊 I think the fact that your (very reasonable) feature request couldn't be implemented previously showed that our fundamental abstractions weren't right. So we should fix them.