Record extra query parameters during upstream callback

And make them available in the templates.
This is useful to get the user display name for Sign-in with Apple

Author: sandhose

Cargo.lock

@@ -15,6 +15,7 @@ workspace = true
 chrono.workspace = true
 thiserror.workspace = true
 serde.workspace = true
+serde_json.workspace = true
 url.workspace = true
 crc = "3.2.1"
 ulid.workspace = true

crates/data-model/Cargo.toml

@@ -19,12 +19,14 @@ pub enum UpstreamOAuthAuthorizationSessionState {
         completed_at: DateTime<Utc>,
         link_id: Ulid,
         id_token: Option<String>,
+        extra_callback_parameters: Option<serde_json::Value>,
     },
     Consumed {
         completed_at: DateTime<Utc>,
         consumed_at: DateTime<Utc>,
         link_id: Ulid,
         id_token: Option<String>,
+        extra_callback_parameters: Option<serde_json::Value>,
     },
 }
 
@@ -42,12 +44,14 @@ impl UpstreamOAuthAuthorizationSessionState {
         completed_at: DateTime<Utc>,
         link: &UpstreamOAuthLink,
         id_token: Option<String>,
+        extra_callback_parameters: Option<serde_json::Value>,
     ) -> Result<Self, InvalidTransitionError> {
         match self {
             Self::Pending => Ok(Self::Completed {
                 completed_at,
                 link_id: link.id,
                 id_token,
+                extra_callback_parameters,
             }),
             Self::Completed { .. } | Self::Consumed { .. } => Err(InvalidTransitionError),
         }
@@ -67,11 +71,13 @@ impl UpstreamOAuthAuthorizationSessionState {
                 completed_at,
                 link_id,
                 id_token,
+                extra_callback_parameters,
             } => Ok(Self::Consumed {
                 completed_at,
                 link_id,
                 consumed_at,
                 id_token,
+                extra_callback_parameters,
             }),
             Self::Pending | Self::Consumed { .. } => Err(InvalidTransitionError),
         }
@@ -124,6 +130,27 @@ impl UpstreamOAuthAuthorizationSessionState {
         }
     }
 
+    /// Get the extra query parameters that were sent to the upstream provider.
+    ///
+    /// Returns `None` if the upstream OAuth 2.0 authorization session state is
+    /// not [`Pending`].
+    ///
+    /// [`Pending`]: UpstreamOAuthAuthorizationSessionState::Pending
+    #[must_use]
+    pub fn extra_callback_parameters(&self) -> Option<&serde_json::Value> {
+        match self {
+            Self::Pending => None,
+            Self::Completed {
+                extra_callback_parameters,
+                ..
+            }
+            | Self::Consumed {
+                extra_callback_parameters,
+                ..
+            } => extra_callback_parameters.as_ref(),
+        }
+    }
+
     /// Get the time at which the upstream OAuth 2.0 authorization session was
     /// consumed.
     ///
@@ -201,8 +228,11 @@ impl UpstreamOAuthAuthorizationSession {
         completed_at: DateTime<Utc>,
         link: &UpstreamOAuthLink,
         id_token: Option<String>,
+        extra_callback_parameters: Option<serde_json::Value>,
     ) -> Result<Self, InvalidTransitionError> {
-        self.state = self.state.complete(completed_at, link, id_token)?;
+        self.state =
+            self.state
+                .complete(completed_at, link, id_token, extra_callback_parameters)?;
         Ok(self)
     }
 

crates/data-model/src/upstream_oauth2/session.rs

@@ -48,6 +48,9 @@ pub struct Params {
 enum CodeOrError {
     Code {
         code: String,
+
+        #[serde(flatten)]
+        extra_callback_parameters: Option<serde_json::Value>,
     },
     Error {
         error: ClientErrorCode,
@@ -201,7 +204,7 @@ pub(crate) async fn handler(
     }
 
     // Let's extract the code from the params, and return if there was an error
-    let code = match params.code_or_error {
+    let (code, extra_callback_parameters) = match params.code_or_error {
         CodeOrError::Error {
             error,
             error_description,
@@ -212,7 +215,10 @@ pub(crate) async fn handler(
                 error_description,
             })
         }
-        CodeOrError::Code { code } => code,
+        CodeOrError::Code {
+            code,
+            extra_callback_parameters,
+        } => (code, extra_callback_parameters),
     };
 
     let mut lazy_metadata = LazyProviderInfos::new(&metadata_cache, &provider, &client);
@@ -266,6 +272,10 @@ pub(crate) async fn handler(
     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
     };
 
@@ -299,7 +309,13 @@ pub(crate) async fn handler(
 
     let session = repo
         .upstream_oauth_session()
-        .complete_with_link(&clock, session, &link, response.id_token)
+        .complete_with_link(
+            &clock,
+            session,
+            &link,
+            response.id_token,
+            extra_callback_parameters,
+        )
         .await?;
 
     let cookie_jar = sessions_cookie

crates/handlers/src/upstream_oauth2/callback.rs

@@ -340,6 +340,10 @@ pub(crate) async fn get(
             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
             };
 
@@ -582,6 +586,10 @@ pub(crate) async fn post(
             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
             };
 
@@ -945,7 +953,13 @@ mod tests {
 
         let session = repo
             .upstream_oauth_session()
-            .complete_with_link(&state.clock, session, &link, Some(id_token.into_string()))
+            .complete_with_link(
+                &state.clock,
+                session,
+                &link,
+                Some(id_token.into_string()),
+                None,
+            )
             .await
             .unwrap();
 

crates/handlers/src/upstream_oauth2/link.rs

@@ -68,6 +68,18 @@ fn string(value: &Value) -> String {
     value.to_string()
 }
 
+fn from_json(value: &str) -> Result<Value, minijinja::Error> {
+    let value: serde_json::Value = serde_json::from_str(value).map_err(|e| {
+        minijinja::Error::new(
+            minijinja::ErrorKind::InvalidOperation,
+            "Failed to decode JSON",
+        )
+        .with_source(e)
+    })?;
+
+    Ok(Value::from_serialize(value))
+}
+
 pub fn environment() -> Environment<'static> {
     let mut env = Environment::new();
 
@@ -77,6 +89,7 @@ pub fn environment() -> Environment<'static> {
     env.add_filter("b64encode", b64encode);
     env.add_filter("tlvdecode", tlvdecode);
     env.add_filter("string", string);
+    env.add_filter("from_json", from_json);
 
     env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
 

crates/handlers/src/upstream_oauth2/template.rs

@@ -0,0 +1,9 @@
+-- Copyright 2024 New Vector Ltd.
+--
+-- SPDX-License-Identifier: AGPL-3.0-only
+-- Please see LICENSE in the repository root for full details.
+
+-- Add a column to the upstream_oauth_authorization_sessions table to store
+-- extra query parameters
+ALTER TABLE "upstream_oauth_authorization_sessions"
+    ADD COLUMN "extra_callback_parameters" JSONB;

…f6e24a91775ceb877735508c1d5b2300d64.json …f966ed71a0ed59be0d48ec66fedf64e707d.jsoncrates/storage-pg/.sqlx/query-b9875a270f7e753e48075ccae233df6e24a91775ceb877735508c1d5b2300d64.json renamed to crates/storage-pg/.sqlx/query-5516235e0983fb64d18e82dbe3e34f966ed71a0ed59be0d48ec66fedf64e707d.json crates/storage-pg/.sqlx/query-b9875a270f7e753e48075ccae233df6e24a91775ceb877735508c1d5b2300d64.json renamed to crates/storage-pg/.sqlx/query-5516235e0983fb64d18e82dbe3e34f966ed71a0ed59be0d48ec66fedf64e707d.json

@@ -145,7 +145,7 @@ mod tests {
 
         let session = repo
             .upstream_oauth_session()
-            .complete_with_link(&clock, session, &link, None)
+            .complete_with_link(&clock, session, &link, None, None)
             .await
             .unwrap();
         // Reload the session