Support State with #[derive(FromRequest[Parts])] (#1391)

* Support `State` with `#[derive(FromRequest[Parts])]`

Fixes https://github.com/tokio-rs/axum/issues/1314

This makes it possible to extract things via `State` in
`#[derive(FromRequet)]`:

```rust
struct Foo {
    state: State<AppState>,
}
```

The state can also be inferred in a lot of cases so you only need to
write:

```rust
struct Foo {
    // since we're using `State<AppState>` we know the state has to be
    // `AppState`
    state: State<AppState>,
}
```

Same for

```rust
struct Foo {
    #[from_request(via(State))]
    state: AppState,
}
```

And

```rust
struct AppState {}
```

I think I've covered all the edge cases but there are (unsurprisingly) a
few.

* make sure things can be combined with other extractors

* main functions in ui tests don't need to be async

* Add test for multiple identicaly state types

* Add failing test for multiple states

axum-macros/tests/from_request/fail/state_infer_multiple_different_types.rs

@@ -0,0 +1,22 @@
+use axum_macros::FromRequest;
+use axum::extract::State;
+
+#[derive(FromRequest)]
+struct Extractor {
+    inner_state: State<AppState>,
+    other_state: State<OtherState>,
+}
+
+#[derive(Clone)]
+struct AppState {}
+
+#[derive(Clone)]
+struct OtherState {}
+
+fn assert_from_request()
+where
+    Extractor: axum::extract::FromRequest<AppState, axum::body::Body, Rejection = axum::response::Response>,
+{
+}
+
+fn main() {}

axum-macros/tests/from_request/fail/state_infer_multiple_different_types.stderr

@@ -0,0 +1,23 @@
+error[E0277]: the trait bound `AppState: FromRef<S>` is not satisfied
+ --> tests/from_request/fail/state_infer_multiple_different_types.rs:6:18
+  |
+6 |     inner_state: State<AppState>,
+  |                  ^^^^^ the trait `FromRef<S>` is not implemented for `AppState`
+  |
+  = note: required because of the requirements on the impl of `FromRequestParts<S>` for `State<AppState>`
+help: consider extending the `where` clause, but there might be an alternative better way to express this requirement
+  |
+4 | #[derive(FromRequest, AppState: FromRef<S>)]
+  |                     ++++++++++++++++++++++
+
+error[E0277]: the trait bound `OtherState: FromRef<S>` is not satisfied
+ --> tests/from_request/fail/state_infer_multiple_different_types.rs:7:18
+  |
+7 |     other_state: State<OtherState>,
+  |                  ^^^^^ the trait `FromRef<S>` is not implemented for `OtherState`
+  |
+  = note: required because of the requirements on the impl of `FromRequestParts<S>` for `State<OtherState>`
+help: consider extending the `where` clause, but there might be an alternative better way to express this requirement
+  |
+4 | #[derive(FromRequest, OtherState: FromRef<S>)]
+  |                     ++++++++++++++++++++++++

axum-macros/tests/from_request/fail/unknown_attr_container.stderr

@@ -1,4 +1,4 @@
-error: expected `via` or `rejection`
+error: expected one of: `via`, `rejection`, `state`
  --> tests/from_request/fail/unknown_attr_container.rs:4:16
   |
 4 | #[from_request(foo)]

axum-macros/tests/from_request/pass/state_cookie.rs

@@ -0,0 +1,27 @@
+use axum_macros::FromRequest;
+use axum::extract::FromRef;
+use axum_extra::extract::cookie::{PrivateCookieJar, Key};
+
+#[derive(FromRequest)]
+#[from_request(state(AppState))]
+struct Extractor {
+    cookies: PrivateCookieJar,
+}
+
+struct AppState {
+    key: Key,
+}
+
+impl FromRef<AppState> for Key {
+    fn from_ref(input: &AppState) -> Self {
+        input.key.clone()
+    }
+}
+
+fn assert_from_request()
+where
+    Extractor: axum::extract::FromRequest<AppState, axum::body::Body, Rejection = axum::response::Response>,
+{
+}
+
+fn main() {}

axum-macros/tests/from_request/pass/state_enum_via.rs

@@ -0,0 +1,34 @@
+use axum::{
+    extract::{State, FromRef},
+    routing::get,
+    Router,
+};
+use axum_macros::FromRequest;
+
+fn main() {
+    let _: Router<AppState> = Router::with_state(AppState::default())
+        .route("/a", get(|_: AppState| async {}))
+        .route("/b", get(|_: InnerState| async {}));
+}
+
+#[derive(Clone, FromRequest)]
+#[from_request(via(State))]
+enum AppState {
+    One,
+}
+
+impl Default for AppState {
+    fn default() -> AppState {
+        Self::One
+    }
+}
+
+#[derive(FromRequest)]
+#[from_request(via(State), state(AppState))]
+enum InnerState {}
+
+impl FromRef<AppState> for InnerState {
+    fn from_ref(_: &AppState) -> Self {
+        todo!(":shrug:")
+    }
+}

axum-macros/tests/from_request/pass/state_enum_via_parts.rs

@@ -0,0 +1,35 @@
+use axum::{
+    extract::{State, FromRef},
+    routing::get,
+    Router,
+};
+use axum_macros::FromRequestParts;
+
+fn main() {
+    let _: Router<AppState> = Router::with_state(AppState::default())
+        .route("/a", get(|_: AppState| async {}))
+        .route("/b", get(|_: InnerState| async {}))
+        .route("/c", get(|_: AppState, _: InnerState| async {}));
+}
+
+#[derive(Clone, FromRequestParts)]
+#[from_request(via(State))]
+enum AppState {
+    One,
+}
+
+impl Default for AppState {
+    fn default() -> AppState {
+        Self::One
+    }
+}
+
+#[derive(FromRequestParts)]
+#[from_request(via(State), state(AppState))]
+enum InnerState {}
+
+impl FromRef<AppState> for InnerState {
+    fn from_ref(_: &AppState) -> Self {
+        todo!(":shrug:")
+    }
+}

axum-macros/tests/from_request/pass/state_explicit.rs

@@ -0,0 +1,44 @@
+use axum_macros::FromRequest;
+use axum::{
+    extract::{FromRef, State},
+    Router,
+    routing::get,
+};
+
+fn main() {
+    let _: Router<AppState> = Router::with_state(AppState::default())
+        .route("/b", get(|_: Extractor| async {}));
+}
+
+#[derive(FromRequest)]
+#[from_request(state(AppState))]
+struct Extractor {
+    app_state: State<AppState>,
+    one: State<One>,
+    two: State<Two>,
+    other_extractor: String,
+}
+
+#[derive(Clone, Default)]
+struct AppState {
+    one: One,
+    two: Two,
+}
+
+#[derive(Clone, Default)]
+struct One {}
+
+impl FromRef<AppState> for One {
+    fn from_ref(input: &AppState) -> Self {
+        input.one.clone()
+    }
+}
+
+#[derive(Clone, Default)]
+struct Two {}
+
+impl FromRef<AppState> for Two {
+    fn from_ref(input: &AppState) -> Self {
+        input.two.clone()
+    }
+}

axum-macros/tests/from_request/pass/state_explicit_parts.rs

@@ -0,0 +1,33 @@
+use axum_macros::FromRequestParts;
+use axum::{
+    extract::{FromRef, State, Query},
+    Router,
+    routing::get,
+};
+use std::collections::HashMap;
+
+fn main() {
+    let _: Router<AppState> = Router::with_state(AppState::default())
+        .route("/b", get(|_: Extractor| async {}));
+}
+
+#[derive(FromRequestParts)]
+#[from_request(state(AppState))]
+struct Extractor {
+    inner_state: State<InnerState>,
+    other: Query<HashMap<String, String>>,
+}
+
+#[derive(Default)]
+struct AppState {
+    inner: InnerState,
+}
+
+#[derive(Clone, Default)]
+struct InnerState {}
+
+impl FromRef<AppState> for InnerState {
+    fn from_ref(input: &AppState) -> Self {
+        input.inner.clone()
+    }
+}

axum-macros/tests/from_request/pass/state_field_explicit.rs

@@ -0,0 +1,34 @@
+use axum::{
+    extract::{State, FromRef},
+    routing::get,
+    Router,
+};
+use axum_macros::FromRequest;
+
+fn main() {
+    let _: Router<AppState> = Router::with_state(AppState::default())
+        .route("/", get(|_: Extractor| async {}));
+}
+
+#[derive(FromRequest)]
+#[from_request(state(AppState))]
+struct Extractor {
+    #[from_request(via(State))]
+    state: AppState,
+    #[from_request(via(State))]
+    inner: InnerState,
+}
+
+#[derive(Clone, Default)]
+struct AppState {
+    inner: InnerState,
+}
+
+#[derive(Clone, Default)]
+struct InnerState {}
+
+impl FromRef<AppState> for InnerState {
+    fn from_ref(input: &AppState) -> Self {
+        input.inner.clone()
+    }
+}

axum-macros/tests/from_request/pass/state_field_infer.rs

@@ -0,0 +1,20 @@
+use axum::{
+    extract::State,
+    routing::get,
+    Router,
+};
+use axum_macros::FromRequest;
+
+fn main() {
+    let _: Router<AppState> = Router::with_state(AppState::default())
+        .route("/", get(|_: Extractor| async {}));
+}
+
+#[derive(FromRequest)]
+struct Extractor {
+    #[from_request(via(State))]
+    state: AppState,
+}
+
+#[derive(Clone, Default)]
+struct AppState {}

axum-macros/tests/from_request/pass/state_infer.rs

@@ -0,0 +1,18 @@
+use axum_macros::FromRequest;
+use axum::extract::State;
+
+#[derive(FromRequest)]
+struct Extractor {
+    inner_state: State<AppState>,
+}
+
+#[derive(Clone)]
+struct AppState {}
+
+fn assert_from_request()
+where
+    Extractor: axum::extract::FromRequest<AppState, axum::body::Body, Rejection = axum::response::Response>,
+{
+}
+
+fn main() {}

axum-macros/tests/from_request/pass/state_infer_multiple.rs

@@ -0,0 +1,19 @@
+use axum_macros::FromRequest;
+use axum::extract::State;
+
+#[derive(FromRequest)]
+struct Extractor {
+    inner_state: State<AppState>,
+    also_inner_state: State<AppState>,
+}
+
+#[derive(Clone)]
+struct AppState {}
+
+fn assert_from_request()
+where
+    Extractor: axum::extract::FromRequest<AppState, axum::body::Body, Rejection = axum::response::Response>,
+{
+}
+
+fn main() {}

axum-macros/tests/from_request/pass/state_infer_parts.rs

@@ -0,0 +1,18 @@
+use axum_macros::FromRequestParts;
+use axum::extract::State;
+
+#[derive(FromRequestParts)]
+struct Extractor {
+    inner_state: State<AppState>,
+}
+
+#[derive(Clone)]
+struct AppState {}
+
+fn assert_from_request()
+where
+    Extractor: axum::extract::FromRequestParts<AppState, Rejection = axum::response::Response>,
+{
+}
+
+fn main() {}

axum-macros/tests/from_request/pass/state_via.rs

@@ -0,0 +1,28 @@
+use axum::{
+    extract::{FromRef, State},
+    routing::get,
+    Router,
+};
+use axum_macros::FromRequest;
+
+fn main() {
+    let _: Router<AppState> = Router::with_state(AppState::default())
+        .route("/b", get(|_: (), _: AppState| async {}))
+        .route("/c", get(|_: (), _: InnerState| async {}));
+}
+
+#[derive(Clone, Default, FromRequest)]
+#[from_request(via(State), state(AppState))]
+struct AppState {
+    inner: InnerState,
+}
+
+#[derive(Clone, Default, FromRequest)]
+#[from_request(via(State), state(AppState))]
+struct InnerState {}
+
+impl FromRef<AppState> for InnerState {
+    fn from_ref(input: &AppState) -> Self {
+        input.inner.clone()
+    }
+}

axum-macros/tests/from_request/pass/state_via_infer.rs

@@ -0,0 +1,17 @@
+use axum::{
+    extract::State,
+    routing::get,
+    Router,
+};
+use axum_macros::FromRequest;
+
+fn main() {
+    let _: Router<AppState> = Router::with_state(AppState::default())
+        .route("/b", get(|_: AppState| async {}));
+}
+
+// if we're extract "via" `State<AppState>` and not specifying state
+// assume `AppState` is the state
+#[derive(Clone, Default, FromRequest)]
+#[from_request(via(State))]
+struct AppState {}

axum-macros/tests/from_request/pass/state_via_parts.rs

@@ -0,0 +1,29 @@
+use axum::{
+    extract::{FromRef, State},
+    routing::get,
+    Router,
+};
+use axum_macros::FromRequestParts;
+
+fn main() {
+    let _: Router<AppState> = Router::with_state(AppState::default())
+        .route("/a", get(|_: AppState, _: InnerState, _: String| async {}))
+        .route("/b", get(|_: AppState, _: String| async {}))
+        .route("/c", get(|_: InnerState, _: String| async {}));
+}
+
+#[derive(Clone, Default, FromRequestParts)]
+#[from_request(via(State))]
+struct AppState {
+    inner: InnerState,
+}
+
+#[derive(Clone, Default, FromRequestParts)]
+#[from_request(via(State), state(AppState))]
+struct InnerState {}
+
+impl FromRef<AppState> for InnerState {
+    fn from_ref(input: &AppState) -> Self {
+        input.inner.clone()
+    }
+}

axum-macros/tests/from_request/pass/state_with_rejection.rs

@@ -0,0 +1,36 @@
+use std::convert::Infallible;
+use axum::{
+    extract::State,
+    response::{IntoResponse, Response},
+    routing::get,
+    Router,
+};
+use axum_macros::FromRequest;
+
+fn main() {
+    let _: Router<AppState> =
+        Router::with_state(AppState::default()).route("/a", get(|_: Extractor| async {}));
+}
+
+#[derive(Clone, Default, FromRequest)]
+#[from_request(rejection(MyRejection))]
+struct Extractor {
+    state: State<AppState>,
+}
+
+#[derive(Clone, Default)]
+struct AppState {}
+
+struct MyRejection {}
+
+impl From<Infallible> for MyRejection {
+    fn from(err: Infallible) -> Self {
+        match err {}
+    }
+}
+
+impl IntoResponse for MyRejection {
+    fn into_response(self) -> Response {
+        ().into_response()
+    }
+}