Fix IntoResponse for tuples overriding error response codes (#3603)

Co-authored-by: David Pedersen <[email protected]>
Co-authored-by: Yann Simon <[email protected]>

axum-core/src/response/into_response.rs

@@ -1,4 +1,4 @@
-use super::{IntoResponseParts, Response, ResponseParts};
+use super::{ForceStatusCode, IntoResponseFailed, IntoResponseParts, Response, ResponseParts};
 use crate::{body::Body, BoxError};
 use bytes::{buf::Chain, Buf, Bytes, BytesMut};
 use http::{
@@ -329,7 +329,9 @@ where
 {
     fn into_response(self) -> Response {
         let mut res = self.1.into_response();
-        *res.status_mut() = self.0;
+        if res.extensions().get::<IntoResponseFailed>().is_none() {
+            *res.status_mut() = self.0;
+        }
         res
     }
 }
@@ -405,18 +407,16 @@ macro_rules! impl_into_response {
                 let ($($ty),*, res) = self;
 
                 let res = res.into_response();
-                let parts = ResponseParts { res };
-
-                $(
-                    let parts = match $ty.into_response_parts(parts) {
+                if res.extensions().get::<IntoResponseFailed>().is_none() {
+                    let parts = ResponseParts { res };
+                    let parts = match ($($ty,)*).into_response_parts(parts) {
                         Ok(parts) => parts,
-                        Err(err) => {
-                            return err.into_response();
-                        }
+                        Err(err) => return err.into_response(),
                     };
-                )*
-
-                parts.res
+                    parts.res
+                } else {
+                    res
+                }
             }
         }
 
@@ -430,16 +430,40 @@ macro_rules! impl_into_response {
                 let (status, $($ty),*, res) = self;
 
                 let res = res.into_response();
-                let parts = ResponseParts { res };
-
-                $(
-                    let parts = match $ty.into_response_parts(parts) {
+                if res.extensions().get::<IntoResponseFailed>().is_none() {
+                    let parts = ResponseParts { res };
+                    let mut parts = match ($($ty,)*).into_response_parts(parts) {
                         Ok(parts) => parts,
-                        Err(err) => {
-                            return err.into_response();
-                        }
+                        Err(err) => return err.into_response(),
                     };
-                )*
+
+                    // Don't call `(status, parts.res).into_response()` since that checks for
+                    // `IntoResponseFailed` and skips setting the status. We've already done that
+                    // check here so overriding the status is required if returning
+                    // `(IntoResponseFailed, StatusCode::INTERNAL_SERVER_ERROR)`
+                    *parts.res.status_mut() = status;
+                    parts.res
+                } else {
+                    res
+                }
+            }
+        }
+
+        #[allow(non_snake_case)]
+        impl<R, $($ty,)*> IntoResponse for (ForceStatusCode, $($ty),*, R)
+        where
+            $( $ty: IntoResponseParts, )*
+            R: IntoResponse,
+        {
+            fn into_response(self) -> Response {
+                let (status, $($ty),*, res) = self;
+
+                let res = res.into_response();
+                let parts = ResponseParts { res };
+                let parts = match ($($ty,)*).into_response_parts(parts) {
+                    Ok(parts) => parts,
+                    Err(err) => return err.into_response(),
+                };
 
                 (status, parts.res).into_response()
             }
@@ -455,17 +479,22 @@ macro_rules! impl_into_response {
                 let (outer_parts, $($ty),*, res) = self;
 
                 let res = res.into_response();
-                let parts = ResponseParts { res };
-                $(
-                    let parts = match $ty.into_response_parts(parts) {
+                if res.extensions().get::<IntoResponseFailed>().is_none() {
+                    let parts = ResponseParts { res };
+                    let mut parts = match ($($ty,)*).into_response_parts(parts) {
                         Ok(parts) => parts,
-                        Err(err) => {
-                            return err.into_response();
-                        }
+                        Err(err) => return err.into_response(),
                     };
-                )*
 
-                (outer_parts, parts.res).into_response()
+                    // Don't call `(outer_parts, parts.res).into_response()` for the same reason we
+                    // don't call `(status, parts.res).into_response()` in the above impl.
+                    *parts.res.status_mut() = outer_parts.status;
+                    parts.res.headers_mut().extend(outer_parts.headers);
+                    parts.res.extensions_mut().extend(outer_parts.extensions);
+                    parts.res
+                } else {
+                    res
+                }
             }
         }
 

axum-core/src/response/into_response_parts.rs

@@ -241,7 +241,9 @@ macro_rules! impl_into_response_parts {
                     let res = match $ty.into_response_parts(res) {
                         Ok(res) => res,
                         Err(err) => {
-                            return Err(err.into_response());
+                            let mut err_res = err.into_response();
+                            err_res.extensions_mut().insert(super::IntoResponseFailed);
+                            return Err(err_res);
                         }
                     };
                 )*
@@ -270,3 +272,19 @@ impl IntoResponseParts for () {
         Ok(res)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use http::StatusCode;
+
+    use crate::response::IntoResponse;
+
+    #[test]
+    fn failed_into_response_parts() {
+        let response = (StatusCode::CREATED, [("\n", "\n")]).into_response();
+        assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
+
+        let response = (StatusCode::CREATED, [("\n", "\n")], ()).into_response();
+        assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
+    }
+}

axum-core/src/response/mod.rs

@@ -4,6 +4,10 @@
 //!
 //! [`axum::response`]: https://docs.rs/axum/0.8/axum/response/index.html
 
+use std::convert::Infallible;
+
+use http::StatusCode;
+
 use crate::body::Body;
 
 mod append_headers;
@@ -128,3 +132,87 @@ where
         Self(value.into_response())
     }
 }
+
+/// Response part that stops status code overrides.
+///
+/// This type should be used by types implementing [`IntoResponseParts`] or
+/// [`IntoResponse`] when they fail to produce the response usually expected of
+/// them and return some sort of error response instead.
+///
+/// It is checked used by the tuple impls of [`IntoResponse`] that have a
+/// [`StatusCode`] as their first element to ignore that status code.
+/// Consider the following example:
+///
+/// ```no_run
+/// # use axum::Json;
+/// # use http::StatusCode;
+/// # #[derive(serde::Serialize)]
+/// # struct CreatedResponse { }
+/// fn my_handler(/* ... */) -> (StatusCode, Json<CreatedResponse>) {
+///     // This response type's serialization may fail
+///     let response = CreatedResponse { /* ... */ };
+///     (StatusCode::CREATED, Json(response))
+/// }
+/// ```
+///
+/// When `response` serialization succeeds, the server responds with a status
+/// code of 201 Created (overwriting `Json`s default status code of 200 OK),
+/// and the expected JSON payload.
+///
+/// When `response` serialization fails hoewever, `impl IntoResponse for Json`
+/// return a response with status code 500 Internal Server Error, and
+/// `IntoResponseFailed` as a response extension, and the 201 Created override
+/// is ignored.
+///
+/// This is a behavior introduced with axum 0.9.\
+/// To force a status code override even when an inner [`IntoResponseParts`] /
+/// [`IntoResponse`] failed, use [`ForceStatusCode`].
+#[derive(Copy, Clone, Debug)]
+pub struct IntoResponseFailed;
+
+impl IntoResponseParts for IntoResponseFailed {
+    type Error = Infallible;
+
+    fn into_response_parts(self, mut res: ResponseParts) -> Result<ResponseParts, Self::Error> {
+        res.extensions_mut().insert(self);
+        Ok(res)
+    }
+}
+
+/// Not sure it makes sense to return `IntoResponseFailed` as the whole response. You should
+/// probably at least combine it with a status code.
+///
+/// ```compile_fail
+/// fn foo()
+/// where
+///     axum_core::response::IntoResponseFailed: axum_core::response::IntoResponse,
+/// {}
+/// ```
+#[allow(dead_code)]
+fn into_response_failed_doesnt_impl_into_response() {}
+
+/// Set the status code regardless of whether [`IntoResponseFailed`] is used or not.
+///
+/// See the docs for [`IntoResponseFailed`] for more details.
+#[derive(Debug, Copy, Clone, Default)]
+pub struct ForceStatusCode(pub StatusCode);
+
+impl IntoResponse for ForceStatusCode {
+    fn into_response(self) -> Response {
+        let mut res = ().into_response();
+        *res.status_mut() = self.0;
+        res
+    }
+}
+
+impl<R> IntoResponse for (ForceStatusCode, R)
+where
+    R: IntoResponse,
+{
+    fn into_response(self) -> Response {
+        let (ForceStatusCode(status), res) = self;
+        let mut res = res.into_response();
+        *res.status_mut() = status;
+        res
+    }
+}