Save and handle post login redirect

Cargo.toml

@@ -33,7 +33,7 @@ web-sys = { version = "0.3", features = [
 ] }
 
 openidconnect = { version = "3.0", optional = true }
-yew-nested-router = { version = ">=0.5, <0.7", optional = true }
+yew-nested-router = { version = ">=0.6.3, <0.7", optional = true }
 
 [features]
 # Enable for Yew nested router support

src/agent/error.rs

@@ -31,3 +31,9 @@ impl From<OAuth2Error> for OAuth2Context {
         OAuth2Context::Failed(err.to_string())
     }
 }
+
+impl OAuth2Error {
+    pub(crate) fn storage_key_empty(key: impl Display) -> Self {
+        Self::Storage(format!("Missing value for key: {key}"))
+    }
+}

src/agent/mod.rs

@@ -4,21 +4,23 @@ pub mod client;
 mod config;
 mod error;
 mod ops;
-mod state;
+pub mod state;
 
 pub use client::*;
 pub use config::*;
 pub use error::*;
 pub use ops::*;
 
 use crate::context::{Authentication, OAuth2Context, Reason};
-use gloo_storage::{SessionStorage, Storage};
+use gloo_storage::{errors::StorageError, SessionStorage, Storage};
 use gloo_timers::callback::Timeout;
 use gloo_utils::{history, window};
 use js_sys::Date;
+use log::error;
 use num_traits::cast::ToPrimitive;
 use reqwest::Url;
 use state::*;
+use std::fmt::Display;
 use std::{collections::HashMap, fmt::Debug, time::Duration};
 use tokio::sync::mpsc::{channel, Receiver, Sender};
 use wasm_bindgen::JsValue;
@@ -39,7 +41,7 @@ use yew::Callback;
 /// # let url = Url::parse("https://example.com").unwrap();
 /// let opts = LoginOptions::default().with_redirect_url(url);
 /// ```
-#[derive(Debug, Clone, Default, PartialEq, Eq)]
+#[derive(Debug, Clone, Default)]
 #[non_exhaustive]
 pub struct LoginOptions {
     pub(crate) query: HashMap<String, String>,
@@ -48,6 +50,11 @@ pub struct LoginOptions {
     ///
     /// If this field is empty, the current URL is used as a redirect URL.
     pub(crate) redirect_url: Option<Url>,
+
+    /// Defines callback used for post-login redirect.
+    ///
+    /// If None, disables post-login redirect
+    pub(crate) post_login_redirect_callback: Option<Callback<String>>,
 }
 
 impl LoginOptions {
@@ -68,10 +75,30 @@ impl LoginOptions {
         self
     }
 
+    /// Define the redirect URL
     pub fn with_redirect_url(mut self, redirect_url: impl Into<Url>) -> Self {
         self.redirect_url = Some(redirect_url.into());
         self
     }
+
+    /// Define callback for post-login redirect
+    pub fn with_redirect_callback(mut self, redirect_callback: Callback<String>) -> Self {
+        self.post_login_redirect_callback = Some(redirect_callback);
+        self
+    }
+
+    /// Use `yew-nested-route` history api for post-login redirect callback
+    #[cfg(feature = "yew-nested-router")]
+    pub fn with_nested_router_redirect(mut self) -> Self {
+        let callback = Callback::from(|url: String| {
+            if yew_nested_router::History::push_state(JsValue::null(), &url).is_err() {
+                error!("Unable to redirect");
+            }
+        });
+
+        self.post_login_redirect_callback = Some(callback);
+        self
+    }
 }
 
 /// Options for the logout process
@@ -255,7 +282,11 @@ where
                     let detected = self.detect_state().await;
                     log::debug!("Detected state: {detected:?}");
                     match detected {
-                        Ok(true) => {}
+                        Ok(true) => {
+                            if let Err(e) = self.post_login_redirect() {
+                                error!("Post-login redirect failed: {e}");
+                            }
+                        }
                         Ok(false) => {
                             self.update_state(
                                 OAuth2Context::NotAuthenticated {
@@ -362,6 +393,24 @@ where
         }
     }
 
+    fn post_login_redirect(&self) -> Result<(), OAuth2Error> {
+        let config = self.config.as_ref().ok_or(OAuth2Error::NotInitialized)?;
+        let Some(redirect_callback) = config
+            .options
+            .as_ref()
+            .and_then(|opts| opts.post_login_redirect_callback.clone())
+        else {
+            return Ok(());
+        };
+        let Some(url) = Self::get_from_store_optional(STORAGE_KEY_POST_LOGIN_URL)? else {
+            return Ok(());
+        };
+        SessionStorage::delete(STORAGE_KEY_POST_LOGIN_URL);
+        redirect_callback.emit(url);
+
+        Ok(())
+    }
+
     fn update_state_from_result(
         &mut self,
         result: Result<(OAuth2Context, C::SessionState), OAuth2Error>,
@@ -405,17 +454,18 @@ where
         }
     }
 
-    fn get_from_store<K: AsRef<str>>(key: K) -> Result<String, OAuth2Error> {
-        let value: String = SessionStorage::get(key.as_ref())
-            .map_err(|err| OAuth2Error::Storage(err.to_string()))?;
+    fn get_from_store<K: AsRef<str> + Display>(key: K) -> Result<String, OAuth2Error> {
+        Self::get_from_store_optional(&key)?.ok_or_else(|| OAuth2Error::storage_key_empty(key))
+    }
 
-        if value.is_empty() {
-            Err(OAuth2Error::Storage(format!(
-                "Missing value for key: {}",
-                key.as_ref()
-            )))
-        } else {
-            Ok(value)
+    fn get_from_store_optional<K: AsRef<str> + Display>(
+        key: K,
+    ) -> Result<Option<String>, OAuth2Error> {
+        match SessionStorage::get::<String>(key.as_ref()) {
+            Err(StorageError::KeyNotFound(_)) => Ok(None),
+            Err(err) => Err(OAuth2Error::Storage(err.to_string())),
+            Ok(value) if value.is_empty() => Err(OAuth2Error::storage_key_empty(key)),
+            Ok(value) => Ok(Some(value)),
         }
     }
 
@@ -460,16 +510,23 @@ where
         let client = self.client.as_ref().ok_or(OAuth2Error::NotInitialized)?;
         let config = self.config.as_ref().ok_or(OAuth2Error::NotInitialized)?;
 
+        let current_url = Self::current_url().map_err(OAuth2Error::StartLogin)?;
+
         // take the parameter value first, then the agent configured value, then fall back to the default
-        let redirect_url = match options.redirect_url.or_else(|| {
-            config
-                .options
-                .as_ref()
-                .and_then(|opts| opts.redirect_url.clone())
-        }) {
-            Some(redirect_url) => redirect_url,
-            None => Self::current_url().map_err(OAuth2Error::StartLogin)?,
-        };
+        let redirect_url = options
+            .redirect_url
+            .or_else(|| {
+                config
+                    .options
+                    .as_ref()
+                    .and_then(|opts| opts.redirect_url.clone())
+            })
+            .unwrap_or_else(|| current_url.clone());
+
+        if redirect_url != current_url {
+            SessionStorage::set(STORAGE_KEY_POST_LOGIN_URL, current_url)
+                .map_err(|err| OAuth2Error::StartLogin(err.to_string()))?;
+        }
 
         let login_context = client.make_login_context(config, redirect_url.clone())?;
 

src/agent/state.rs

@@ -1,6 +1,7 @@
 pub const STORAGE_KEY_CSRF_TOKEN: &str = "ctron/oauth2/csrfToken";
 pub const STORAGE_KEY_LOGIN_STATE: &str = "ctron/oauth2/loginState";
 pub const STORAGE_KEY_REDIRECT_URL: &str = "ctron/oauth2/redirectUrl";
+pub const STORAGE_KEY_POST_LOGIN_URL: &str = "ctron/oauth2/postLoginUrl";
 
 #[derive(Debug)]
 pub struct State {