Add type safe state extractor (#1155)

* begin threading the state through

* Pass state to extractors

* make state extractor work

* make sure nesting with different states work

* impl Service for MethodRouter<()>

* Fix some of axum-macro's tests

* Implement more traits for `State`

* Update examples to use `State`

* consistent naming of request body param

* swap type params

* Default the state param to ()

* fix docs references

* Docs and handler state refactoring

* docs clean ups

* more consistent naming

* when does MethodRouter implement Service?

* add missing docs

* use `Router`'s default state type param

* changelog

* don't use default type param for FromRequest and RequestParts

probably safer for library authors so you don't accidentally forget

* fix examples

* minor docs tweaks

* clarify how to convert handlers into services

* group methods in one impl block

* make sure merged `MethodRouter`s can access state

* fix docs link

* test merge with same state type

* Document how to access state from middleware

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

* Updates ECOSYSTEM with a new sample project (#1252)

* Avoid unhelpful compiler suggestion (#1251)

* fix docs typo

* document how library authors should access state

* Add `RequestParts::with_state`

* fix example

* apply suggestions from review

* add relevant changes to axum-extra and axum-core changelogs

* Add `route_service_with_tsr`

* fix trybuild expectations

* make sure `SpaRouter` works with routers that have state

* Change order of type params on FromRequest and RequestParts

* reverse order of `RequestParts::with_state` args to match type params

* Add `FromRef` trait (#1268)

* Add `FromRef` trait

* Remove unnecessary type params

* format

* fix docs link

* format examples

* Avoid unnecessary `MethodRouter`

* apply suggestions from review

Co-authored-by: Dani Pardo <[email protected]>
Co-authored-by: Jonas Platte <[email protected]>

axum-core/src/extract/mod.rs

@@ -12,9 +12,12 @@ use std::convert::Infallible;
 
 pub mod rejection;
 
+mod from_ref;
 mod request_parts;
 mod tuple;
 
+pub use self::from_ref::FromRef;
+
 /// Types that can be created from requests.
 ///
 /// See [`axum::extract`] for more details.
@@ -42,13 +45,15 @@ mod tuple;
 /// struct MyExtractor;
 ///
 /// #[async_trait]
-/// impl<B> FromRequest<B> for MyExtractor
+/// impl<S, B> FromRequest<S, B> for MyExtractor
 /// where
-///     B: Send, // required by `async_trait`
+///     // these bounds are required by `async_trait`
+///     B: Send,
+///     S: Send,
 /// {
 ///     type Rejection = http::StatusCode;
 ///
-///     async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
+///     async fn from_request(req: &mut RequestParts<S, B>) -> Result<Self, Self::Rejection> {
 ///         // ...
 ///         # unimplemented!()
 ///     }
@@ -60,20 +65,21 @@ mod tuple;
 /// [`http::Request<B>`]: http::Request
 /// [`axum::extract`]: https://docs.rs/axum/latest/axum/extract/index.html
 #[async_trait]
-pub trait FromRequest<B>: Sized {
+pub trait FromRequest<S, B>: Sized {
     /// If the extractor fails it'll use this "rejection" type. A rejection is
     /// a kind of error that can be converted into a response.
     type Rejection: IntoResponse;
 
     /// Perform the extraction.
-    async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection>;
+    async fn from_request(req: &mut RequestParts<S, B>) -> Result<Self, Self::Rejection>;
 }
 
 /// The type used with [`FromRequest`] to extract data from requests.
 ///
 /// Has several convenience methods for getting owned parts of the request.
 #[derive(Debug)]
-pub struct RequestParts<B> {
+pub struct RequestParts<S, B> {
+    state: S,
     method: Method,
     uri: Uri,
     version: Version,
@@ -82,8 +88,8 @@ pub struct RequestParts<B> {
     body: Option<B>,
 }
 
-impl<B> RequestParts<B> {
-    /// Create a new `RequestParts`.
+impl<B> RequestParts<(), B> {
+    /// Create a new `RequestParts` without any state.
     ///
     /// You generally shouldn't need to construct this type yourself, unless
     /// using extractors outside of axum for example to implement a
@@ -91,6 +97,19 @@ impl<B> RequestParts<B> {
     ///
     /// [`tower::Service`]: https://docs.rs/tower/lastest/tower/trait.Service.html
     pub fn new(req: Request<B>) -> Self {
+        Self::with_state((), req)
+    }
+}
+
+impl<S, B> RequestParts<S, B> {
+    /// Create a new `RequestParts` with the given state.
+    ///
+    /// You generally shouldn't need to construct this type yourself, unless
+    /// using extractors outside of axum for example to implement a
+    /// [`tower::Service`].
+    ///
+    /// [`tower::Service`]: https://docs.rs/tower/lastest/tower/trait.Service.html
+    pub fn with_state(state: S, req: Request<B>) -> Self {
         let (
             http::request::Parts {
                 method,
@@ -104,6 +123,7 @@ impl<B> RequestParts<B> {
         ) = req.into_parts();
 
         RequestParts {
+            state,
             method,
             uri,
             version,
@@ -130,10 +150,14 @@ impl<B> RequestParts<B> {
     /// use http::{Method, Uri};
     ///
     /// #[async_trait]
-    /// impl<B: Send> FromRequest<B> for MyExtractor {
+    /// impl<S, B> FromRequest<S, B> for MyExtractor
+    /// where
+    ///     B: Send,
+    ///     S: Send,
+    /// {
     ///     type Rejection = Infallible;
     ///
-    ///     async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Infallible> {
+    ///     async fn from_request(req: &mut RequestParts<S, B>) -> Result<Self, Infallible> {
     ///         let method = req.extract::<Method>().await?;
     ///         let path = req.extract::<Uri>().await?.path().to_owned();
     ///
@@ -141,7 +165,10 @@ impl<B> RequestParts<B> {
     ///     }
     /// }
     /// ```
-    pub async fn extract<E: FromRequest<B>>(&mut self) -> Result<E, E::Rejection> {
+    pub async fn extract<E>(&mut self) -> Result<E, E::Rejection>
+    where
+        E: FromRequest<S, B>,
+    {
         E::from_request(self).await
     }
 
@@ -153,6 +180,7 @@ impl<B> RequestParts<B> {
     /// [`take_body`]: RequestParts::take_body
     pub fn try_into_request(self) -> Result<Request<B>, BodyAlreadyExtracted> {
         let Self {
+            state: _,
             method,
             uri,
             version,
@@ -245,30 +273,37 @@ impl<B> RequestParts<B> {
     pub fn take_body(&mut self) -> Option<B> {
         self.body.take()
     }
+
+    /// Get a reference to the state.
+    pub fn state(&self) -> &S {
+        &self.state
+    }
 }
 
 #[async_trait]
-impl<T, B> FromRequest<B> for Option<T>
+impl<S, T, B> FromRequest<S, B> for Option<T>
 where
-    T: FromRequest<B>,
+    T: FromRequest<S, B>,
     B: Send,
+    S: Send,
 {
     type Rejection = Infallible;
 
-    async fn from_request(req: &mut RequestParts<B>) -> Result<Option<T>, Self::Rejection> {
+    async fn from_request(req: &mut RequestParts<S, B>) -> Result<Option<T>, Self::Rejection> {
         Ok(T::from_request(req).await.ok())
     }
 }
 
 #[async_trait]
-impl<T, B> FromRequest<B> for Result<T, T::Rejection>
+impl<S, T, B> FromRequest<S, B> for Result<T, T::Rejection>
 where
-    T: FromRequest<B>,
+    T: FromRequest<S, B>,
     B: Send,
+    S: Send,
 {
     type Rejection = Infallible;
 
-    async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
+    async fn from_request(req: &mut RequestParts<S, B>) -> Result<Self, Self::Rejection> {
         Ok(T::from_request(req).await)
     }
 }