Pass claims as context instead of globals during attr mapping

This is in preparation for also fetching the claims from the userinfo
endpoint

Author: sandhose

crates/handlers/src/upstream_oauth2/callback.rs

@@ -30,7 +30,9 @@ use thiserror::Error;
 use ulid::Ulid;
 
 use super::{
-    cache::LazyProviderInfos, client_credentials_for_provider, template::environment,
+    cache::LazyProviderInfos,
+    client_credentials_for_provider,
+    template::{environment, AttributeMappingContext},
     UpstreamSessionsCookie,
 };
 use crate::{impl_from_error_for_route, upstream_oauth2::cache::MetadataCache};
@@ -269,15 +271,13 @@ pub(crate) async fn handler(
 
     let (_header, id_token) = id_token.ok_or(RouteError::MissingIDToken)?.into_parts();
 
-    let env = {
-        let mut env = environment();
-        env.add_global("user", minijinja::Value::from_serialize(&id_token));
-        env.add_global(
-            "extra_callback_parameters",
-            minijinja::Value::from_serialize(&extra_callback_parameters),
-        );
-        env
-    };
+    let mut context = AttributeMappingContext::new().with_id_token_claims(id_token);
+    if let Some(extra_callback_parameters) = extra_callback_parameters.clone() {
+        context = context.with_extra_callback_parameters(extra_callback_parameters);
+    }
+    let context = context.build();
+
+    let env = environment();
 
     let template = provider
         .claims_imports
@@ -286,7 +286,7 @@ pub(crate) async fn handler(
         .as_deref()
         .unwrap_or("{{ user.sub }}");
     let subject = env
-        .render_str(template, ())
+        .render_str(template, context)
         .map_err(RouteError::ExtractSubject)?;
 
     if subject.is_empty() {

crates/handlers/src/upstream_oauth2/link.rs

@@ -38,7 +38,10 @@ use thiserror::Error;
 use tracing::warn;
 use ulid::Ulid;
 
-use super::{template::environment, UpstreamSessionsCookie};
+use super::{
+    template::{environment, AttributeMappingContext},
+    UpstreamSessionsCookie,
+};
 use crate::{
     impl_from_error_for_route, views::shared::OptionalPostAuthAction, PreferredLanguage, SiteConfig,
 };
@@ -130,9 +133,10 @@ impl IntoResponse for RouteError {
 fn render_attribute_template(
     environment: &Environment,
     template: &str,
+    context: &minijinja::Value,
     required: bool,
 ) -> Result<Option<String>, RouteError> {
-    match environment.render_str(template, ()) {
+    match environment.render_str(template, context) {
         Ok(value) if value.is_empty() => {
             if required {
                 return Err(RouteError::RequiredAttributeEmpty {
@@ -320,32 +324,27 @@ 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 id_token = upstream_session
-                .id_token()
-                .map(Jwt::<'_, minijinja::Value>::try_from)
-                .transpose()?;
+            let id_token = upstream_session.id_token().map(Jwt::try_from).transpose()?;
 
             let provider = repo
                 .upstream_oauth_provider()
                 .lookup(link.provider_id)
                 .await?
                 .ok_or(RouteError::ProviderNotFound)?;
 
-            let payload = id_token
-                .map(|id_token| id_token.into_parts().1)
-                .unwrap_or_default();
-
             let ctx = UpstreamRegister::default();
 
-            let env = {
-                let mut e = environment();
-                e.add_global("user", payload);
-                e.add_global(
-                    "extra_callback_parameters",
-                    minijinja::Value::from_serialize(upstream_session.extra_callback_parameters()),
-                );
-                e
-            };
+            let env = environment();
+
+            let mut context = AttributeMappingContext::new();
+            if let Some(id_token) = id_token {
+                let (_, payload) = id_token.into_parts();
+                context = context.with_id_token_claims(payload);
+            }
+            if let Some(extra_callback_parameters) = upstream_session.extra_callback_parameters() {
+                context = context.with_extra_callback_parameters(extra_callback_parameters.clone());
+            }
+            let context = context.build();
 
             let ctx = if provider.claims_imports.displayname.ignore() {
                 ctx
@@ -360,6 +359,7 @@ pub(crate) async fn get(
                 match render_attribute_template(
                     &env,
                     template,
+                    &context,
                     provider.claims_imports.displayname.is_required(),
                 )? {
                     Some(value) => ctx
@@ -381,6 +381,7 @@ pub(crate) async fn get(
                 match render_attribute_template(
                     &env,
                     template,
+                    &context,
                     provider.claims_imports.email.is_required(),
                 )? {
                     Some(value) => ctx.with_email(value, provider.claims_imports.email.is_forced()),
@@ -401,6 +402,7 @@ pub(crate) async fn get(
                 match render_attribute_template(
                     &env,
                     template,
+                    &context,
                     provider.claims_imports.localpart.is_required(),
                 )? {
                     Some(localpart) => {
@@ -561,37 +563,31 @@ pub(crate) async fn post(
             let import_display_name = import_display_name.is_some();
             let accept_terms = accept_terms.is_some();
 
-            let id_token = upstream_session
-                .id_token()
-                .map(Jwt::<'_, minijinja::Value>::try_from)
-                .transpose()?;
+            let id_token = upstream_session.id_token().map(Jwt::try_from).transpose()?;
 
             let provider = repo
                 .upstream_oauth_provider()
                 .lookup(link.provider_id)
                 .await?
                 .ok_or(RouteError::ProviderNotFound)?;
 
-            let payload = id_token
-                .map(|id_token| id_token.into_parts().1)
-                .unwrap_or_default();
+            // Let's try to import the claims from the ID token
+            let env = environment();
 
-            // Is the email verified according to the upstream provider?
-            let provider_email_verified = payload
-                .get_item(&minijinja::Value::from("email_verified"))
-                .map(|v| v.is_true())
-                .unwrap_or(false);
+            let mut context = AttributeMappingContext::new();
+            if let Some(id_token) = id_token {
+                let (_, payload) = id_token.into_parts();
+                context = context.with_id_token_claims(payload);
+            }
+            if let Some(extra_callback_parameters) = upstream_session.extra_callback_parameters() {
+                context = context.with_extra_callback_parameters(extra_callback_parameters.clone());
+            }
+            let context = context.build();
 
-            // Let's try to import the claims from the ID token
-            let env = {
-                let mut e = environment();
-                e.add_global("user", payload);
-                e.add_global(
-                    "extra_callback_parameters",
-                    minijinja::Value::from_serialize(upstream_session.extra_callback_parameters()),
-                );
-                e
-            };
+            // Is the email verified according to the upstream provider?
+            let provider_email_verified = env
+                .render_str("{{ user.email_verified | string }}", &context)
+                .map_or(false, |v| v == "true");
 
             // Create a template context in case we need to re-render because of an error
             let ctx = UpstreamRegister::default();
@@ -611,6 +607,7 @@ pub(crate) async fn post(
                 render_attribute_template(
                     &env,
                     template,
+                    &context,
                     provider.claims_imports.displayname.is_required(),
                 )?
             } else {
@@ -637,6 +634,7 @@ pub(crate) async fn post(
                 render_attribute_template(
                     &env,
                     template,
+                    &context,
                     provider.claims_imports.email.is_required(),
                 )?
             } else {
@@ -660,6 +658,7 @@ pub(crate) async fn post(
                 render_attribute_template(
                     &env,
                     template,
+                    &context,
                     provider.claims_imports.email.is_required(),
                 )?
             } else {

crates/handlers/src/upstream_oauth2/template.rs

@@ -7,7 +7,76 @@
 use std::{collections::HashMap, sync::Arc};
 
 use base64ct::{Base64, Base64Unpadded, Base64Url, Base64UrlUnpadded, Encoding};
-use minijinja::{Environment, Error, ErrorKind, Value};
+use minijinja::{
+    value::{Enumerator, Object},
+    Environment, Error, ErrorKind, Value,
+};
+
+/// Context passed to the attribute mapping template
+///
+/// The variables available in the template are:
+/// - `user`: claims for the user, currently from the ID token. Later, we'll
+///   also allow importing from the userinfo endpoint
+/// - `id_token_claims`: claims from the ID token
+/// - `extra_callback_parameters`: extra parameters passed to the callback
+#[derive(Debug, Default)]
+pub(crate) struct AttributeMappingContext {
+    id_token_claims: Option<HashMap<String, serde_json::Value>>,
+    extra_callback_parameters: Option<serde_json::Value>,
+}
+
+impl AttributeMappingContext {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn with_id_token_claims(
+        mut self,
+        id_token_claims: HashMap<String, serde_json::Value>,
+    ) -> Self {
+        self.id_token_claims = Some(id_token_claims);
+        self
+    }
+
+    pub fn with_extra_callback_parameters(
+        mut self,
+        extra_callback_parameters: serde_json::Value,
+    ) -> Self {
+        self.extra_callback_parameters = Some(extra_callback_parameters);
+        self
+    }
+
+    pub fn build(self) -> Value {
+        Value::from_object(self)
+    }
+}
+
+impl Object for AttributeMappingContext {
+    fn get_value(self: &Arc<Self>, name: &Value) -> Option<Value> {
+        match name.as_str()? {
+            "user" | "id_token_claims" => self.id_token_claims.as_ref().map(Value::from_serialize),
+            "extra_callback_parameters" => self
+                .extra_callback_parameters
+                .as_ref()
+                .map(Value::from_serialize),
+            _ => None,
+        }
+    }
+
+    fn enumerate(self: &Arc<Self>) -> Enumerator {
+        match (
+            self.id_token_claims.is_some(),
+            self.extra_callback_parameters.is_some(),
+        ) {
+            (true, true) => {
+                Enumerator::Str(&["user", "id_token_claims", "extra_callback_parameters"])
+            }
+            (true, false) => Enumerator::Str(&["user", "id_token_claims"]),
+            (false, true) => Enumerator::Str(&["extra_callback_parameters"]),
+            (false, false) => Enumerator::Str(&["user"]),
+        }
+    }
+}
 
 fn b64decode(value: &str) -> Result<Value, Error> {
     // We're not too concerned about the performance of this filter, so we'll just