Port cookie extractors to use state to extract keys (#1250)

axum-extra/CHANGELOG.md

@@ -17,6 +17,8 @@ and this project adheres to [Semantic Versioning].
   literal `Response`
 - **added:** Support chaining handlers with `HandlerCallWithExtractors::or` ([#1170])
 - **change:** axum-extra's MSRV is now 1.60 ([#1239])
+- **breaking:** `SignedCookieJar` and `PrivateCookieJar` now extracts the keys
+  from the router's state, rather than extensions
 
 [#1086]: https://github.com/tokio-rs/axum/pull/1086
 [#1119]: https://github.com/tokio-rs/axum/pull/1119

axum-extra/src/extract/cookie/mod.rs

@@ -208,7 +208,7 @@ fn set_cookies(jar: cookie::CookieJar, headers: &mut HeaderMap) {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use axum::{body::Body, http::Request, routing::get, Extension, Router};
+    use axum::{body::Body, http::Request, routing::get, Router};
     use tower::ServiceExt;
 
     macro_rules! cookie_test {
@@ -227,12 +227,15 @@ mod tests {
                     jar.remove(Cookie::named("key"))
                 }
 
-                let app = Router::<_, Body>::new()
+                let state = AppState {
+                    key: Key::generate(),
+                    custom_key: CustomKey(Key::generate()),
+                };
+
+                let app = Router::<_, Body>::with_state(state)
                     .route("/set", get(set_cookie))
                     .route("/get", get(get_cookie))
-                    .route("/remove", get(remove_cookie))
-                    .layer(Extension(Key::generate()))
-                    .layer(Extension(CustomKey(Key::generate())));
+                    .route("/remove", get(remove_cookie));
 
                 let res = app
                     .clone()
@@ -280,6 +283,24 @@ mod tests {
     cookie_test!(private_cookies, PrivateCookieJar);
     cookie_test!(private_cookies_with_custom_key, PrivateCookieJar<CustomKey>);
 
+    #[derive(Clone)]
+    struct AppState {
+        key: Key,
+        custom_key: CustomKey,
+    }
+
+    impl From<AppState> for Key {
+        fn from(state: AppState) -> Key {
+            state.key
+        }
+    }
+
+    impl From<AppState> for CustomKey {
+        fn from(state: AppState) -> CustomKey {
+            state.custom_key
+        }
+    }
+
     #[derive(Clone)]
     struct CustomKey(Key);
 
@@ -295,9 +316,12 @@ mod tests {
             format!("{:?}", jar.get("key"))
         }
 
-        let app = Router::<_, Body>::new()
-            .route("/get", get(get_cookie))
-            .layer(Extension(Key::generate()));
+        let state = AppState {
+            key: Key::generate(),
+            custom_key: CustomKey(Key::generate()),
+        };
+
+        let app = Router::<_, Body>::with_state(state).route("/get", get(get_cookie));
 
         let res = app
             .clone()

axum-extra/src/extract/cookie/private.rs

@@ -3,7 +3,6 @@ use axum::{
     async_trait,
     extract::{FromRequest, RequestParts},
     response::{IntoResponse, IntoResponseParts, Response, ResponseParts},
-    Extension,
 };
 use cookie::PrivateJar;
 use std::{convert::Infallible, fmt, marker::PhantomData};
@@ -22,7 +21,6 @@ use std::{convert::Infallible, fmt, marker::PhantomData};
 /// ```rust
 /// use axum::{
 ///     Router,
-///     Extension,
 ///     routing::{post, get},
 ///     extract::TypedHeader,
 ///     response::{IntoResponse, Redirect},
@@ -44,22 +42,36 @@ use std::{convert::Infallible, fmt, marker::PhantomData};
 ///     }
 /// }
 ///
-/// // Generate a secure key
-/// //
-/// // You probably don't wanna generate a new one each time the app starts though
-/// let key = Key::generate();
+/// // our application state
+/// #[derive(Clone)]
+/// struct AppState {
+///     // that holds the key used to sign cookies
+///     key: Key,
+/// }
+///
+/// // this impl tells `SignedCookieJar` how to access the key from our state
+/// impl From<AppState> for Key {
+///     fn from(state: AppState) -> Self {
+///         state.key
+///     }
+/// }
+///
+/// let state = AppState {
+///     // Generate a secure key
+///     //
+///     // You probably don't wanna generate a new one each time the app starts though
+///     key: Key::generate(),
+/// };
 ///
-/// let app = Router::new()
+/// let app = Router::with_state(state)
 ///     .route("/set", post(set_secret))
-///     .route("/get", get(get_secret))
-///     // add extension with the key so `PrivateCookieJar` can access it
-///     .layer(Extension(key));
-/// # let app: Router = app;
+///     .route("/get", get(get_secret));
+/// # let app: Router<_> = app;
 /// ```
 pub struct PrivateCookieJar<K = Key> {
     jar: cookie::CookieJar,
     key: Key,
-    // The key used to extract the key extension. Allows users to use multiple keys for different
+    // The key used to extract the key. Allows users to use multiple keys for different
     // jars. Maybe a library wants its own key.
     _marker: PhantomData<K>,
 }
@@ -77,13 +89,15 @@ impl<K> fmt::Debug for PrivateCookieJar<K> {
 impl<S, B, K> FromRequest<B, S> for PrivateCookieJar<K>
 where
     B: Send,
-    S: Send,
+    S: Into<K> + Clone + Send,
     K: Into<Key> + Clone + Send + Sync + 'static,
 {
-    type Rejection = <axum::Extension<K> as FromRequest<B, S>>::Rejection;
+    type Rejection = Infallible;
 
     async fn from_request(req: &mut RequestParts<B, S>) -> Result<Self, Self::Rejection> {
-        let key = Extension::<K>::from_request(req).await?.0.into();
+        let state = req.state().clone();
+        let key: K = state.into();
+        let key: Key = key.into();
 
         let mut jar = cookie::CookieJar::new();
         let mut private_jar = jar.private_mut(&key);

axum-extra/src/extract/cookie/signed.rs

@@ -3,7 +3,6 @@ use axum::{
     async_trait,
     extract::{FromRequest, RequestParts},
     response::{IntoResponse, IntoResponseParts, Response, ResponseParts},
-    Extension,
 };
 use cookie::SignedJar;
 use cookie::{Cookie, Key};
@@ -23,7 +22,6 @@ use std::{convert::Infallible, fmt, marker::PhantomData};
 /// ```rust
 /// use axum::{
 ///     Router,
-///     Extension,
 ///     routing::{post, get},
 ///     extract::TypedHeader,
 ///     response::{IntoResponse, Redirect},
@@ -62,22 +60,36 @@ use std::{convert::Infallible, fmt, marker::PhantomData};
 ///     # todo!()
 /// }
 ///
-/// // Generate a secure key
-/// //
-/// // You probably don't wanna generate a new one each time the app starts though
-/// let key = Key::generate();
+/// // our application state
+/// #[derive(Clone)]
+/// struct AppState {
+///     // that holds the key used to sign cookies
+///     key: Key,
+/// }
+///
+/// // this impl tells `SignedCookieJar` how to access the key from our state
+/// impl From<AppState> for Key {
+///     fn from(state: AppState) -> Self {
+///         state.key
+///     }
+/// }
+///
+/// let state = AppState {
+///     // Generate a secure key
+///     //
+///     // You probably don't wanna generate a new one each time the app starts though
+///     key: Key::generate(),
+/// };
 ///
-/// let app = Router::new()
+/// let app = Router::with_state(state)
 ///     .route("/sessions", post(create_session))
-///     .route("/me", get(me))
-///     // add extension with the key so `SignedCookieJar` can access it
-///     .layer(Extension(key));
-/// # let app: Router = app;
+///     .route("/me", get(me));
+/// # let app: Router<_> = app;
 /// ```
 pub struct SignedCookieJar<K = Key> {
     jar: cookie::CookieJar,
     key: Key,
-    // The key used to extract the key extension. Allows users to use multiple keys for different
+    // The key used to extract the key. Allows users to use multiple keys for different
     // jars. Maybe a library wants its own key.
     _marker: PhantomData<K>,
 }
@@ -95,13 +107,15 @@ impl<K> fmt::Debug for SignedCookieJar<K> {
 impl<S, B, K> FromRequest<B, S> for SignedCookieJar<K>
 where
     B: Send,
-    S: Send,
+    S: Into<K> + Clone + Send,
     K: Into<Key> + Clone + Send + Sync + 'static,
 {
-    type Rejection = <axum::Extension<K> as FromRequest<B, S>>::Rejection;
+    type Rejection = Infallible;
 
     async fn from_request(req: &mut RequestParts<B, S>) -> Result<Self, Self::Rejection> {
-        let key = Extension::<K>::from_request(req).await?.0.into();
+        let state = req.state().clone();
+        let key: K = state.into();
+        let key: Key = key.into();
 
         let mut jar = cookie::CookieJar::new();
         let mut signed_jar = jar.signed_mut(&key);