Add Redirect response (#192)

* Add `Redirect` response * Add `Redirect::found`
tokio-rs · Aug 16, 2021 · b4cbd7f · b4cbd7f
1 parent dda6257
commit b4cbd7f
Show file tree

Hide file tree

Showing 4 changed files with 107 additions and 31 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `NestedUri` for extracting request URI in nested services ([#161](/~https://github.com/tokio-rs/axum/pull/161))
 - Implement `FromRequest` for `http::Extensions`
 - Implement SSE as an `IntoResponse` instead of a service ([#98](/~https://github.com/tokio-rs/axum/pull/98))
+- Add `Redirect` response.
 
 ## Breaking changes
 

diff --git a/examples/oauth.rs b/examples/oauth.rs
@@ -9,14 +9,13 @@
 use async_session::{MemoryStore, Session, SessionStore};
 use axum::{
     async_trait,
+    body::{Bytes, Empty},
     extract::{Extension, FromRequest, Query, RequestParts, TypedHeader},
     prelude::*,
-    response::IntoResponse,
+    response::{IntoResponse, Redirect},
     AddExtensionLayer,
 };
-use http::header::SET_COOKIE;
-use http::StatusCode;
-use hyper::Body;
+use http::{header::SET_COOKIE, HeaderMap};
 use oauth2::{
     basic::BasicClient, reqwest::async_http_client, AuthUrl, AuthorizationCode, ClientId,
     ClientSecret, CsrfToken, RedirectUrl, Scope, TokenResponse, TokenUrl,
@@ -118,7 +117,7 @@ async fn discord_auth(Extension(client): Extension<BasicClient>) -> impl IntoRes
         .url();
 
     // Redirect to Discord's oauth service
-    Redirect(auth_url.into())
+    Redirect::found(auth_url.to_string().parse().unwrap())
 }
 
 // Valid user session required. If there is none, redirect to the auth page
@@ -137,12 +136,12 @@ async fn logout(
     let session = match store.load_session(cookie.to_string()).await.unwrap() {
         Some(s) => s,
         // No session active, just redirect
-        None => return Redirect("/".to_string()),
+        None => return Redirect::found("/".parse().unwrap()),
     };
 
     store.destroy_session(session).await.unwrap();
 
-    Redirect("/".to_string())
+    Redirect::found("/".parse().unwrap())
 }
 
 #[derive(Debug, Deserialize)]
@@ -187,35 +186,20 @@ async fn login_authorized(
     let cookie = format!("{}={}; SameSite=Lax; Path=/", COOKIE_NAME, cookie);
 
     // Set cookie
-    let r = http::Response::builder()
-        .header("Location", "/")
-        .header(SET_COOKIE, cookie)
-        .status(302);
+    let mut headers = HeaderMap::new();
+    headers.insert(SET_COOKIE, cookie.parse().unwrap());
 
-    r.body(Body::empty()).unwrap()
-}
-
-// Utility to save some lines of code
-struct Redirect(String);
-impl IntoResponse for Redirect {
-    type Body = Body;
-    type BodyError = hyper::Error;
-
-    fn into_response(self) -> http::Response<Body> {
-        let builder = http::Response::builder()
-            .header("Location", self.0)
-            .status(StatusCode::FOUND);
-        builder.body(Body::empty()).unwrap()
-    }
+    (headers, Redirect::found("/".parse().unwrap()))
 }
 
 struct AuthRedirect;
+
 impl IntoResponse for AuthRedirect {
-    type Body = Body;
-    type BodyError = hyper::Error;
+    type Body = Empty<Bytes>;
+    type BodyError = <Self::Body as axum::body::HttpBody>::Error;
 
-    fn into_response(self) -> http::Response<Body> {
-        Redirect("/auth/discord".to_string()).into_response()
+    fn into_response(self) -> http::Response<Self::Body> {
+        Redirect::found("/auth/discord".parse().unwrap()).into_response()
     }
 }
 

diff --git a/src/response/mod.rs b/src/response/mod.rs
@@ -16,8 +16,11 @@ use tower::{util::Either, BoxError};
 #[doc(no_inline)]
 pub use crate::Json;
 
-pub mod sse;
+mod redirect;
+
+pub use self::redirect::Redirect;
 
+pub mod sse;
 pub use sse::{sse, Sse};
 
 /// Trait for generating responses.

diff --git a/src/response/redirect.rs b/src/response/redirect.rs
@@ -0,0 +1,88 @@
+use super::IntoResponse;
+use bytes::Bytes;
+use http::{header::LOCATION, HeaderValue, Response, StatusCode, Uri};
+use http_body::{Body, Empty};
+use std::convert::TryFrom;
+
+/// Response that redirects the request to another location.
+///
+/// # Example
+///
+/// ```rust
+/// use axum::{prelude::*, response::Redirect};
+///
+/// let app = route("/old", get(|| async { Redirect::permanent("/new".parse().unwrap()) }))
+///     .route("/new", get(|| async { "Hello!" }));
+/// # async {
+/// # hyper::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap();
+/// # };
+/// ```
+#[derive(Debug, Clone)]
+pub struct Redirect {
+    status_code: StatusCode,
+    location: HeaderValue,
+}
+
+impl Redirect {
+    /// Create a new [`Redirect`] that uses a [`307 Temporary Redirect`][mdn] status code.
+    ///
+    /// # Panics
+    ///
+    /// If `uri` isn't a valid [`HeaderValue`].
+    ///
+    /// [mdn]: https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/307
+    pub fn temporary(uri: Uri) -> Self {
+        Self::with_status_code(StatusCode::TEMPORARY_REDIRECT, uri)
+    }
+
+    /// Create a new [`Redirect`] that uses a [`308 Permanent Redirect`][mdn] status code.
+    ///
+    /// # Panics
+    ///
+    /// If `uri` isn't a valid [`HeaderValue`].
+    ///
+    /// [mdn]: https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/308
+    pub fn permanent(uri: Uri) -> Self {
+        Self::with_status_code(StatusCode::PERMANENT_REDIRECT, uri)
+    }
+
+    /// Create a new [`Redirect`] that uses a [`302 Found`][mdn] status code.
+    ///
+    /// # Panics
+    ///
+    /// If `uri` isn't a valid [`HeaderValue`].
+    ///
+    /// [mdn]: https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/302
+    pub fn found(uri: Uri) -> Self {
+        Self::with_status_code(StatusCode::FOUND, uri)
+    }
+
+    // This is intentionally not public since other kinds of redirects might not
+    // use the `Location` header, namely `304 Not Modified`.
+    //
+    // We're open to adding more constructors upon request, if they make sense :)
+    fn with_status_code(status_code: StatusCode, uri: Uri) -> Self {
+        assert!(
+            status_code.is_redirection(),
+            "not a redirection status code"
+        );
+
+        Self {
+            status_code,
+            location: HeaderValue::try_from(uri.to_string())
+                .expect("URI isn't a valid header value"),
+        }
+    }
+}
+
+impl IntoResponse for Redirect {
+    type Body = Empty<Bytes>;
+    type BodyError = <Self::Body as Body>::Error;
+
+    fn into_response(self) -> Response<Self::Body> {
+        let mut res = Response::new(Empty::new());
+        *res.status_mut() = self.status_code;
+        res.headers_mut().insert(LOCATION, self.location);
+        res
+    }
+}