Record the decoded ID token claims on upstream auth sessions

Author: sandhose

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

@@ -19,6 +19,7 @@ pub enum UpstreamOAuthAuthorizationSessionState {
         completed_at: DateTime<Utc>,
         link_id: Ulid,
         id_token: Option<String>,
+        id_token_claims: Option<serde_json::Value>,
         extra_callback_parameters: Option<serde_json::Value>,
         userinfo: Option<serde_json::Value>,
     },
@@ -27,6 +28,7 @@ pub enum UpstreamOAuthAuthorizationSessionState {
         consumed_at: DateTime<Utc>,
         link_id: Ulid,
         id_token: Option<String>,
+        id_token_claims: Option<serde_json::Value>,
         extra_callback_parameters: Option<serde_json::Value>,
         userinfo: Option<serde_json::Value>,
     },
@@ -35,6 +37,7 @@ pub enum UpstreamOAuthAuthorizationSessionState {
         consumed_at: Option<DateTime<Utc>>,
         unlinked_at: DateTime<Utc>,
         id_token: Option<String>,
+        id_token_claims: Option<serde_json::Value>,
     },
 }
 
@@ -52,6 +55,7 @@ impl UpstreamOAuthAuthorizationSessionState {
         completed_at: DateTime<Utc>,
         link: &UpstreamOAuthLink,
         id_token: Option<String>,
+        id_token_claims: Option<serde_json::Value>,
         extra_callback_parameters: Option<serde_json::Value>,
         userinfo: Option<serde_json::Value>,
     ) -> Result<Self, InvalidTransitionError> {
@@ -60,6 +64,7 @@ impl UpstreamOAuthAuthorizationSessionState {
                 completed_at,
                 link_id: link.id,
                 id_token,
+                id_token_claims,
                 extra_callback_parameters,
                 userinfo,
             }),
@@ -83,13 +88,15 @@ impl UpstreamOAuthAuthorizationSessionState {
                 completed_at,
                 link_id,
                 id_token,
+                id_token_claims,
                 extra_callback_parameters,
                 userinfo,
             } => Ok(Self::Consumed {
                 completed_at,
                 link_id,
                 consumed_at,
                 id_token,
+                id_token_claims,
                 extra_callback_parameters,
                 userinfo,
             }),
@@ -146,6 +153,29 @@ impl UpstreamOAuthAuthorizationSessionState {
         }
     }
 
+    /// Get the ID token claims for the upstream OAuth 2.0 authorization
+    /// session.
+    ///
+    /// Returns `None` if the upstream OAuth 2.0 authorization session state is
+    /// not [`Pending`].
+    ///
+    /// [`Pending`]: UpstreamOAuthAuthorizationSessionState::Pending
+    #[must_use]
+    pub fn id_token_claims(&self) -> Option<&serde_json::Value> {
+        match self {
+            Self::Pending => None,
+            Self::Completed {
+                id_token_claims, ..
+            }
+            | Self::Consumed {
+                id_token_claims, ..
+            }
+            | Self::Unlinked {
+                id_token_claims, ..
+            } => id_token_claims.as_ref(),
+        }
+    }
+
     /// Get the extra query parameters that were sent to the upstream provider.
     ///
     /// Returns `None` if the upstream OAuth 2.0 authorization session state is
@@ -277,13 +307,15 @@ impl UpstreamOAuthAuthorizationSession {
         completed_at: DateTime<Utc>,
         link: &UpstreamOAuthLink,
         id_token: Option<String>,
+        id_token_claims: Option<serde_json::Value>,
         extra_callback_parameters: Option<serde_json::Value>,
         userinfo: Option<serde_json::Value>,
     ) -> Result<Self, InvalidTransitionError> {
         self.state = self.state.complete(
             completed_at,
             link,
             id_token,
+            id_token_claims,
             extra_callback_parameters,
             userinfo,
         )?;

crates/handlers/src/admin/v1/upstream_oauth_links/delete.rs

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

crates/handlers/src/upstream_oauth2/callback.rs

@@ -312,6 +312,7 @@ pub(crate) async fn handler(
     .await?;
 
     let mut jwks = None;
+    let mut id_token_claims = None;
 
     let mut context = AttributeMappingContext::new();
     if let Some(id_token) = token_response.id_token.as_ref() {
@@ -337,6 +338,14 @@ pub(crate) async fn handler(
 
         let (_headers, mut claims) = id_token.into_parts();
 
+        // Save a copy of the claims for later; the claims extract methods
+        // remove them from the map, and we want to store the original claims.
+        // We anyway need this to be a serde_json::Value
+        id_token_claims = Some(
+            serde_json::to_value(&claims)
+                .expect("serializing a HashMap<String, Value> into a Value should never fail"),
+        );
+
         // Access token hash must match.
         mas_jose::claims::AT_HASH
             .extract_optional_with_options(
@@ -472,6 +481,7 @@ pub(crate) async fn handler(
             session,
             &link,
             token_response.id_token,
+            id_token_claims,
             params.extra_callback_parameters,
             userinfo,
         )

crates/handlers/src/upstream_oauth2/link.rs

@@ -934,7 +934,7 @@ mod tests {
             ..UpstreamOAuthProviderClaimsImports::default()
         };
 
-        let id_token = serde_json::json!({
+        let id_token_claims = serde_json::json!({
             "preferred_username": "john",
             "email": "[email protected]",
             "email_verified": true,
@@ -953,7 +953,8 @@ mod tests {
             .signing_key_for_alg(&JsonWebSignatureAlg::Rs256)
             .unwrap();
         let header = JsonWebSignatureHeader::new(JsonWebSignatureAlg::Rs256);
-        let id_token = Jwt::sign_with_rng(&mut rng, header, id_token, &signer).unwrap();
+        let id_token =
+            Jwt::sign_with_rng(&mut rng, header, id_token_claims.clone(), &signer).unwrap();
 
         // Provision a provider and a link
         let mut repo = state.repository().await.unwrap();
@@ -1022,6 +1023,7 @@ mod tests {
                 session,
                 &link,
                 Some(id_token.into_string()),
+                Some(id_token_claims),
                 None,
                 None,
             )

crates/storage-pg/.sqlx/query-5f5245ace61b896f92be78ab4fef701b37c9e3c2f4a332f418b9fb2625a0fe3f.json

@@ -0,0 +1,8 @@
+-- Copyright 2025 New Vector Ltd.
+--
+-- SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
+-- Please see LICENSE in the repository root for full details.
+
+-- This is the decoded claims from the ID token stored as JSONB
+ALTER TABLE upstream_oauth_authorization_sessions
+    ADD COLUMN id_token_claims JSONB;

crates/storage-pg/.sqlx/query-37a124678323380357fa9d1375fd125fb35476ac3008e5adbd04a761d5edcd42.json renamed to crates/storage-pg/.sqlx/query-e62d043f86e7232e6e9433631f8273e7ed0770c81071cf1f17516d3a45881ae9.json

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