1 files changed, 21 insertions, 22 deletions
diff --git a/atuin-server/src/router.rs b/atuin-server/src/router.rs
index 08eea996..c4c15f1f 100644
--- a/atuin-server/src/router.rs
+++ b/atuin-server/src/router.rs
@@ -1,12 +1,12 @@
 use async_trait::async_trait;
 use axum::{
-    extract::{FromRequest, RequestParts},
-    handler::Handler,
+    extract::FromRequestParts,
     response::IntoResponse,
     routing::{get, post},
-    Extension, Router,
+    Router,
 };
 use eyre::Result;
+use http::request::Parts;
 use tower::ServiceBuilder;
 use tower_http::trace::TraceLayer;
 
@@ -17,20 +17,15 @@ use super::{
 use crate::{models::User, settings::Settings};
 
 #[async_trait]
-impl<B> FromRequest<B> for User
-where
-    B: Send,
-{
+impl FromRequestParts<AppState> for User {
     type Rejection = http::StatusCode;
 
-    async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
-        let postgres = req
-            .extensions()
-            .get::<Postgres>()
-            .ok_or(http::StatusCode::INTERNAL_SERVER_ERROR)?;
-
+    async fn from_request_parts(
+        req: &mut Parts,
+        state: &AppState,
+    ) -> Result<Self, Self::Rejection> {
         let auth_header = req
-            .headers()
+            .headers
             .get(http::header::AUTHORIZATION)
             .ok_or(http::StatusCode::FORBIDDEN)?;
         let auth_header = auth_header
@@ -44,7 +39,8 @@ where
             return Err(http::StatusCode::FORBIDDEN);
         }
 
-        let user = postgres
+        let user = state
+            .postgres
             .get_session_user(token)
             .await
             .map_err(|_| http::StatusCode::FORBIDDEN)?;
@@ -56,6 +52,13 @@ where
 async fn teapot() -> impl IntoResponse {
     (http::StatusCode::IM_A_TEAPOT, "☕")
 }
+
+#[derive(Clone)]
+pub struct AppState {
+    pub postgres: Postgres,
+    pub settings: Settings,
+}
+
 pub fn router(postgres: Postgres, settings: Settings) -> Router {
     let routes = Router::new()
         .route("/", get(handlers::index))
@@ -73,11 +76,7 @@ pub fn router(postgres: Postgres, settings: Settings) -> Router {
     } else {
         Router::new().nest(path, routes)
     }
-    .fallback(teapot.into_service())
-    .layer(
-        ServiceBuilder::new()
-            .layer(TraceLayer::new_for_http())
-            .layer(Extension(postgres))
-            .layer(Extension(settings)),
-    )
+    .fallback(teapot)
+    .with_state(AppState { postgres, settings })
+    .layer(ServiceBuilder::new().layer(TraceLayer::new_for_http()))
 }