Upgrade axum to 0.6.0-rc.1

crates/handlers/Cargo.toml

@@ -20,9 +20,9 @@ anyhow = "1.0.64"
 hyper = { version = "0.14.20", features = ["full"] }
 tower = "0.4.13"
 tower-http = { version = "0.3.4", features = ["cors"] }
-axum = "0.5.15"
+axum = "0.6.0-rc.1"
 axum-macros = "0.2.3"
-axum-extra = { version = "0.3.7", features = ["cookie-private"] }
+axum-extra = { version = "0.4.0-rc.1", features = ["cookie-private"] }
 
 # Emails
 lettre = { version = "0.10.1", default-features = false, features = ["builder"] }

crates/handlers/src/app_state.rs

@@ -0,0 +1,85 @@
+// Copyright 2022 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use axum::extract::FromRef;
+use mas_email::Mailer;
+use mas_keystore::{Encrypter, Keystore};
+use mas_policy::PolicyFactory;
+use mas_router::UrlBuilder;
+use mas_templates::Templates;
+use sqlx::PgPool;
+
+use crate::MatrixHomeserver;
+
+#[derive(Clone)]
+pub struct AppState {
+    pub pool: PgPool,
+    pub templates: Templates,
+    pub key_store: Keystore,
+    pub encrypter: Encrypter,
+    pub url_builder: UrlBuilder,
+    pub mailer: Mailer,
+    pub homeserver: MatrixHomeserver,
+    pub policy_factory: Arc<PolicyFactory>,
+}
+
+impl FromRef<AppState> for PgPool {
+    fn from_ref(input: &AppState) -> Self {
+        input.pool.clone()
+    }
+}
+
+impl FromRef<AppState> for Templates {
+    fn from_ref(input: &AppState) -> Self {
+        input.templates.clone()
+    }
+}
+
+impl FromRef<AppState> for Keystore {
+    fn from_ref(input: &AppState) -> Self {
+        input.key_store.clone()
+    }
+}
+
+impl FromRef<AppState> for Encrypter {
+    fn from_ref(input: &AppState) -> Self {
+        input.encrypter.clone()
+    }
+}
+
+impl FromRef<AppState> for UrlBuilder {
+    fn from_ref(input: &AppState) -> Self {
+        input.url_builder.clone()
+    }
+}
+
+impl FromRef<AppState> for Mailer {
+    fn from_ref(input: &AppState) -> Self {
+        input.mailer.clone()
+    }
+}
+
+impl FromRef<AppState> for MatrixHomeserver {
+    fn from_ref(input: &AppState) -> Self {
+        input.homeserver.clone()
+    }
+}
+
+impl FromRef<AppState> for Arc<PolicyFactory> {
+    fn from_ref(input: &AppState) -> Self {
+        input.policy_factory.clone()
+    }
+}

crates/handlers/src/compat/login.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use axum::{response::IntoResponse, Extension, Json};
+use axum::{extract::State, response::IntoResponse, Json};
 use chrono::{Duration, Utc};
 use hyper::StatusCode;
 use mas_data_model::{CompatSession, CompatSsoLoginState, Device, TokenType};
@@ -197,8 +197,8 @@ impl IntoResponse for RouteError {
 
 #[tracing::instrument(skip_all, err)]
 pub(crate) async fn post(
-    Extension(pool): Extension<PgPool>,
-    Extension(homeserver): Extension<MatrixHomeserver>,
+    State(pool): State<PgPool>,
+    State(homeserver): State<MatrixHomeserver>,
     Json(input): Json<RequestBody>,
 ) -> Result<impl IntoResponse, RouteError> {
     let mut txn = pool.begin().await?;

crates/handlers/src/compat/login_sso_complete.rs

@@ -16,9 +16,8 @@
 use std::collections::HashMap;
 
 use axum::{
-    extract::{Form, Path, Query},
+    extract::{Form, Path, Query, State},
     response::{Html, IntoResponse, Redirect, Response},
-    Extension,
 };
 use axum_extra::extract::PrivateCookieJar;
 use chrono::{Duration, Utc};
@@ -50,8 +49,8 @@ pub struct Params {
 }
 
 pub async fn get(
-    Extension(pool): Extension<PgPool>,
-    Extension(templates): Extension<Templates>,
+    State(pool): State<PgPool>,
+    State(templates): State<Templates>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Path(id): Path<i64>,
     Query(params): Query<Params>,
@@ -114,12 +113,12 @@ pub async fn get(
 }
 
 pub async fn post(
-    Extension(pool): Extension<PgPool>,
-    Extension(templates): Extension<Templates>,
+    State(pool): State<PgPool>,
+    State(templates): State<Templates>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Path(id): Path<i64>,
-    Form(form): Form<ProtectedForm<()>>,
     Query(params): Query<Params>,
+    Form(form): Form<ProtectedForm<()>>,
 ) -> Result<Response, FancyError> {
     let mut txn = pool.begin().await?;
 

crates/handlers/src/compat/login_sso_redirect.rs

@@ -13,7 +13,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use axum::{extract::Query, response::IntoResponse, Extension};
+use axum::{
+    extract::{Query, State},
+    response::IntoResponse,
+};
 use hyper::StatusCode;
 use mas_router::{CompatLoginSsoAction, CompatLoginSsoComplete, UrlBuilder};
 use mas_storage::compat::insert_compat_sso_login;
@@ -63,8 +66,8 @@ impl IntoResponse for RouteError {
 
 #[tracing::instrument(skip(pool, url_builder), err)]
 pub async fn get(
-    Extension(pool): Extension<PgPool>,
-    Extension(url_builder): Extension<UrlBuilder>,
+    State(pool): State<PgPool>,
+    State(url_builder): State<UrlBuilder>,
     Query(params): Query<Params>,
 ) -> Result<impl IntoResponse, RouteError> {
     // Check the redirectUrl parameter

crates/handlers/src/compat/logout.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use axum::{response::IntoResponse, Extension, Json, TypedHeader};
+use axum::{extract::State, response::IntoResponse, Json, TypedHeader};
 use headers::{authorization::Bearer, Authorization};
 use hyper::StatusCode;
 use mas_data_model::{TokenFormatError, TokenType};
@@ -64,7 +64,7 @@ impl From<TokenFormatError> for RouteError {
 }
 
 pub(crate) async fn post(
-    Extension(pool): Extension<PgPool>,
+    State(pool): State<PgPool>,
     maybe_authorization: Option<TypedHeader<Authorization<Bearer>>>,
 ) -> Result<impl IntoResponse, RouteError> {
     let mut conn = pool.acquire().await?;

crates/handlers/src/compat/refresh.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use axum::{response::IntoResponse, Extension, Json};
+use axum::{extract::State, response::IntoResponse, Json};
 use chrono::Duration;
 use hyper::StatusCode;
 use mas_data_model::{TokenFormatError, TokenType};
@@ -96,7 +96,7 @@ pub struct ResponseBody {
 }
 
 pub(crate) async fn post(
-    Extension(pool): Extension<PgPool>,
+    State(pool): State<PgPool>,
     Json(input): Json<RequestBody>,
 ) -> Result<impl IntoResponse, RouteError> {
     let mut txn = pool.begin().await?;

crates/handlers/src/health.rs

@@ -12,12 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use axum::{extract::Extension, response::IntoResponse};
+use axum::{extract::State, response::IntoResponse};
 use mas_axum_utils::FancyError;
 use sqlx::PgPool;
 use tracing::{info_span, Instrument};
 
-pub async fn get(Extension(pool): Extension<PgPool>) -> Result<impl IntoResponse, FancyError> {
+pub async fn get(State(pool): State<PgPool>) -> Result<impl IntoResponse, FancyError> {
     let mut conn = pool.acquire().await?;
 
     sqlx::query("SELECT $1")
@@ -38,7 +38,9 @@ mod tests {
 
     #[sqlx::test(migrator = "mas_storage::MIGRATOR")]
     async fn test_get_health(pool: PgPool) -> Result<(), anyhow::Error> {
-        let app = crate::test_router(&pool).await?;
+        let state = crate::test_state(pool).await?;
+        let app = crate::api_router(state);
+
         let request = Request::builder().uri("/health").body(Body::empty())?;
 
         let response = app.oneshot(request).await?;

crates/handlers/src/lib.rs

@@ -23,7 +23,7 @@ use std::{convert::Infallible, sync::Arc, time::Duration};
 
 use axum::{
     body::HttpBody,
-    extract::Extension,
+    extract::FromRef,
     response::{Html, IntoResponse},
     routing::{get, on, post, MethodFilter},
     Router,
@@ -37,9 +37,10 @@ use mas_policy::PolicyFactory;
 use mas_router::{Route, UrlBuilder};
 use mas_templates::{ErrorContext, Templates};
 use sqlx::PgPool;
-use tower::util::ThenLayer;
+use tower::util::AndThenLayer;
 use tower_http::cors::{Any, CorsLayer};
 
+mod app_state;
 mod compat;
 mod health;
 mod oauth2;
@@ -47,30 +48,24 @@ mod views;
 
 pub use compat::MatrixHomeserver;
 
+pub use self::app_state::AppState;
+
 #[must_use]
-#[allow(
-    clippy::too_many_lines,
-    clippy::missing_panics_doc,
-    clippy::too_many_arguments,
-    clippy::trait_duplication_in_bounds
-)]
-pub fn router<B>(
-    pool: &PgPool,
-    templates: &Templates,
-    key_store: &Keystore,
-    encrypter: &Encrypter,
-    mailer: &Mailer,
-    url_builder: &UrlBuilder,
-    homeserver: &MatrixHomeserver,
-    policy_factory: &Arc<PolicyFactory>,
-) -> Router<B>
+#[allow(clippy::trait_duplication_in_bounds)]
+pub fn api_router<S, B>(state: Arc<S>) -> Router<S, B>
 where
     B: HttpBody + Send + 'static,
     <B as HttpBody>::Data: Send,
     <B as HttpBody>::Error: std::error::Error + Send + Sync,
+    S: Send + Sync + 'static,
+    Keystore: FromRef<S>,
+    UrlBuilder: FromRef<S>,
+    Arc<PolicyFactory>: FromRef<S>,
+    PgPool: FromRef<S>,
+    Encrypter: FromRef<S>,
 {
     // All those routes are API-like, with a common CORS layer
-    let api_router = Router::new()
+    Router::with_state_arc(state)
         .route(
             mas_router::ChangePasswordDiscovery::route(),
             get(|| async { mas_router::AccountPassword.go() }),
@@ -118,9 +113,21 @@ where
                     CONTENT_TYPE,
                 ])
                 .max_age(Duration::from_secs(60 * 60)),
-        );
-
-    let compat_router = Router::new()
+        )
+}
+#[must_use]
+#[allow(clippy::trait_duplication_in_bounds)]
+pub fn compat_router<S, B>(state: Arc<S>) -> Router<S, B>
+where
+    B: HttpBody + Send + 'static,
+    <B as HttpBody>::Data: Send,
+    <B as HttpBody>::Error: std::error::Error + Send + Sync,
+    S: Send + Sync + 'static,
+    UrlBuilder: FromRef<S>,
+    PgPool: FromRef<S>,
+    MatrixHomeserver: FromRef<S>,
+{
+    Router::with_state_arc(state)
         .route(
             mas_router::CompatLogin::route(),
             get(self::compat::login::get).post(self::compat::login::post),
@@ -146,106 +153,131 @@ where
                     HeaderName::from_static("x-requested-with"),
                 ])
                 .max_age(Duration::from_secs(60 * 60)),
-        );
+        )
+}
 
-    let human_router = {
-        let templates = templates.clone();
-        Router::new()
-            .route(mas_router::Index::route(), get(self::views::index::get))
-            .route(mas_router::Healthcheck::route(), get(self::health::get))
-            .route(
-                mas_router::Login::route(),
-                get(self::views::login::get).post(self::views::login::post),
-            )
-            .route(mas_router::Logout::route(), post(self::views::logout::post))
-            .route(
-                mas_router::Reauth::route(),
-                get(self::views::reauth::get).post(self::views::reauth::post),
-            )
-            .route(
-                mas_router::Register::route(),
-                get(self::views::register::get).post(self::views::register::post),
-            )
-            .route(mas_router::Account::route(), get(self::views::account::get))
-            .route(
-                mas_router::AccountPassword::route(),
-                get(self::views::account::password::get).post(self::views::account::password::post),
-            )
-            .route(
-                mas_router::AccountEmails::route(),
-                get(self::views::account::emails::get).post(self::views::account::emails::post),
-            )
-            .route(
-                mas_router::AccountVerifyEmail::route(),
-                get(self::views::account::emails::verify::get)
-                    .post(self::views::account::emails::verify::post),
-            )
-            .route(
-                mas_router::AccountAddEmail::route(),
-                get(self::views::account::emails::add::get)
-                    .post(self::views::account::emails::add::post),
-            )
-            .route(
-                mas_router::OAuth2AuthorizationEndpoint::route(),
-                get(self::oauth2::authorization::get),
-            )
-            .route(
-                mas_router::ContinueAuthorizationGrant::route(),
-                get(self::oauth2::authorization::complete::get),
-            )
-            .route(
-                mas_router::Consent::route(),
-                get(self::oauth2::consent::get).post(self::oauth2::consent::post),
-            )
-            .route(
-                mas_router::CompatLoginSsoRedirect::route(),
-                get(self::compat::login_sso_redirect::get),
-            )
-            .route(
-                mas_router::CompatLoginSsoRedirectIdp::route(),
-                get(self::compat::login_sso_redirect::get),
-            )
-            .route(
-                mas_router::CompatLoginSsoComplete::route(),
-                get(self::compat::login_sso_complete::get)
-                    .post(self::compat::login_sso_complete::post),
-            )
-            .layer(ThenLayer::new(
-                move |result: Result<axum::response::Response, Infallible>| async move {
-                    let response = result.unwrap();
-
-                    if response.status().is_server_error() {
-                        // Error responses should have an ErrorContext attached to them
-                        let ext = response.extensions().get::<ErrorContext>();
-                        if let Some(ctx) = ext {
-                            if let Ok(res) = templates.render_error(ctx).await {
-                                let (mut parts, _original_body) = response.into_parts();
-                                parts.headers.remove(CONTENT_TYPE);
-                                return Ok((parts, Html(res)).into_response());
-                            }
+#[must_use]
+#[allow(clippy::trait_duplication_in_bounds)]
+pub fn human_router<S, B>(state: Arc<S>) -> Router<S, B>
+where
+    B: HttpBody + Send + 'static,
+    <B as HttpBody>::Data: Send,
+    <B as HttpBody>::Error: std::error::Error + Send + Sync,
+    S: Send + Sync + 'static,
+    UrlBuilder: FromRef<S>,
+    Arc<PolicyFactory>: FromRef<S>,
+    PgPool: FromRef<S>,
+    Encrypter: FromRef<S>,
+    Templates: FromRef<S>,
+    Mailer: FromRef<S>,
+{
+    let templates = Templates::from_ref(&state);
+    Router::with_state_arc(state)
+        .route(mas_router::Index::route(), get(self::views::index::get))
+        .route(mas_router::Healthcheck::route(), get(self::health::get))
+        .route(
+            mas_router::Login::route(),
+            get(self::views::login::get).post(self::views::login::post),
+        )
+        .route(mas_router::Logout::route(), post(self::views::logout::post))
+        .route(
+            mas_router::Reauth::route(),
+            get(self::views::reauth::get).post(self::views::reauth::post),
+        )
+        .route(
+            mas_router::Register::route(),
+            get(self::views::register::get).post(self::views::register::post),
+        )
+        .route(mas_router::Account::route(), get(self::views::account::get))
+        .route(
+            mas_router::AccountPassword::route(),
+            get(self::views::account::password::get).post(self::views::account::password::post),
+        )
+        .route(
+            mas_router::AccountEmails::route(),
+            get(self::views::account::emails::get).post(self::views::account::emails::post),
+        )
+        .route(
+            mas_router::AccountVerifyEmail::route(),
+            get(self::views::account::emails::verify::get)
+                .post(self::views::account::emails::verify::post),
+        )
+        .route(
+            mas_router::AccountAddEmail::route(),
+            get(self::views::account::emails::add::get)
+                .post(self::views::account::emails::add::post),
+        )
+        .route(
+            mas_router::OAuth2AuthorizationEndpoint::route(),
+            get(self::oauth2::authorization::get),
+        )
+        .route(
+            mas_router::ContinueAuthorizationGrant::route(),
+            get(self::oauth2::authorization::complete::get),
+        )
+        .route(
+            mas_router::Consent::route(),
+            get(self::oauth2::consent::get).post(self::oauth2::consent::post),
+        )
+        .route(
+            mas_router::CompatLoginSsoRedirect::route(),
+            get(self::compat::login_sso_redirect::get),
+        )
+        .route(
+            mas_router::CompatLoginSsoRedirectIdp::route(),
+            get(self::compat::login_sso_redirect::get),
+        )
+        .route(
+            mas_router::CompatLoginSsoComplete::route(),
+            get(self::compat::login_sso_complete::get).post(self::compat::login_sso_complete::post),
+        )
+        .layer(AndThenLayer::new(
+            move |response: axum::response::Response| async move {
+                if response.status().is_server_error() {
+                    // Error responses should have an ErrorContext attached to them
+                    let ext = response.extensions().get::<ErrorContext>();
+                    if let Some(ctx) = ext {
+                        if let Ok(res) = templates.render_error(ctx).await {
+                            let (mut parts, _original_body) = response.into_parts();
+                            parts.headers.remove(CONTENT_TYPE);
+                            return Ok((parts, Html(res)).into_response());
                         }
                     }
+                }
 
-                    Ok(response)
-                },
-            ))
-    };
+                Ok::<_, Infallible>(response)
+            },
+        ))
+}
 
-    human_router
-        .merge(api_router)
-        .merge(compat_router)
-        .layer(Extension(pool.clone()))
-        .layer(Extension(templates.clone()))
-        .layer(Extension(key_store.clone()))
-        .layer(Extension(encrypter.clone()))
-        .layer(Extension(url_builder.clone()))
-        .layer(Extension(mailer.clone()))
-        .layer(Extension(homeserver.clone()))
-        .layer(Extension(policy_factory.clone()))
+#[must_use]
+#[allow(clippy::trait_duplication_in_bounds)]
+pub fn router<S, B>(state: S) -> Router<S, B>
+where
+    B: HttpBody + Send + 'static,
+    <B as HttpBody>::Data: Send,
+    <B as HttpBody>::Error: std::error::Error + Send + Sync,
+    S: Send + Sync + 'static,
+    Keystore: FromRef<S>,
+    UrlBuilder: FromRef<S>,
+    Arc<PolicyFactory>: FromRef<S>,
+    PgPool: FromRef<S>,
+    Encrypter: FromRef<S>,
+    Templates: FromRef<S>,
+    Mailer: FromRef<S>,
+    MatrixHomeserver: FromRef<S>,
+{
+    let state = Arc::new(state);
+
+    let api_router = api_router(state.clone());
+    let compat_router = compat_router(state.clone());
+    let human_router = human_router(state);
+
+    human_router.merge(api_router).merge(compat_router)
 }
 
 #[cfg(test)]
-async fn test_router(pool: &PgPool) -> Result<Router, anyhow::Error> {
+async fn test_state(pool: PgPool) -> Result<Arc<AppState>, anyhow::Error> {
     use mas_email::MailTransport;
 
     let templates = Templates::load(None, true).await?;
@@ -265,14 +297,14 @@ async fn test_router(pool: &PgPool) -> Result<Router, anyhow::Error> {
     let policy_factory = PolicyFactory::load_default(serde_json::json!({})).await?;
     let policy_factory = Arc::new(policy_factory);
 
-    Ok(router(
+    Ok(Arc::new(AppState {
         pool,
-        &templates,
-        &key_store,
-        &encrypter,
-        &mailer,
-        &url_builder,
-        &homeserver,
-        &policy_factory,
-    ))
+        templates,
+        key_store,
+        encrypter,
+        url_builder,
+        mailer,
+        homeserver,
+        policy_factory,
+    }))
 }

crates/handlers/src/oauth2/authorization/complete.rs

@@ -16,9 +16,8 @@ use std::sync::Arc;
 
 use anyhow::anyhow;
 use axum::{
-    extract::Path,
+    extract::{Path, State},
     response::{IntoResponse, Response},
-    Extension,
 };
 use axum_extra::extract::PrivateCookieJar;
 use hyper::StatusCode;
@@ -104,9 +103,9 @@ impl From<CallbackDestinationError> for RouteError {
 }
 
 pub(crate) async fn get(
-    Extension(policy_factory): Extension<Arc<PolicyFactory>>,
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(policy_factory): State<Arc<PolicyFactory>>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Path(grant_id): Path<i64>,
 ) -> Result<Response, RouteError> {

crates/handlers/src/oauth2/authorization/mod.rs

@@ -16,7 +16,7 @@ use std::sync::Arc;
 
 use anyhow::{anyhow, Context};
 use axum::{
-    extract::{Extension, Form},
+    extract::{Form, State},
     response::{IntoResponse, Response},
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -156,9 +156,9 @@ fn resolve_response_mode(
 
 #[allow(clippy::too_many_lines)]
 pub(crate) async fn get(
-    Extension(policy_factory): Extension<Arc<PolicyFactory>>,
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(policy_factory): State<Arc<PolicyFactory>>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Form(params): Form<Params>,
 ) -> Result<Response, RouteError> {

crates/handlers/src/oauth2/consent.rs

@@ -16,7 +16,7 @@ use std::sync::Arc;
 
 use anyhow::Context;
 use axum::{
-    extract::{Extension, Form, Path},
+    extract::{Form, Path, State},
     response::{Html, IntoResponse, Response},
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -50,9 +50,9 @@ impl IntoResponse for RouteError {
 }
 
 pub(crate) async fn get(
-    Extension(policy_factory): Extension<Arc<PolicyFactory>>,
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(policy_factory): State<Arc<PolicyFactory>>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Path(grant_id): Path<i64>,
 ) -> Result<Response, RouteError> {
@@ -112,8 +112,8 @@ pub(crate) async fn get(
 }
 
 pub(crate) async fn post(
-    Extension(policy_factory): Extension<Arc<PolicyFactory>>,
-    Extension(pool): Extension<PgPool>,
+    State(policy_factory): State<Arc<PolicyFactory>>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Path(grant_id): Path<i64>,
     Form(form): Form<ProtectedForm<()>>,

crates/handlers/src/oauth2/discovery.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use axum::{extract::Extension, response::IntoResponse, Json};
+use axum::{extract::State, response::IntoResponse, Json};
 use mas_iana::{
     jose::JsonWebSignatureAlg,
     oauth::{
@@ -30,8 +30,8 @@ use oauth2_types::{
 
 #[allow(clippy::too_many_lines)]
 pub(crate) async fn get(
-    Extension(key_store): Extension<Keystore>,
-    Extension(url_builder): Extension<UrlBuilder>,
+    State(key_store): State<Keystore>,
+    State(url_builder): State<UrlBuilder>,
 ) -> impl IntoResponse {
     // This is how clients can authenticate
     let client_auth_methods_supported = Some(vec![

crates/handlers/src/oauth2/introspection.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use axum::{extract::Extension, response::IntoResponse, Json};
+use axum::{extract::State, response::IntoResponse, Json};
 use hyper::StatusCode;
 use mas_axum_utils::client_authorization::{ClientAuthorization, CredentialsVerificationError};
 use mas_data_model::{TokenFormatError, TokenType};
@@ -154,8 +154,8 @@ const INACTIVE: IntrospectionResponse = IntrospectionResponse {
 
 #[tracing::instrument(skip_all, err)]
 pub(crate) async fn post(
-    Extension(pool): Extension<PgPool>,
-    Extension(encrypter): Extension<Encrypter>,
+    State(pool): State<PgPool>,
+    State(encrypter): State<Encrypter>,
     client_authorization: ClientAuthorization<IntrospectionRequest>,
 ) -> Result<impl IntoResponse, RouteError> {
     let mut conn = pool.acquire().await?;

crates/handlers/src/oauth2/keys.rs

@@ -12,10 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use axum::{extract::Extension, response::IntoResponse, Json};
+use axum::{extract::State, response::IntoResponse, Json};
 use mas_keystore::Keystore;
 
-pub(crate) async fn get(Extension(key_store): Extension<Keystore>) -> impl IntoResponse {
+pub(crate) async fn get(State(key_store): State<Keystore>) -> impl IntoResponse {
     let jwks = key_store.public_jwks();
     Json(jwks)
 }

crates/handlers/src/oauth2/registration.rs

@@ -14,7 +14,7 @@
 
 use std::sync::Arc;
 
-use axum::{response::IntoResponse, Extension, Json};
+use axum::{extract::State, response::IntoResponse, Json};
 use hyper::StatusCode;
 use mas_policy::{PolicyFactory, Violation};
 use mas_storage::oauth2::client::insert_client;
@@ -105,8 +105,8 @@ impl IntoResponse for RouteError {
 
 #[tracing::instrument(skip_all, err)]
 pub(crate) async fn post(
-    Extension(pool): Extension<PgPool>,
-    Extension(policy_factory): Extension<Arc<PolicyFactory>>,
+    State(pool): State<PgPool>,
+    State(policy_factory): State<Arc<PolicyFactory>>,
     Json(body): Json<ClientMetadata>,
 ) -> Result<impl IntoResponse, RouteError> {
     info!(?body, "Client registration");

crates/handlers/src/oauth2/token.rs

@@ -15,7 +15,7 @@
 use std::collections::HashMap;
 
 use anyhow::Context;
-use axum::{extract::Extension, response::IntoResponse, Json};
+use axum::{extract::State, response::IntoResponse, Json};
 use chrono::{DateTime, Duration, Utc};
 use data_encoding::BASE64URL_NOPAD;
 use headers::{CacheControl, HeaderMap, HeaderMapExt, Pragma};
@@ -188,11 +188,11 @@ impl From<JwtSignatureError> for RouteError {
 
 #[tracing::instrument(skip_all, err)]
 pub(crate) async fn post(
+    State(key_store): State<Keystore>,
+    State(url_builder): State<UrlBuilder>,
+    State(pool): State<PgPool>,
+    State(encrypter): State<Encrypter>,
     client_authorization: ClientAuthorization<AccessTokenRequest>,
-    Extension(key_store): Extension<Keystore>,
-    Extension(url_builder): Extension<UrlBuilder>,
-    Extension(pool): Extension<PgPool>,
-    Extension(encrypter): Extension<Encrypter>,
 ) -> Result<impl IntoResponse, RouteError> {
     let mut txn = pool.begin().await?;
 

crates/handlers/src/oauth2/userinfo.rs

@@ -14,7 +14,7 @@
 
 use anyhow::Context;
 use axum::{
-    extract::Extension,
+    extract::State,
     response::{IntoResponse, Response},
     Json,
 };
@@ -48,9 +48,9 @@ struct SignedUserInfo {
 }
 
 pub async fn get(
-    Extension(url_builder): Extension<UrlBuilder>,
-    Extension(pool): Extension<PgPool>,
-    Extension(key_store): Extension<Keystore>,
+    State(url_builder): State<UrlBuilder>,
+    State(pool): State<PgPool>,
+    State(key_store): State<Keystore>,
     user_authorization: UserAuthorization,
 ) -> Result<Response, FancyError> {
     // TODO: error handling

crates/handlers/src/oauth2/webfinger.rs

@@ -12,7 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use axum::{extract::Query, response::IntoResponse, Extension, Json, TypedHeader};
+use axum::{
+    extract::{Query, State},
+    response::IntoResponse,
+    Json, TypedHeader,
+};
 use headers::ContentType;
 use mas_router::UrlBuilder;
 use oauth2_types::webfinger::WebFingerResponse;
@@ -33,7 +37,7 @@ fn jrd() -> mime::Mime {
 
 pub(crate) async fn get(
     Query(params): Query<Params>,
-    Extension(url_builder): Extension<UrlBuilder>,
+    State(url_builder): State<UrlBuilder>,
 ) -> impl IntoResponse {
     // TODO: should we validate the subject?
     let subject = params.resource;

crates/handlers/src/views/account/emails/add.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use axum::{
-    extract::{Extension, Form, Query},
+    extract::{Form, Query, State},
     response::{Html, IntoResponse, Response},
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -38,8 +38,8 @@ pub struct EmailForm {
 }
 
 pub(crate) async fn get(
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
 ) -> Result<Response, FancyError> {
     let mut conn = pool.begin().await?;
@@ -66,8 +66,8 @@ pub(crate) async fn get(
 }
 
 pub(crate) async fn post(
-    Extension(pool): Extension<PgPool>,
-    Extension(mailer): Extension<Mailer>,
+    State(pool): State<PgPool>,
+    State(mailer): State<Mailer>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Query(query): Query<OptionalPostAuthAction>,
     Form(form): Form<ProtectedForm<EmailForm>>,

crates/handlers/src/views/account/emails/mod.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use axum::{
-    extract::{Extension, Form},
+    extract::{Form, State},
     response::{Html, IntoResponse, Response},
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -52,8 +52,8 @@ pub enum ManagementForm {
 }
 
 pub(crate) async fn get(
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
 ) -> Result<Response, FancyError> {
     let mut conn = pool.acquire().await?;
@@ -118,9 +118,9 @@ async fn start_email_verification(
 }
 
 pub(crate) async fn post(
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
-    Extension(mailer): Extension<Mailer>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
+    State(mailer): State<Mailer>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Form(form): Form<ProtectedForm<ManagementForm>>,
 ) -> Result<Response, FancyError> {

crates/handlers/src/views/account/emails/verify.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use axum::{
-    extract::{Extension, Form, Path, Query},
+    extract::{Form, Path, Query, State},
     response::{Html, IntoResponse, Response},
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -40,8 +40,8 @@ pub struct CodeForm {
 }
 
 pub(crate) async fn get(
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     Query(query): Query<OptionalPostAuthAction>,
     Path(id): Path<i64>,
     cookie_jar: PrivateCookieJar<Encrypter>,
@@ -78,7 +78,7 @@ pub(crate) async fn get(
 }
 
 pub(crate) async fn post(
-    Extension(pool): Extension<PgPool>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Query(query): Query<OptionalPostAuthAction>,
     Path(id): Path<i64>,

crates/handlers/src/views/account/mod.rs

@@ -16,7 +16,7 @@ pub mod emails;
 pub mod password;
 
 use axum::{
-    extract::Extension,
+    extract::State,
     response::{Html, IntoResponse, Response},
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -28,8 +28,8 @@ use mas_templates::{AccountContext, TemplateContext, Templates};
 use sqlx::PgPool;
 
 pub(crate) async fn get(
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
 ) -> Result<Response, FancyError> {
     let mut conn = pool.acquire().await?;

crates/handlers/src/views/account/password.rs

@@ -14,7 +14,7 @@
 
 use argon2::Argon2;
 use axum::{
-    extract::{Extension, Form},
+    extract::{Form, State},
     response::{Html, IntoResponse, Response},
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -41,8 +41,8 @@ pub struct ChangeForm {
 }
 
 pub(crate) async fn get(
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
 ) -> Result<Response, FancyError> {
     let mut conn = pool.acquire().await?;
@@ -76,8 +76,8 @@ async fn render(
 }
 
 pub(crate) async fn post(
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Form(form): Form<ProtectedForm<ChangeForm>>,
 ) -> Result<Response, FancyError> {

crates/handlers/src/views/index.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use axum::{
-    extract::Extension,
+    extract::State,
     response::{Html, IntoResponse},
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -24,9 +24,9 @@ use mas_templates::{IndexContext, TemplateContext, Templates};
 use sqlx::PgPool;
 
 pub async fn get(
-    Extension(templates): Extension<Templates>,
-    Extension(url_builder): Extension<UrlBuilder>,
-    Extension(pool): Extension<PgPool>,
+    State(templates): State<Templates>,
+    State(url_builder): State<UrlBuilder>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
 ) -> Result<impl IntoResponse, FancyError> {
     let mut conn = pool.acquire().await?;

crates/handlers/src/views/login.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use axum::{
-    extract::{Extension, Form, Query},
+    extract::{Form, Query, State},
     response::{Html, IntoResponse, Response},
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -44,8 +44,8 @@ impl ToFormState for LoginForm {
 
 #[tracing::instrument(skip(templates, pool, cookie_jar))]
 pub(crate) async fn get(
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     Query(query): Query<OptionalPostAuthAction>,
     cookie_jar: PrivateCookieJar<Encrypter>,
 ) -> Result<Response, FancyError> {
@@ -74,8 +74,8 @@ pub(crate) async fn get(
 }
 
 pub(crate) async fn post(
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     Query(query): Query<OptionalPostAuthAction>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Form(form): Form<ProtectedForm<LoginForm>>,

crates/handlers/src/views/logout.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use axum::{
-    extract::{Extension, Form},
+    extract::{Form, State},
     response::IntoResponse,
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -27,7 +27,7 @@ use mas_storage::user::end_session;
 use sqlx::PgPool;
 
 pub(crate) async fn post(
-    Extension(pool): Extension<PgPool>,
+    State(pool): State<PgPool>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Form(form): Form<ProtectedForm<Option<PostAuthAction>>>,
 ) -> Result<impl IntoResponse, FancyError> {

crates/handlers/src/views/reauth.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use axum::{
-    extract::{Extension, Form, Query},
+    extract::{Form, Query, State},
     response::{Html, IntoResponse, Response},
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -36,8 +36,8 @@ pub(crate) struct ReauthForm {
 }
 
 pub(crate) async fn get(
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     Query(query): Query<OptionalPostAuthAction>,
     cookie_jar: PrivateCookieJar<Encrypter>,
 ) -> Result<Response, FancyError> {
@@ -75,7 +75,7 @@ pub(crate) async fn get(
 }
 
 pub(crate) async fn post(
-    Extension(pool): Extension<PgPool>,
+    State(pool): State<PgPool>,
     Query(query): Query<OptionalPostAuthAction>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Form(form): Form<ProtectedForm<ReauthForm>>,

crates/handlers/src/views/register.rs

@@ -18,7 +18,7 @@ use std::{str::FromStr, sync::Arc};
 
 use argon2::Argon2;
 use axum::{
-    extract::{Extension, Form, Query},
+    extract::{Form, Query, State},
     response::{Html, IntoResponse, Response},
 };
 use axum_extra::extract::PrivateCookieJar;
@@ -57,8 +57,8 @@ impl ToFormState for RegisterForm {
 }
 
 pub(crate) async fn get(
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     Query(query): Query<OptionalPostAuthAction>,
     cookie_jar: PrivateCookieJar<Encrypter>,
 ) -> Result<Response, FancyError> {
@@ -87,10 +87,10 @@ pub(crate) async fn get(
 }
 
 pub(crate) async fn post(
-    Extension(mailer): Extension<Mailer>,
-    Extension(policy_factory): Extension<Arc<PolicyFactory>>,
-    Extension(templates): Extension<Templates>,
-    Extension(pool): Extension<PgPool>,
+    State(mailer): State<Mailer>,
+    State(policy_factory): State<Arc<PolicyFactory>>,
+    State(templates): State<Templates>,
+    State(pool): State<PgPool>,
     Query(query): Query<OptionalPostAuthAction>,
     cookie_jar: PrivateCookieJar<Encrypter>,
     Form(form): Form<ProtectedForm<RegisterForm>>,