axum-login updated (lazy sessions)

GrantSparks · Jan 1, 2024 · eded705 · eded705
1 parent a79eb30
commit eded705
Show file tree

Hide file tree

Showing 7 changed files with 49 additions and 26 deletions.
diff --git a/build.rs b/build.rs
@@ -1,5 +1,7 @@
 // generated by `sqlx migrate build-script`
 fn main() {
-    // trigger recompilation when a new migration is added
-    println!("cargo:rerun-if-changed=grafton-server/migrations");
+    // trigger recompilation when a sql migration is added or changed.
+    // this is only useful if you have a persistent database.
+    // We currently don't have a persistent database, so we don't need this.
+    // println!("cargo:rerun-if-changed=grafton-server/migrations");
 }
diff --git a/examples/chatgpt-plugin/public/www/index.html b/examples/chatgpt-plugin/public/www/index.html
@@ -10,7 +10,9 @@
   </head>
 
   <body>
-    <div class="container">&nbsp;</div>
+    <div class="container">
+      <a href="/protected">Login to portal</a>
+    </div>
     <script src="script.js"></script>
   </body>
 </html>
diff --git a/grafton-server/Cargo.toml b/grafton-server/Cargo.toml
@@ -15,9 +15,10 @@ include = [
     "LICENSE",
     "README.md",
     "build.rs",
-    "src/**/*",
+    "examples/**/*",
     "migrations/**/*",
     "templates/**/*",
+    "src/**/*",
 ]
 
 [profile.dev.build-override]
@@ -54,7 +55,7 @@ version = "*"
 features = ["macros"] # Hack to enable macros in axum-login crate dependency
 
 [dependencies.axum-login]
-version = "0.10.2" # Pin version until next update
+version = ">=0.11.2"
 
 [dependencies.reqwest]
 features = ["json"]

diff --git a/grafton-server/src/web/oauth2/callback.rs b/grafton-server/src/web/oauth2/callback.rs
@@ -50,6 +50,7 @@ mod get {
 
         let old_state = session
             .get(CSRF_STATE_KEY)
+            .await
             .map_err(|_| AppError::SessionStateError("Failed to retrieve CSRF state".to_string()))?
             .ok_or(AppError::MissingCSRFState)?;
 
@@ -97,7 +98,7 @@ mod get {
             ));
         }
 
-        match session.remove::<String>(NEXT_URL_KEY) {
+        match session.remove::<String>(NEXT_URL_KEY).await {
             Ok(Some(next)) if !next.is_empty() => Ok(Redirect::to(&next).into_response()),
             Ok(Some(_)) | Ok(None) => Ok(Redirect::to("/").into_response()),
             Err(e) => {

diff --git a/grafton-server/src/web/oauth2/login.rs b/grafton-server/src/web/oauth2/login.rs
@@ -37,6 +37,7 @@ mod post {
         axum::{response::Redirect, Form},
         tower_sessions::Session,
     };
+    use tracing::warn;
 
     use crate::{web::oauth2::CSRF_STATE_KEY, AppError, AuthSession};
 
@@ -51,11 +52,17 @@ mod post {
     ) -> Result<impl IntoResponse, AppError> {
         match auth_session.backend.authorize_url(provider.clone()) {
             Ok((url, token)) => {
-                if let Err(e) = session.insert(CSRF_STATE_KEY, token.secret()) {
+                if let Err(e) = session.insert(CSRF_STATE_KEY, token.secret()).await {
                     error!("Error serializing CSRF token: {:?}", e);
                     return Err(AppError::SerializationError(e.to_string()));
                 }
-                if let Err(e) = session.insert(NEXT_URL_KEY, next) {
+
+                // Check if 'next' is None or an empty String
+                if next.as_ref().map(|s| s.is_empty()).unwrap_or(false) {
+                    warn!("NEXT_URL_KEY is empty or null");
+                }
+
+                if let Err(e) = session.insert(NEXT_URL_KEY, next).await {
                     error!("Error serializing next URL: {:?}", e);
                     return Err(AppError::SerializationError(e.to_string()));
                 }

diff --git a/grafton-server/src/web/oauth2/logout.rs b/grafton-server/src/web/oauth2/logout.rs
@@ -15,7 +15,7 @@ mod get {
     use crate::AuthSession;
 
     pub async fn logout(mut auth_session: AuthSession) -> impl IntoResponse {
-        match auth_session.logout() {
+        match auth_session.logout().await {
             Ok(_) => Redirect::to("/").into_response(),
             Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
         }

diff --git a/grafton-server/src/web/protected_app.rs b/grafton-server/src/web/protected_app.rs
@@ -4,16 +4,16 @@ use {
     askama_axum::IntoResponse,
     axum_login::{
         axum::{
-            error_handling::HandleErrorLayer, http::StatusCode, middleware::from_fn,
-            middleware::Next, response::Redirect, BoxError,
+            error_handling::HandleErrorLayer, extract::OriginalUri, http::StatusCode,
+            middleware::from_fn, middleware::Next, response::Redirect, BoxError,
         },
         tower_sessions::{MemoryStore, SessionManagerLayer},
-        urlencoding, AuthManagerLayerBuilder,
+        url_with_redirect_query, AuthManagerLayerBuilder,
     },
     oauth2::{basic::BasicClient, AuthUrl, TokenUrl},
     sqlx::SqlitePool,
     tower::ServiceBuilder,
-    tracing::{debug, error, info, warn},
+    tracing::{debug, error, info},
 };
 
 use crate::{
@@ -126,20 +126,30 @@ impl ProtectedApp {
             .layer(AuthManagerLayerBuilder::new(backend, self.session_layer).build());
 
         let login_url = Arc::new(self.login_url);
-        let auth_middleware = from_fn(move |auth_session: AuthSession, req, next: Next| {
-            let login_url_clone = login_url.clone();
-            async move {
-                if auth_session.user.is_some() {
-                    next.run(req).await
-                } else {
-                    warn!("Unauthorized access attempt, redirecting to login");
-                    let uri = req.uri().to_string();
-                    let next = urlencoding::encode(&uri);
-                    let redirect_url = format!("{}?next={}", login_url_clone, next);
-                    Redirect::temporary(&redirect_url).into_response()
+        let auth_middleware = from_fn(
+            move |auth_session: AuthSession,
+                  OriginalUri(original_uri): OriginalUri,
+                  req,
+                  next: Next| {
+                let login_url_clone = login_url.clone();
+                async move {
+                    if auth_session.user.is_some() {
+                        next.run(req).await
+                    } else {
+                        match url_with_redirect_query(&login_url_clone, "next", original_uri) {
+                            Ok(login_url) => {
+                                Redirect::temporary(&login_url.to_string()).into_response()
+                            }
+
+                            Err(err) => {
+                                error!(err = %err);
+                                StatusCode::INTERNAL_SERVER_ERROR.into_response()
+                            }
+                        }
+                    }
                 }
-            }
-        });
+            },
+        );
         info!("Auth middleware created");
 
         let router = match self.protected_router {