feat: with return url

Signed-off-by: kjuulh <contact@kjuulh.io>

crates/nefarious-login/src/auth.rs

@@ -14,9 +14,14 @@ use crate::{
 
 #[async_trait]
 pub trait Auth {
-    async fn login(&self) -> anyhow::Result<Url>;
+    async fn login(&self, return_url: Option<String>) -> anyhow::Result<Url>;
     async fn login_token(&self, user: &str, password: &str) -> anyhow::Result<IdToken>;
-    async fn login_authorized(&self, code: &str, state: &str) -> anyhow::Result<(HeaderMap, Url)>;
+    async fn login_authorized(
+        &self,
+        code: &str,
+        state: &str,
+        return_url: Option<String>,
+    ) -> anyhow::Result<(HeaderMap, Url)>;
     async fn get_user_from_session(&self, cookie: &str) -> anyhow::Result<User>;
 }
 
@@ -87,12 +92,17 @@ pub static COOKIE_NAME: &str = "SESSION";
 
 #[async_trait]
 impl Auth for ZitadelAuthService {
-    async fn login(&self) -> anyhow::Result<Url> {
-        let authorize_url = self.oauth.authorize_url().await?;
+    async fn login(&self, return_url: Option<String>) -> anyhow::Result<Url> {
+        let authorize_url = self.oauth.authorize_url(return_url).await?;
 
         Ok(authorize_url)
     }
-    async fn login_authorized(&self, code: &str, _state: &str) -> anyhow::Result<(HeaderMap, Url)> {
+    async fn login_authorized(
+        &self,
+        code: &str,
+        _state: &str,
+        return_path: Option<String>,
+    ) -> anyhow::Result<(HeaderMap, Url)> {
         let token = self.oauth.exchange(code).await?;
         let id_token = self.introspection.get_id_token(token.as_str()).await?;
         let cookie_value = self.session.insert_user("user", id_token).await?;
@@ -102,9 +112,14 @@ impl Auth for ZitadelAuthService {
         let mut headers = HeaderMap::new();
         headers.insert(SET_COOKIE, cookie.parse().unwrap());
 
+        let mut return_url = self.config.return_url.clone();
+        if let Some(return_path) = return_path {
+            return_url.push_str(&format!("?returnPath={return_path}"));
+        }
+
         Ok((
             headers,
-            Url::parse(&self.config.return_url)
+            Url::parse(&return_url)
                 .context("failed to parse login_authorized zitadel return url")?,
         ))
     }
@@ -126,7 +141,7 @@ pub struct NoopAuthService {
 
 #[async_trait]
 impl Auth for NoopAuthService {
-    async fn login(&self) -> anyhow::Result<Url> {
+    async fn login(&self, return_url: Option<String>) -> anyhow::Result<Url> {
         let url = Url::parse(&format!(
             "{}/auth/authorized?code=noop&state=noop",
             self.config
@@ -142,6 +157,7 @@ impl Auth for NoopAuthService {
         &self,
         _code: &str,
         _state: &str,
+        _return_url: Option<String>,
     ) -> anyhow::Result<(HeaderMap, Url)> {
         let cookie_value = self
             .session

crates/nefarious-login/src/axum.rs

@@ -18,7 +18,6 @@ use crate::session::User;
 
 #[derive(Debug, Deserialize)]
 pub struct ZitadelAuthParams {
-    #[allow(dead_code)]
     return_url: Option<String>,
 }
 
@@ -51,7 +50,7 @@ where
 pub async fn zitadel_auth(
     State(auth_service): State<AuthService>,
 ) -> Result<impl IntoResponse, ErrorResponse> {
-    let url = auth_service.login().await.into_response()?;
+    let url = auth_service.login(None).await.into_response()?;
 
     Ok(Redirect::to(url.as_ref()))
 }
@@ -61,6 +60,8 @@ pub async fn zitadel_auth(
 pub struct AuthRequest {
     code: String,
     state: String,
+    #[serde(alias = "returnUrl")]
+    return_url: Option<String>,
 }
 
 pub async fn login_authorized(
@@ -68,7 +69,7 @@ pub async fn login_authorized(
     State(auth_service): State<AuthService>,
 ) -> Result<impl IntoResponse, ErrorResponse> {
     let (headers, url) = auth_service
-        .login_authorized(&query.code, &query.state)
+        .login_authorized(&query.code, &query.state, query.return_url)
         .await
         .into_response()?;
 

crates/nefarious-login/src/oauth.rs

@@ -31,7 +31,7 @@ impl Deref for OAuth {
 #[async_trait]
 pub trait OAuthClient {
     async fn get_token(&self) -> anyhow::Result<()>;
-    async fn authorize_url(&self) -> anyhow::Result<Url>;
+    async fn authorize_url(&self, return_url: Option<String>) -> anyhow::Result<Url>;
     async fn exchange(&self, code: &str) -> anyhow::Result<String>;
 }
 

crates/nefarious-login/src/oauth/noop.rs

@@ -10,7 +10,7 @@ impl OAuthClient for NoopOAuthClient {
     async fn get_token(&self) -> anyhow::Result<()> {
         Ok(())
     }
-    async fn authorize_url(&self) -> anyhow::Result<Url> {
+    async fn authorize_url(&self, return_url: Option<String>) -> anyhow::Result<Url> {
         Ok(Url::parse("http://localhost:3000/auth/zitadel").unwrap())
     }
 

crates/nefarious-login/src/oauth/zitadel.rs

@@ -104,15 +104,28 @@ impl OAuthClient for ZitadelOAuthClient {
     async fn get_token(&self) -> anyhow::Result<()> {
         Ok(())
     }
-    async fn authorize_url(&self) -> anyhow::Result<Url> {
-        let (auth_url, _csrf_token) = self
+    async fn authorize_url(&self, return_url: Option<String>) -> anyhow::Result<Url> {
+        let req = self
             .client
             .authorize_url(CsrfToken::new_random)
             .add_scope(Scope::new("identify".to_string()))
             .add_scope(Scope::new("openid".to_string()))
             .add_scope(Scope::new("email".to_string()))
-            .add_scope(Scope::new("profile".to_string()))
-            .url();
+            .add_scope(Scope::new("profile".to_string()));
+
+        let req = {
+            if let Some(return_url) = return_url {
+                let mut redirect_url = self.client.redirect_url().unwrap().as_str().to_string();
+
+                redirect_url.push_str(&format!("?returnUrl={}", return_url));
+
+                req.set_redirect_uri(std::borrow::Cow::Owned(RedirectUrl::new(redirect_url)?))
+            } else {
+                req
+            }
+        };
+
+        let (auth_url, _csrf_token) = req.url();
 
         Ok(auth_url)
     }