Tidy up upstream linking templates

crates/handlers/src/upstream_oauth2/link.rs

@@ -34,7 +34,10 @@ use mas_storage::{
     },
     LookupResultExt,
 };
-use mas_templates::{EmptyContext, TemplateContext, Templates, UpstreamExistingLinkContext};
+use mas_templates::{
+    EmptyContext, TemplateContext, Templates, UpstreamExistingLinkContext, UpstreamRegister,
+    UpstreamSuggestLink,
+};
 use serde::Deserialize;
 use sqlx::PgPool;
 use thiserror::Error;
@@ -174,7 +177,7 @@ pub(crate) async fn get(
 
         (Some(user_session), None) => {
             // Session not linked, but user logged in: suggest linking account
-            let ctx = EmptyContext
+            let ctx = UpstreamSuggestLink::new(link.id)
                 .with_session(user_session)
                 .with_csrf(csrf_token.form_value());
 
@@ -193,7 +196,7 @@ pub(crate) async fn get(
         (None, None) => {
             // Session not linked and used not logged in: suggest creating an
             // account or logging in an existing user
-            let ctx = EmptyContext.with_csrf(csrf_token.form_value());
+            let ctx = UpstreamRegister::new(link.id).with_csrf(csrf_token.form_value());
 
             templates.render_upstream_oauth2_do_register(&ctx).await?
         }

crates/handlers/src/views/logout.rs

@@ -48,10 +48,10 @@ pub(crate) async fn post(
     txn.commit().await?;
 
     let destination = if let Some(action) = form {
-        mas_router::Login::and_then(action)
+        action.go_next()
     } else {
-        mas_router::Login::default()
+        mas_router::Login::default().go()
     };
 
-    Ok((cookie_jar, destination.go()))
+    Ok((cookie_jar, destination))
 }

crates/handlers/src/views/shared.rs

@@ -47,12 +47,27 @@ impl OptionalPostAuthAction {
                 let grant = Box::new(grant.into());
                 Ok(Some(PostAuthContext::ContinueAuthorizationGrant { grant }))
             }
+
             Some(PostAuthAction::ContinueCompatSsoLogin { data }) => {
                 let login = get_compat_sso_login_by_id(conn, *data).await?;
                 let login = Box::new(login.into());
                 Ok(Some(PostAuthContext::ContinueCompatSsoLogin { login }))
             }
+
             Some(PostAuthAction::ChangePassword) => Ok(Some(PostAuthContext::ChangePassword)),
+
+            Some(PostAuthAction::LinkUpstream { id }) => {
+                let (link, provider_id, _user_id) =
+                    mas_storage::upstream_oauth2::lookup_link(&mut *conn, *id).await?;
+
+                let provider =
+                    mas_storage::upstream_oauth2::lookup_provider(&mut *conn, provider_id).await?;
+
+                let provider = Box::new(provider);
+                let link = Box::new(link);
+                Ok(Some(PostAuthContext::LinkUpstream { provider, link }))
+            }
+
             None => Ok(None),
         }
     }

crates/router/src/endpoints.rs

@@ -31,6 +31,10 @@ pub enum PostAuthAction {
         data: Ulid,
     },
     ChangePassword,
+    LinkUpstream {
+        #[serde_as(as = "DisplayFromStr")]
+        id: Ulid,
+    },
 }
 
 impl PostAuthAction {
@@ -49,6 +53,7 @@ impl PostAuthAction {
             Self::ContinueAuthorizationGrant { data } => ContinueAuthorizationGrant(*data).go(),
             Self::ContinueCompatSsoLogin { data } => CompatLoginSsoComplete::new(*data, None).go(),
             Self::ChangePassword => AccountPassword.go(),
+            Self::LinkUpstream { id } => UpstreamOAuth2Link::new(*id).go(),
         }
     }
 }

crates/templates/src/context.rs

@@ -18,10 +18,10 @@
 
 use chrono::Utc;
 use mas_data_model::{
-    AuthorizationGrant, BrowserSession, CompatSsoLogin, CompatSsoLoginState, StorageBackend, User,
-    UserEmail, UserEmailVerification,
+    AuthorizationGrant, BrowserSession, CompatSsoLogin, CompatSsoLoginState, StorageBackend,
+    UpstreamOAuthLink, UpstreamOAuthProvider, User, UserEmail, UserEmailVerification,
 };
-use mas_router::PostAuthAction;
+use mas_router::{PostAuthAction, Route};
 use serde::{ser::SerializeStruct, Deserialize, Serialize};
 use ulid::Ulid;
 use url::Url;
@@ -263,6 +263,15 @@ pub enum PostAuthContext {
 
     /// Change the account password
     ChangePassword,
+
+    /// Link an upstream account
+    LinkUpstream {
+        /// The upstream provider
+        provider: Box<UpstreamOAuthProvider>,
+
+        /// The link
+        link: Box<UpstreamOAuthLink>,
+    },
 }
 
 /// Context used by the `login.html` template
@@ -779,6 +788,57 @@ impl TemplateContext for UpstreamExistingLinkContext {
     }
 }
 
+/// Context used by the `pages/upstream_oauth2/suggest_link.html`
+/// templates
+#[derive(Serialize)]
+pub struct UpstreamSuggestLink {
+    post_logout_action: PostAuthAction,
+}
+
+impl UpstreamSuggestLink {
+    /// Constructs a new context with an existing linked user
+    #[must_use]
+    pub fn new(link_id: Ulid) -> Self {
+        let post_logout_action = PostAuthAction::LinkUpstream { id: link_id };
+        Self { post_logout_action }
+    }
+}
+
+impl TemplateContext for UpstreamSuggestLink {
+    fn sample(_now: chrono::DateTime<Utc>) -> Vec<Self>
+    where
+        Self: Sized,
+    {
+        vec![Self::new(Ulid::nil())]
+    }
+}
+
+/// Context used by the `pages/upstream_oauth2/do_register.html`
+/// templates
+#[derive(Serialize)]
+pub struct UpstreamRegister {
+    login_link: String,
+}
+
+impl UpstreamRegister {
+    /// Constructs a new context with an existing linked user
+    #[must_use]
+    pub fn new(link_id: Ulid) -> Self {
+        let action = PostAuthAction::LinkUpstream { id: link_id };
+        let login_link = mas_router::Login::and_then(action).relative_url().into();
+        Self { login_link }
+    }
+}
+
+impl TemplateContext for UpstreamRegister {
+    fn sample(_now: chrono::DateTime<Utc>) -> Vec<Self>
+    where
+        Self: Sized,
+    {
+        vec![Self::new(Ulid::nil())]
+    }
+}
+
 /// Context used by the `form_post.html` template
 #[derive(Serialize)]
 pub struct FormPostContext<T> {

crates/templates/src/lib.rs

@@ -49,7 +49,8 @@ pub use self::{
         EmailVerificationContext, EmailVerificationPageContext, EmptyContext, ErrorContext,
         FormPostContext, IndexContext, LoginContext, LoginFormField, PolicyViolationContext,
         PostAuthContext, ReauthContext, ReauthFormField, RegisterContext, RegisterFormField,
-        TemplateContext, UpstreamExistingLinkContext, WithCsrf, WithOptionalSession, WithSession,
+        TemplateContext, UpstreamExistingLinkContext, UpstreamRegister, UpstreamSuggestLink,
+        WithCsrf, WithOptionalSession, WithSession,
     },
     forms::{FieldError, FormError, FormField, FormState, ToFormState},
 };
@@ -233,13 +234,13 @@ register_templates! {
     pub fn render_upstream_oauth2_link_mismatch(WithCsrf<WithSession<UpstreamExistingLinkContext>>) { "pages/upstream_oauth2/link_mismatch.html" }
 
     /// Render the upstream suggest link message
-    pub fn render_upstream_oauth2_suggest_link(WithCsrf<WithSession<EmptyContext>>) { "pages/upstream_oauth2/suggest_link.html" }
+    pub fn render_upstream_oauth2_suggest_link(WithCsrf<WithSession<UpstreamSuggestLink>>) { "pages/upstream_oauth2/suggest_link.html" }
 
     /// Render the upstream login screen
     pub fn render_upstream_oauth2_do_login(WithCsrf<UpstreamExistingLinkContext>) { "pages/upstream_oauth2/do_login.html" }
 
     /// Render the upstream register screen
-    pub fn render_upstream_oauth2_do_register(WithCsrf<EmptyContext>) { "pages/upstream_oauth2/do_register.html" }
+    pub fn render_upstream_oauth2_do_register(WithCsrf<UpstreamRegister>) { "pages/upstream_oauth2/do_register.html" }
 }
 
 impl Templates {